diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 1a3288b9..b2daab71 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,8 +24,24 @@ jobs: fail-fast: false matrix: include: + - os: ubuntu-24.04 + label: "ubuntu-24.04-gcc-14" + cmake_args: >- + -G Ninja + -DCMAKE_C_COMPILER:STRING=gcc-14 + -DCMAKE_CXX_COMPILER:STRING=g++-14 + -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed + tmpdir: "/tmp" + - os: ubuntu-24.04 + label: "ubuntu-24.04-gcc-13" + cmake_args: >- + -G Ninja + -DCMAKE_C_COMPILER:STRING=gcc-13 + -DCMAKE_CXX_COMPILER:STRING=g++-13 + -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed + tmpdir: "/tmp" - os: ubuntu-22.04 - additional_label: "-gcc-12" + label: "ubuntu-22.04-gcc-12" cmake_args: >- -G Ninja -DCMAKE_C_COMPILER:STRING=gcc-12 @@ -33,7 +49,7 @@ jobs: -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed tmpdir: "/tmp" - os: ubuntu-22.04 - additional_label: "-gcc-11" + label: "ubuntu-22.04-gcc-11" cmake_args: >- -G Ninja -DCMAKE_C_COMPILER:STRING=gcc-11 @@ -41,23 +57,47 @@ jobs: -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed tmpdir: "/tmp" - os: ubuntu-22.04 - additional_label: "-gcc-10" + label: "ubuntu-22.04-gcc-10" cmake_args: >- -G Ninja -DCMAKE_C_COMPILER:STRING=gcc-10 -DCMAKE_CXX_COMPILER:STRING=g++-10 -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed tmpdir: "/tmp" - - os: ubuntu-22.04 - additional_label: "-gcc-9" + - os: ubuntu-24.04 + label: "ubuntu-24.04-clang-18" cmake_args: >- -G Ninja - -DCMAKE_C_COMPILER:STRING=gcc-9 - -DCMAKE_CXX_COMPILER:STRING=g++-9 - -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed + -DCMAKE_C_COMPILER:STRING=clang-18 + -DCMAKE_CXX_COMPILER:STRING=clang++-18 + -DMUJOCO_HARDEN:BOOL=ON + tmpdir: "/tmp" + - os: ubuntu-24.04 + label: "ubuntu-24.04-clang-17" + cmake_args: >- + -G Ninja + -DCMAKE_C_COMPILER:STRING=clang-17 + -DCMAKE_CXX_COMPILER:STRING=clang++-17 + -DMUJOCO_HARDEN:BOOL=ON + tmpdir: "/tmp" + - os: ubuntu-24.04 + label: "ubuntu-24.04-clang-16" + cmake_args: >- + -G Ninja + -DCMAKE_C_COMPILER:STRING=clang-16 + -DCMAKE_CXX_COMPILER:STRING=clang++-16 + -DMUJOCO_HARDEN:BOOL=ON tmpdir: "/tmp" - os: ubuntu-22.04 - additional_label: "-clang-14" + label: "ubuntu-22.04-clang-15" + cmake_args: >- + -G Ninja + -DCMAKE_C_COMPILER:STRING=clang-15 + -DCMAKE_CXX_COMPILER:STRING=clang++-15 + -DMUJOCO_HARDEN:BOOL=ON + tmpdir: "/tmp" + - os: ubuntu-22.04 + label: "ubuntu-22.04-clang-14" cmake_args: >- -G Ninja -DCMAKE_C_COMPILER:STRING=clang-14 @@ -65,49 +105,33 @@ jobs: -DMUJOCO_HARDEN:BOOL=ON tmpdir: "/tmp" - os: ubuntu-22.04 - additional_label: "-clang-13" + label: "ubuntu-22.04-clang-13" cmake_args: >- -G Ninja -DCMAKE_C_COMPILER:STRING=clang-13 -DCMAKE_CXX_COMPILER:STRING=clang++-13 -DMUJOCO_HARDEN:BOOL=ON tmpdir: "/tmp" - - os: ubuntu-20.04 - additional_label: "-clang-12" - cmake_args: >- - -G Ninja - -DCMAKE_C_COMPILER:STRING=clang-12 - -DCMAKE_CXX_COMPILER:STRING=clang++-12 - -DMUJOCO_HARDEN:BOOL=ON - tmpdir: "/tmp" - - os: ubuntu-20.04 - additional_label: "-clang-11" - cmake_args: >- - -G Ninja - -DCMAKE_C_COMPILER:STRING=clang-11 - -DCMAKE_CXX_COMPILER:STRING=clang++-11 - -DMUJOCO_HARDEN:BOOL=ON - tmpdir: "/tmp" - - os: ubuntu-20.04 - additional_label: "-clang-10" - cmake_args: >- - -G Ninja - -DCMAKE_C_COMPILER:STRING=clang-10 - -DCMAKE_CXX_COMPILER:STRING=clang++-10 - -DMUJOCO_HARDEN:BOOL=ON - tmpdir: "/tmp" - - os: macos-13 + - os: macos-15 + label: "macos-15-arm64" cmake_args: >- -G Ninja -DMUJOCO_HARDEN:BOOL=ON tmpdir: "/tmp" - - os: windows-2022 + - os: macos-15-large + label: "macos-15-x86_64" cmake_args: >- - -DCMAKE_SYSTEM_VERSION="10.0.22621.0" + -G Ninja + -DMUJOCO_HARDEN:BOOL=ON + tmpdir: "/tmp" + - os: windows-2025 + label: "windows-2025" + cmake_args: >- + -DCMAKE_SYSTEM_VERSION="10.0.26100.0" cmake_build_args: "-- -m" tmpdir: "C:/Temp" - name: "${{ matrix.os }}${{ matrix.additional_label }}" + name: "${{ matrix.label }}" runs-on: ${{ matrix.os }} steps: @@ -179,7 +203,7 @@ jobs: cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && - cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin + cp lib/libsdf_plugin.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin - name: Copy plugins (Windows) if: ${{ runner.os == 'Windows' }} working-directory: build @@ -273,7 +297,7 @@ jobs: CHATMSG_AUTHOR_NAME: ${{ github.event.head_commit.author.name }} CHATMSG_AUTHOR_EMAIL: ${{ github.event.head_commit.author.email }} CHATMSG_COMMIT_MESSAGE: ${{ github.event.head_commit.message }} - CHATMSG_JOB_ID: ${{ matrix.os }}${{ matrix.additional_label }} + CHATMSG_JOB_ID: ${{ matrix.label }} if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API_URL != '' }} run: | CHATMSG="$(cat <<-'EOF' | python3 @@ -288,7 +312,7 @@ jobs: email=env('CHATMSG_AUTHOR_EMAIL'), msg=env('CHATMSG_COMMIT_MESSAGE').replace('```', '') ) - text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name}<{email}>\n\n{msg}```'.format(**data) + text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name} <{email}>\n\n{msg}```'.format(**data) print(json.dumps({'text' : text})) EOF )" && diff --git a/CMakeLists.txt b/CMakeLists.txt index 62a396ea..3cab120d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON) project( mujoco - VERSION 3.3.1 + VERSION 3.3.3 DESCRIPTION "MuJoCo Physics Simulator" HOMEPAGE_URL "https://mujoco.org" ) @@ -42,6 +42,14 @@ option(MUJOCO_BUILD_EXAMPLES "Build samples for MuJoCo" ON) option(MUJOCO_BUILD_SIMULATE "Build simulate library for MuJoCo" ON) option(MUJOCO_BUILD_TESTS "Build tests for MuJoCo" ON) option(MUJOCO_TEST_PYTHON_UTIL "Build and test utility libraries for Python bindings" ON) +option(MUJOCO_BUILD_USD_PLUGINS "Build OpenUSD plugins" OFF) + +# USD libs to compile against. +set(MUJOCO_USD_ALLOWED_TARGET_VALUES "USD" "Houdini") +set(MUJOCO_USD_TARGET "USD" CACHE STRING "Select the USD target for the project.") +set_property(CACHE MUJOCO_USD_TARGET + PROPERTY STRINGS ${MUJOCO_USD_ALLOWED_TARGET_VALUES} +) if(APPLE AND (MUJOCO_BUILD_EXAMPLES OR MUJOCO_BUILD_SIMULATE)) enable_language(OBJC) @@ -181,6 +189,10 @@ if(MUJOCO_BUILD_EXAMPLES) add_subdirectory(sample) endif() +if(MUJOCO_BUILD_USD_PLUGINS) + add_subdirectory(src/experimental/usd/plugins) +endif() + if(BUILD_TESTING AND MUJOCO_BUILD_TESTS) enable_testing() add_subdirectory(test) diff --git a/README.md b/README.md index 935c80b7..62d57c19 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,8 @@ running on Google Colab: - The **introductory** tutorial teaches MuJoCo basics: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb) + - The **Model Editing** tutorial shows how to create and edit models procedurally: + [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb) - The **rollout** tutorial shows how to use the multithreaded `rollout` module: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb) - The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake index 004aedc4..9d5dc01c 100644 --- a/cmake/MujocoDependencies.cmake +++ b/cmake/MujocoDependencies.cmake @@ -39,12 +39,12 @@ set(MUJOCO_DEP_VERSION_qhull CACHE STRING "Version of `qhull` to be fetched." ) set(MUJOCO_DEP_VERSION_Eigen3 - 66f7f51b7e069d0a03a21157fa60b24aece69aeb + 464c1d097891a1462ab28bf8bb763c1683883892 CACHE STRING "Version of `Eigen3` to be fetched." ) set(MUJOCO_DEP_VERSION_abseil - 9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0 + d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1 CACHE STRING "Version of `abseil` to be fetched." ) diff --git a/dist/mujoco.rc b/dist/mujoco.rc index 171b9810..6191980f 100644 --- a/dist/mujoco.rc +++ b/dist/mujoco.rc @@ -1,6 +1,6 @@ 1 VERSIONINFO -FILEVERSION 3,3,1,0 -PRODUCTVERSION 3,3,1,0 +FILEVERSION 3,3,3,0 +PRODUCTVERSION 3,3,3,0 FILEOS 0x4 FILETYPE 0x1 { @@ -9,9 +9,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "3.3.1" + VALUE "ProductVersion", "3.3.3" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "3.3.1" + VALUE "FileVersion", "3.3.3" VALUE "InternalName", "mujoco.dll" VALUE "OriginalFilename", "mujoco.dll" VALUE "CompanyName", "Google DeepMind" diff --git a/dist/simulate.rc b/dist/simulate.rc index 7a2492ad..fca19ae2 100644 --- a/dist/simulate.rc +++ b/dist/simulate.rc @@ -1,8 +1,8 @@ MUJOCO ICON "mujoco.ico" 1 VERSIONINFO -FILEVERSION 3,3,1,0 -PRODUCTVERSION 3,3,1,0 +FILEVERSION 3,3,3,0 +PRODUCTVERSION 3,3,3,0 FILEOS 0x4 FILETYPE 0x1 { @@ -11,9 +11,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "3.3.1" + VALUE "ProductVersion", "3.3.3" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "3.3.1" + VALUE "FileVersion", "3.3.3" VALUE "InternalName", "simulate.exe" VALUE "OriginalFilename", "simulate.exe" VALUE "CompanyName", "Google DeepMind" diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index c116257a..11077782 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -448,7 +448,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - 200 - The number of iterations where solver statistics can be stored in ``mjData.solver``. This array is used to store diagnostic information about each iteration of the constraint solver. - The actual number of iterations is given by ``mjData.solver_iter``. + The actual number of iterations is given by ``mjData.solver_niter``. * - ``mjNISLAND`` - 20 - The number of islands for which solver statistics can be stored in ``mjData.solver``. This array is @@ -517,7 +517,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - Maximum number of UI rectangles. Defined in `mjui.h `_. * - ``mjVERSION_HEADER`` - - 331 + - 333 - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning but for the compiled library. diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index 7e2befa0..f1df0615 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -165,6 +165,17 @@ values are used in ``m->cam_mode`` and ``m->light_mode``. .. mujoco-include:: mjtCamLight +.. _mjtLightType: + +mjtLightType +~~~~~~~~~~~~ + +The type of a light source describing how its position, orientation and other properties will interact with the +objects in the scene. These values are used in ``m->light_type``. + +.. mujoco-include:: mjtLightType + + .. _mjtTexture: mjtTexture @@ -357,6 +368,26 @@ last argument to :ref:`mj_local2global`. .. mujoco-include:: mjtSameFrame +.. _mjtFlexSelf: + +mjtFlexSelf +~~~~~~~~~~~~ + +Types of flex self-collisions midphase. + +.. mujoco-include:: mjtFlexSelf + + +.. _mjtSDFType: + +mjtSDFType +~~~~~~~~~~~ + +Formulas used to combine SDFs when calling mjc_distance and mjc_gradient. + +.. mujoco-include:: mjtSDFType + + .. _tyDataEnums: Data @@ -883,7 +914,7 @@ mjSolverStat This is the data structure holding information about one solver iteration. ``mjData.solver`` is a preallocated array of mjSolverStat data structures, one for each iteration of the solver, up to a maximum of mjNSOLVER. The actual number -of solver iterations is given by ``mjData.solver_iter``. +of solver iterations is given by ``mjData.solver_niter``. .. mujoco-include:: mjSolverStat @@ -967,17 +998,6 @@ This structure contains everything needed to render the 3D scene in OpenGL. .. mujoco-include:: mjvScene -.. _mjvSceneState: - -mjvSceneState -~~~~~~~~~~~~~ - -This structure contains the portions of :ref:`mjModel` and :ref:`mjData` that are required for -various ``mjv_*`` functions. - -.. mujoco-include:: mjvSceneState - - .. _mjvFigure: mjvFigure diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 7ca8f985..2bba2a3c 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -51,6 +51,15 @@ Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multi :ref:`mjModel` instance that takes the edits into account. If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`. +.. _mj_copyBack: + +`mj_copyBack <#mj_copyBack>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mj_copyBack + +Copy real-valued arrays from model to spec, returns 1 on success. + .. _mj_recompile: `mj_recompile <#mj_recompile>`__ @@ -459,7 +468,7 @@ Multiply vector by (inertia matrix)^(1/2). .. mujoco-include:: mj_addM Add inertia matrix to destination matrix. -Destination can be sparse uncompressed, or dense when all int* are NULL +Destination can be sparse or dense when all int* are NULL. .. _mj_applyFT: @@ -748,8 +757,7 @@ Compare forward and inverse dynamics, save results in fwdinv. Sub components ^^^^^^^^^^^^^^ -These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely -that the user will need to call them. +These are sub-components of the simulation pipeline, called internally from the components above. .. _mj_sensorPos: @@ -886,6 +894,18 @@ Compute actuator transmission lengths and moments. Run composite rigid body inertia algorithm (CRB). +.. _mj_makeM: + +`mj_makeM <#mj_makeM>`__ +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mj_makeM + +Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due +to :ref:`tendon armature`. The joint-space inertia matrix is stored in both ``mjData.qM`` and +``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row, +respectively). + .. _mj_factorM: `mj_factorM <#mj_factorM>`__ @@ -1364,6 +1384,15 @@ If the model buffer is unallocated the initial configuration will not be set. Copy mjData. m is only required to contain the size fields from MJMODEL_INTS. +.. _mjv_copyData: + +`mjv_copyData <#mjv_copyData>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjv_copyData + +Copy mjData, skip large arrays not required for visualization. + .. _mj_resetData: `mj_resetData <#mj_resetData>`__ @@ -1995,15 +2024,6 @@ Rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_ Move camera with mouse; action is mjtMouse. -.. _mjv_moveCameraFromState: - -`mjv_moveCameraFromState <#mjv_moveCameraFromState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_moveCameraFromState - -Move camera with mouse given a scene state; action is mjtMouse. - .. _mjv_movePerturb: `mjv_movePerturb <#mjv_movePerturb>`__ @@ -2013,15 +2033,6 @@ Move camera with mouse given a scene state; action is mjtMouse. Move perturb object with mouse; action is mjtMouse. -.. _mjv_movePerturbFromState: - -`mjv_movePerturbFromState <#mjv_movePerturbFromState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_movePerturbFromState - -Move perturb object with mouse given a scene state; action is mjtMouse. - .. _mjv_moveModel: `mjv_moveModel <#mjv_moveModel>`__ @@ -2164,15 +2175,6 @@ Free abstract scene. Update entire scene given model state. -.. _mjv_updateSceneFromState: - -`mjv_updateSceneFromState <#mjv_updateSceneFromState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_updateSceneFromState - -Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings. - .. _mjv_copyModel: `mjv_copyModel <#mjv_copyModel>`__ @@ -2182,42 +2184,6 @@ Update entire scene from a scene state, return the number of new mjWARN_VGEOMFUL Copy mjModel, skip large arrays not required for abstract visualization. -.. _mjv_defaultSceneState: - -`mjv_defaultSceneState <#mjv_defaultSceneState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_defaultSceneState - -Set default scene state. - -.. _mjv_makeSceneState: - -`mjv_makeSceneState <#mjv_makeSceneState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_makeSceneState - -Allocate resources and initialize a scene state object. - -.. _mjv_freeSceneState: - -`mjv_freeSceneState <#mjv_freeSceneState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_freeSceneState - -Free scene state. - -.. _mjv_updateSceneState: - -`mjv_updateSceneState <#mjv_updateSceneState>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjv_updateSceneState - -Update a scene state from model and data. - .. _mjv_addGeoms: `mjv_addGeoms <#mjv_addGeoms>`__ @@ -3808,41 +3774,14 @@ Free all pointers with ``mju_free()``. Attachment ^^^^^^^^^^ -.. _mjs_attachBody: +.. _mjs_attach: -`mjs_attachBody <#mjs_attachBody>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +`mjs_attach <#mjs_attach>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -.. mujoco-include:: mjs_attachBody +.. mujoco-include:: mjs_attach -Attach child body to a parent frame, return the attached body if success or NULL otherwise. - -.. _mjs_attachFrame: - -`mjs_attachFrame <#mjs_attachFrame>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjs_attachFrame - -Attach child frame to a parent body, return the attached frame if success or NULL otherwise. - -.. _mjs_attachToSite: - -`mjs_attachToSite <#mjs_attachToSite>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjs_attachToSite - -Attach child body to a parent site, return the attached body if success or NULL otherwise. - -.. _mjs_attachFrameToSite: - -`mjs_attachFrameToSite <#mjs_attachFrameToSite>`__ -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. mujoco-include:: mjs_attachFrameToSite - -Attach child frame to a parent site, return the attached frame if success or NULL otherwise. +Attach child to a parent, return the attached element if success or NULL otherwise. .. _mjs_detachBody: @@ -4428,6 +4367,15 @@ Get string contents. Get double array contents and optionally its size. +.. _mjs_getPluginAttributes: + +`mjs_getPluginAttributes <#mjs_getPluginAttributes>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_getPluginAttributes + +Get plugin attributes. + .. _SpecUtilities: Spec utilities @@ -4448,7 +4396,7 @@ Set element's default. .. mujoco-include:: mjs_setFrame -Set element's enclosing frame. +Set element's enclosing frame, return 0 on success. .. _mjs_resolveOrientation: @@ -4477,6 +4425,17 @@ Transform body into a frame. Set user payload, overriding the existing value for the specified key if present. +.. _mjs_setUserValueWithCleanup: + +`mjs_setUserValueWithCleanup <#mjs_setUserValueWithCleanup>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_setUserValueWithCleanup + +Set user payload, overriding the existing value for the specified key if +present. This version differs from mjs_setUserValue in that it takes a +cleanup function that will be called when the user payload is deleted. + .. _mjs_getUserValue: `mjs_getUserValue <#mjs_getUserValue>`__ diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index ec72a183..5d00db0a 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -113,8 +113,14 @@ Integrates the simulation state using an implicit-in-velocity integrator (either .. _Subcomponents: -These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely -that the user will need to call them. +These are sub-components of the simulation pipeline, called internally from the components above. + +.. _mj_makeM: + +Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due +to :ref:`tendon armature`. The joint-space inertia matrix is stored in both ``mjData.qM`` and +``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row, +respectively). .. _mj_factorM: diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index ca7258f8..588235ab 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -175,11 +175,10 @@ replicating 200 times, suffixes will be ``000, 001, ...`` etc). All referencing and namespaced appropriately. Detailed examples of models using replicate can be found in the `model/replicate/ `__ directory. -There is a caveat concerning :ref:`keyframes` when using replicate. Since :ref:`mjs_attachFrame` is used to +There is a caveat concerning :ref:`keyframes` when using replicate. Since :ref:`mjs_attach` is used to self-attach multiple times the enclosed kinematic tree, if this tree contains further :ref:`attach` elements, keyframes will not be replicated nor namespaced by :ref:`replicate`, but they will be attached and -namespaced once by the innermost call of :ref:`mjs_attachFrame` or :ref:`mjs_attachBody`. See the limitations discussed -in :ref:`attach`. +namespaced once by the innermost call of :ref:`mjs_attach`. See the limitations discussed in :ref:`attach`. .. _replicate-count: @@ -833,6 +832,10 @@ has any effect. The settings here are global and apply to the entire model. necessary to adjust this attribute and the geom-specific groups so as to exclude world geoms from the inertial computation. +.. _compiler-saveinertial: + +:at:`saveinertial`: :at-val:`[false, true], "false"` + If set to "true", the compiler will save explicit :ref:`inertial ` clauses for all bodies. .. _compiler-lengthrange: @@ -1461,16 +1464,19 @@ still be specified here but this functionality is now deprecated and will be rem :el-prefix:`asset/` |-| **texture** (*) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -| This element creates a texture asset, which is then referenced from a :ref:`material ` asset, which is - finally referenced from a model element that needs to be textured. MuJoCo provides access to the texture mapping - mechanism in OpenGL. Texture coordinates are generated automatically in GL_OBJECT_PLANE mode, using either 2D or cube - mapping. MIP maps are always enabled in GL_LINEAR_MIPMAP_LINEAR mode. The texture color is combined with the object - color in GL_MODULATE mode. The texture data can be loaded from PNG files, with provisions for loading cube and skybox - textures. Alternatively the data can be generated by the compiler as a procedural texture. Because different texture - types require different parameters, only a subset of the attributes below are used for any given texture. -| A second file format is supported for loading textures, in addition to PNG. If the file name extension is - different from .png or .PNG, or if the ``content_type`` attribute is set to ``image/vnd.mujoco.texture``, then MuJoCo - assumes that the texture is in this format. This is a custom binary file format, containing the following data: + This element creates a texture asset, which is then referenced from a :ref:`material ` asset, which + is finally referenced from a model element that needs to be textured. + + The texture data can be loaded from files or can be generated by the compiler as a procedural texture. Because + different texture types require different parameters, only a subset of the attributes below are used for any given + texture. Provisions are provided for loading cube and skybox textures from individual image files. + + Currently, three file formats are supported for loading textures: PNG, KTX, and a custom MuJoCo texture format. The + loader will use the extension of the file name to determine which format to use, defaulting to the custom format if + the extension is not recognized. Alternatively, the content_type attribute can be used to specify the format + explicitly. Only ``image/png``, ``image/ktx``, or ``image/vnd.mujoco.texture`` are supported. + + The custom MuJoCo format is assumed to be a binary file containing the following data: .. code:: Text @@ -1523,13 +1529,19 @@ still be specified here but this functionality is now deprecated and will be rem with the texrepeat attribute of :ref:`material `. The data can be loaded from a single file or created procedurally. +.. _asset-texture-colorspace: + +:at:`colorspace`: :at-val:`[auto, linear, sRGB], "auto"` + This attribute determines the color space of the texture. The default value ``auto`` means that the color space will + be determined from the image file itself. If no color space is defined in the file, then ``linear`` is assumed. + .. _asset-texture-content_type: :at:`content_type`: :at-val:`string, optional` If the file attribute is specified, then this sets the `Media Type `_ (formerly known as MIME types) of the - file to be loaded. Any filename extensions will be ignored. Currently ``image/png`` and ``image/vnd.mujoco.texture`` - are supported. + file to be loaded. Any filename extensions will be ignored. Currently ``image/png``, ``image/ktx``, and + ``image/vnd.mujoco.texture`` are supported. .. _asset-texture-file: @@ -1777,11 +1789,11 @@ properties are grouped together. .. _asset-material-rgba: :at:`rgba`: :at-val:`real(4), "1 1 1 1"` - Color and transparency of the material. All components should be in the range [0 1]. Note that textures are applied - in GL_MODULATE mode, meaning that the texture color and the color specified here are multiplied component-wise. Thus - the default value of "1 1 1 1" has the effect of leaving the texture unchanged. When the material is applied to a - model element which defines its own local rgba attribute, the local definition has precedence. Note that this "local" - definition could in fact come from a defaults class. The remaining material properties always apply. + Color and transparency of the material. All components should be in the range [0 1]. Note that the texture color (if + assigned) and the color specified here are multiplied component-wise. Thus the default value of "1 1 1 1" has the + effect of leaving the texture unchanged. When the material is applied to a model element which defines its own local + rgba attribute, the local definition has precedence. Note that this "local" definition could in fact come from a + defaults class. The remaining material properties always apply. .. _material-layer: @@ -2151,14 +2163,28 @@ rotations as unit quaternions. corresponding to mjModel.qpos_spring is also used to compute the spring reference lengths of all tendons, stored in mjModel.tendon_lengthspring. This is because :ref:`tendons ` can also have springs. +.. image:: images/XMLreference/armature.gif + :width: 40% + :align: right + :class: only-light + :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml +.. image:: images/XMLreference/armature_dark.gif + :width: 40% + :align: right + :class: only-dark + :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml + .. _body-joint-armature: :at:`armature`: :at-val:`real, "0"` Additional inertia associated with movement of the joint that is not due to body mass. This added inertia is usually due to a rotor (a.k.a `armature `__) spinning faster than the - joint itself due to a geared transmission; in this case the added inertia is known as "reflected inertia" and its - value is the rotational inertia of the spinning element multiplied by the square of the gear ratio. The value applies - to all degrees of freedom created by this joint. + joint itself due to a geared transmission. In the illustration, we compare (*left*) a 2-dof system with an armature + body (purple box), coupled with a gear ratio of :math:`3` to the pendulum using a :ref:`joint + equality` constraint, and (*right*) a simple 1-dof pendulum with an equivalent :at:`armature`. + Because the gear ratio appears twice, multiplying both forces and lengths, the effect is known as "reflected + inertia" and the equivalent value is the inertia of the spinning body multiplied by the *square of the gear ratio*, + in this case :math:`9=3^2`. The value applies to all degrees of freedom created by this joint. Besides increasing the realism of joints with geared transmission, positive :at:`armature` significantly improves simulation stability, even for small values, and is a recommended possible fix when encountering stability issues. @@ -2849,11 +2875,15 @@ and the +Y axis points up. Thus the frame position and orientation are the key a ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ This element creates a light, which moves with the body where it is defined. To create a fixed light, define it in the -world body. The lights created here are in addition to the default headlight which is always defined and is adjusted via -the :ref:`visual ` element. MuJoCo relies on the standard lighting model in OpenGL (fixed functionality) -augmented with shadow mapping. The effects of lights are additive, thus adding a light always makes the scene brighter. -The maximum number of lights that can be active simultaneously is 8, counting the headlight. The light is shining along -the direction specified by the dir attribute. It does not have a full spatial frame with three orthogonal axes. +world body. The lights created here are in addition to the headlight which is always defined and is configured via the +:ref:`visual ` element. Lights shine along the direction specified by the dir attribute. They do not have +a full spatial frame with three orthogonal axes. + +By default, MuJoCo uses the standard OpenGL (fixed functional) Phong lighting model for its rendering, with augmented +with shadow mapping. (See the OpenGL documentation for more information, including details about various attributes.) + +MJCF also supports alternative lighting models (e.g. physically-based rendering) by providing additional attributes. +Attributes may be applied or ignored depending on the lighting model being used. .. _body-light-name: @@ -2879,10 +2909,17 @@ the direction specified by the dir attribute. It does not have a full spatial fr This is identical to the target attribute of :ref:`camera ` above. It specifies which body should be targeted in "targetbody" and "targetbodycom" modes. +.. _body-light-type: + +:at:`type`: :at-val:`[spot, directional, point, image], "spot"` + Determines the type of light. Note that some light types may not be supported by some renderers (e.g. only spot and + directional lights are supported by the default native renderer). + .. _body-light-directional: :at:`directional`: :at-val:`[false, true], "false"` - The light is directional if this attribute is "true", otherwise it is a spotlight. + This is a deprecated legacy attribute. Please use :ref:`light ` type instead. If set to "true", and + no type is specified, this will change the light type to be directional. .. _body-light-castshadow: @@ -2902,12 +2939,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr these clipping planes bound the cone or box shadow volume in the light direction. As a result, some shadows (especially those very close to the light) may be clipped. -.. _body-light-bulbradius: - -:at:`radius`: :at-val:`real, "0.02"` - Radius of the light, affects shadow softness. This attribute has no effect in MuJoCo's native renderer, but it can be - useful when rendering scenes with an external renderer. - .. _body-light-active: :at:`active`: :at-val:`[false, true], "true"` @@ -2924,11 +2955,50 @@ the direction specified by the dir attribute. It does not have a full spatial fr :at:`dir`: :at-val:`real(3), "0 0 -1"` Direction of the light. +.. _body-light-diffuse: + +:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"` + The color of the light. For the Phong (default) lighting model, this defines the diffuse color of + the light. + +.. _body-light-texture: + +:at:`texture`: :at-val:`string, optional` + The texture to use for image-based lighting. This is unused by the default Phong lighting model. + +.. _body-light-intensity: + +:at:`intensity`: :at-val:`real, "0.0"` + The intensity of the light source, measured in candela, used for physically-based lighting models. + This is unused by the default Phong lighting model. + +.. _body-light-ambient: + +:at:`ambient`: :at-val:`real(3), "0 0 0"` + The ambient color of the light, used by the default Phong lighting model. + +.. _body-light-specular: + +:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"` + The specular color of the light, used by the default Phong lighting model. + +.. _body-light-range: + +:at:`range`: :at-val:`real, "10.0"` + The effective range of the light. Objects further than this distance from the light position + will not be illuminated by this light. This only applies to spotlights. + +.. _body-light-bulbradius: + +:at:`bulbradius`: :at-val:`real, "0.02"` + The radius of the light source which can affect shadow softness depending on the + renderer. This only applies to spotlights. + .. _body-light-attenuation: :at:`attenuation`: :at-val:`real(3), "1 0 0"` - These are the constant, linear and quadratic attenuation coefficients in OpenGL. The default corresponds to no - attenuation. See the OpenGL documentation for more information on this and all other OpenGL-related properties. + These are the constant, linear and quadratic attenuation coefficients for Phong lighting. + The default corresponds to no attenuation. .. _body-light-cutoff: @@ -2940,22 +3010,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr :at:`exponent`: :at-val:`real, "10"` Exponent for spotlights. This setting controls the softness of the spotlight cutoff. -.. _body-light-ambient: - -:at:`ambient`: :at-val:`real(3), "0 0 0"` - The ambient color of the light. - -.. _body-light-diffuse: - -:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"` - The diffuse color of the light. - -.. _body-light-specular: - -:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"` - The specular color of the light. - - .. _body-composite: :el-prefix:`body/` |-| **composite** (*) @@ -3003,8 +3057,13 @@ cable, which produces an inextensible chain of bodies connected with ball joints .. _body-composite-offset: :at:`offset`: :at-val:`real(3), "0 0 0"` - It specifies a 3D offset from the center of the parent body to the center of the grid of elements. The offset is - expressed in the local coordinate frame of the parent body. + It specifies a 3D offset from the center of the parent body to the center of the first body of the cable. The offset + is expressed in the local coordinate frame of the parent body. + +.. _body-composite-quat: + +:at:`quat`: :at-val:`real(4), "1 0 0 0"` + It specifies a quaternion that rotates the first body frame. The quaternion is expressed in the parent body frame. .. _body-composite-vertex: @@ -3532,6 +3591,7 @@ saving the XML: .. _flexcomp-contact-internal: .. _flexcomp-contact-selfcollide: +.. _flexcomp-contact-vertcollide: .. _flexcomp-contact-activelayers: .. _flexcomp-contact-contype: .. _flexcomp-contact-conaffinity: @@ -3545,8 +3605,8 @@ saving the XML: .. _flexcomp-contact-gap: .. |body/flexcomp/contact attrib list| replace:: - :at:`internal`, :at:`selfcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`, :at:`condim`, - :at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap` + :at:`internal`, :at:`selfcollide`, :at:`vertcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`, + :at:`condim`, :at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap` |body/flexcomp/contact attrib list| Same as in :ref:`flex/contact`. All attributes are passed through to the automatically-generated flex. @@ -3588,9 +3648,10 @@ element is used to adjust the properties of all edges in the flex. .. _flexcomp-elasticity-poisson: .. _flexcomp-elasticity-damping: .. _flexcomp-elasticity-thickness: +.. _flexcomp-elasticity-elastic2d: .. |body/flexcomp/elasticity attrib list| replace:: - :at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness` + :at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness`, :at:`elastic2d` |body/flexcomp/elasticity attrib list| Same as in :ref:`flex/elasticity`. @@ -3698,8 +3759,9 @@ all attachments will appear in the saved XML file. .. _body-attach-body: -:at:`body`: :at-val:`string, required` - Name of the body in the sub-model to attach here. The body and its subtree will be attached. +:at:`body`: :at-val:`string, optional` + Name of the body in the sub-model to attach here. The body and its subtree will be attached. If this attribute is not + specified, the contents of the world body will be attached in a new :ref:`frame`. .. _body-attach-prefix: @@ -4030,6 +4092,11 @@ stress-strain relationship.. See also :ref:`deformable ` objects. This thickness can be set equal to 2 times the :ref:`radius ` in order to match the geometry, but is exposed separately since the radius might be constrained by considerations related to collision detection. +.. _flex-elasticity-elastic2d: + +:at:`elastic2d`: :at-val:`[none, bend, stretch, both], "none"` + Elastic contribution to passive forces of 2D flexes. "none": none, "bend": bending only, "stretch": stretching only, + "both": bending and stretching. .. _flex-contact: @@ -4063,6 +4130,13 @@ extensions specific to flexes. **sap** in 1D and 2D, and **bvh** in 3D. Which strategy performs better depends on the specifics of the model. The automatic setting is just a simple rule which we have found to perform well in general. +.. _flex-contact-vertcollide: + +:at:`vertcollide`: :at-val:`[true, false], "false"` + Enables or disables vertex collisions. if **true**, spherical geoms are added at the vertices of flex, with radius + equal to the radius of the flex. These geoms can collide with other geoms and are not visible by default. If + **false**, no additional geoms are added. + .. _flex-contact-activelayers: :at:`activelayers`: :at-val:`int(1), "1"` @@ -4607,12 +4681,28 @@ length X, as in the clip on the right of `this example model solver. If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler `, length limits will be enabled if range is defined. +.. _tendon-spatial-actuatorfrclimited: + +:at:`actuatorfrclimited`: :at-val:`[false, true, auto], "auto"` + This attribute specifies whether actuator forces acting on the tendon should be clamped. See :ref:`CForceRange` for + details. This attribute interacts with the :ref:`actuatorfrcrange` attribute. If + this attribute is "false", actuator force clamping is disabled. If it is "true", actuator force clamping is enabled. + If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler `, actuator force clamping will + be enabled if :at:`actuatorfrcrange` is defined. + .. _tendon-spatial-range: :at:`range`: :at-val:`real(2), "0 0"` Range of allowed tendon lengths. Setting this attribute without specifying :at:`limited` is an error, unless :at:`autolimits` is set in :ref:`compiler `. +.. _tendon-spatial-actuatorfrcrange: + +:at:`actuatorfrcrange`: :at-val:`real(2), "0 0"` + Range for clamping total actuator forces acting on this tendon. See :ref:`CForceRange` for details. The compiler + expects the lower bound to be nonpositive and the upper bound to be nonnegative. |br| Setting this attribute without + specifying :at:`actuatorfrclimited` is an error if :at:`compiler-autolimits` is "false". + .. _tendon-spatial-solreflimit: .. _tendon-spatial-solimplimit: @@ -4683,6 +4773,31 @@ length X, as in the clip on the right of `this example model joint damping which is integrated implicitly by the Euler method, tendon damping is not integrated implicitly, thus joint damping should be used if possible. +.. image:: images/XMLreference/tendon_armature.gif + :width: 30% + :align: right + :class: only-light + :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml +.. image:: images/XMLreference/tendon_armature_dark.gif + :width: 30% + :align: right + :class: only-dark + :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml + +.. _tendon-spatial-armature: + +:at:`armature`: :at-val:`real, "0"` + Inertia associated with changes in tendon length. Setting this attribute to a positive value :math:`m` adds a kinetic + energy term :math:`\frac{1}{2}mv^2`, where :math:`v` is the tendon velocity. Tendon inertia is most valuable + when modeling the :ref:`armature` inertia in a linear actuator which contains a spinning element + or the inertial motion of a fluid in a linear hydraulic actuator. In the illustration, we compare (*left*) a 3-dof + system with a "tendon" implemented with a rotational joint and a slider joint with + :ref:`armature`, attached to the world with a :ref:`connect` constraint and + (*right*) an equivalent 1-dof model with an armature-bearing tendon. Like joint :ref:`armature`, + this added inertia is only associated with changes in tendon length, and would not affect the dynamics of a moving + fixed-length tendon. Because the tendon Jacobian :math:`J` is position-dependent, tendon armature leads to an + additional bias-force term :math:`c = m J \dot{J}^T \dot{q}`. + .. _tendon-spatial-user: :at:`user`: :at-val:`real(nuser_tendon), "0 0 ..."` @@ -4772,8 +4887,12 @@ as above. .. _tendon-fixed-limited: +.. _tendon-fixed-actuatorfrclimited: + .. _tendon-fixed-range: +.. _tendon-fixed-actuatorfrcrange: + .. _tendon-fixed-solreflimit: .. _tendon-fixed-solimplimit: @@ -4792,6 +4911,8 @@ as above. .. _tendon-fixed-damping: +.. _tendon-fixed-armature: + .. _tendon-fixed-user: .. |tendon/fixed attrib list| replace:: @@ -6335,6 +6456,33 @@ joint or when a single actuator acts on multiple joints. See :ref:`CForceRange` The joint where actuator forces will be sensed. The sensor output is copied from ``mjData.qfrc_actuator``. +.. _sensor-tendonactuatorfrc: + +:el-prefix:`sensor/` |-| **tendonactuatorfrc** (*) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This element creates an actuator force sensor, measured at a tendon. The quantity being sensed is the total force +contributed by all actuators to a single tendon. This type of sensor is important when multiple actuators act on a +single tendon. See :ref:`CForceRange` for details. + + +.. _sensor-tendonactuatorfrc-name: + +.. _sensor-tendonactuatorfrc-noise: + +.. _sensor-tendonactuatorfrc-cutoff: + +.. _sensor-tendonactuatorfrc-user: + +:at:`name`, :at:`noise`, :at:`cutoff`, :at:`user` + See :ref:`CSensor`. + +.. _sensor-tendonactuatorfrc-tendon: + +:at:`tendon`: :at-val:`string, required` + The tendon where actuator forces will be sensed. + + .. _sensor-ballquat: :el-prefix:`sensor/` |-| **ballquat** (*) @@ -8152,7 +8300,7 @@ if omitted. .. _default-light-dir: -.. _default-light-bulbradius: +.. _default-light-type: .. _default-light-directional: @@ -8160,18 +8308,24 @@ if omitted. .. _default-light-active: +.. _default-light-diffuse: + +.. _default-light-intensity: + +.. _default-light-ambient: + +.. _default-light-specular: + +.. _default-light-bulbradius: + +.. _default-light-range: + .. _default-light-attenuation: .. _default-light-cutoff: .. _default-light-exponent: -.. _default-light-ambient: - -.. _default-light-diffuse: - -.. _default-light-specular: - .. _default-light-mode: :el-prefix:`default/` |-| **light** (?) @@ -8227,8 +8381,12 @@ if omitted. .. _default-tendon-limited: +.. _default-tendon-actuatorfrclimited: + .. _default-tendon-range: +.. _default-tendon-actuatorfrcrange: + .. _default-tendon-solreflimit: .. _default-tendon-solimplimit: diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 0f71f6a3..21272c90 100644 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -54,7 +54,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`discardvisual` | :ref:`usethread` | :ref:`fusestatic` | :ref:`inertiafromgeom` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`inertiagrouprange` | :ref:`assetdir` | :ref:`alignfree` | | | +| | | | :ref:`inertiagrouprange` | :ref:`saveinertial` | :ref:`assetdir` | :ref:`alignfree` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| compiler |br| |_| |L| | | .. table:: | @@ -153,17 +153,17 @@ | :ref:`texture | \* | :class: mjcf-attributes | | ` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`name` | :ref:`type` | :ref:`content_type` | :ref:`file` | | +| | | | :ref:`name` | :ref:`type` | :ref:`colorspace` | :ref:`content_type` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`gridsize` | :ref:`gridlayout` | :ref:`fileright` | :ref:`fileleft` | | +| | | | :ref:`file` | :ref:`gridsize` | :ref:`gridlayout` | :ref:`fileright` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`fileup` | :ref:`filedown` | :ref:`filefront` | :ref:`fileback` | | +| | | | :ref:`fileleft` | :ref:`fileup` | :ref:`filedown` | :ref:`filefront` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`builtin` | :ref:`rgb1` | :ref:`rgb2` | :ref:`mark` | | +| | | | :ref:`fileback` | :ref:`builtin` | :ref:`rgb1` | :ref:`rgb2` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`markrgb` | :ref:`random` | :ref:`width` | :ref:`height` | | +| | | | :ref:`mark` | :ref:`markrgb` | :ref:`random` | :ref:`width` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`hflip` | :ref:`vflip` | :ref:`nchannel` | | | +| | | | :ref:`height` | :ref:`hflip` | :ref:`vflip` | :ref:`nchannel` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| asset |br| |_| |L| | | .. table:: | @@ -313,13 +313,15 @@ | :ref:`light | \* | :class: mjcf-attributes | | ` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`name` | :ref:`class` | :ref:`directional` | :ref:`castshadow` | | +| | | | :ref:`name` | :ref:`class` | :ref:`directional` | :ref:`type` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`active` | :ref:`pos` | :ref:`dir` | :ref:`bulbradius` | | +| | | | :ref:`castshadow` | :ref:`active` | :ref:`pos` | :ref:`dir` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`attenuation` | :ref:`cutoff` | :ref:`exponent` | :ref:`ambient` | | +| | | | :ref:`bulbradius` | :ref:`intensity` | :ref:`range` | :ref:`attenuation` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`diffuse` | :ref:`specular` | :ref:`mode` | :ref:`target` | | +| | | | :ref:`cutoff` | :ref:`exponent` | :ref:`ambient` | :ref:`diffuse` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`specular` | :ref:`mode` | :ref:`target` | :ref:`texture` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| body |br| |_| |L| | | .. table:: | @@ -344,6 +346,8 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`vertex` | :ref:`initial` | :ref:`curve` | :ref:`size` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`quat` | | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| composite |br| |_2| |L| | | .. table:: | | :ref:`joint | \* | :class: mjcf-attributes | @@ -439,6 +443,8 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`young` | :ref:`poisson` | :ref:`damping` | :ref:`thickness` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`elastic2d` | | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flexcomp |br| |_2| |L| | | .. table:: | | :ref:`contact | ? | :class: mjcf-attributes | @@ -450,7 +456,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`margin` | :ref:`gap` | :ref:`internal` | :ref:`selfcollide` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`activelayers` | | | | | +| | | | :ref:`activelayers` | :ref:`vertcollide` | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flexcomp |br| |_2| |L| | | .. table:: | @@ -500,7 +506,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`margin` | :ref:`gap` | :ref:`internal` | :ref:`selfcollide` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`activelayers` | | | | | +| | | | :ref:`activelayers` | :ref:`vertcollide` | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flex |br| |_2| |L| | | .. table:: | @@ -516,6 +522,8 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`young` | :ref:`poisson` | :ref:`damping` | :ref:`thickness` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`elastic2d` | | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| deformable |br| |_| |L| | | .. table:: | | :ref:`skin | \* | :class: mjcf-attributes | @@ -619,13 +627,15 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`name` | :ref:`class` | :ref:`group` | :ref:`limited` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`range` | :ref:`solreflimit` | :ref:`solimplimit` | :ref:`solreffriction` | | +| | | | :ref:`actuatorfrclimited` | :ref:`range` | :ref:`actuatorfrcrange` | :ref:`solreflimit` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`solimpfriction` | :ref:`frictionloss` | :ref:`springlength` | :ref:`width` | | +| | | | :ref:`solimplimit` | :ref:`solreffriction` | :ref:`solimpfriction` | :ref:`frictionloss` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`material` | :ref:`margin` | :ref:`stiffness` | :ref:`damping` | | +| | | | :ref:`springlength` | :ref:`width` | :ref:`material` | :ref:`margin` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`rgba` | :ref:`user` | | | | +| | | | :ref:`stiffness` | :ref:`damping` | :ref:`armature` | :ref:`rgba` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`user` | | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| spatial |br| |_2| |L| | | .. table:: | @@ -655,11 +665,13 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`name` | :ref:`class` | :ref:`group` | :ref:`limited` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`range` | :ref:`solreflimit` | :ref:`solimplimit` | :ref:`solreffriction` | | +| | | | :ref:`actuatorfrclimited` | :ref:`range` | :ref:`actuatorfrcrange` | :ref:`solreflimit` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`solimpfriction` | :ref:`frictionloss` | :ref:`springlength` | :ref:`margin` | | +| | | | :ref:`solimplimit` | :ref:`solreffriction` | :ref:`solimpfriction` | :ref:`frictionloss` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`stiffness` | :ref:`damping` | :ref:`user` | | | +| | | | :ref:`springlength` | :ref:`margin` | :ref:`stiffness` | :ref:`damping` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`armature` | :ref:`user` | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| fixed |br| |_2| |L| | | .. table:: | @@ -1002,6 +1014,15 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| sensor |br| |_| |L| | | .. table:: | +| :ref:`tendonactuatorfrc | \* | :class: mjcf-attributes | +| ` | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`name` | :ref:`tendon` | :ref:`cutoff` | :ref:`noise` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`user` | | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | ++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| |_| sensor |br| |_| |L| | | .. table:: | | :ref:`ballquat | \* | :class: mjcf-attributes | | ` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | @@ -1442,13 +1463,13 @@ | :ref:`light | ? | :class: mjcf-attributes | | ` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`pos` | :ref:`dir` | :ref:`bulbradius` | :ref:`directional` | | +| | | | :ref:`pos` | :ref:`dir` | :ref:`bulbradius` | :ref:`intensity` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`castshadow` | :ref:`active` | :ref:`attenuation` | :ref:`cutoff` | | +| | | | :ref:`range` | :ref:`directional` | :ref:`type` | :ref:`castshadow` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`exponent` | :ref:`ambient` | :ref:`diffuse` | :ref:`specular` | | +| | | | :ref:`active` | :ref:`attenuation` | :ref:`cutoff` | :ref:`exponent` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`mode` | | | | | +| | | | :ref:`ambient` | :ref:`diffuse` | :ref:`specular` | :ref:`mode` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| default |br| |_| |L| | | .. table:: | diff --git a/doc/changelog.rst b/doc/changelog.rst index 05d68306..324d41e7 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,26 +2,100 @@ Changelog ========= -Upcoming version (not yet released) ------------------------------------ +Upcoming version (not yet release) +---------------------------------- + +General +^^^^^^^ +- Refactored island implementation so that island data is memory-contiguous. This speeds up island processing in the + solver and clears the way for the addition of the Newton and PGS solvers (currently only CG is supported). +- Removed the :at:`shell` plugin. This is now supported by :ref:`flexcomp` and is active depending on + the :ref:`elastic2d` attribute (off by default). +- Replaced the :ref:`directional` (boolean) field for lights with a + :ref:`type` field (of type :ref:`mjtLightType`) to allow for additional lighting + types. +- Added new sub-component :ref:`mj_makeM` which combines the :ref:`mj_crb` call with additional logic to support the + introduction in 3.3.1 of :ref:`tendon armature`. In addition to the traditional + ``mjData.qM``, :ref:`mj_makeM` also computes ``mjData.M``, a CSR representation of the same matrix. +- Added a new function :ref:`mj_copyBack` to copy real-valued arrays in an mjModel to a compatible mjSpec. + +Simulate +^^^^^^^^ +- The struct ``mjv_sceneState`` has been removed. This struct was used for partial synchronization of ``mjModel`` and + ``mjData`` when the Python viewer is used in passive mode. This functionality is now provided by :ref:`mjv_copyModel` + and :ref:`mjv_copyData`, which don't copy arrays which are not required for visualization. + +.. image:: images/changelog/procedural_terrain_generation.png + :width: 25% + :align: right + +Python bindings +^^^^^^^^^^^^^^^ + +- Added examples of procedural terrain generation to the Model Editing tutorial: |mjspec_colab| + +Version 3.3.2 (April 28, 2025) +------------------------------ + +MJX +^^^ +1. Added inverse dynamics. +2. Added tendon actuator force sensor. +3. Fix :github:issue:`2606` such that ``make_data`` copies over ``mocap_pos`` and ``mocap_quat`` + from ``body_pos`` and ``body_quat``. + +Version 3.3.1 (Apr 9, 2025) +---------------------------- .. admonition:: Breaking API changes :class: attention - - The default value of the flag for toggling :ref:`internal flex contacts` was changed from - "true" to "false". This feature has proven to be counterintuitive for users. + 1. The default value of the flag for toggling :ref:`internal flex contacts` was changed from + "true" to "false". This feature has proven to be counterintuitive for users. + 2. All of the attach functions (``mjs_attachBody``, ``mjs_attachFrame``, ``mjs_attachToSite``, + ``mjs_attachFrameToSite``) have been removed and replaced by a single function :ref:`mjs_attach`. + +General +^^^^^^^ +3. Added :ref:`tendon armature`: inertia associated with changes in tendon length. +4. Added the :ref:`compiler/saveinertial` flag, writing explicit inertial clauses for all + bodies when saving to XML. +5. Added :ref:`orientation` attribute to :ref:`composite`. Moreover, allow the + composite to be the direct child of a frame. +6. Added :ref:`tendon actuator force limits` and + :ref:`tendon actuator force sensor`. + +MJX +^^^ +7. Added tendon actuator force limits. Bug fixes ^^^^^^^^^ -- :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to - which the Jacobian is computed, now fixed. -- Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an - mjSpec to a frame or a site. +8. :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to + which the Jacobian is computed, now fixed. +9. Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an + mjSpec to a frame or a site. +10. Fixed a bug that caused shadow rendering to flicker on platforms (e.g., MacOS) that do not support + ARB_clip_control. Fixed in collaboration with :github:user:`aftersomemath`. + +Python bindings +^^^^^^^^^^^^^^^ + +.. youtube:: LbANnKMDOHg + :aspect: 16:7 + :align: right + :width: 240px + +11. Added examples of procedural model creation to the Model Editing tutorial: |mjspec_colab| +12. Added support for nameless :ref:`mjSpec` objects in the ``bind`` method, see the corresponding + :ref:`section` in the documentation. + +.. |mjspec_colab| image:: https://colab.research.google.com/assets/colab-badge.svg + :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb Version 3.3.0 (Feb 26, 2025) ---------------------------- - Feature promotion ^^^^^^^^^^^^^^^^^ .. youtube:: qJFbx-FR7Bc diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 1021067a..ff87af05 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -1692,7 +1692,7 @@ The stages below compute quantities that depend on the generalized positions ``m 4. Compute quantities related to :ref:`flex` objects: :ref:`mj_flex` 5. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial tendons: :ref:`mj_tendon` -6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb` +6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_makeM` 7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM` 8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection: :ref:`mj_collision` diff --git a/doc/images/XMLreference/armature.gif b/doc/images/XMLreference/armature.gif new file mode 100644 index 00000000..93f0e583 Binary files /dev/null and b/doc/images/XMLreference/armature.gif differ diff --git a/doc/images/XMLreference/armature_dark.gif b/doc/images/XMLreference/armature_dark.gif new file mode 100644 index 00000000..a478ff5c Binary files /dev/null and b/doc/images/XMLreference/armature_dark.gif differ diff --git a/doc/images/XMLreference/tendon_armature.gif b/doc/images/XMLreference/tendon_armature.gif new file mode 100644 index 00000000..5a1c6d98 Binary files /dev/null and b/doc/images/XMLreference/tendon_armature.gif differ diff --git a/doc/images/XMLreference/tendon_armature_dark.gif b/doc/images/XMLreference/tendon_armature_dark.gif new file mode 100644 index 00000000..9860860d Binary files /dev/null and b/doc/images/XMLreference/tendon_armature_dark.gif differ diff --git a/doc/images/changelog/procedural_terrain_generation.png b/doc/images/changelog/procedural_terrain_generation.png new file mode 100644 index 00000000..7c083dff Binary files /dev/null and b/doc/images/changelog/procedural_terrain_generation.png differ diff --git a/doc/includes/references.h b/doc/includes/references.h index 39eece88..3f725ad9 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -154,7 +154,6 @@ struct mjData_ { // solver statistics mjSolverStat solver[mjNISLAND*mjNSOLVER]; // solver statistics per island, per iteration - int solver_nisland; // number of islands processed by solver int solver_niter[mjNISLAND]; // number of solver iterations, per island int solver_nnz[mjNISLAND]; // number of nonzeros in Hessian or efc_AR, per island mjtNum solver_fwdinv[2]; // forward-inverse comparison: qfrc, efc @@ -172,6 +171,7 @@ struct mjData_ { int nJ; // number of non-zeros in constraint Jacobian int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands + int nidof; // number of dofs in all islands // global properties mjtNum time; // simulation time @@ -267,12 +267,13 @@ struct mjData_ { int* moment_colind; // column indices in sparse Jacobian (nJmom x 1) mjtNum* actuator_moment; // actuator moments (nJmom x 1) - // computed by mj_fwdPosition/mj_crb + // computed by mj_fwdPosition/mj_makeM mjtNum* crb; // com-based composite inertia and mass (nbody x 10) - mjtNum* qM; // total inertia (sparse) (nM x 1) + mjtNum* qM; // inertia (sparse) (nM x 1) + mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1) // computed by mj_fwdPosition/mj_factorM - mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) // computed by mj_collisionTree @@ -305,27 +306,23 @@ struct mjData_ { mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) // computed by mj_Euler or mj_implicit - mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) + mjtNum* qH; // L'*D*L factorization of modified M (nC x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) // computed by mj_resetData int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1) int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1) int* B_colind; // body-dof: column indices of non-zeros (nB x 1) - int* M_rownnz; // inertia: non-zeros in each row (nv x 1) - int* M_rowadr; // inertia: address of each row in M_colind (nv x 1) - int* M_colind; // inertia: column indices of non-zeros (nM x 1) - int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1) - int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1) - int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1) - int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1) - int* mapM2C; // index mapping from M to C (nC x 1) - int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1) - int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1) - int* D_diag; // dof-dof: index of diagonal element (nv x 1) - int* D_colind; // dof-dof: column indices of non-zeros (nD x 1) - int* mapM2D; // index mapping from M to D (nD x 1) - int* mapD2M; // index mapping from D to M (nM x 1) + int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1) + int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1) + int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1) + int* mapM2M; // index mapping from qM to M (nC x 1) + int* D_rownnz; // full inertia: non-zeros in each row (nv x 1) + int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1) + int* D_diag; // full inertia: index of diagonal element (nv x 1) + int* D_colind; // full inertia: column indices of non-zeros (nD x 1) + int* mapM2D; // index mapping from qM to D (nD x 1) + int* mapD2M; // index mapping from D to qM (nM x 1) // computed by mj_implicit/mj_derivative mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1) @@ -347,8 +344,8 @@ struct mjData_ { mjtNum* qfrc_constraint; // constraint force (nv x 1) // computed by mj_inverse - mjtNum* qfrc_inverse; // net external force; should equal: (nv x 1) - // qfrc_applied + J'*xfrc_applied + qfrc_actuator + mjtNum* qfrc_inverse; // net external force; should equal: + // qfrc_applied + J'*xfrc_applied + qfrc_actuator (nv x 1) // computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format mjtNum* cacc; // com-based acceleration (nbody x 6) @@ -382,16 +379,50 @@ struct mjData_ { mjtNum* efc_R; // inverse constraint mass (nefc x 1) int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1) - // computed by mj_island + // computed by mj_island (island dof structure) int* dof_island; // island id of this dof; -1: none (nv x 1) - int* island_dofnum; // number of dofs in island (nisland x 1) - int* island_dofadr; // start address in island_dofind (nisland x 1) - int* island_dofind; // island dof indices; -1: none (nv x 1) - int* dof_islandind; // dof island indices; -1: none (nv x 1) + int* island_nv; // number of dofs in this island (nisland x 1) + int* island_idofadr; // island start address in idof vector (nisland x 1) + int* island_dofadr; // island start address in dof vector (nisland x 1) + int* map_dof2idof; // map from dof to idof (nv x 1) + int* map_idof2dof; // map from idof to dof; >= nidof: unconstrained (nv x 1) + + // computed by mj_island (dofs sorted by island) + mjtNum* ifrc_smooth; // net unconstrained force (nidof x 1) + mjtNum* iacc_smooth; // unconstrained acceleration (nidof x 1) + int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) + int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) + int* iM_colind; // inertia: column indices of non-zeros (nC x 1) + mjtNum* iM; // total inertia (sparse) (nC x 1) + mjtNum* iLD; // L'*D*L factorization of M (sparse) (nC x 1) + mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) + mjtNum* iacc; // acceleration (nidof x 1) + + // computed by mj_island (island constraint structure) int* efc_island; // island id of this constraint (nefc x 1) - int* island_efcnum; // number of constraints in island (nisland x 1) - int* island_efcadr; // start address in island_efcind (nisland x 1) - int* island_efcind; // island constraint indices (nefc x 1) + int* island_ne; // number of equality constraints in island (nisland x 1) + int* island_nf; // number of friction constraints in island (nisland x 1) + int* island_nefc; // number of constraints in island (nisland x 1) + int* island_iefcadr; // start address in iefc vector (nisland x 1) + int* map_efc2iefc; // map from efc to iefc (nefc x 1) + int* map_iefc2efc; // map from iefc to efc (nefc x 1) + + // computed by mj_island (constraints sorted by island) + int* iefc_type; // constraint type (mjtConstraint) (nefc x 1) + int* iefc_id; // id of object of specified type (nefc x 1) + int* iefc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* iefc_J_rowadr; // row start address in colind array (nefc x 1) + int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) + int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) + int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) + int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) + int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) + mjtNum* iefc_J; // constraint Jacobian (nJ x 1) + mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) + mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* iefc_D; // constraint mass (nefc x 1) + mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) @@ -409,11 +440,18 @@ struct mjData_ { // computed by mj_fwdConstraint/mj_inverse mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1) + int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* iefc_force; // constraint force in constraint space (nefc x 1) int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* ifrc_constraint; // constraint force (nidof x 1) // thread pool pointer uintptr_t threadpool; + + // compilation signature + uint64_t signature; // also held by the mjSpec that compiled the model }; typedef struct mjData_ mjData; typedef enum mjtDisableBit_ { // disable default feature bitflags @@ -488,6 +526,12 @@ typedef enum mjtCamLight_ { // tracking mode for camera and light mjCAMLIGHT_TARGETBODY, // pos fixed in body, rot tracks target body mjCAMLIGHT_TARGETBODYCOM // pos fixed in body, rot tracks target subtree com } mjtCamLight; +typedef enum mjtLightType_ { // type of light + mjLIGHT_SPOT = 0, // spot + mjLIGHT_DIRECTIONAL, // directional + mjLIGHT_POINT, // point + mjLIGHT_IMAGE, // image-based +} mjtLightType; typedef enum mjtTexture_ { // type of texture mjTEXTURE_2D = 0, // 2d texture, suitable for planes and hfields mjTEXTURE_CUBE, // cube texture, suitable for all other geom types @@ -506,6 +550,11 @@ typedef enum mjtTextureRole_ { // role of texture map in rendering mjTEXROLE_ORM, // occlusion, roughness, metallic mjNTEXROLE } mjtTextureRole; +typedef enum mjtColorSpace_ { // type of color space encoding + mjCOLORSPACE_AUTO = 0, // attempts to autodetect color space, defaults to linear + mjCOLORSPACE_LINEAR, // linear color space + mjCOLORSPACE_SRGB // standard RGB color space +} mjtColorSpace; typedef enum mjtIntegrator_ { // integrator mode mjINT_EULER = 0, // semi-implicit Euler mjINT_RK4, // 4th-order Runge Kutta @@ -604,7 +653,8 @@ typedef enum mjtObj_ { // type of MujoCo object // meta elements, do not appear in mjModel mjOBJ_FRAME = 100, // frame - mjOBJ_DEFAULT // default + mjOBJ_DEFAULT, // default + mjOBJ_MODEL // entire model } mjtObj; typedef enum mjtConstraint_ { // type of constraint @@ -645,6 +695,7 @@ typedef enum mjtSensor_ { // type of sensor mjSENS_ACTUATORVEL, // scalar actuator velocity mjSENS_ACTUATORFRC, // scalar actuator force mjSENS_JOINTACTFRC, // scalar actuator force, measured at the joint + mjSENS_TENDONACTFRC, // scalar actuator force, measured at the tendon // sensors related to ball joints mjSENS_BALLQUAT, // 4D ball joint quaternion @@ -722,6 +773,12 @@ typedef enum mjtFlexSelf_ { // mode for flex selfcollide mjFLEXSELF_SAP, // use SAP in midphase mjFLEXSELF_AUTO // choose between BVH and SAP automatically } mjtFlexSelf; +typedef enum mjtSDFType_ { // signed distance function (SDF) type + mjSDFTYPE_SINGLE = 0, // single SDF + mjSDFTYPE_INTERSECTION, // max(A, B) + mjSDFTYPE_MIDSURFACE, // A - B + mjSDFTYPE_COLLISION, // A + B + abs(max(A, B)) +} mjtSDFType; struct mjLROpt_ { // options for mj_setLengthRange() // flags int mode; // which actuators to process (mjtLRMode) @@ -907,7 +964,7 @@ struct mjModel_ { int ncam; // number of cameras int nlight; // number of lights int nflex; // number of flexes - int nflexnode; // number of dofs in all flexes + int nflexnode; // number of dofs in all flexes int nflexvert; // number of vertices in all flexes int nflexedge; // number of edges in all flexes int nflexelem; // number of elements in all flexes @@ -1128,9 +1185,12 @@ struct mjModel_ { int* light_mode; // light tracking mode (mjtCamLight) (nlight x 1) int* light_bodyid; // id of light's body (nlight x 1) int* light_targetbodyid; // id of targeted body; -1: none (nlight x 1) - mjtByte* light_directional; // directional light (nlight x 1) + int* light_type; // spot, directional, etc. (mjtLightType) (nlight x 1) + int* light_texid; // texture id for image lights (nlight x 1) mjtByte* light_castshadow; // does light cast shadows (nlight x 1) float* light_bulbradius; // light radius for soft shadows (nlight x 1) + float* light_intensity; // intensity, in candela (nlight x 1) + float* light_range; // range of effectiveness (nlight x 1) mjtByte* light_active; // is light on (nlight x 1) mjtNum* light_pos; // position rel. to body frame (nlight x 3) mjtNum* light_dir; // direction rel. to body frame (nlight x 3) @@ -1182,6 +1242,7 @@ struct mjModel_ { int* flex_nodebodyid; // node body ids (nflexnode x 1) int* flex_vertbodyid; // vertex body ids (nflexvert x 1) int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2) + int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2) int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1) int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1) int* flex_elemedge; // element edge ids (nflexelemedge x 1) @@ -1196,6 +1257,7 @@ struct mjModel_ { mjtNum* flexedge_invweight0; // edge inv. weight in qpos0 (nflexedge x 1) mjtNum* flex_radius; // radius around primitive element (nflex x 1) mjtNum* flex_stiffness; // finite element stiffness matrix (nflexelem x 21) + mjtNum* flex_bending; // bending stiffness (nflexedge x 16) mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1) mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1) mjtNum* flex_edgedamping; // edge damping (nflex x 1) @@ -1276,6 +1338,7 @@ struct mjModel_ { // textures int* tex_type; // texture type (mjtTexture) (ntex x 1) + int* tex_colorspace; // texture colorspace (mjtColorSpace) (ntex x 1) int* tex_height; // number of rows in texture image (ntex x 1) int* tex_width; // number of columns in texture image (ntex x 1) int* tex_nchannel; // number of channels in texture image (ntex x 1) @@ -1326,15 +1389,18 @@ struct mjModel_ { int* tendon_matid; // material id for rendering (ntendon x 1) int* tendon_group; // group for visibility (ntendon x 1) mjtByte* tendon_limited; // does tendon have length limits (ntendon x 1) + mjtByte* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1) mjtNum* tendon_width; // width for rendering (ntendon x 1) mjtNum* tendon_solref_lim; // constraint solver reference: limit (ntendon x mjNREF) mjtNum* tendon_solimp_lim; // constraint solver impedance: limit (ntendon x mjNIMP) mjtNum* tendon_solref_fri; // constraint solver reference: friction (ntendon x mjNREF) mjtNum* tendon_solimp_fri; // constraint solver impedance: friction (ntendon x mjNIMP) mjtNum* tendon_range; // tendon length limits (ntendon x 2) + mjtNum* tendon_actfrcrange; // range of total actuator force (ntendon x 2) mjtNum* tendon_margin; // min distance for limit detection (ntendon x 1) mjtNum* tendon_stiffness; // stiffness coefficient (ntendon x 1) mjtNum* tendon_damping; // damping coefficient (ntendon x 1) + mjtNum* tendon_armature; // inertia associated with tendon velocity (ntendon x 1) mjtNum* tendon_frictionloss; // loss due to friction (ntendon x 1) mjtNum* tendon_lengthspring; // spring resting length range (ntendon x 2) mjtNum* tendon_length0; // tendon length in qpos0 (ntendon x 1) @@ -1451,6 +1517,9 @@ struct mjModel_ { // paths char* paths; // paths to assets, 0-terminated (npaths x 1) + + // compilation signature + uint64_t signature; // also held by the mjSpec that compiled this model }; typedef struct mjModel_ mjModel; struct mjResource_ { @@ -1536,6 +1605,15 @@ struct mjpPlugin_ { void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes); }; typedef struct mjpPlugin_ mjpPlugin; +struct mjSDF_ { + const mjpPlugin** plugin; + int* id; + mjtSDFType type; + mjtNum* relpos; + mjtNum* relmat; + mjtGeom* geomtype; +}; +typedef struct mjSDF_ mjSDF; typedef enum mjtGridPos_ { // grid position for overlay mjGRID_TOPLEFT = 0, // top left mjGRID_TOPRIGHT, // top right @@ -1713,6 +1791,7 @@ typedef enum mjtOrientation_ { // type of orientation specifier } mjtOrientation; typedef struct mjsElement_ { // element type, do not modify mjtObj elemtype; // element type + uint64_t signature; // compilation signature } mjsElement; typedef struct mjsCompiler_ { // compiler options mjtByte autolimits; // infer "limited" attribute based on range @@ -1728,6 +1807,7 @@ typedef struct mjsCompiler_ { // compiler options mjtByte fusestatic; // fuse static bodies with parent int inertiafromgeom; // use geom inertias (mjtInertiaFromGeom) int inertiagrouprange[2]; // range of geom groups used to compute inertia + mjtByte saveinertial; // save explicit inertial clause for all bodies to XML int alignfree; // align free joints with inertial frame mjLROpt LRopt; // options for lengthrange computation } mjsCompiler; @@ -1963,9 +2043,12 @@ typedef struct mjsLight_ { // light specification // intrinsics mjtByte active; // is light active - mjtByte directional; // is light directional or spot + mjtLightType type; // type of light + mjString* texture; // texture name for image lights mjtByte castshadow; // does light cast shadows - double bulbradius; // bulb radius, for soft shadows + float bulbradius; // bulb radius, for soft shadows + float intensity; // intensity, in candelas + float range; // range of effectiveness float attenuation[3]; // OpenGL attenuation (quadratic model) float cutoff; // OpenGL cutoff float exponent; // OpenGL exponent @@ -1997,7 +2080,8 @@ typedef struct mjsFlex_ { // flex specification double radius; // radius around primitive element mjtByte internal; // enable internal collisions mjtByte flatskin; // render flex skin with flat shading - int selfcollide; // mode for flex self colllision + int selfcollide; // mode for flex self collision + int vertcollide; // mode for vertex collision int activelayers; // number of active element layers in 3D int group; // group for visualizatioh double edgestiffness; // edge stiffness @@ -2008,6 +2092,7 @@ typedef struct mjsFlex_ { // flex specification double poisson; // Poisson's ratio double damping; // Rayleigh's damping double thickness; // thickness (2D only) + int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both // mesh properties mjStringVec* nodebody; // node body names @@ -2079,6 +2164,7 @@ typedef struct mjsTexture_ { // texture specification mjsElement* element; // element type mjString* name; // name mjtTexture type; // texture type + mjtColorSpace colorspace; // colorspace // method 1: builtin int builtin; // builtin type (mjtBuiltin) @@ -2165,17 +2251,20 @@ typedef struct mjsTendon_ { // tendon specification mjsElement* element; // element type mjString* name; // name - // stiffness, damping, friction + // stiffness, damping, friction, armature double stiffness; // stiffness coefficient double springlength[2]; // spring resting length; {-1, -1}: use qpos_spring double damping; // damping coefficient double frictionloss; // friction loss mjtNum solref_friction[mjNREF]; // solver reference: tendon friction mjtNum solimp_friction[mjNIMP]; // solver impedance: tendon friction + double armature; // inertia associated with tendon velocity // length range int limited; // does tendon have limits (mjtLimited) + int actfrclimited; // does tendon have actuator force limits double range[2]; // length limits + double actfrcrange[2]; // actuator force limits double margin; // margin value for tendon limit detection mjtNum solref_limit[mjNREF]; // solver reference: tendon limits mjtNum solimp_limit[mjNIMP]; // solver impedance: tendon limits @@ -2728,7 +2817,7 @@ typedef struct mjvGLCamera_ mjvGLCamera; struct mjvGeom_ { // abstract geom // type info int type; // geom type (mjtGeom) - int dataid; // mesh, hfield or plane id; -1: none + int dataid; // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull) int objtype; // mujoco object type; mjOBJ_UNKNOWN for decor int objid; // mujoco object id; -1 for decor int category; // visual category @@ -2759,6 +2848,8 @@ typedef struct mjvGeom_ mjvGeom; struct mjvLight_ { // OpenGL light float pos[3]; // position rel. to body frame float dir[3]; // direction rel. to body frame + int type; // type (mjtLightType) + int texid; // texture id for image lights float attenuation[3]; // OpenGL attenuation (quadratic model) float cutoff; // OpenGL cutoff float exponent; // OpenGL exponent @@ -2766,9 +2857,10 @@ struct mjvLight_ { // OpenGL light float diffuse[3]; // diffuse rgb (alpha=1) float specular[3]; // specular rgb (alpha=1) mjtByte headlight; // headlight - mjtByte directional; // directional light mjtByte castshadow; // does light cast shadows float bulbradius; // bulb radius for soft shadows + float intensity; // intensity, in candelas + float range; // range of effectiveness }; typedef struct mjvLight_ mjvLight; struct mjvOption_ { // abstract visualization options @@ -2889,287 +2981,6 @@ struct mjvFigure_ { // abstract 2D figure passed to OpenGL rendere float yaxisdata[2]; // range of y-axis in data units }; typedef struct mjvFigure_ mjvFigure; -struct mjvSceneState_ { - int nbuffer; // size of the buffer in bytes - void* buffer; // heap-allocated memory for all arrays in this struct - int maxgeom; // maximum number of mjvGeom supported by this state object - mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins - - // fields in mjModel that are necessary to re-render a scene - struct { - int nv; - int nu; - int na; - int nbody; - int nbvh; - int nbvhstatic; - int njnt; - int ngeom; - int nsite; - int ncam; - int nlight; - int nmesh; - int nskin; - int nflex; - int nflexvert; - int nflextexcoord; - int nskinvert; - int nskinface; - int nskinbone; - int nskinbonevert; - int nmat; - int neq; - int ntendon; - int ntree; - int nwrap; - int nsensor; - int nnames; - int npaths; - int nsensordata; - int narena; - - mjOption opt; - mjVisual vis; - mjStatistic stat; - - int* body_parentid; - int* body_rootid; - int* body_weldid; - int* body_mocapid; - int* body_jntnum; - int* body_jntadr; - int* body_dofnum; - int* body_dofadr; - int* body_geomnum; - int* body_geomadr; - mjtNum* body_iquat; - mjtNum* body_mass; - mjtNum* body_inertia; - int* body_bvhadr; - int* body_bvhnum; - - int* bvh_depth; - int* bvh_child; - int* bvh_nodeid; - mjtNum* bvh_aabb; - - int* jnt_type; - int* jnt_bodyid; - int* jnt_group; - - int* geom_type; - int* geom_bodyid; - int* geom_contype; - int* geom_conaffinity; - int* geom_dataid; - int* geom_matid; - int* geom_group; - mjtNum* geom_size; - mjtNum* geom_aabb; - mjtNum* geom_rbound; - float* geom_rgba; - - int* site_type; - int* site_bodyid; - int* site_matid; - int* site_group; - mjtNum* site_size; - float* site_rgba; - - int* cam_orthographic; - mjtNum* cam_fovy; - mjtNum* cam_ipd; - int* cam_resolution; - float* cam_sensorsize; - float* cam_intrinsic; - - mjtByte* light_directional; - mjtByte* light_castshadow; - float* light_bulbradius; - mjtByte* light_active; - float* light_attenuation; - float* light_cutoff; - float* light_exponent; - float* light_ambient; - float* light_diffuse; - float* light_specular; - - mjtByte* flex_flatskin; - int* flex_dim; - int* flex_matid; - int* flex_group; - int* flex_interp; - int* flex_nodeadr; - int* flex_nodenum; - int* flex_nodebodyid; - int* flex_vertadr; - int* flex_vertnum; - int* flex_elem; - int* flex_elemtexcoord; - int* flex_elemlayer; - int* flex_elemadr; - int* flex_elemnum; - int* flex_elemdataadr; - int* flex_shell; - int* flex_shellnum; - int* flex_shelldataadr; - int* flex_texcoordadr; - int* flex_bvhadr; - int* flex_bvhnum; - mjtByte* flex_centered; - mjtNum* flex_node; - mjtNum* flex_radius; - float* flex_rgba; - float* flex_texcoord; - - int* hfield_pathadr; - - int* mesh_bvhadr; - int* mesh_bvhnum; - int* mesh_texcoordadr; - int* mesh_graphadr; - int* mesh_pathadr; - - int* skin_matid; - int* skin_group; - float* skin_rgba; - float* skin_inflate; - int* skin_vertadr; - int* skin_vertnum; - int* skin_texcoordadr; - int* skin_faceadr; - int* skin_facenum; - int* skin_boneadr; - int* skin_bonenum; - float* skin_vert; - int* skin_face; - int* skin_bonevertadr; - int* skin_bonevertnum; - float* skin_bonebindpos; - float* skin_bonebindquat; - int* skin_bonebodyid; - int* skin_bonevertid; - float* skin_bonevertweight; - int* skin_pathadr; - - int* tex_pathadr; - - int* mat_texid; - mjtByte* mat_texuniform; - float* mat_texrepeat; - float* mat_emission; - float* mat_specular; - float* mat_shininess; - float* mat_reflectance; - float* mat_metallic; - float* mat_roughness; - float* mat_rgba; - - int* eq_type; - int* eq_obj1id; - int* eq_obj2id; - int* eq_objtype; - mjtNum* eq_data; - - int* tendon_num; - int* tendon_matid; - int* tendon_group; - mjtByte* tendon_limited; - mjtNum* tendon_width; - mjtNum* tendon_range; - mjtNum* tendon_stiffness; - mjtNum* tendon_damping; - mjtNum* tendon_frictionloss; - mjtNum* tendon_lengthspring; - float* tendon_rgba; - - int* actuator_trntype; - int* actuator_dyntype; - int* actuator_trnid; - int* actuator_actadr; - int* actuator_actnum; - int* actuator_group; - mjtByte* actuator_ctrllimited; - mjtByte* actuator_actlimited; - mjtNum* actuator_ctrlrange; - mjtNum* actuator_actrange; - mjtNum* actuator_cranklength; - - int* sensor_type; - int* sensor_objid; - int* sensor_adr; - - int* name_bodyadr; - int* name_jntadr; - int* name_geomadr; - int* name_siteadr; - int* name_camadr; - int* name_lightadr; - int* name_eqadr; - int* name_tendonadr; - int* name_actuatoradr; - char* names; - char* paths; - } model; - - // fields in mjData that are necessary to re-render a scene - struct { - mjWarningStat warning[mjNWARNING]; - - int nefc; - int ncon; - int nisland; - - mjtNum time; - - mjtNum* act; - - mjtNum* ctrl; - mjtNum* xfrc_applied; - mjtByte* eq_active; - - mjtNum* sensordata; - - mjtNum* xpos; - mjtNum* xquat; - mjtNum* xmat; - mjtNum* xipos; - mjtNum* ximat; - mjtNum* xanchor; - mjtNum* xaxis; - mjtNum* geom_xpos; - mjtNum* geom_xmat; - mjtNum* site_xpos; - mjtNum* site_xmat; - mjtNum* cam_xpos; - mjtNum* cam_xmat; - mjtNum* light_xpos; - mjtNum* light_xdir; - - mjtNum* subtree_com; - - int* ten_wrapadr; - int* ten_wrapnum; - int* wrap_obj; - mjtNum* ten_length; - mjtNum* wrap_xpos; - - mjtNum* bvh_aabb_dyn; - mjtByte* bvh_active; - int* island_dofadr; - int* island_dofind; - int* dof_island; - int* efc_island; - int* tendon_efcadr; - - mjtNum* flexvert_xpos; - - mjContact* contact; - mjtNum* efc_force; - void* arena; - } data; -}; -typedef struct mjvSceneState_ mjvSceneState; //----------------------------- MJAPI FUNCTIONS -------------------------------- void mj_defaultVFS(mjVFS* vfs); @@ -3181,6 +2992,7 @@ mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int err mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz); mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz); mjModel* mj_compile(mjSpec* s, const mjVFS* vfs); +int mj_copyBack(mjSpec* s, const mjModel* m); int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz); void mj_freeLastXML(void); @@ -3204,6 +3016,7 @@ void mj_deleteModel(mjModel* m); int mj_sizeModel(const mjModel* m); mjData* mj_makeData(const mjModel* m); mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src); +mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src); void mj_resetData(const mjModel* m, mjData* d); void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value); void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key); @@ -3259,6 +3072,7 @@ void mj_flex(const mjModel* m, mjData* d); void mj_tendon(const mjModel* m, mjData* d); void mj_transmission(const mjModel* m, mjData* d); void mj_crb(const mjModel* m, mjData* d); +void mj_makeM(const mjModel* m, mjData* d); void mj_factorM(const mjModel* m, mjData* d); void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, @@ -3357,14 +3171,8 @@ mjtNum mjv_frustumHeight(const mjvScene* scn); void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum forward[3]); void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy, const mjvScene* scn, mjvCamera* cam); -void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvCamera* cam); void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx, mjtNum reldy, const mjvScene* scn, mjvPerturb* pert); -void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvPerturb* pert); void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy, const mjtNum roomup[3], mjvScene* scn); void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPerturb* pert); @@ -3387,16 +3195,7 @@ void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom); void mjv_freeScene(mjvScene* scn); void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt, const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn); -int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt, - const mjvPerturb* pert, mjvCamera* cam, int catmask, - mjvScene* scn); void mjv_copyModel(mjModel* dest, const mjModel* src); -void mjv_defaultSceneState(mjvSceneState* scnstate); -void mjv_makeSceneState(const mjModel* m, const mjData* d, - mjvSceneState* scnstate, int maxgeom); -void mjv_freeSceneState(mjvSceneState* scnstate); -void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt, - mjvSceneState* scnstate); void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt, const mjvPerturb* pert, int catmask, mjvScene* scn); void mjv_makeLights(const mjModel* m, const mjData* d, mjvScene* scn); @@ -3580,6 +3379,10 @@ void mju_insertionSortInt(int* list, int n); mjtNum mju_Halton(int index, int base); char* mju_strncpy(char *dst, const char *src, int n); mjtNum mju_sigmoid(mjtNum x); +const mjpPlugin* mjc_getSDF(const mjModel* m, int id); +mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]); +void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3], + const mjtNum x[3]); void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered, mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D); void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation, @@ -3605,14 +3408,8 @@ void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task); void mju_threadPoolDestroy(mjThreadPool* thread_pool); void mju_defaultTask(mjTask* task); void mju_taskJoin(mjTask* task); -mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child, - const char* prefix, const char* suffix); -mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child, - const char* prefix, const char* suffix); -mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child, - const char* prefix, const char* suffix); -mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child, - const char* prefix, const char* suffix); +mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child, + const char* prefix, const char* suffix); int mjs_detachBody(mjSpec* s, mjsBody* b); int mjs_detachDefault(mjSpec* s, mjsDefault* d); mjsBody* mjs_addBody(mjsBody* body, const mjsDefault* def); @@ -3675,12 +3472,16 @@ void mjs_setDouble(mjDoubleVec* dest, const double* array, int size); void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes); const char* mjs_getString(const mjString* source); const double* mjs_getDouble(const mjDoubleVec* source, int* size); +const void* mjs_getPluginAttributes(const mjsPlugin* plugin); void mjs_setDefault(mjsElement* element, const mjsDefault* def); -void mjs_setFrame(mjsElement* dest, mjsFrame* frame); +int mjs_setFrame(mjsElement* dest, mjsFrame* frame); const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence, const mjsOrientation* orientation); mjsFrame* mjs_bodyToFrame(mjsBody** body); void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); const void* mjs_getUserValue(mjsElement* element, const char* key); void mjs_deleteUserValue(mjsElement* element, const char* key); void mjs_defaultSpec(mjSpec* spec); diff --git a/doc/mjx.rst b/doc/mjx.rst index 98fb36b3..e0106b5f 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -235,6 +235,8 @@ The following features are **fully supported** in MJX: - 1, 3, 4, 6 (1 is not supported with ``ELLIPTIC``) * - :ref:`Solver ` - ``CG``, ``NEWTON`` + * - Dynamics + - :ref:`Inverse ` * - Fluid Model - :ref:`flInertia` * - :ref:`Tendons ` @@ -244,7 +246,7 @@ The following features are **fully supported** in MJX: ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``, ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``TENDONVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, - ``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC`` + ``ACTUATORFRC``, ``JOINTACTFRC``, ``TENDONACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC`` (``ACCELEROMETER``, ``FORCE``, ``TORQUE`` not supported with connect or weld equality constraints) The following features are **in development** and coming soon: @@ -262,8 +264,6 @@ The following features are **in development** and coming soon: (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``. * - :ref:`Integrator ` - ``IMPLICIT`` - * - Dynamics - - :ref:`Inverse ` * - Fluid Model - :ref:`flEllipsoid` * - :ref:`Sensors ` diff --git a/doc/modeling.rst b/doc/modeling.rst index 4ef331a4..61d79c13 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -724,7 +724,10 @@ Force clamping at joint input with :ref:`joint/actuatorfrcrange` sensor to report the total actuator force acting on a joint. The standard :ref:`actuatorfrc` sensor will continue to report the pre-clamped actuator force. -The three clamping options above are non-exclusive and can be combined as required. +Force clamping at tendon input with :ref:`tendon/actuatorfrcrange`: + This tendon attribute clamps input forces from all actuators acting on the tendon. + +The clamping options above are non-exclusive and can be combined as required. .. _CLengthRange: diff --git a/doc/programming/modeledit.rst b/doc/programming/modeledit.rst index cec8f03d..ac74efdf 100644 --- a/doc/programming/modeledit.rst +++ b/doc/programming/modeledit.rst @@ -111,9 +111,10 @@ This framework introduces a powerful new feature: attaching and detaching model to power the :ref:`attach` an :ref:`replicate` meta-elements in MJCF. Attachment allows the user to move or copy a subtree from one model into another, while also copying or moving related referenced assets and referencing elements from outside the kinematic tree (e.g., actuators and sensors). Similarly, detaching a subtree will -remove all associated elements from the model. The default behavior is to move during attach. The user can select to -instead copy by passing the corresponding flag to ``mjs_setDeepCopy``. This flag is temporary set to true while parsing -XMLs. It is possible to :ref:`attach a body to a frame`: +remove all associated elements from the model. The default behavior is to move the child into the parent while +attaching, so subsequent changes to the child will also change the parent. Alternatively, the user can choose to make an +entirely new copy during attach using :ref:`mjs_setDeepCopy`. This flag is temporarily set to true while parsing XMLs. +It is possible to :ref:`attach a body to a frame`: .. code-block:: C @@ -121,29 +122,29 @@ XMLs. It is possible to :ref:`attach a body to a frame`: mjSpec* child = mj_makeSpec(); parent->compiler.degree = 0; child->compiler.degree = 1; - mjsFrame* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL); - mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL); - mjsBody* attached_body_1 = mjs_attachBody(frame, body, "attached-", "-1"); + mjsElement* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL)->element; + mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element; + mjsBody* attached_body_1 = mjs_asBody(mjs_attach(frame, body, "attached-", "-1")); -or :ref:`attach a body to a site`: +or :ref:`attach a body to a site`: .. code-block:: C mjSpec* parent = mj_makeSpec(); mjSpec* child = mj_makeSpec(); - mjsSite* site = mjs_addSite(mjs_findBody(parent, "world"), NULL); - mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL); - mjsBody* attached_body_2 = mjs_attachToSite(site, body, "attached-", "-2"); + mjsElement* site = mjs_addSite(mjs_findBody(parent, "world"), NULL)->element; + mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element; + mjsBody* attached_body_2 = mjs_asBody(mjs_attach(site, body, "attached-", "-2")); -or :ref:`attach a frame to a body`: +or :ref:`attach a frame to a body`: .. code-block:: C mjSpec* parent = mj_makeSpec(); mjSpec* child = mj_makeSpec(); - mjsBody* body = mjs_addBody(mjs_findBody(parent, "world"), NULL); - mjsFrame* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL); - mjsFrame* attached_frame = mjs_attachFrame(body, frame, "attached-", "-1"); + mjsElement* body = mjs_addBody(mjs_findBody(parent, "world"), NULL)->element; + mjsElement* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL)->element; + mjsFrame* attached_frame = mjs_asFrame(mjs_attach(body, frame, "attached-", "-1")); Note that in the above examples, the parent and child models have different values for ``compiler.degree``, corresponding to the :ref:`compiler/angle` attribute, specifying the units in which angles are diff --git a/doc/programming/simulation.rst b/doc/programming/simulation.rst index 3739f548..26bf3fe7 100644 --- a/doc/programming/simulation.rst +++ b/doc/programming/simulation.rst @@ -807,7 +807,7 @@ to implement high-resolution timers in C without bringing in additional dependen does not need timing, and in that case there is no reason to call timing functions. One part of the simulation pipeline that needs to be monitored closely is the iterative constraint solver. The -simplest diagnostic here is ``mjData.solver_iter`` which shows how many iterations the solver took on the last call to +simplest diagnostic here is ``mjData.solver_niter`` which shows how many iterations the solver took on the last call to mj_step or ``mj_forward``. Note that the solver has tolerance parameters for early termination, so this number is usually smaller than the maximum number of iterations allowed. The array ``mjData.solver`` contains one :ref:`mjSolverStat` data structure per iteration of the constraint solver, with information about the constraint state @@ -913,7 +913,7 @@ elliptic, depending on which solver is selected in ``mjModel.opt``. The function can be used to determine which friction cone type is used. For pyramidal cones, the interpretation of the contact force (whose address we computed above) is non-trivial, because the components are forces along redundant non-orthogonal axes corresponding to the edges of the pyramid. The function :ref:`mj_contactForce` can be -used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D toque. +used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D torque. The torque component will be zero when :at:`condim` is 1 or 3, and non-zero otherwise. This force and torque are expressed in the contact frame given by mjContact.frame. Unlike all other matrices in mjData, this matrix is stored in transposed form. Normally a 3-by-3 matrix corresponding to a coordinate frame would have the frame axes along the diff --git a/doc/python.rst b/doc/python.rst index edb2067b..61ebcdc8 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -601,6 +601,14 @@ Lists of all elements in a spec can be accessed using named properties, using th ``equalities``, ``tendons``, ``actuators``, ``skins``, ``textures``, ``texts``, ``tuples``, ``flexes``, ``hfields``, ``keys``, ``numerics``, ``excludes``, ``sensors``, ``plugins``. +Element removal +^^^^^^^^^^^^^^^ +For elements that can have children (bodies and defaults), the methods ``spec.detach_body(body)`` and +``spec.detach_default(def)`` remove, respectively, ``body`` and ``def`` from the spec, together with all of their +children. When detaching body subtrees, all elements which reference elements in the subtree, will also be removed. For +all other elements, the method ``delete()`` removes the corresponding element from the spec, e.g. +``spec.geom('my_geom').delete()`` will remove the geom named "my_geom" and all of the elements that reference it. + Tree traversal ^^^^^^^^^^^^^^ Traversal of the kinematic tree is aided by the following methods which return tree-related lists of elements: @@ -620,8 +628,16 @@ Parent: The parent body of a given element -- including bodies and frames -- can be accessed via the ``parent`` property. For example, the parent of a site can be accessed via ``site.parent``. -Relationship to ``PyMJCF`` --------------------------- +Serialization +^^^^^^^^^^^^^ +The ``MjSpec`` object can be serialized with all of its assets using the function ``spec.to_zip(file)``, where ``file`` +can be either a path to a file or a file object. In order to load the spec from a zip file, use ``spec = +MjSpec.from_zip(file)``, where ``file`` is a path to a zip file or a zip file object. + +.. _PyMJCF: + +Relationship to ``PyMJCF`` and ``bind`` +--------------------------------------- `dm_control `__'s `PyMJCF `__ module provides similar @@ -637,9 +653,8 @@ includes a reimplementation of the ``PyMJCF`` example in the ``dm_control`` ``PyMJCF`` provides a notion of "binding", giving access to :ref:`mjModel` and :ref:`mjData` values via a helper class. In the native API, the helper class is not needed, so it is possible to directly bind an ``mjs`` object to -:ref:`mjModel` and :ref:`mjData`. This requires the objects to have a non-empty name. For example, say we have multiple -geoms containing the string "torso" in their name. We want to get their Cartesian positions in the XY plane from -``mjData``. This can be done as follows: +:ref:`mjModel` and :ref:`mjData`. For example, say we have multiple geoms containing the string "torso" in their name. +We want to get their Cartesian positions in the XY plane from ``mjData``. This can be done as follows: .. code-block:: python @@ -647,6 +662,9 @@ geoms containing the string "torso" in their name. We want to get their Cartesia pos_x = [torso.xpos[0] for torso in torsos] pos_y = [torso.xpos[1] for torso in torsos] +Using the ``bind`` method requires the :ref:`mjModel` and :ref:`mjData` to be compiled from the :`ref:`mjSpec`. If +objects are added or removed from the :ref:`mjSpec` since the last compilation, an error is raised. + Notes ----- diff --git a/doc/requirements.txt b/doc/requirements.txt index f5bddde1..aa0c90a4 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -1,6 +1,6 @@ Sphinx==5.3.0 furo==2022.9.29 -sphinxcontrib-bibtex==2.6.1 +sphinxcontrib-bibtex==2.6.3 sphinxcontrib-katex==0.9.4 sphinxcontrib-youtube==1.2.0 sphinx-copybutton==0.5.2 @@ -10,7 +10,7 @@ sphinx-toolbox==3.8.2 nbsphinx==0.9.1 pandoc==1.1.0 pygments==2.15.0 -jq==1.4.1 +jq==1.8.0 Jinja2~=3.0 wheel # see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe diff --git a/doc/unity.rst b/doc/unity.rst index 1f499b8f..876a96a6 100644 --- a/doc/unity.rst +++ b/doc/unity.rst @@ -37,14 +37,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.3.3.1.dylib`` (it can be +``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.3.3.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-3.3.1/lib/libmujoco.so.3.3.1`` and rename it as ``libmujoco.so``. +``~/.mujoco/mujoco-3.3.3/lib/libmujoco.so.3.3.3`` and rename it as ``libmujoco.so``. Windows _______ diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index ddb42215..1d894fd4 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -182,7 +182,6 @@ struct mjData_ { // solver statistics mjSolverStat solver[mjNISLAND*mjNSOLVER]; // solver statistics per island, per iteration - int solver_nisland; // number of islands processed by solver int solver_niter[mjNISLAND]; // number of solver iterations, per island int solver_nnz[mjNISLAND]; // number of nonzeros in Hessian or efc_AR, per island mjtNum solver_fwdinv[2]; // forward-inverse comparison: qfrc, efc @@ -200,6 +199,7 @@ struct mjData_ { int nJ; // number of non-zeros in constraint Jacobian int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands + int nidof; // number of dofs in all islands // global properties mjtNum time; // simulation time @@ -295,12 +295,13 @@ struct mjData_ { int* moment_colind; // column indices in sparse Jacobian (nJmom x 1) mjtNum* actuator_moment; // actuator moments (nJmom x 1) - // computed by mj_fwdPosition/mj_crb + // computed by mj_fwdPosition/mj_makeM mjtNum* crb; // com-based composite inertia and mass (nbody x 10) - mjtNum* qM; // total inertia (sparse) (nM x 1) + mjtNum* qM; // inertia (sparse) (nM x 1) + mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1) // computed by mj_fwdPosition/mj_factorM - mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) // computed by mj_collisionTree @@ -333,27 +334,23 @@ struct mjData_ { mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) // computed by mj_Euler or mj_implicit - mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) + mjtNum* qH; // L'*D*L factorization of modified M (nC x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) // computed by mj_resetData int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1) int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1) int* B_colind; // body-dof: column indices of non-zeros (nB x 1) - int* M_rownnz; // inertia: non-zeros in each row (nv x 1) - int* M_rowadr; // inertia: address of each row in M_colind (nv x 1) - int* M_colind; // inertia: column indices of non-zeros (nM x 1) - int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1) - int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1) - int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1) - int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1) - int* mapM2C; // index mapping from M to C (nC x 1) - int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1) - int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1) - int* D_diag; // dof-dof: index of diagonal element (nv x 1) - int* D_colind; // dof-dof: column indices of non-zeros (nD x 1) - int* mapM2D; // index mapping from M to D (nD x 1) - int* mapD2M; // index mapping from D to M (nM x 1) + int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1) + int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1) + int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1) + int* mapM2M; // index mapping from qM to M (nC x 1) + int* D_rownnz; // full inertia: non-zeros in each row (nv x 1) + int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1) + int* D_diag; // full inertia: index of diagonal element (nv x 1) + int* D_colind; // full inertia: column indices of non-zeros (nD x 1) + int* mapM2D; // index mapping from qM to D (nD x 1) + int* mapD2M; // index mapping from D to qM (nM x 1) // computed by mj_implicit/mj_derivative mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1) @@ -375,8 +372,8 @@ struct mjData_ { mjtNum* qfrc_constraint; // constraint force (nv x 1) // computed by mj_inverse - mjtNum* qfrc_inverse; // net external force; should equal: (nv x 1) - // qfrc_applied + J'*xfrc_applied + qfrc_actuator + mjtNum* qfrc_inverse; // net external force; should equal: + // qfrc_applied + J'*xfrc_applied + qfrc_actuator (nv x 1) // computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format mjtNum* cacc; // com-based acceleration (nbody x 6) @@ -410,16 +407,50 @@ struct mjData_ { mjtNum* efc_R; // inverse constraint mass (nefc x 1) int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1) - // computed by mj_island + // computed by mj_island (island dof structure) int* dof_island; // island id of this dof; -1: none (nv x 1) - int* island_dofnum; // number of dofs in island (nisland x 1) - int* island_dofadr; // start address in island_dofind (nisland x 1) - int* island_dofind; // island dof indices; -1: none (nv x 1) - int* dof_islandind; // dof island indices; -1: none (nv x 1) + int* island_nv; // number of dofs in this island (nisland x 1) + int* island_idofadr; // island start address in idof vector (nisland x 1) + int* island_dofadr; // island start address in dof vector (nisland x 1) + int* map_dof2idof; // map from dof to idof (nv x 1) + int* map_idof2dof; // map from idof to dof; >= nidof: unconstrained (nv x 1) + + // computed by mj_island (dofs sorted by island) + mjtNum* ifrc_smooth; // net unconstrained force (nidof x 1) + mjtNum* iacc_smooth; // unconstrained acceleration (nidof x 1) + int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) + int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) + int* iM_colind; // inertia: column indices of non-zeros (nC x 1) + mjtNum* iM; // total inertia (sparse) (nC x 1) + mjtNum* iLD; // L'*D*L factorization of M (sparse) (nC x 1) + mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) + mjtNum* iacc; // acceleration (nidof x 1) + + // computed by mj_island (island constraint structure) int* efc_island; // island id of this constraint (nefc x 1) - int* island_efcnum; // number of constraints in island (nisland x 1) - int* island_efcadr; // start address in island_efcind (nisland x 1) - int* island_efcind; // island constraint indices (nefc x 1) + int* island_ne; // number of equality constraints in island (nisland x 1) + int* island_nf; // number of friction constraints in island (nisland x 1) + int* island_nefc; // number of constraints in island (nisland x 1) + int* island_iefcadr; // start address in iefc vector (nisland x 1) + int* map_efc2iefc; // map from efc to iefc (nefc x 1) + int* map_iefc2efc; // map from iefc to efc (nefc x 1) + + // computed by mj_island (constraints sorted by island) + int* iefc_type; // constraint type (mjtConstraint) (nefc x 1) + int* iefc_id; // id of object of specified type (nefc x 1) + int* iefc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* iefc_J_rowadr; // row start address in colind array (nefc x 1) + int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) + int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) + int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) + int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) + int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) + mjtNum* iefc_J; // constraint Jacobian (nJ x 1) + mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) + mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* iefc_D; // constraint mass (nefc x 1) + mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) @@ -437,11 +468,18 @@ struct mjData_ { // computed by mj_fwdConstraint/mj_inverse mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1) + int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* iefc_force; // constraint force in constraint space (nefc x 1) int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* ifrc_constraint; // constraint force (nidof x 1) // thread pool pointer uintptr_t threadpool; + + // compilation signature + uint64_t signature; // also held by the mjSpec that compiled the model }; typedef struct mjData_ mjData; diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 0099cf78..75eb3cd0 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -129,6 +129,14 @@ typedef enum mjtCamLight_ { // tracking mode for camera and light } mjtCamLight; +typedef enum mjtLightType_ { // type of light + mjLIGHT_SPOT = 0, // spot + mjLIGHT_DIRECTIONAL, // directional + mjLIGHT_POINT, // point + mjLIGHT_IMAGE, // image-based +} mjtLightType; + + typedef enum mjtTexture_ { // type of texture mjTEXTURE_2D = 0, // 2d texture, suitable for planes and hfields mjTEXTURE_CUBE, // cube texture, suitable for all other geom types @@ -151,6 +159,13 @@ typedef enum mjtTextureRole_ { // role of texture map in rendering } mjtTextureRole; +typedef enum mjtColorSpace_ { // type of color space encoding + mjCOLORSPACE_AUTO = 0, // attempts to autodetect color space, defaults to linear + mjCOLORSPACE_LINEAR, // linear color space + mjCOLORSPACE_SRGB // standard RGB color space +} mjtColorSpace; + + typedef enum mjtIntegrator_ { // integrator mode mjINT_EULER = 0, // semi-implicit Euler mjINT_RK4, // 4th-order Runge Kutta @@ -269,7 +284,8 @@ typedef enum mjtObj_ { // type of MujoCo object // meta elements, do not appear in mjModel mjOBJ_FRAME = 100, // frame - mjOBJ_DEFAULT // default + mjOBJ_DEFAULT, // default + mjOBJ_MODEL // entire model } mjtObj; @@ -316,6 +332,7 @@ typedef enum mjtSensor_ { // type of sensor mjSENS_ACTUATORVEL, // scalar actuator velocity mjSENS_ACTUATORFRC, // scalar actuator force mjSENS_JOINTACTFRC, // scalar actuator force, measured at the joint + mjSENS_TENDONACTFRC, // scalar actuator force, measured at the tendon // sensors related to ball joints mjSENS_BALLQUAT, // 4D ball joint quaternion @@ -405,6 +422,14 @@ typedef enum mjtFlexSelf_ { // mode for flex selfcollide } mjtFlexSelf; +typedef enum mjtSDFType_ { // signed distance function (SDF) type + mjSDFTYPE_SINGLE = 0, // single SDF + mjSDFTYPE_INTERSECTION, // max(A, B) + mjSDFTYPE_MIDSURFACE, // A - B + mjSDFTYPE_COLLISION, // A + B + abs(max(A, B)) +} mjtSDFType; + + //---------------------------------- mjLROpt ------------------------------------------------------- struct mjLROpt_ { // options for mj_setLengthRange() @@ -611,7 +636,7 @@ struct mjModel_ { int ncam; // number of cameras int nlight; // number of lights int nflex; // number of flexes - int nflexnode; // number of dofs in all flexes + int nflexnode; // number of dofs in all flexes int nflexvert; // number of vertices in all flexes int nflexedge; // number of edges in all flexes int nflexelem; // number of elements in all flexes @@ -832,9 +857,12 @@ struct mjModel_ { int* light_mode; // light tracking mode (mjtCamLight) (nlight x 1) int* light_bodyid; // id of light's body (nlight x 1) int* light_targetbodyid; // id of targeted body; -1: none (nlight x 1) - mjtByte* light_directional; // directional light (nlight x 1) + int* light_type; // spot, directional, etc. (mjtLightType) (nlight x 1) + int* light_texid; // texture id for image lights (nlight x 1) mjtByte* light_castshadow; // does light cast shadows (nlight x 1) float* light_bulbradius; // light radius for soft shadows (nlight x 1) + float* light_intensity; // intensity, in candela (nlight x 1) + float* light_range; // range of effectiveness (nlight x 1) mjtByte* light_active; // is light on (nlight x 1) mjtNum* light_pos; // position rel. to body frame (nlight x 3) mjtNum* light_dir; // direction rel. to body frame (nlight x 3) @@ -886,6 +914,7 @@ struct mjModel_ { int* flex_nodebodyid; // node body ids (nflexnode x 1) int* flex_vertbodyid; // vertex body ids (nflexvert x 1) int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2) + int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2) int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1) int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1) int* flex_elemedge; // element edge ids (nflexelemedge x 1) @@ -900,6 +929,7 @@ struct mjModel_ { mjtNum* flexedge_invweight0; // edge inv. weight in qpos0 (nflexedge x 1) mjtNum* flex_radius; // radius around primitive element (nflex x 1) mjtNum* flex_stiffness; // finite element stiffness matrix (nflexelem x 21) + mjtNum* flex_bending; // bending stiffness (nflexedge x 16) mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1) mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1) mjtNum* flex_edgedamping; // edge damping (nflex x 1) @@ -980,6 +1010,7 @@ struct mjModel_ { // textures int* tex_type; // texture type (mjtTexture) (ntex x 1) + int* tex_colorspace; // texture colorspace (mjtColorSpace) (ntex x 1) int* tex_height; // number of rows in texture image (ntex x 1) int* tex_width; // number of columns in texture image (ntex x 1) int* tex_nchannel; // number of channels in texture image (ntex x 1) @@ -1030,15 +1061,18 @@ struct mjModel_ { int* tendon_matid; // material id for rendering (ntendon x 1) int* tendon_group; // group for visibility (ntendon x 1) mjtByte* tendon_limited; // does tendon have length limits (ntendon x 1) + mjtByte* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1) mjtNum* tendon_width; // width for rendering (ntendon x 1) mjtNum* tendon_solref_lim; // constraint solver reference: limit (ntendon x mjNREF) mjtNum* tendon_solimp_lim; // constraint solver impedance: limit (ntendon x mjNIMP) mjtNum* tendon_solref_fri; // constraint solver reference: friction (ntendon x mjNREF) mjtNum* tendon_solimp_fri; // constraint solver impedance: friction (ntendon x mjNIMP) mjtNum* tendon_range; // tendon length limits (ntendon x 2) + mjtNum* tendon_actfrcrange; // range of total actuator force (ntendon x 2) mjtNum* tendon_margin; // min distance for limit detection (ntendon x 1) mjtNum* tendon_stiffness; // stiffness coefficient (ntendon x 1) mjtNum* tendon_damping; // damping coefficient (ntendon x 1) + mjtNum* tendon_armature; // inertia associated with tendon velocity (ntendon x 1) mjtNum* tendon_frictionloss; // loss due to friction (ntendon x 1) mjtNum* tendon_lengthspring; // spring resting length range (ntendon x 2) mjtNum* tendon_length0; // tendon length in qpos0 (ntendon x 1) @@ -1155,6 +1189,9 @@ struct mjModel_ { // paths char* paths; // paths to assets, 0-terminated (npaths x 1) + + // compilation signature + uint64_t signature; // also held by the mjSpec that compiled this model }; typedef struct mjModel_ mjModel; diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h index 0fc31a6c..44f2af20 100644 --- a/include/mujoco/mjplugin.h +++ b/include/mujoco/mjplugin.h @@ -135,6 +135,16 @@ struct mjpPlugin_ { }; typedef struct mjpPlugin_ mjpPlugin; +struct mjSDF_ { + const mjpPlugin** plugin; + int* id; + mjtSDFType type; + mjtNum* relpos; + mjtNum* relmat; + mjtGeom* geomtype; +}; +typedef struct mjSDF_ mjSDF; + #if defined(__has_attribute) #if __has_attribute(constructor) diff --git a/include/mujoco/mjspec.h b/include/mujoco/mjspec.h index 52e96f41..3c00f478 100644 --- a/include/mujoco/mjspec.h +++ b/include/mujoco/mjspec.h @@ -23,6 +23,7 @@ // this is a C-API #ifdef __cplusplus #include +#include #include #include @@ -119,6 +120,7 @@ typedef enum mjtOrientation_ { // type of orientation specifier typedef struct mjsElement_ { // element type, do not modify mjtObj elemtype; // element type + uint64_t signature; // compilation signature } mjsElement; @@ -136,6 +138,7 @@ typedef struct mjsCompiler_ { // compiler options mjtByte fusestatic; // fuse static bodies with parent int inertiafromgeom; // use geom inertias (mjtInertiaFromGeom) int inertiagrouprange[2]; // range of geom groups used to compute inertia + mjtByte saveinertial; // save explicit inertial clause for all bodies to XML int alignfree; // align free joints with inertial frame mjLROpt LRopt; // options for lengthrange computation } mjsCompiler; @@ -391,9 +394,12 @@ typedef struct mjsLight_ { // light specification // intrinsics mjtByte active; // is light active - mjtByte directional; // is light directional or spot + mjtLightType type; // type of light + mjString* texture; // texture name for image lights mjtByte castshadow; // does light cast shadows - double bulbradius; // bulb radius, for soft shadows + float bulbradius; // bulb radius, for soft shadows + float intensity; // intensity, in candelas + float range; // range of effectiveness float attenuation[3]; // OpenGL attenuation (quadratic model) float cutoff; // OpenGL cutoff float exponent; // OpenGL exponent @@ -427,7 +433,8 @@ typedef struct mjsFlex_ { // flex specification double radius; // radius around primitive element mjtByte internal; // enable internal collisions mjtByte flatskin; // render flex skin with flat shading - int selfcollide; // mode for flex self colllision + int selfcollide; // mode for flex self collision + int vertcollide; // mode for vertex collision int activelayers; // number of active element layers in 3D int group; // group for visualizatioh double edgestiffness; // edge stiffness @@ -438,6 +445,7 @@ typedef struct mjsFlex_ { // flex specification double poisson; // Poisson's ratio double damping; // Rayleigh's damping double thickness; // thickness (2D only) + int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both // mesh properties mjStringVec* nodebody; // node body names @@ -518,6 +526,7 @@ typedef struct mjsTexture_ { // texture specification mjsElement* element; // element type mjString* name; // name mjtTexture type; // texture type + mjtColorSpace colorspace; // colorspace // method 1: builtin int builtin; // builtin type (mjtBuiltin) @@ -614,17 +623,20 @@ typedef struct mjsTendon_ { // tendon specification mjsElement* element; // element type mjString* name; // name - // stiffness, damping, friction + // stiffness, damping, friction, armature double stiffness; // stiffness coefficient double springlength[2]; // spring resting length; {-1, -1}: use qpos_spring double damping; // damping coefficient double frictionloss; // friction loss mjtNum solref_friction[mjNREF]; // solver reference: tendon friction mjtNum solimp_friction[mjNIMP]; // solver impedance: tendon friction + double armature; // inertia associated with tendon velocity // length range int limited; // does tendon have limits (mjtLimited) + int actfrclimited; // does tendon have actuator force limits double range[2]; // length limits + double actfrcrange[2]; // actuator force limits double margin; // margin value for tendon limit detection mjtNum solref_limit[mjNREF]; // solver reference: tendon limits mjtNum solimp_limit[mjNIMP]; // solver impedance: tendon limits diff --git a/include/mujoco/mjvisualize.h b/include/mujoco/mjvisualize.h index 19dad80c..efff4d03 100644 --- a/include/mujoco/mjvisualize.h +++ b/include/mujoco/mjvisualize.h @@ -226,7 +226,7 @@ typedef struct mjvGLCamera_ mjvGLCamera; struct mjvGeom_ { // abstract geom // type info int type; // geom type (mjtGeom) - int dataid; // mesh, hfield or plane id; -1: none + int dataid; // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull) int objtype; // mujoco object type; mjOBJ_UNKNOWN for decor int objid; // mujoco object id; -1 for decor int category; // visual category @@ -261,6 +261,8 @@ typedef struct mjvGeom_ mjvGeom; struct mjvLight_ { // OpenGL light float pos[3]; // position rel. to body frame float dir[3]; // direction rel. to body frame + int type; // type (mjtLightType) + int texid; // texture id for image lights float attenuation[3]; // OpenGL attenuation (quadratic model) float cutoff; // OpenGL cutoff float exponent; // OpenGL exponent @@ -268,9 +270,10 @@ struct mjvLight_ { // OpenGL light float diffuse[3]; // diffuse rgb (alpha=1) float specular[3]; // specular rgb (alpha=1) mjtByte headlight; // headlight - mjtByte directional; // directional light mjtByte castshadow; // does light cast shadows float bulbradius; // bulb radius for soft shadows + float intensity; // intensity, in candelas + float range; // range of effectiveness }; typedef struct mjvLight_ mjvLight; @@ -404,289 +407,4 @@ struct mjvFigure_ { // abstract 2D figure passed to OpenGL rendere }; typedef struct mjvFigure_ mjvFigure; - -//---------------------------------- mjvSceneState ------------------------------------------------- - -struct mjvSceneState_ { - int nbuffer; // size of the buffer in bytes - void* buffer; // heap-allocated memory for all arrays in this struct - int maxgeom; // maximum number of mjvGeom supported by this state object - mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins - - // fields in mjModel that are necessary to re-render a scene - struct { - int nv; - int nu; - int na; - int nbody; - int nbvh; - int nbvhstatic; - int njnt; - int ngeom; - int nsite; - int ncam; - int nlight; - int nmesh; - int nskin; - int nflex; - int nflexvert; - int nflextexcoord; - int nskinvert; - int nskinface; - int nskinbone; - int nskinbonevert; - int nmat; - int neq; - int ntendon; - int ntree; - int nwrap; - int nsensor; - int nnames; - int npaths; - int nsensordata; - int narena; - - mjOption opt; - mjVisual vis; - mjStatistic stat; - - int* body_parentid; - int* body_rootid; - int* body_weldid; - int* body_mocapid; - int* body_jntnum; - int* body_jntadr; - int* body_dofnum; - int* body_dofadr; - int* body_geomnum; - int* body_geomadr; - mjtNum* body_iquat; - mjtNum* body_mass; - mjtNum* body_inertia; - int* body_bvhadr; - int* body_bvhnum; - - int* bvh_depth; - int* bvh_child; - int* bvh_nodeid; - mjtNum* bvh_aabb; - - int* jnt_type; - int* jnt_bodyid; - int* jnt_group; - - int* geom_type; - int* geom_bodyid; - int* geom_contype; - int* geom_conaffinity; - int* geom_dataid; - int* geom_matid; - int* geom_group; - mjtNum* geom_size; - mjtNum* geom_aabb; - mjtNum* geom_rbound; - float* geom_rgba; - - int* site_type; - int* site_bodyid; - int* site_matid; - int* site_group; - mjtNum* site_size; - float* site_rgba; - - int* cam_orthographic; - mjtNum* cam_fovy; - mjtNum* cam_ipd; - int* cam_resolution; - float* cam_sensorsize; - float* cam_intrinsic; - - mjtByte* light_directional; - mjtByte* light_castshadow; - float* light_bulbradius; - mjtByte* light_active; - float* light_attenuation; - float* light_cutoff; - float* light_exponent; - float* light_ambient; - float* light_diffuse; - float* light_specular; - - mjtByte* flex_flatskin; - int* flex_dim; - int* flex_matid; - int* flex_group; - int* flex_interp; - int* flex_nodeadr; - int* flex_nodenum; - int* flex_nodebodyid; - int* flex_vertadr; - int* flex_vertnum; - int* flex_elem; - int* flex_elemtexcoord; - int* flex_elemlayer; - int* flex_elemadr; - int* flex_elemnum; - int* flex_elemdataadr; - int* flex_shell; - int* flex_shellnum; - int* flex_shelldataadr; - int* flex_texcoordadr; - int* flex_bvhadr; - int* flex_bvhnum; - mjtByte* flex_centered; - mjtNum* flex_node; - mjtNum* flex_radius; - float* flex_rgba; - float* flex_texcoord; - - int* hfield_pathadr; - - int* mesh_bvhadr; - int* mesh_bvhnum; - int* mesh_texcoordadr; - int* mesh_graphadr; - int* mesh_pathadr; - - int* skin_matid; - int* skin_group; - float* skin_rgba; - float* skin_inflate; - int* skin_vertadr; - int* skin_vertnum; - int* skin_texcoordadr; - int* skin_faceadr; - int* skin_facenum; - int* skin_boneadr; - int* skin_bonenum; - float* skin_vert; - int* skin_face; - int* skin_bonevertadr; - int* skin_bonevertnum; - float* skin_bonebindpos; - float* skin_bonebindquat; - int* skin_bonebodyid; - int* skin_bonevertid; - float* skin_bonevertweight; - int* skin_pathadr; - - int* tex_pathadr; - - int* mat_texid; - mjtByte* mat_texuniform; - float* mat_texrepeat; - float* mat_emission; - float* mat_specular; - float* mat_shininess; - float* mat_reflectance; - float* mat_metallic; - float* mat_roughness; - float* mat_rgba; - - int* eq_type; - int* eq_obj1id; - int* eq_obj2id; - int* eq_objtype; - mjtNum* eq_data; - - int* tendon_num; - int* tendon_matid; - int* tendon_group; - mjtByte* tendon_limited; - mjtNum* tendon_width; - mjtNum* tendon_range; - mjtNum* tendon_stiffness; - mjtNum* tendon_damping; - mjtNum* tendon_frictionloss; - mjtNum* tendon_lengthspring; - float* tendon_rgba; - - int* actuator_trntype; - int* actuator_dyntype; - int* actuator_trnid; - int* actuator_actadr; - int* actuator_actnum; - int* actuator_group; - mjtByte* actuator_ctrllimited; - mjtByte* actuator_actlimited; - mjtNum* actuator_ctrlrange; - mjtNum* actuator_actrange; - mjtNum* actuator_cranklength; - - int* sensor_type; - int* sensor_objid; - int* sensor_adr; - - int* name_bodyadr; - int* name_jntadr; - int* name_geomadr; - int* name_siteadr; - int* name_camadr; - int* name_lightadr; - int* name_eqadr; - int* name_tendonadr; - int* name_actuatoradr; - char* names; - char* paths; - } model; - - // fields in mjData that are necessary to re-render a scene - struct { - mjWarningStat warning[mjNWARNING]; - - int nefc; - int ncon; - int nisland; - - mjtNum time; - - mjtNum* act; - - mjtNum* ctrl; - mjtNum* xfrc_applied; - mjtByte* eq_active; - - mjtNum* sensordata; - - mjtNum* xpos; - mjtNum* xquat; - mjtNum* xmat; - mjtNum* xipos; - mjtNum* ximat; - mjtNum* xanchor; - mjtNum* xaxis; - mjtNum* geom_xpos; - mjtNum* geom_xmat; - mjtNum* site_xpos; - mjtNum* site_xmat; - mjtNum* cam_xpos; - mjtNum* cam_xmat; - mjtNum* light_xpos; - mjtNum* light_xdir; - - mjtNum* subtree_com; - - int* ten_wrapadr; - int* ten_wrapnum; - int* wrap_obj; - mjtNum* ten_length; - mjtNum* wrap_xpos; - - mjtNum* bvh_aabb_dyn; - mjtByte* bvh_active; - int* island_dofadr; - int* island_dofind; - int* dof_island; - int* efc_island; - int* tendon_efcadr; - - mjtNum* flexvert_xpos; - - mjContact* contact; - mjtNum* efc_force; - void* arena; - } data; -}; -typedef struct mjvSceneState_ mjvSceneState; - #endif // MUJOCO_MJVISUALIZE_H_ diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 34a6dee4..93542a88 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -66,92 +66,92 @@ //-------------------------------- mjModel --------------------------------------------------------- // int fields of mjModel -#define MJMODEL_INTS \ - X ( nq ) \ - XMJV( nv ) \ - XMJV( nu ) \ - XMJV( na ) \ - XMJV( nbody ) \ - XMJV( nbvh ) \ - XMJV( nbvhstatic ) \ - X ( nbvhdynamic ) \ - XMJV( njnt ) \ - XMJV( ngeom ) \ - XMJV( nsite ) \ - XMJV( ncam ) \ - XMJV( nlight ) \ - XMJV( nflex ) \ - X ( nflexnode ) \ - XMJV( nflexvert ) \ - X ( nflexedge ) \ - X ( nflexelem ) \ - X ( nflexelemdata ) \ - X ( nflexelemedge ) \ - X ( nflexshelldata ) \ - X ( nflexevpair ) \ - XMJV( nflextexcoord ) \ - XMJV( nmesh ) \ - X ( nmeshvert ) \ - X ( nmeshnormal ) \ - X ( nmeshtexcoord ) \ - X ( nmeshface ) \ - X ( nmeshgraph ) \ - X ( nmeshpoly ) \ - X ( nmeshpolyvert ) \ - X ( nmeshpolymap ) \ - XMJV( nskin ) \ - XMJV( nskinvert ) \ - X ( nskintexvert ) \ - XMJV( nskinface ) \ - XMJV( nskinbone ) \ - XMJV( nskinbonevert ) \ - X ( nhfield ) \ - X ( nhfielddata ) \ - X ( ntex ) \ - X ( ntexdata ) \ - XMJV( nmat ) \ - X ( npair ) \ - X ( nexclude ) \ - XMJV( neq ) \ - XMJV( ntendon ) \ - XMJV( nwrap ) \ - XMJV( nsensor ) \ - X ( nnumeric ) \ - X ( nnumericdata ) \ - X ( ntext ) \ - X ( ntextdata ) \ - X ( ntuple ) \ - X ( ntupledata ) \ - X ( nkey ) \ - X ( nmocap ) \ - X ( nplugin ) \ - X ( npluginattr ) \ - X ( nuser_body ) \ - X ( nuser_jnt ) \ - X ( nuser_geom ) \ - X ( nuser_site ) \ - X ( nuser_cam ) \ - X ( nuser_tendon ) \ - X ( nuser_actuator ) \ - X ( nuser_sensor ) \ - XMJV( nnames ) \ - XMJV( npaths ) \ - X ( nnames_map ) \ - X ( nM ) \ - X ( nB ) \ - X ( nC ) \ - X ( nD ) \ - X ( nJmom ) \ - XMJV( ntree ) \ - X ( ngravcomp ) \ - X ( nemax ) \ - X ( njmax ) \ - X ( nconmax ) \ - X ( nuserdata ) \ - XMJV( nsensordata ) \ - X ( npluginstate ) \ - X ( narena ) \ - X ( nbuffer ) +#define MJMODEL_INTS \ + X( nq ) \ + X( nv ) \ + X( nu ) \ + X( na ) \ + X( nbody ) \ + X( nbvh ) \ + X( nbvhstatic ) \ + X( nbvhdynamic ) \ + X( njnt ) \ + X( ngeom ) \ + X( nsite ) \ + X( ncam ) \ + X( nlight ) \ + X( nflex ) \ + X( nflexnode ) \ + X( nflexvert ) \ + X( nflexedge ) \ + X( nflexelem ) \ + X( nflexelemdata ) \ + X( nflexelemedge ) \ + X( nflexshelldata ) \ + X( nflexevpair ) \ + X( nflextexcoord ) \ + X( nmesh ) \ + X( nmeshvert ) \ + X( nmeshnormal ) \ + X( nmeshtexcoord ) \ + X( nmeshface ) \ + X( nmeshgraph ) \ + X( nmeshpoly ) \ + X( nmeshpolyvert ) \ + X( nmeshpolymap ) \ + X( nskin ) \ + X( nskinvert ) \ + X( nskintexvert ) \ + X( nskinface ) \ + X( nskinbone ) \ + X( nskinbonevert ) \ + X( nhfield ) \ + X( nhfielddata ) \ + X( ntex ) \ + X( ntexdata ) \ + X( nmat ) \ + X( npair ) \ + X( nexclude ) \ + X( neq ) \ + X( ntendon ) \ + X( nwrap ) \ + X( nsensor ) \ + X( nnumeric ) \ + X( nnumericdata ) \ + X( ntext ) \ + X( ntextdata ) \ + X( ntuple ) \ + X( ntupledata ) \ + X( nkey ) \ + X( nmocap ) \ + X( nplugin ) \ + X( npluginattr ) \ + X( nuser_body ) \ + X( nuser_jnt ) \ + X( nuser_geom ) \ + X( nuser_site ) \ + X( nuser_cam ) \ + X( nuser_tendon ) \ + X( nuser_actuator ) \ + X( nuser_sensor ) \ + X( nnames ) \ + X( npaths ) \ + X( nnames_map ) \ + X( nM ) \ + X( nB ) \ + X( nC ) \ + X( nD ) \ + X( nJmom ) \ + X( ntree ) \ + X( ngravcomp ) \ + X( nemax ) \ + X( njmax ) \ + X( nconmax ) \ + X( nuserdata ) \ + X( nsensordata ) \ + X( npluginstate ) \ + X( narena ) \ + X( nbuffer ) /* nbuffer needs to be the final field */ @@ -178,31 +178,31 @@ // pointer fields of mjModel -// XMJV means that the field is required to construct mjvScene -// (by default we define XMJV to be the same as X) +// XNV means that the field is not required to construct mjvScene +// (by default we define XNV to be the same as X) #define MJMODEL_POINTERS \ X ( mjtNum, qpos0, nq, 1 ) \ X ( mjtNum, qpos_spring, nq, 1 ) \ - XMJV( int, body_parentid, nbody, 1 ) \ - XMJV( int, body_rootid, nbody, 1 ) \ - XMJV( int, body_weldid, nbody, 1 ) \ - XMJV( int, body_mocapid, nbody, 1 ) \ - XMJV( int, body_jntnum, nbody, 1 ) \ - XMJV( int, body_jntadr, nbody, 1 ) \ - XMJV( int, body_dofnum, nbody, 1 ) \ - XMJV( int, body_dofadr, nbody, 1 ) \ + X ( int, body_parentid, nbody, 1 ) \ + X ( int, body_rootid, nbody, 1 ) \ + X ( int, body_weldid, nbody, 1 ) \ + X ( int, body_mocapid, nbody, 1 ) \ + X ( int, body_jntnum, nbody, 1 ) \ + X ( int, body_jntadr, nbody, 1 ) \ + X ( int, body_dofnum, nbody, 1 ) \ + X ( int, body_dofadr, nbody, 1 ) \ X ( int, body_treeid, nbody, 1 ) \ - XMJV( int, body_geomnum, nbody, 1 ) \ - XMJV( int, body_geomadr, nbody, 1 ) \ + X ( int, body_geomnum, nbody, 1 ) \ + X ( int, body_geomadr, nbody, 1 ) \ X ( mjtByte, body_simple, nbody, 1 ) \ X ( mjtByte, body_sameframe, nbody, 1 ) \ X ( mjtNum, body_pos, nbody, 3 ) \ X ( mjtNum, body_quat, nbody, 4 ) \ X ( mjtNum, body_ipos, nbody, 3 ) \ - XMJV( mjtNum, body_iquat, nbody, 4 ) \ - XMJV( mjtNum, body_mass, nbody, 1 ) \ + X ( mjtNum, body_iquat, nbody, 4 ) \ + X ( mjtNum, body_mass, nbody, 1 ) \ X ( mjtNum, body_subtreemass, nbody, 1 ) \ - XMJV( mjtNum, body_inertia, nbody, 3 ) \ + X ( mjtNum, body_inertia, nbody, 3 ) \ X ( mjtNum, body_invweight0, nbody, 2 ) \ X ( mjtNum, body_gravcomp, nbody, 1 ) \ X ( mjtNum, body_margin, nbody, 1 ) \ @@ -210,17 +210,17 @@ X ( int, body_plugin, nbody, 1 ) \ X ( int, body_contype, nbody, 1 ) \ X ( int, body_conaffinity, nbody, 1 ) \ - XMJV( int, body_bvhadr, nbody, 1 ) \ - XMJV( int, body_bvhnum, nbody, 1 ) \ - XMJV( int, bvh_depth, nbvh, 1 ) \ - XMJV( int, bvh_child, nbvh, 2 ) \ - XMJV( int, bvh_nodeid, nbvh, 1 ) \ - XMJV( mjtNum, bvh_aabb, nbvhstatic, 6 ) \ - XMJV( int, jnt_type, njnt, 1 ) \ + X ( int, body_bvhadr, nbody, 1 ) \ + X ( int, body_bvhnum, nbody, 1 ) \ + X ( int, bvh_depth, nbvh, 1 ) \ + X ( int, bvh_child, nbvh, 2 ) \ + X ( int, bvh_nodeid, nbvh, 1 ) \ + X ( mjtNum, bvh_aabb, nbvhstatic, 6 ) \ + X ( int, jnt_type, njnt, 1 ) \ X ( int, jnt_qposadr, njnt, 1 ) \ X ( int, jnt_dofadr, njnt, 1 ) \ - XMJV( int, jnt_bodyid, njnt, 1 ) \ - XMJV( int, jnt_group, njnt, 1 ) \ + X ( int, jnt_bodyid, njnt, 1 ) \ + X ( int, jnt_group, njnt, 1 ) \ X ( mjtByte, jnt_limited, njnt, 1 ) \ X ( mjtByte, jnt_actfrclimited, njnt, 1 ) \ X ( mjtByte, jnt_actgravcomp, njnt, 1 ) \ @@ -246,41 +246,41 @@ X ( mjtNum, dof_damping, nv, 1 ) \ X ( mjtNum, dof_invweight0, nv, 1 ) \ X ( mjtNum, dof_M0, nv, 1 ) \ - XMJV( int, geom_type, ngeom, 1 ) \ - XMJV( int, geom_contype, ngeom, 1 ) \ - XMJV( int, geom_conaffinity, ngeom, 1 ) \ + X ( int, geom_type, ngeom, 1 ) \ + X ( int, geom_contype, ngeom, 1 ) \ + X ( int, geom_conaffinity, ngeom, 1 ) \ X ( int, geom_condim, ngeom, 1 ) \ - XMJV( int, geom_bodyid, ngeom, 1 ) \ - XMJV( int, geom_dataid, ngeom, 1 ) \ - XMJV( int, geom_matid, ngeom, 1 ) \ - XMJV( int, geom_group, ngeom, 1 ) \ + X ( int, geom_bodyid, ngeom, 1 ) \ + X ( int, geom_dataid, ngeom, 1 ) \ + X ( int, geom_matid, ngeom, 1 ) \ + X ( int, geom_group, ngeom, 1 ) \ X ( int, geom_priority, ngeom, 1 ) \ X ( int, geom_plugin, ngeom, 1 ) \ X ( mjtByte, geom_sameframe, ngeom, 1 ) \ X ( mjtNum, geom_solmix, ngeom, 1 ) \ X ( mjtNum, geom_solref, ngeom, mjNREF ) \ X ( mjtNum, geom_solimp, ngeom, mjNIMP ) \ - XMJV( mjtNum, geom_size, ngeom, 3 ) \ - XMJV( mjtNum, geom_aabb, ngeom, 6 ) \ - XMJV( mjtNum, geom_rbound, ngeom, 1 ) \ + X ( mjtNum, geom_size, ngeom, 3 ) \ + X ( mjtNum, geom_aabb, ngeom, 6 ) \ + X ( mjtNum, geom_rbound, ngeom, 1 ) \ X ( mjtNum, geom_pos, ngeom, 3 ) \ X ( mjtNum, geom_quat, ngeom, 4 ) \ X ( mjtNum, geom_friction, ngeom, 3 ) \ X ( mjtNum, geom_margin, ngeom, 1 ) \ X ( mjtNum, geom_gap, ngeom, 1 ) \ - X ( mjtNum, geom_fluid, ngeom, mjNFLUID ) \ + XNV ( mjtNum, geom_fluid, ngeom, mjNFLUID ) \ X ( mjtNum, geom_user, ngeom, MJ_M(nuser_geom) ) \ - XMJV( float, geom_rgba, ngeom, 4 ) \ - XMJV( int, site_type, nsite, 1 ) \ - XMJV( int, site_bodyid, nsite, 1 ) \ - XMJV( int, site_matid, nsite, 1 ) \ - XMJV( int, site_group, nsite, 1 ) \ + X ( float, geom_rgba, ngeom, 4 ) \ + X ( int, site_type, nsite, 1 ) \ + X ( int, site_bodyid, nsite, 1 ) \ + X ( int, site_matid, nsite, 1 ) \ + X ( int, site_group, nsite, 1 ) \ X ( mjtByte, site_sameframe, nsite, 1 ) \ - XMJV( mjtNum, site_size, nsite, 3 ) \ + X ( mjtNum, site_size, nsite, 3 ) \ X ( mjtNum, site_pos, nsite, 3 ) \ X ( mjtNum, site_quat, nsite, 4 ) \ X ( mjtNum, site_user, nsite, MJ_M(nuser_site) ) \ - XMJV( float, site_rgba, nsite, 4 ) \ + X ( float, site_rgba, nsite, 4 ) \ X ( int, cam_mode, ncam, 1 ) \ X ( int, cam_bodyid, ncam, 1 ) \ X ( int, cam_targetbodyid, ncam, 1 ) \ @@ -289,31 +289,34 @@ X ( mjtNum, cam_poscom0, ncam, 3 ) \ X ( mjtNum, cam_pos0, ncam, 3 ) \ X ( mjtNum, cam_mat0, ncam, 9 ) \ - XMJV( int, cam_orthographic, ncam, 1 ) \ - XMJV( mjtNum, cam_fovy, ncam, 1 ) \ - XMJV( mjtNum, cam_ipd, ncam, 1 ) \ - XMJV( int, cam_resolution, ncam, 2 ) \ - XMJV( float, cam_sensorsize, ncam, 2 ) \ - XMJV( float, cam_intrinsic, ncam, 4 ) \ + X ( int, cam_orthographic, ncam, 1 ) \ + X ( mjtNum, cam_fovy, ncam, 1 ) \ + X ( mjtNum, cam_ipd, ncam, 1 ) \ + X ( int, cam_resolution, ncam, 2 ) \ + X ( float, cam_sensorsize, ncam, 2 ) \ + X ( float, cam_intrinsic, ncam, 4 ) \ X ( mjtNum, cam_user, ncam, MJ_M(nuser_cam) ) \ X ( int, light_mode, nlight, 1 ) \ X ( int, light_bodyid, nlight, 1 ) \ X ( int, light_targetbodyid, nlight, 1 ) \ - XMJV( mjtByte, light_directional, nlight, 1 ) \ - XMJV( mjtByte, light_castshadow, nlight, 1 ) \ - XMJV( float, light_bulbradius, nlight, 1 ) \ - XMJV( mjtByte, light_active, nlight, 1 ) \ + X ( int, light_type, nlight, 1 ) \ + X ( int, light_texid, nlight, 1 ) \ + X ( mjtByte, light_castshadow, nlight, 1 ) \ + X ( float, light_bulbradius, nlight, 1 ) \ + X ( float, light_intensity, nlight, 1 ) \ + X ( float, light_range, nlight, 1 ) \ + X ( mjtByte, light_active, nlight, 1 ) \ X ( mjtNum, light_pos, nlight, 3 ) \ X ( mjtNum, light_dir, nlight, 3 ) \ X ( mjtNum, light_poscom0, nlight, 3 ) \ X ( mjtNum, light_pos0, nlight, 3 ) \ X ( mjtNum, light_dir0, nlight, 3 ) \ - XMJV( float, light_attenuation, nlight, 3 ) \ - XMJV( float, light_cutoff, nlight, 1 ) \ - XMJV( float, light_exponent, nlight, 1 ) \ - XMJV( float, light_ambient, nlight, 3 ) \ - XMJV( float, light_diffuse, nlight, 3 ) \ - XMJV( float, light_specular, nlight, 3 ) \ + X ( float, light_attenuation, nlight, 3 ) \ + X ( float, light_cutoff, nlight, 1 ) \ + X ( float, light_exponent, nlight, 1 ) \ + X ( float, light_ambient, nlight, 3 ) \ + X ( float, light_diffuse, nlight, 3 ) \ + X ( float, light_specular, nlight, 3 ) \ X ( int, flex_contype, nflex, 1 ) \ X ( int, flex_conaffinity, nflex, 1 ) \ X ( int, flex_condim, nflex, 1 ) \ @@ -327,65 +330,67 @@ X ( mjtByte, flex_internal, nflex, 1 ) \ X ( int, flex_selfcollide, nflex, 1 ) \ X ( int, flex_activelayers, nflex, 1 ) \ - XMJV( int, flex_dim, nflex, 1 ) \ - XMJV( int, flex_matid, nflex, 1 ) \ - XMJV( int, flex_group, nflex, 1 ) \ - XMJV( int, flex_interp, nflex, 1 ) \ - XMJV( int, flex_nodeadr, nflex, 1 ) \ - XMJV( int, flex_nodenum, nflex, 1 ) \ - XMJV( int, flex_vertadr, nflex, 1 ) \ - XMJV( int, flex_vertnum, nflex, 1 ) \ + X ( int, flex_dim, nflex, 1 ) \ + X ( int, flex_matid, nflex, 1 ) \ + X ( int, flex_group, nflex, 1 ) \ + X ( int, flex_interp, nflex, 1 ) \ + X ( int, flex_nodeadr, nflex, 1 ) \ + X ( int, flex_nodenum, nflex, 1 ) \ + X ( int, flex_vertadr, nflex, 1 ) \ + X ( int, flex_vertnum, nflex, 1 ) \ X ( int, flex_edgeadr, nflex, 1 ) \ X ( int, flex_edgenum, nflex, 1 ) \ - XMJV( int, flex_elemadr, nflex, 1 ) \ - XMJV( int, flex_elemnum, nflex, 1 ) \ - XMJV( int, flex_elemdataadr, nflex, 1 ) \ + X ( int, flex_elemadr, nflex, 1 ) \ + X ( int, flex_elemnum, nflex, 1 ) \ + X ( int, flex_elemdataadr, nflex, 1 ) \ X ( int, flex_elemedgeadr, nflex, 1 ) \ - XMJV( int, flex_shellnum, nflex, 1 ) \ - XMJV( int, flex_shelldataadr, nflex, 1 ) \ + X ( int, flex_shellnum, nflex, 1 ) \ + X ( int, flex_shelldataadr, nflex, 1 ) \ X ( int, flex_evpairadr, nflex, 1 ) \ X ( int, flex_evpairnum, nflex, 1 ) \ - XMJV( int, flex_texcoordadr, nflex, 1 ) \ - XMJV( int, flex_nodebodyid, nflexnode, 1 ) \ + X ( int, flex_texcoordadr, nflex, 1 ) \ + X ( int, flex_nodebodyid, nflexnode, 1 ) \ X ( int, flex_vertbodyid, nflexvert, 1 ) \ X ( int, flex_edge, nflexedge, 2 ) \ - XMJV( int, flex_elem, nflexelemdata, 1 ) \ - XMJV( int, flex_elemtexcoord, nflexelemdata, 1 ) \ + X ( int, flex_edgeflap, nflexedge, 2 ) \ + X ( int, flex_elem, nflexelemdata, 1 ) \ + X ( int, flex_elemtexcoord, nflexelemdata, 1 ) \ X ( int, flex_elemedge, nflexelemedge, 1 ) \ - XMJV( int, flex_elemlayer, nflexelem, 1 ) \ - XMJV( int, flex_shell, nflexshelldata,1 ) \ + X ( int, flex_elemlayer, nflexelem, 1 ) \ + X ( int, flex_shell, nflexshelldata,1 ) \ X ( int, flex_evpair, nflexevpair, 2 ) \ X ( mjtNum, flex_vert, nflexvert, 3 ) \ X ( mjtNum, flex_vert0, nflexvert, 3 ) \ - XMJV( mjtNum, flex_node, nflexnode, 3 ) \ + X ( mjtNum, flex_node, nflexnode, 3 ) \ X ( mjtNum, flex_node0, nflexnode, 3 ) \ X ( mjtNum, flexedge_length0, nflexedge, 1 ) \ X ( mjtNum, flexedge_invweight0, nflexedge, 1 ) \ - XMJV( mjtNum, flex_radius, nflex, 1 ) \ + X ( mjtNum, flex_radius, nflex, 1 ) \ X ( mjtNum, flex_stiffness, nflexelem, 21 ) \ + X ( mjtNum, flex_bending, nflexedge, 16 ) \ X ( mjtNum, flex_damping, nflex, 1 ) \ X ( mjtNum, flex_edgestiffness, nflex, 1 ) \ X ( mjtNum, flex_edgedamping, nflex, 1 ) \ X ( mjtByte, flex_edgeequality, nflex, 1 ) \ X ( mjtByte, flex_rigid, nflex, 1 ) \ X ( mjtByte, flexedge_rigid, nflexedge, 1 ) \ - XMJV( mjtByte, flex_centered, nflex, 1 ) \ - XMJV( mjtByte, flex_flatskin, nflex, 1 ) \ - XMJV( int, flex_bvhadr, nflex, 1 ) \ - XMJV( int, flex_bvhnum, nflex, 1 ) \ - XMJV( float, flex_rgba, nflex, 4 ) \ - XMJV( float, flex_texcoord, nflextexcoord, 2 ) \ + X ( mjtByte, flex_centered, nflex, 1 ) \ + X ( mjtByte, flex_flatskin, nflex, 1 ) \ + X ( int, flex_bvhadr, nflex, 1 ) \ + X ( int, flex_bvhnum, nflex, 1 ) \ + X ( float, flex_rgba, nflex, 4 ) \ + X ( float, flex_texcoord, nflextexcoord, 2 ) \ X ( int, mesh_vertadr, nmesh, 1 ) \ X ( int, mesh_vertnum, nmesh, 1 ) \ X ( int, mesh_normaladr, nmesh, 1 ) \ X ( int, mesh_normalnum, nmesh, 1 ) \ - XMJV( int, mesh_texcoordadr, nmesh, 1 ) \ + X ( int, mesh_texcoordadr, nmesh, 1 ) \ X ( int, mesh_texcoordnum, nmesh, 1 ) \ X ( int, mesh_faceadr, nmesh, 1 ) \ X ( int, mesh_facenum, nmesh, 1 ) \ - XMJV( int, mesh_bvhadr, nmesh, 1 ) \ - XMJV( int, mesh_bvhnum, nmesh, 1 ) \ - XMJV( int, mesh_graphadr, nmesh, 1 ) \ + X ( int, mesh_bvhadr, nmesh, 1 ) \ + X ( int, mesh_bvhnum, nmesh, 1 ) \ + X ( int, mesh_graphadr, nmesh, 1 ) \ X ( mjtNum, mesh_scale, nmesh, 3 ) \ X ( mjtNum, mesh_pos, nmesh, 3 ) \ X ( mjtNum, mesh_quat, nmesh, 4 ) \ @@ -396,61 +401,62 @@ XNV ( int, mesh_facenormal, nmeshface, 3 ) \ XNV ( int, mesh_facetexcoord, nmeshface, 3 ) \ XNV ( int, mesh_graph, nmeshgraph, 1 ) \ - XMJV( int, mesh_pathadr, nmesh, 1 ) \ - X ( int, mesh_polynum, nmesh, 1 ) \ - X ( int, mesh_polyadr, nmesh, 1 ) \ - X ( mjtNum, mesh_polynormal, nmeshpoly, 3 ) \ - X ( int, mesh_polyvertadr, nmeshpoly, 1 ) \ - X ( int, mesh_polyvertnum, nmeshpoly, 1 ) \ - X ( int, mesh_polyvert, nmeshpolyvert, 1 ) \ - X ( int, mesh_polymapadr, nmeshvert, 1 ) \ - X ( int, mesh_polymapnum, nmeshvert, 1 ) \ - X ( int, mesh_polymap, nmeshpolymap, 1 ) \ - XMJV( int, skin_matid, nskin, 1 ) \ - XMJV( int, skin_group, nskin, 1 ) \ - XMJV( float, skin_rgba, nskin, 4 ) \ - XMJV( float, skin_inflate, nskin, 1 ) \ - XMJV( int, skin_vertadr, nskin, 1 ) \ - XMJV( int, skin_vertnum, nskin, 1 ) \ - XMJV( int, skin_texcoordadr, nskin, 1 ) \ - XMJV( int, skin_faceadr, nskin, 1 ) \ - XMJV( int, skin_facenum, nskin, 1 ) \ - XMJV( int, skin_boneadr, nskin, 1 ) \ - XMJV( int, skin_bonenum, nskin, 1 ) \ - XMJV( float, skin_vert, nskinvert, 3 ) \ + X ( int, mesh_pathadr, nmesh, 1 ) \ + XNV ( int, mesh_polynum, nmesh, 1 ) \ + XNV ( int, mesh_polyadr, nmesh, 1 ) \ + XNV ( mjtNum, mesh_polynormal, nmeshpoly, 3 ) \ + XNV ( int, mesh_polyvertadr, nmeshpoly, 1 ) \ + XNV ( int, mesh_polyvertnum, nmeshpoly, 1 ) \ + XNV ( int, mesh_polyvert, nmeshpolyvert, 1 ) \ + XNV ( int, mesh_polymapadr, nmeshvert, 1 ) \ + XNV ( int, mesh_polymapnum, nmeshvert, 1 ) \ + XNV ( int, mesh_polymap, nmeshpolymap, 1 ) \ + X ( int, skin_matid, nskin, 1 ) \ + X ( int, skin_group, nskin, 1 ) \ + X ( float, skin_rgba, nskin, 4 ) \ + X ( float, skin_inflate, nskin, 1 ) \ + X ( int, skin_vertadr, nskin, 1 ) \ + X ( int, skin_vertnum, nskin, 1 ) \ + X ( int, skin_texcoordadr, nskin, 1 ) \ + X ( int, skin_faceadr, nskin, 1 ) \ + X ( int, skin_facenum, nskin, 1 ) \ + X ( int, skin_boneadr, nskin, 1 ) \ + X ( int, skin_bonenum, nskin, 1 ) \ + X ( float, skin_vert, nskinvert, 3 ) \ X ( float, skin_texcoord, nskintexvert, 2 ) \ - XMJV( int, skin_face, nskinface, 3 ) \ - XMJV( int, skin_bonevertadr, nskinbone, 1 ) \ - XMJV( int, skin_bonevertnum, nskinbone, 1 ) \ - XMJV( float, skin_bonebindpos, nskinbone, 3 ) \ - XMJV( float, skin_bonebindquat, nskinbone, 4 ) \ - XMJV( int, skin_bonebodyid, nskinbone, 1 ) \ - XMJV( int, skin_bonevertid, nskinbonevert, 1 ) \ - XMJV( float, skin_bonevertweight, nskinbonevert, 1 ) \ - XMJV( int, skin_pathadr, nskin, 1 ) \ + X ( int, skin_face, nskinface, 3 ) \ + X ( int, skin_bonevertadr, nskinbone, 1 ) \ + X ( int, skin_bonevertnum, nskinbone, 1 ) \ + X ( float, skin_bonebindpos, nskinbone, 3 ) \ + X ( float, skin_bonebindquat, nskinbone, 4 ) \ + X ( int, skin_bonebodyid, nskinbone, 1 ) \ + X ( int, skin_bonevertid, nskinbonevert, 1 ) \ + X ( float, skin_bonevertweight, nskinbonevert, 1 ) \ + X ( int, skin_pathadr, nskin, 1 ) \ X ( mjtNum, hfield_size, nhfield, 4 ) \ X ( int, hfield_nrow, nhfield, 1 ) \ X ( int, hfield_ncol, nhfield, 1 ) \ X ( int, hfield_adr, nhfield, 1 ) \ XNV ( float, hfield_data, nhfielddata, 1 ) \ - XMJV( int, hfield_pathadr, nhfield, 1 ) \ + X ( int, hfield_pathadr, nhfield, 1 ) \ X ( int, tex_type, ntex, 1 ) \ + X ( int, tex_colorspace, ntex, 1 ) \ X ( int, tex_height, ntex, 1 ) \ X ( int, tex_width, ntex, 1 ) \ X ( int, tex_nchannel, ntex, 1 ) \ X ( int, tex_adr, ntex, 1 ) \ XNV ( mjtByte, tex_data, ntexdata, 1 ) \ - XMJV( int, tex_pathadr, ntex, 1 ) \ - XMJV( int, mat_texid, nmat, mjNTEXROLE ) \ - XMJV( mjtByte, mat_texuniform, nmat, 1 ) \ - XMJV( float, mat_texrepeat, nmat, 2 ) \ - XMJV( float, mat_emission, nmat, 1 ) \ - XMJV( float, mat_specular, nmat, 1 ) \ - XMJV( float, mat_shininess, nmat, 1 ) \ - XMJV( float, mat_reflectance, nmat, 1 ) \ - XMJV( float, mat_metallic, nmat, 1 ) \ - XMJV( float, mat_roughness, nmat, 1 ) \ - XMJV( float, mat_rgba, nmat, 4 ) \ + X ( int, tex_pathadr, ntex, 1 ) \ + X ( int, mat_texid, nmat, mjNTEXROLE ) \ + X ( mjtByte, mat_texuniform, nmat, 1 ) \ + X ( float, mat_texrepeat, nmat, 2 ) \ + X ( float, mat_emission, nmat, 1 ) \ + X ( float, mat_specular, nmat, 1 ) \ + X ( float, mat_shininess, nmat, 1 ) \ + X ( float, mat_reflectance, nmat, 1 ) \ + X ( float, mat_metallic, nmat, 1 ) \ + X ( float, mat_roughness, nmat, 1 ) \ + X ( float, mat_rgba, nmat, 4 ) \ X ( int, pair_dim, npair, 1 ) \ X ( int, pair_geom1, npair, 1 ) \ X ( int, pair_geom2, npair, 1 ) \ @@ -462,71 +468,74 @@ X ( mjtNum, pair_gap, npair, 1 ) \ X ( mjtNum, pair_friction, npair, 5 ) \ X ( int, exclude_signature, nexclude, 1 ) \ - XMJV( int, eq_type, neq, 1 ) \ - XMJV( int, eq_obj1id, neq, 1 ) \ - XMJV( int, eq_obj2id, neq, 1 ) \ - XMJV( int, eq_objtype, neq, 1 ) \ + X ( int, eq_type, neq, 1 ) \ + X ( int, eq_obj1id, neq, 1 ) \ + X ( int, eq_obj2id, neq, 1 ) \ + X ( int, eq_objtype, neq, 1 ) \ X ( mjtByte, eq_active0, neq, 1 ) \ X ( mjtNum, eq_solref, neq, mjNREF ) \ X ( mjtNum, eq_solimp, neq, mjNIMP ) \ - XMJV( mjtNum, eq_data, neq, mjNEQDATA ) \ + X ( mjtNum, eq_data, neq, mjNEQDATA ) \ X ( int, tendon_adr, ntendon, 1 ) \ - XMJV( int, tendon_num, ntendon, 1 ) \ - XMJV( int, tendon_matid, ntendon, 1 ) \ - XMJV( int, tendon_group, ntendon, 1 ) \ - XMJV( mjtByte, tendon_limited, ntendon, 1 ) \ - XMJV( mjtNum, tendon_width, ntendon, 1 ) \ + X ( int, tendon_num, ntendon, 1 ) \ + X ( int, tendon_matid, ntendon, 1 ) \ + X ( int, tendon_group, ntendon, 1 ) \ + X ( mjtByte, tendon_limited, ntendon, 1 ) \ + X ( mjtByte, tendon_actfrclimited, ntendon, 1 ) \ + X ( mjtNum, tendon_width, ntendon, 1 ) \ X ( mjtNum, tendon_solref_lim, ntendon, mjNREF ) \ X ( mjtNum, tendon_solimp_lim, ntendon, mjNIMP ) \ X ( mjtNum, tendon_solref_fri, ntendon, mjNREF ) \ X ( mjtNum, tendon_solimp_fri, ntendon, mjNIMP ) \ - XMJV( mjtNum, tendon_range, ntendon, 2 ) \ + X ( mjtNum, tendon_range, ntendon, 2 ) \ + X ( mjtNum, tendon_actfrcrange, ntendon, 2 ) \ X ( mjtNum, tendon_margin, ntendon, 1 ) \ - XMJV( mjtNum, tendon_stiffness, ntendon, 1 ) \ - XMJV( mjtNum, tendon_damping, ntendon, 1 ) \ - XMJV( mjtNum, tendon_frictionloss, ntendon, 1 ) \ - XMJV( mjtNum, tendon_lengthspring, ntendon, 2 ) \ + X ( mjtNum, tendon_stiffness, ntendon, 1 ) \ + X ( mjtNum, tendon_damping, ntendon, 1 ) \ + X ( mjtNum, tendon_armature, ntendon, 1 ) \ + X ( mjtNum, tendon_frictionloss, ntendon, 1 ) \ + X ( mjtNum, tendon_lengthspring, ntendon, 2 ) \ X ( mjtNum, tendon_length0, ntendon, 1 ) \ X ( mjtNum, tendon_invweight0, ntendon, 1 ) \ X ( mjtNum, tendon_user, ntendon, MJ_M(nuser_tendon) ) \ - XMJV( float, tendon_rgba, ntendon, 4 ) \ + X ( float, tendon_rgba, ntendon, 4 ) \ X ( int, wrap_type, nwrap, 1 ) \ X ( int, wrap_objid, nwrap, 1 ) \ X ( mjtNum, wrap_prm, nwrap, 1 ) \ - XMJV( int, actuator_trntype, nu, 1 ) \ - XMJV( int, actuator_dyntype, nu, 1 ) \ + X ( int, actuator_trntype, nu, 1 ) \ + X ( int, actuator_dyntype, nu, 1 ) \ X ( int, actuator_gaintype, nu, 1 ) \ X ( int, actuator_biastype, nu, 1 ) \ - XMJV( int, actuator_trnid, nu, 2 ) \ - XMJV( int, actuator_actadr, nu, 1 ) \ - XMJV( int, actuator_actnum, nu, 1 ) \ - XMJV( int, actuator_group, nu, 1 ) \ - XMJV( mjtByte, actuator_ctrllimited, nu, 1 ) \ + X ( int, actuator_trnid, nu, 2 ) \ + X ( int, actuator_actadr, nu, 1 ) \ + X ( int, actuator_actnum, nu, 1 ) \ + X ( int, actuator_group, nu, 1 ) \ + X ( mjtByte, actuator_ctrllimited, nu, 1 ) \ X ( mjtByte, actuator_forcelimited, nu, 1 ) \ - XMJV( mjtByte, actuator_actlimited, nu, 1 ) \ + X ( mjtByte, actuator_actlimited, nu, 1 ) \ X ( mjtNum, actuator_dynprm, nu, mjNDYN ) \ X ( mjtNum, actuator_gainprm, nu, mjNGAIN ) \ X ( mjtNum, actuator_biasprm, nu, mjNBIAS ) \ X ( mjtByte, actuator_actearly, nu, 1 ) \ - XMJV( mjtNum, actuator_ctrlrange, nu, 2 ) \ + X ( mjtNum, actuator_ctrlrange, nu, 2 ) \ X ( mjtNum, actuator_forcerange, nu, 2 ) \ - XMJV( mjtNum, actuator_actrange, nu, 2 ) \ + X ( mjtNum, actuator_actrange, nu, 2 ) \ X ( mjtNum, actuator_gear, nu, 6 ) \ - XMJV( mjtNum, actuator_cranklength, nu, 1 ) \ + X ( mjtNum, actuator_cranklength, nu, 1 ) \ X ( mjtNum, actuator_acc0, nu, 1 ) \ X ( mjtNum, actuator_length0, nu, 1 ) \ X ( mjtNum, actuator_lengthrange, nu, 2 ) \ X ( mjtNum, actuator_user, nu, MJ_M(nuser_actuator) ) \ X ( int, actuator_plugin, nu, 1 ) \ - XMJV( int, sensor_type, nsensor, 1 ) \ + X ( int, sensor_type, nsensor, 1 ) \ X ( int, sensor_datatype, nsensor, 1 ) \ X ( int, sensor_needstage, nsensor, 1 ) \ X ( int, sensor_objtype, nsensor, 1 ) \ - XMJV( int, sensor_objid, nsensor, 1 ) \ + X ( int, sensor_objid, nsensor, 1 ) \ X ( int, sensor_reftype, nsensor, 1 ) \ X ( int, sensor_refid, nsensor, 1 ) \ X ( int, sensor_dim, nsensor, 1 ) \ - XMJV( int, sensor_adr, nsensor, 1 ) \ + X ( int, sensor_adr, nsensor, 1 ) \ X ( mjtNum, sensor_cutoff, nsensor, 1 ) \ X ( mjtNum, sensor_noise, nsensor, 1 ) \ X ( mjtNum, sensor_user, nsensor, MJ_M(nuser_sensor) ) \ @@ -554,12 +563,12 @@ X ( mjtNum, key_mpos, nkey, MJ_M(nmocap)*3 ) \ X ( mjtNum, key_mquat, nkey, MJ_M(nmocap)*4 ) \ X ( mjtNum, key_ctrl, nkey, MJ_M(nu) ) \ - XMJV( int, name_bodyadr, nbody, 1 ) \ - XMJV( int, name_jntadr, njnt, 1 ) \ - XMJV( int, name_geomadr, ngeom, 1 ) \ - XMJV( int, name_siteadr, nsite, 1 ) \ - XMJV( int, name_camadr, ncam, 1 ) \ - XMJV( int, name_lightadr, nlight, 1 ) \ + X ( int, name_bodyadr, nbody, 1 ) \ + X ( int, name_jntadr, njnt, 1 ) \ + X ( int, name_geomadr, ngeom, 1 ) \ + X ( int, name_siteadr, nsite, 1 ) \ + X ( int, name_camadr, ncam, 1 ) \ + X ( int, name_lightadr, nlight, 1 ) \ X ( int, name_flexadr, nflex, 1 ) \ X ( int, name_meshadr, nmesh, 1 ) \ X ( int, name_skinadr, nskin, 1 ) \ @@ -568,18 +577,18 @@ X ( int, name_matadr, nmat, 1 ) \ X ( int, name_pairadr, npair, 1 ) \ X ( int, name_excludeadr, nexclude, 1 ) \ - XMJV( int, name_eqadr, neq, 1 ) \ - XMJV( int, name_tendonadr, ntendon, 1 ) \ - XMJV( int, name_actuatoradr, nu, 1 ) \ + X ( int, name_eqadr, neq, 1 ) \ + X ( int, name_tendonadr, ntendon, 1 ) \ + X ( int, name_actuatoradr, nu, 1 ) \ X ( int, name_sensoradr, nsensor, 1 ) \ X ( int, name_numericadr, nnumeric, 1 ) \ X ( int, name_textadr, ntext, 1 ) \ X ( int, name_tupleadr, ntuple, 1 ) \ X ( int, name_keyadr, nkey, 1 ) \ X ( int, name_pluginadr, nplugin, 1 ) \ - XMJV( char, names, nnames, 1 ) \ + X ( char, names, nnames, 1 ) \ X ( int, names_map, nnames_map, 1 ) \ - XMJV( char, paths, npaths, 1 ) \ + X ( char, paths, npaths, 1 ) \ //-------------------------------- mjData ---------------------------------------------------------- @@ -589,71 +598,72 @@ // pointer fields of mjData -// XMJV means that the field is required to construct mjvScene -// (by default we define XMJV to be the same as X) +// XNV means that the field is not required to construct mjvScene +// (by default we define XNV to be the same as X) #define MJDATA_POINTERS \ X ( mjtNum, qpos, nq, 1 ) \ X ( mjtNum, qvel, nv, 1 ) \ - XMJV( mjtNum, act, na, 1 ) \ + X ( mjtNum, act, na, 1 ) \ X ( mjtNum, qacc_warmstart, nv, 1 ) \ X ( mjtNum, plugin_state, npluginstate, 1 ) \ - XMJV( mjtNum, ctrl, nu, 1 ) \ + X ( mjtNum, ctrl, nu, 1 ) \ X ( mjtNum, qfrc_applied, nv, 1 ) \ - XMJV( mjtNum, xfrc_applied, nbody, 6 ) \ - XMJV( mjtByte, eq_active, neq, 1 ) \ + X ( mjtNum, xfrc_applied, nbody, 6 ) \ + X ( mjtByte, eq_active, neq, 1 ) \ X ( mjtNum, mocap_pos, nmocap, 3 ) \ X ( mjtNum, mocap_quat, nmocap, 4 ) \ X ( mjtNum, qacc, nv, 1 ) \ X ( mjtNum, act_dot, na, 1 ) \ X ( mjtNum, userdata, nuserdata, 1 ) \ - XMJV( mjtNum, sensordata, nsensordata, 1 ) \ + X ( mjtNum, sensordata, nsensordata, 1 ) \ X ( int, plugin, nplugin, 1 ) \ X ( uintptr_t, plugin_data, nplugin, 1 ) \ - XMJV( mjtNum, xpos, nbody, 3 ) \ - XMJV( mjtNum, xquat, nbody, 4 ) \ - XMJV( mjtNum, xmat, nbody, 9 ) \ - XMJV( mjtNum, xipos, nbody, 3 ) \ - XMJV( mjtNum, ximat, nbody, 9 ) \ - XMJV( mjtNum, xanchor, njnt, 3 ) \ - XMJV( mjtNum, xaxis, njnt, 3 ) \ - XMJV( mjtNum, geom_xpos, ngeom, 3 ) \ - XMJV( mjtNum, geom_xmat, ngeom, 9 ) \ - XMJV( mjtNum, site_xpos, nsite, 3 ) \ - XMJV( mjtNum, site_xmat, nsite, 9 ) \ - XMJV( mjtNum, cam_xpos, ncam, 3 ) \ - XMJV( mjtNum, cam_xmat, ncam, 9 ) \ - XMJV( mjtNum, light_xpos, nlight, 3 ) \ - XMJV( mjtNum, light_xdir, nlight, 3 ) \ - XMJV( mjtNum, subtree_com, nbody, 3 ) \ + X ( mjtNum, xpos, nbody, 3 ) \ + X ( mjtNum, xquat, nbody, 4 ) \ + X ( mjtNum, xmat, nbody, 9 ) \ + X ( mjtNum, xipos, nbody, 3 ) \ + X ( mjtNum, ximat, nbody, 9 ) \ + X ( mjtNum, xanchor, njnt, 3 ) \ + X ( mjtNum, xaxis, njnt, 3 ) \ + X ( mjtNum, geom_xpos, ngeom, 3 ) \ + X ( mjtNum, geom_xmat, ngeom, 9 ) \ + X ( mjtNum, site_xpos, nsite, 3 ) \ + X ( mjtNum, site_xmat, nsite, 9 ) \ + X ( mjtNum, cam_xpos, ncam, 3 ) \ + X ( mjtNum, cam_xmat, ncam, 9 ) \ + X ( mjtNum, light_xpos, nlight, 3 ) \ + X ( mjtNum, light_xdir, nlight, 3 ) \ + X ( mjtNum, subtree_com, nbody, 3 ) \ X ( mjtNum, cdof, nv, 6 ) \ X ( mjtNum, cinert, nbody, 10 ) \ - XMJV( mjtNum, flexvert_xpos, nflexvert, 3 ) \ + X ( mjtNum, flexvert_xpos, nflexvert, 3 ) \ X ( mjtNum, flexelem_aabb, nflexelem, 6 ) \ X ( int, flexedge_J_rownnz, nflexedge, 1 ) \ X ( int, flexedge_J_rowadr, nflexedge, 1 ) \ X ( int, flexedge_J_colind, nflexedge, MJ_M(nv) ) \ X ( mjtNum, flexedge_J, nflexedge, MJ_M(nv) ) \ X ( mjtNum, flexedge_length, nflexedge, 1 ) \ - XMJV( int, ten_wrapadr, ntendon, 1 ) \ - XMJV( int, ten_wrapnum, ntendon, 1 ) \ + X ( int, ten_wrapadr, ntendon, 1 ) \ + X ( int, ten_wrapnum, ntendon, 1 ) \ X ( int, ten_J_rownnz, ntendon, 1 ) \ X ( int, ten_J_rowadr, ntendon, 1 ) \ X ( int, ten_J_colind, ntendon, MJ_M(nv) ) \ - XMJV( mjtNum, ten_length, ntendon, 1 ) \ + X ( mjtNum, ten_length, ntendon, 1 ) \ X ( mjtNum, ten_J, ntendon, MJ_M(nv) ) \ - XMJV( int, wrap_obj, nwrap, 2 ) \ - XMJV( mjtNum, wrap_xpos, nwrap, 6 ) \ + X ( int, wrap_obj, nwrap, 2 ) \ + X ( mjtNum, wrap_xpos, nwrap, 6 ) \ X ( mjtNum, actuator_length, nu, 1 ) \ X ( int, moment_rownnz, nu, 1 ) \ X ( int, moment_rowadr, nu, 1 ) \ X ( int, moment_colind, nJmom, 1 ) \ X ( mjtNum, actuator_moment, nJmom, 1 ) \ - X ( mjtNum, crb, nbody, 10 ) \ - X ( mjtNum, qM, nM, 1 ) \ - X ( mjtNum, qLD, nM, 1 ) \ + XNV ( mjtNum, crb, nbody, 10 ) \ + XNV ( mjtNum, qM, nM, 1 ) \ + XNV ( mjtNum, M, nC, 1 ) \ + XNV ( mjtNum, qLD, nC, 1 ) \ X ( mjtNum, qLDiagInv, nv, 1 ) \ - XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ - XMJV( mjtByte, bvh_active, nbvh, 1 ) \ + X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ + X ( mjtByte, bvh_active, nbvh, 1 ) \ X ( mjtNum, flexedge_velocity, nflexedge, 1 ) \ X ( mjtNum, ten_velocity, ntendon, 1 ) \ X ( mjtNum, actuator_velocity, nu, 1 ) \ @@ -667,27 +677,23 @@ X ( mjtNum, qfrc_passive, nv, 1 ) \ X ( mjtNum, subtree_linvel, nbody, 3 ) \ X ( mjtNum, subtree_angmom, nbody, 3 ) \ - X ( mjtNum, qH, nM, 1 ) \ + XNV ( mjtNum, qH, nC, 1 ) \ X ( mjtNum, qHDiagInv, nv, 1 ) \ - X ( int, B_rownnz, nbody, 1 ) \ - X ( int, B_rowadr, nbody, 1 ) \ - X ( int, B_colind, nB, 1 ) \ - X ( int, M_rownnz, nv, 1 ) \ - X ( int, M_rowadr, nv, 1 ) \ - X ( int, M_colind, nM, 1 ) \ - X ( int, mapM2M, nM, 1 ) \ - X ( int, C_rownnz, nv, 1 ) \ - X ( int, C_rowadr, nv, 1 ) \ - X ( int, C_colind, nC, 1 ) \ - X ( int, mapM2C, nC, 1 ) \ - X ( int, D_rownnz, nv, 1 ) \ - X ( int, D_rowadr, nv, 1 ) \ - X ( int, D_diag, nv, 1 ) \ - X ( int, D_colind, nD, 1 ) \ - X ( int, mapM2D, nD, 1 ) \ - X ( int, mapD2M, nM, 1 ) \ - X ( mjtNum, qDeriv, nD, 1 ) \ - X ( mjtNum, qLU, nD, 1 ) \ + XNV ( int, B_rownnz, nbody, 1 ) \ + XNV ( int, B_rowadr, nbody, 1 ) \ + XNV ( int, B_colind, nB, 1 ) \ + XNV ( int, M_rownnz, nv, 1 ) \ + XNV ( int, M_rowadr, nv, 1 ) \ + XNV ( int, M_colind, nC, 1 ) \ + XNV ( int, mapM2M, nC, 1 ) \ + XNV ( int, D_rownnz, nv, 1 ) \ + XNV ( int, D_rowadr, nv, 1 ) \ + XNV ( int, D_diag, nv, 1 ) \ + XNV ( int, D_colind, nD, 1 ) \ + XNV ( int, mapM2D, nD, 1 ) \ + XNV ( int, mapD2M, nM, 1 ) \ + XNV ( mjtNum, qDeriv, nD, 1 ) \ + XNV ( mjtNum, qLU, nD, 1 ) \ X ( mjtNum, actuator_force, nu, 1 ) \ X ( mjtNum, qfrc_actuator, nv, 1 ) \ X ( mjtNum, qfrc_smooth, nv, 1 ) \ @@ -708,51 +714,83 @@ X( mjContact, contact, MJ_D(ncon), 1 ) // array fields of mjData that are used in the primal problem -#define MJDATA_ARENA_POINTERS_SOLVER \ - X( int, efc_type, MJ_D(nefc), 1 ) \ - X( int, efc_id, MJ_D(nefc), 1 ) \ - X( int, efc_J_rownnz, MJ_D(nefc), 1 ) \ - X( int, efc_J_rowadr, MJ_D(nefc), 1 ) \ - X( int, efc_J_rowsuper, MJ_D(nefc), 1 ) \ - X( int, efc_J_colind, MJ_D(nJ), 1 ) \ - X( int, efc_JT_rownnz, MJ_M(nv), 1 ) \ - X( int, efc_JT_rowadr, MJ_M(nv), 1 ) \ - X( int, efc_JT_rowsuper, MJ_M(nv), 1 ) \ - X( int, efc_JT_colind, MJ_D(nJ), 1 ) \ - X( mjtNum, efc_J, MJ_D(nJ), 1 ) \ - X( mjtNum, efc_JT, MJ_D(nJ), 1 ) \ - X( mjtNum, efc_pos, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_margin, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_frictionloss, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_diagApprox, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_KBIP, MJ_D(nefc), 4 ) \ - X( mjtNum, efc_D, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_R, MJ_D(nefc), 1 ) \ - X( int, tendon_efcadr, MJ_M(ntendon), 1 ) \ - X( mjtNum, efc_vel, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_aref, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_b, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_force, MJ_D(nefc), 1 ) \ - X( int, efc_state, MJ_D(nefc), 1 ) +#define MJDATA_ARENA_POINTERS_SOLVER \ + X ( int, efc_type, MJ_D(nefc), 1 ) \ + X ( int, efc_id, MJ_D(nefc), 1 ) \ + XNV( int, efc_J_rownnz, MJ_D(nefc), 1 ) \ + XNV( int, efc_J_rowadr, MJ_D(nefc), 1 ) \ + XNV( int, efc_J_rowsuper, MJ_D(nefc), 1 ) \ + XNV( int, efc_J_colind, MJ_D(nJ), 1 ) \ + XNV( int, efc_JT_rownnz, MJ_M(nv), 1 ) \ + XNV( int, efc_JT_rowadr, MJ_M(nv), 1 ) \ + XNV( int, efc_JT_rowsuper, MJ_M(nv), 1 ) \ + XNV( int, efc_JT_colind, MJ_D(nJ), 1 ) \ + XNV( mjtNum, efc_J, MJ_D(nJ), 1 ) \ + XNV( mjtNum, efc_JT, MJ_D(nJ), 1 ) \ + X ( mjtNum, efc_pos, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_margin, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_frictionloss, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_diagApprox, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_KBIP, MJ_D(nefc), 4 ) \ + X ( mjtNum, efc_D, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_R, MJ_D(nefc), 1 ) \ + X ( int, tendon_efcadr, MJ_M(ntendon), 1 ) \ + X ( mjtNum, efc_vel, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_aref, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_b, MJ_D(nefc), 1 ) \ + X ( mjtNum, efc_force, MJ_D(nefc), 1 ) \ + X ( int, efc_state, MJ_D(nefc), 1 ) // array fields of mjData that are used in the dual problem -#define MJDATA_ARENA_POINTERS_DUAL \ - X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ - X( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ - X( int, efc_AR_colind, MJ_D(nA), 1 ) \ - X( mjtNum, efc_AR, MJ_D(nA), 1 ) +#define MJDATA_ARENA_POINTERS_DUAL \ + XNV( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ + XNV( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ + XNV( int, efc_AR_colind, MJ_D(nA), 1 ) \ + XNV( mjtNum, efc_AR, MJ_D(nA), 1 ) // array fields of mjData that are used for constraint islands -#define MJDATA_ARENA_POINTERS_ISLAND \ - X( int, dof_island, MJ_M(nv), 1 ) \ - X( int, island_dofnum, MJ_D(nisland), 1 ) \ - X( int, island_dofadr, MJ_D(nisland), 1 ) \ - X( int, island_dofind, MJ_M(nv), 1 ) \ - X( int, dof_islandind, MJ_M(nv), 1 ) \ - X( int, efc_island, MJ_D(nefc), 1 ) \ - X( int, island_efcnum, MJ_D(nisland), 1 ) \ - X( int, island_efcadr, MJ_D(nisland), 1 ) \ - X( int, island_efcind, MJ_D(nefc), 1 ) +#define MJDATA_ARENA_POINTERS_ISLAND \ + X ( int, dof_island, MJ_M(nv), 1 ) \ + X ( int, island_nv, MJ_D(nisland), 1 ) \ + X ( int, island_idofadr, MJ_D(nisland), 1 ) \ + X ( int, island_dofadr, MJ_D(nisland), 1 ) \ + X ( int, map_dof2idof, MJ_M(nv), 1 ) \ + X ( int, map_idof2dof, MJ_M(nv), 1 ) \ + X ( mjtNum, ifrc_smooth, MJ_D(nidof), 1 ) \ + X ( mjtNum, iacc_smooth, MJ_D(nidof), 1 ) \ + XNV( int, iM_rownnz, MJ_D(nidof), 1 ) \ + XNV( int, iM_rowadr, MJ_D(nidof), 1 ) \ + XNV( int, iM_colind, MJ_M(nC), 1 ) \ + XNV( mjtNum, iM, MJ_M(nC), 1 ) \ + XNV( mjtNum, iLD, MJ_M(nC), 1 ) \ + X ( mjtNum, iLDiagInv, MJ_D(nidof), 1 ) \ + X ( mjtNum, iacc, MJ_D(nidof), 1 ) \ + X ( int, efc_island, MJ_D(nefc), 1 ) \ + X ( int, island_ne, MJ_D(nisland), 1 ) \ + X ( int, island_nf, MJ_D(nisland), 1 ) \ + X ( int, island_nefc, MJ_D(nisland), 1 ) \ + X ( int, island_iefcadr, MJ_D(nisland), 1 ) \ + X ( int, map_efc2iefc, MJ_D(nefc), 1 ) \ + X ( int, map_iefc2efc, MJ_D(nefc), 1 ) \ + X ( int, iefc_type, MJ_D(nefc), 1 ) \ + X ( int, iefc_id, MJ_D(nefc), 1 ) \ + XNV( int, iefc_J_rownnz, MJ_D(nefc), 1 ) \ + XNV( int, iefc_J_rowadr, MJ_D(nefc), 1 ) \ + XNV( int, iefc_J_rowsuper, MJ_D(nefc), 1 ) \ + XNV( int, iefc_J_colind, MJ_D(nJ), 1 ) \ + XNV( int, iefc_JT_rownnz, MJ_D(nidof), 1 ) \ + XNV( int, iefc_JT_rowadr, MJ_D(nidof), 1 ) \ + XNV( int, iefc_JT_rowsuper, MJ_D(nidof), 1 ) \ + XNV( int, iefc_JT_colind, MJ_D(nJ), 1 ) \ + XNV( mjtNum, iefc_J, MJ_D(nJ), 1 ) \ + XNV( mjtNum, iefc_JT, MJ_D(nJ), 1 ) \ + X ( mjtNum, iefc_frictionloss, MJ_D(nefc), 1 ) \ + X ( mjtNum, iefc_D, MJ_D(nefc), 1 ) \ + X ( mjtNum, iefc_R, MJ_D(nefc), 1 ) \ + X ( mjtNum, iefc_aref, MJ_D(nefc), 1 ) \ + X ( int, iefc_state, MJ_D(nefc), 1 ) \ + X ( mjtNum, iefc_force, MJ_D(nefc), 1 ) \ + X ( mjtNum, ifrc_constraint, MJ_D(nidof), 1 ) // array fields of mjData that live in d->arena #define MJDATA_ARENA_POINTERS \ @@ -774,7 +812,6 @@ X( size_t, maxuse_arena ) \ X( int, maxuse_con ) \ X( int, maxuse_efc ) \ - X( int, solver_nisland ) \ X( int, ncon ) \ X( int, ne ) \ X( int, nf ) \ @@ -783,6 +820,7 @@ X( int, nJ ) \ X( int, nA ) \ X( int, nisland ) \ + X( int, nidof ) \ X( mjtNum, time ) \ X( uintptr_t, threadpool ) @@ -798,12 +836,6 @@ X( mjtNum, solver_fwdinv, 2, 1 ) \ X( mjtNum, energy, 2, 1 ) - -// alias XMJV to be the same as X -// to obtain only X macros for fields that are relevant for mjvScene creation, -// redefine X to expand to nothing, and XMJV to do what's required -#define XMJV X - // alias XNV to be the same as X // to obtain only X macros for fields that are relevant for mjvScene creation, // redefine XNV to expand to nothing diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 5158e1f8..59283756 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -16,7 +16,7 @@ #define MUJOCO_MUJOCO_H_ // header version; should match the library version as returned by mj_version() -#define mjVERSION_HEADER 331 +#define mjVERSION_HEADER 333 // needed to define size_t, fabs and log10 #include @@ -107,6 +107,9 @@ MJAPI mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, // Compile spec to model. MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs); +// Copy real-valued arrays from model to spec, returns 1 on success. +MJAPI int mj_copyBack(mjSpec* s, const mjModel* m); + // Recompile spec to model, preserving the state, return 0 on success. MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); @@ -187,6 +190,9 @@ MJAPI mjData* mj_makeData(const mjModel* m); // m is only required to contain the size fields from MJMODEL_INTS. MJAPI mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src); +// Copy mjData, skip large arrays not required for visualization. +MJAPI mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src); + // Reset data to defaults. MJAPI void mj_resetData(const mjModel* m, mjData* d); @@ -361,6 +367,9 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d); // Run composite rigid body inertia algorithm (CRB). MJAPI void mj_crb(const mjModel* m, mjData* d); +// Make inertia matrix. +MJAPI void mj_makeM(const mjModel* m, mjData* d); + // Compute sparse L'*D*L factorizaton of inertia matrix. MJAPI void mj_factorM(const mjModel* m, mjData* d); @@ -485,7 +494,7 @@ MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); // Add inertia matrix to destination matrix. -// Destination can be sparse uncompressed, or dense when all int* are NULL +// Destination can be sparse or dense when all int* are NULL. MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind); // Apply Cartesian force and torque (outside xfrc_applied mechanism). @@ -622,20 +631,10 @@ MJAPI void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum fo MJAPI void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy, const mjvScene* scn, mjvCamera* cam); -// Move camera with mouse given a scene state; action is mjtMouse. -MJAPI void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvCamera* cam); - // Move perturb object with mouse; action is mjtMouse. MJAPI void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx, mjtNum reldy, const mjvScene* scn, mjvPerturb* pert); -// Move perturb object with mouse given a scene state; action is mjtMouse. -MJAPI void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvPerturb* pert); - // Move model with mouse; action is mjtMouse. MJAPI void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy, const mjtNum roomup[3], mjvScene* scn); @@ -692,28 +691,9 @@ MJAPI void mjv_freeScene(mjvScene* scn); MJAPI void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt, const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn); -// Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings. -MJAPI int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt, - const mjvPerturb* pert, mjvCamera* cam, int catmask, - mjvScene* scn); - // Copy mjModel, skip large arrays not required for abstract visualization. MJAPI void mjv_copyModel(mjModel* dest, const mjModel* src); -// Set default scene state. -MJAPI void mjv_defaultSceneState(mjvSceneState* scnstate); - -// Allocate resources and initialize a scene state object. -MJAPI void mjv_makeSceneState(const mjModel* m, const mjData* d, - mjvSceneState* scnstate, int maxgeom); - -// Free scene state. -MJAPI void mjv_freeSceneState(mjvSceneState* scnstate); - -// Update a scene state from model and data. -MJAPI void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt, - mjvSceneState* scnstate); - // Add geoms from selected categories. MJAPI void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt, const mjvPerturb* pert, int catmask, mjvScene* scn); @@ -1308,6 +1288,19 @@ MJAPI char* mju_strncpy(char *dst, const char *src, int n); MJAPI mjtNum mju_sigmoid(mjtNum x); +//---------------------------------- Signed Distance Function -------------------------------------- + +// get sdf from geom id +MJAPI const mjpPlugin* mjc_getSDF(const mjModel* m, int id); + +// signed distance function +MJAPI mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]); + +// gradient of sdf +MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3], + const mjtNum x[3]); + + //---------------------------------- Derivatives --------------------------------------------------- // Finite differenced transition matrices (control theory notation) @@ -1414,21 +1407,9 @@ MJAPI void mju_taskJoin(mjTask* task); //---------------------------------- Attachment ---------------------------------------------------- -// Attach child body to a parent frame, return the attached body if success or NULL otherwise. -MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child, - const char* prefix, const char* suffix); - -// Attach child frame to a parent body, return the attached frame if success or NULL otherwise. -MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child, - const char* prefix, const char* suffix); - -// Attach child body to a parent site, return the attached body if success or NULL otherwise. -MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child, - const char* prefix, const char* suffix); - -// Attach child frame to a parent site, return the attached frame if success or NULL otherwise. -MJAPI mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child, - const char* prefix, const char* suffix); +// Attach child to a parent, return the attached element if success or NULL otherwise. +MJAPI mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child, + const char* prefix, const char* suffix); // Delete body and descendants from mjSpec, remove all references, return 0 on success. MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b); @@ -1634,14 +1615,17 @@ MJAPI const char* mjs_getString(const mjString* source); // Get double array contents and optionally its size. MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size); +// Get plugin attributes. +MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin); + //---------------------------------- Spec utilities ------------------------------------------------ // Set element's default. MJAPI void mjs_setDefault(mjsElement* element, const mjsDefault* def); -// Set element's enclosing frame. -MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame); +// Set element's enclosing frame, return 0 on success. +MJAPI int mjs_setFrame(mjsElement* dest, mjsFrame* frame); // Resolve alternative orientations to quat, return error if any. MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence, @@ -1653,6 +1637,13 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload, overriding the existing value for the specified key if present. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload, overriding the existing value for the specified key if +// present. This version differs from mjs_setUserValue in that it takes a +// cleanup function that will be called when the user payload is deleted. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 6a031b57..c60c785c 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -17,6 +17,7 @@ # 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.derivative import deriv_smooth_vel from mujoco.mjx._src.forward import euler from mujoco.mjx._src.forward import forward from mujoco.mjx._src.forward import fwd_acceleration @@ -26,6 +27,7 @@ from mujoco.mjx._src.forward import fwd_velocity from mujoco.mjx._src.forward import implicit from mujoco.mjx._src.forward import rungekutta4 from mujoco.mjx._src.forward import step +from mujoco.mjx._src.inverse import inverse from mujoco.mjx._src.io import get_data from mujoco.mjx._src.io import get_data_into from mujoco.mjx._src.io import make_data diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 31ce2057..90a709ab 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -28,8 +28,10 @@ from mujoco.mjx._src.collision_types import FunctionKey from mujoco.mjx._src.collision_types import GeomInfo from mujoco.mjx._src.collision_types import HFieldInfo from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import GeomType from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX # pylint: enable=g-importing-member _GeomInfo = Union[GeomInfo, ConvexInfo] @@ -42,6 +44,9 @@ def collider(ncon: int): def collide( m: Model, d: Data, key: FunctionKey, geom: jax.Array ) -> Collision: + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('collider requires JAX backend implementation.') + g1, g2 = geom.T infos = [ GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]), @@ -56,7 +61,7 @@ def collider(ncon: int): pos=0, mat=0, size=0, face=0, vert=0 ) elif key.types[i] == GeomType.MESH: - c, cm = infos[i], m.mesh_convex[key.data_ids[i]] + c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]] infos[i] = ConvexInfo(**vars(c), **vars(cm)) in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( pos=0, mat=0, size=0 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 7b947ce1..2f45bca8 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -71,9 +71,11 @@ from mujoco.mjx._src.collision_sdf import sphere_ellipsoid from mujoco.mjx._src.collision_types import FunctionKey from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import GeomType from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX # pylint: enable=g-importing-member import numpy as np @@ -227,7 +229,7 @@ def _geom_groups( if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD: # add static grid bounds to the grouping key for hfield collisions geom_rbound_hfield = ( - m.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound + m._impl.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound # pytype: disable=attribute-error ) nrow, ncol = m.hfield_nrow[data_ids[0]], m.hfield_ncol[data_ids[0]] xsize, ysize = m.hfield_size[data_ids[0]][:2] @@ -323,11 +325,11 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]: solref=solref, solreffriction=solreffriction, solimp=solimp, - dim=d.contact.dim, + dim=d._impl.contact.dim, # pytype: disable=attribute-error geom1=jp.array(geom[:, 0]), geom2=jp.array(geom[:, 1]), geom=jp.array(geom[:, :2]), - efc_address=d.contact.efc_address, + efc_address=d._impl.contact.efc_address, # pytype: disable=attribute-error ) return groups @@ -361,8 +363,12 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray: condim_counts = {} for k, v in group_counts.items(): - func = _COLLISION_FUNC[k.types] - num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v # pytype: disable=attribute-error + if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF: + ncon = m.opt.sdf_initpoints + else: + func = _COLLISION_FUNC[k.types] + ncon = func.ncon # pytype: disable=attribute-error + num_contacts = condim_counts.get(k.condim, 0) + ncon * v if max_contact_points > -1: num_contacts = min(max_contact_points, num_contacts) condim_counts[k.condim] = num_contacts @@ -374,7 +380,10 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray: def collision(m: Model, d: Data) -> Data: """Collides geometries.""" - if d.ncon == 0: + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('collision requires JAX backend implementation.') + + if d._impl.ncon == 0: # pytype: disable=attribute-error return d max_geom_pairs = _numeric(m, 'max_geom_pairs') @@ -424,4 +433,4 @@ def collision(m: Model, d: Data) -> Data: contacts = sum([condim_groups[k] for k in sorted(condim_groups)], []) contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts) - return d.replace(contact=contact) + return d.tree_replace({'_impl.contact': contact}) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index f33cbb07..ab41bc26 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -119,7 +119,7 @@ class SphereCollisionTest(parameterized.TestCase): 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) + _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-5) _SPHERE_CONVEX = """ @@ -143,7 +143,7 @@ class SphereCollisionTest(parameterized.TestCase): ) d, dx = _collide(xml) self.assertEmpty(d.contact.dist) - self.assertGreater(dx.contact.dist, 0) + self.assertGreater(dx._impl.contact.dist, 0) # face contact xml = self._SPHERE_CONVEX.replace( @@ -151,38 +151,38 @@ class SphereCollisionTest(parameterized.TestCase): ) d, dx = _collide(xml) for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'face', 1e-4) + _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'face', 1e-4) # deep face contact xml = self._SPHERE_CONVEX.replace( '', '' ) d, dx = _collide(xml) - self.assertTrue((dx.contact.dist < 0).all()) + self.assertTrue((dx._impl.contact.dist < 0).all()) self.assertTrue((d.contact.dist < 0).all()) - np.testing.assert_allclose(dx.contact.dist, [-0.07], atol=1e-5) - np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos) + np.testing.assert_allclose(dx._impl.contact.dist, [-0.07], atol=1e-5) + np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos) np.testing.assert_array_almost_equal( - dx.contact.frame, d.contact.frame.reshape((-1, 3, 3)) + dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3)) ) def test_sphere_convex_edge(self): # edge contact d, dx = _collide(self._SPHERE_CONVEX) for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'edge', 1e-4) + _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'edge', 1e-4) # deep edge penetration xml = self._SPHERE_CONVEX.replace( '', '' ) d, dx = _collide(xml) - self.assertTrue((dx.contact.dist < 0).all()) + self.assertTrue((dx._impl.contact.dist < 0).all()) self.assertTrue((d.contact.dist < 0).all()) - np.testing.assert_allclose(dx.contact.dist, [-0.06], atol=1e-5) - np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos) + np.testing.assert_allclose(dx._impl.contact.dist, [-0.06], atol=1e-5) + np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos) np.testing.assert_array_almost_equal( - dx.contact.frame, d.contact.frame.reshape((-1, 3, 3)) + dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3)) ) # vertex contact @@ -191,7 +191,7 @@ class SphereCollisionTest(parameterized.TestCase): ) d, dx = _collide(xml) for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex', 1e-4) + _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'vertex', 1e-4) # sphere center on vertex xml = self._SPHERE_CONVEX.replace( @@ -199,7 +199,9 @@ class SphereCollisionTest(parameterized.TestCase): ) d, dx = _collide(xml) for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex_center', 1e-4) + _assert_attr_eq( + dx._impl.contact, d.contact, field.name, 'vertex_center', 1e-4 + ) class EllipsoidCollisionTest(parameterized.TestCase): @@ -219,10 +221,10 @@ class EllipsoidCollisionTest(parameterized.TestCase): def test_plane_ellipsoid(self): """Tests ellipsoid plane contact.""" d, dx = _collide(self._ELLIPSOID_PLANE) - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5 + dx._impl.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5 ) _ELLIPSOID_ELLIPSOID = """ @@ -242,10 +244,10 @@ class EllipsoidCollisionTest(parameterized.TestCase): def test_ellipsoid_ellipsoid(self): """Tests ellipsoid ellipsoid contact.""" d, dx = _collide(self._ELLIPSOID_ELLIPSOID) - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2 + dx._impl.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2 ) _ELLIPSOID_SPHERE = """ @@ -266,10 +268,10 @@ class EllipsoidCollisionTest(parameterized.TestCase): """Tests ellipsoid capsule contact.""" d, dx = _collide(self._ELLIPSOID_SPHERE) d.contact.pos[0][2] = 0.03 # MJX finds the point on the surface - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4 + dx._impl.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4 ) _ELLIPSOID_CAPSULE = """ @@ -289,10 +291,10 @@ class EllipsoidCollisionTest(parameterized.TestCase): def test_capsule_ellipsoid(self): """Tests ellipsoid capsule contact.""" d, dx = _collide(self._ELLIPSOID_CAPSULE) - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5 + dx._impl.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5 ) _ELLIPSOID_CYLINDER = """ @@ -313,10 +315,10 @@ class EllipsoidCollisionTest(parameterized.TestCase): """Tests ellipsoid cylinder contact.""" d, dx = _collide(self._ELLIPSOID_CYLINDER) d.contact.pos[0][2] = 0.04 # MJX finds the deepest point on the surface - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4 + dx._impl.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4 ) @@ -357,7 +359,7 @@ class CapsuleCollisionTest(parameterized.TestCase): 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) + _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-4) _PARALLEL_CAP = """ @@ -378,14 +380,14 @@ class CapsuleCollisionTest(parameterized.TestCase): """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._impl.contact.dist, -0.05) np.testing.assert_allclose( - dx.contact.pos[0], + dx._impl.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 + dx._impl.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5 ) _CAP_BOX = """ @@ -408,30 +410,32 @@ class CapsuleCollisionTest(parameterized.TestCase): d, dx = _collide(self._CAP_BOX) # sort positions for comparison - idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1])) - dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]}) + idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1])) + dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]}) idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1])) d.contact.pos[:] = d.contact.pos[idx] d.contact.frame[:] = d.contact.frame[idx] d.contact.dist[:] = d.contact.dist[idx] for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4) + _assert_attr_eq( + dx._impl.contact, d.contact, field.name, 'capsule_convex', 1e-4 + ) def test_capsule_convex_face_deep(self): """Tests deep face penetration.""" xml = self._CAP_BOX.replace('', '') _, dx = _collide(xml) - self.assertTrue((dx.contact.dist < 0).all()) + self.assertTrue((dx._impl.contact.dist < 0).all()) np.testing.assert_array_almost_equal( - dx.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]]) + dx._impl.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]]) ) np.testing.assert_array_almost_equal( - dx.contact.dist, np.array([-0.15, -0.15]) + dx._impl.contact.dist, np.array([-0.15, -0.15]) ) np.testing.assert_array_almost_equal( - dx.contact.frame[:, 0], np.array([[0, 0, -1]] * 2) + dx._impl.contact.frame[:, 0], np.array([[0, 0, -1]] * 2) ) _CAP_EDGE_BOX = """ @@ -453,11 +457,11 @@ class CapsuleCollisionTest(parameterized.TestCase): """Tests edge contact.""" d, dx = _collide(self._CAP_EDGE_BOX) - c = dx.contact + c = dx._impl.contact self.assertEqual(c.pos.shape[0], 2) self.assertGreater(c.dist[1], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact) c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) @@ -469,13 +473,17 @@ class CapsuleCollisionTest(parameterized.TestCase): ) _, dx = _collide(xml) - np.testing.assert_array_equal(dx.contact.dist < 0, np.array([True, False])) - np.testing.assert_array_almost_equal(dx.contact.dist[0], np.array([-0.13])) - np.testing.assert_array_almost_equal( - dx.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3 + np.testing.assert_array_equal( + dx._impl.contact.dist < 0, np.array([True, False]) ) np.testing.assert_array_almost_equal( - dx.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3 + dx._impl.contact.dist[0], np.array([-0.13]) + ) + np.testing.assert_array_almost_equal( + dx._impl.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3 + ) + np.testing.assert_array_almost_equal( + dx._impl.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3 ) def test_capsule_convex_edge_shallow_tip(self): @@ -489,16 +497,16 @@ class CapsuleCollisionTest(parameterized.TestCase): xml = xml.replace(' @@ -586,10 +596,10 @@ class CylinderTest(absltest.TestCase): """Tests sphere cylinder contact.""" d, dx = _collide(self._SPHERE_CYLINDER) d.contact.pos[0][2] = 0.05 # MJX finds the deepest point on the surface - self.assertLess(dx.contact.dist[0], 0) + self.assertLess(dx._impl.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4 + dx._impl.contact, d.contact, field.name, 'sphere-cylinder', 1e-4 ) @@ -612,11 +622,11 @@ class ConvexTest(absltest.TestCase): """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) + np.testing.assert_array_less(dx._impl.contact.dist[:2], 0) + np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0) # extract the contact points with penetration c = jax.tree_util.tree_map( - lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact + lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx._impl.contact ) c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]]) for field in dataclasses.fields(Contact): @@ -638,16 +648,18 @@ class ConvexTest(absltest.TestCase): """Tests box collision with a plane.""" d, dx = _collide(self._FLAT_BOX_PLANE) - np.testing.assert_array_less(dx.contact.dist, 0) + np.testing.assert_array_less(dx._impl.contact.dist, 0) # sort positions for comparison - idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1])) - dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]}) + idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1])) + dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]}) idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1])) d.contact.pos[:] = d.contact.pos[idx] for field in dataclasses.fields(Contact): - _assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5) + _assert_attr_eq( + dx._impl.contact, d.contact, field.name, 'flat_box_plane', 1e-5 + ) _BOX_BOX = """ @@ -669,7 +681,7 @@ class ConvexTest(absltest.TestCase): def test_box_box(self): """Tests a face contact for a box-box collision.""" d, dx = _collide(self._BOX_BOX, keyframe=0) - c = dx.contact + c = dx._impl.contact self.assertEqual(c.pos.shape[0], 4) np.testing.assert_array_less(c.dist, 0) @@ -701,10 +713,10 @@ class ConvexTest(absltest.TestCase): 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) + np.testing.assert_array_less(dx._impl.contact.dist[:1], 0) + np.testing.assert_array_less(-dx._impl.contact.dist[1:], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact) c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) @@ -736,7 +748,7 @@ class ConvexTest(absltest.TestCase): ).read_bytes(), } _, dx = _collide(self._CONVEX_CONVEX, assets=assets) - c = dx.contact + c = dx._impl.contact # Only one contact point for an edge contact. self.assertLess(c.dist[0], 0) @@ -766,7 +778,7 @@ class ConvexTest(absltest.TestCase): def test_convex_convex_edge(self): """Tests convex-convex collisions with edge contact via _sat_gaussmap.""" _, dx = _collide(self._CONVEX_CONVEX_THIN) - c = dx.contact + c = dx._impl.contact # Only one contact point for an edge contact. self.assertLess(c.dist[0], 0) @@ -781,7 +793,7 @@ class ConvexTest(absltest.TestCase): 'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"' ) ) - c = dx.contact + c = dx._impl.contact self.assertTrue((c.dist > 0).all()) @@ -837,8 +849,8 @@ class HFieldTest(absltest.TestCase): # check that all geoms are colliding with the hfield for geom_id in [1, 2, 3]: - mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1) - c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact) + mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1) + c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact) self.assertTrue((c.dist < 0).any()) self.assertTrue((c.dist > -1e-3).any()) # all contact normals are roughly pointing in the right direction @@ -850,7 +862,7 @@ class HFieldTest(absltest.TestCase): for p in positions: xml = self._HFIELD.replace('= 0).all()) + self.assertTrue((dx._impl.contact.dist >= 0).all()) def test_hfield_deep(self): """Tests that objects with deep penetration do not get stuck.""" @@ -868,8 +880,8 @@ class HFieldTest(absltest.TestCase): # check that all geoms are colliding with the hfield for geom_id in [1, 2, 3]: - mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1) - c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact) + mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1) + c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact) # all contact normals are in the top half-face of the hfield self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all()) @@ -892,8 +904,8 @@ class BodyPairFilterTest(absltest.TestCase): 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) + self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx._impl.contact.pos.shape[0], 0) _PARENT_CHILD = """ @@ -923,8 +935,8 @@ class BodyPairFilterTest(absltest.TestCase): 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) + self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx._impl.contact.pos.shape[0], 0) def test_disable_filter_parent_child(self): """Tests that filterparent flag disables parent-child filtering.""" @@ -941,8 +953,8 @@ class BodyPairFilterTest(absltest.TestCase): 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) + self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx._impl.contact.pos.shape[0], 1) class DimTest(parameterized.TestCase): @@ -1027,8 +1039,8 @@ class TopKContactTest(absltest.TestCase): dx_all = collision_jit_fn(mx_all, dx) dx_top_k = collision_jit_fn(mx_top_k, dx) - self.assertEqual(dx_all.contact.dist.shape, (3,)) - self.assertEqual(dx_top_k.contact.dist.shape, (2,)) + self.assertEqual(dx_all._impl.contact.dist.shape, (3,)) + self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,)) _CAPSULES_MAX_PAIR = """ @@ -1073,9 +1085,9 @@ class TopKContactTest(absltest.TestCase): dx_all = collision_jit_fn(mx_all, dx) dx_top_k = collision_jit_fn(mx_top_k, dx) - self.assertEqual(dx_all.contact.dist.shape, (6,)) - self.assertEqual(dx_top_k.contact.dist.shape, (2,)) - self.assertTrue((dx_top_k.contact.dist < 0).all()) + self.assertEqual(dx_all._impl.contact.dist.shape, (6,)) + self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,)) + self.assertTrue((dx_top_k._impl.contact.dist < 0).all()) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py index 1874211d..3bf8b824 100644 --- a/mjx/mujoco/mjx/_src/collision_sdf.py +++ b/mjx/mujoco/mjx/_src/collision_sdf.py @@ -32,7 +32,9 @@ from mujoco.mjx._src.collision_types import Collision from mujoco.mjx._src.collision_types import GeomInfo from mujoco.mjx._src.dataclasses import PyTreeNode from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX # pylint: enable=g-importing-member # the SDF function takes position in, and returns a distance or objective diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index e778bcd8..9b41daa7 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -28,10 +28,12 @@ from mujoco.mjx._src.types import ConeType from mujoco.mjx._src.types import ConstraintType from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX 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 +from mujoco.mjx._src.types import ModelJAX from mujoco.mjx._src.types import ObjType # pylint: enable=g-importing-member import numpy as np @@ -281,6 +283,10 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]: def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for tendon equality constraints.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError( + '_efc_equality_tendon requires JAX backend implementation.' + ) eq_id = np.nonzero(m.eq_type == EqType.TENDON)[0] @@ -314,9 +320,9 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]: return jax.tree_util.tree_map(lambda x: x * active, efc) inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id] - jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id] - pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id] - pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id] + jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id] + pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id] + pos2 = d._impl.ten_length[obj2id] - m.tendon_length0[obj2id] invweight = inv1 + inv2 * (obj2id > -1) return rows( @@ -326,8 +332,11 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]: def _efc_friction(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for dof frictionloss.""" - dof_id = np.nonzero(m.dof_hasfrictionloss)[0] - tendon_id = np.nonzero(m.tendon_hasfrictionloss)[0] + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('_efc_friction requires JAX backend implementation.') + + dof_id = np.nonzero(m._impl.dof_hasfrictionloss)[0] + tendon_id = np.nonzero(m._impl.tendon_hasfrictionloss)[0] size = dof_id.size + tendon_id.size if (m.opt.disableflags & DisableBit.FRICTIONLOSS) or (size == 0): @@ -337,7 +346,7 @@ def _efc_friction(m: Model, d: Data) -> Optional[_Efc]: args_dof += (m.dof_solimp,) args_dof = jax.tree_util.tree_map(lambda x: x[dof_id], args_dof) - args_ten = (d.ten_J, m.tendon_frictionloss, m.tendon_invweight0) + args_ten = (d._impl.ten_J, m.tendon_frictionloss, m.tendon_invweight0) args_ten += (m.tendon_solref_fri, m.tendon_solimp_fri) args_ten = jax.tree_util.tree_map(lambda x: x[tendon_id], args_ten) @@ -415,6 +424,9 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for tendon limits.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('_efc_limit_tendon requires JAX backend implementation.') + tendon_id = np.nonzero(m.tendon_limited)[0] if (m.opt.disableflags & DisableBit.LIMIT) or tendon_id.size == 0: @@ -423,8 +435,8 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]: length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map( lambda x: x[tendon_id], ( - d.ten_length, - d.ten_J, + d._impl.ten_length, + d._impl.ten_J, m.tendon_range, m.tendon_margin, m.tendon_invweight0, @@ -446,8 +458,12 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]: def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for frictionless contacts.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError( + '_efc_contact_frictionless requires JAX backend implementation.' + ) - con_id = np.nonzero(d.contact.dim == 1)[0] + con_id = np.nonzero(d._impl.contact.dim == 1)[0] if con_id.size == 0: return None @@ -473,15 +489,19 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]: jp.zeros_like(pos), ) - contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) + contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact) return rows(contact) def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]: """Calculates constraint rows for frictional pyramidal contacts.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError( + '_efc_contact_pyramidal requires JAX backend implementation.' + ) - con_id = np.nonzero(d.contact.dim == condim)[0] + con_id = np.nonzero(d._impl.contact.dim == condim)[0] if con_id.size == 0: return None @@ -518,15 +538,19 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]: jp.zeros_like(pos), ) - contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) + contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact) # concatenate to drop row grouping return jax.tree_util.tree_map(jp.concatenate, rows(contact)) def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]: """Calculates constraint rows for frictional elliptic contacts.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError( + '_efc_contact_elliptic requires JAX backend implementation.' + ) - con_id = np.nonzero(d.contact.dim == condim)[0] + con_id = np.nonzero(d._impl.contact.dim == condim)[0] if con_id.size == 0: return None @@ -563,7 +587,7 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]: jp.zeros_like(pos), ) - contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) + contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact) # concatenate to drop row grouping return jax.tree_util.tree_map(jp.concatenate, rows(contact)) @@ -602,14 +626,14 @@ def make_efc_type( if not m.opt.disableflags & DisableBit.FRICTIONLOSS: nf_dof = ( - m.dof_hasfrictionloss.sum() - if isinstance(m, Model) + m._impl.dof_hasfrictionloss.sum() + if isinstance(m, Model) and isinstance(m._impl, ModelJAX) else (m.dof_frictionloss > 0).sum() ) efc_types += [ConstraintType.FRICTION_DOF] * nf_dof nf_tendon = ( - m.tendon_hasfrictionloss.sum() - if isinstance(m, Model) + m._impl.tendon_hasfrictionloss.sum() + if isinstance(m, Model) and isinstance(m._impl, ModelJAX) else (m.tendon_frictionloss > 0).sum() ) efc_types += [ConstraintType.FRICTION_TENDON] * nf_tendon @@ -683,10 +707,14 @@ def make_constraint(m: Model, d: Data) -> Data: if not efcs: z = jp.empty(0) - d = d.replace(efc_J=jp.empty((0, m.nv))) - d = d.replace( - efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z, efc_margin=z - ) + d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))}) + d = d.tree_replace({ + '_impl.efc_D': z, + '_impl.efc_aref': z, + '_impl.efc_frictionloss': z, + '_impl.efc_pos': z, + '_impl.efc_margin': z, + }) return d efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs) @@ -699,9 +727,13 @@ def make_constraint(m: Model, d: Data) -> Data: return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss aref, r, pos, margin, frictionloss = fn(efc) - d = d.replace( - efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin - ) - d = d.replace(efc_frictionloss=frictionloss) + d = d.tree_replace({ + '_impl.efc_J': efc.J, + '_impl.efc_D': 1 / r, + '_impl.efc_aref': aref, + '_impl.efc_pos': pos, + '_impl.efc_margin': margin, + }) + d = d.tree_replace({'_impl.efc_frictionloss': frictionloss}) return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 4e53fb13..f94873ff 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -69,16 +69,16 @@ class ConstraintTest(parameterized.TestCase): order = test_util.efc_order(m, d, dx) d_efc_j = d.efc_J.reshape((-1, m.nv)) - _assert_eq(d_efc_j, dx.efc_J[order][: d.nefc], 'efc_J') - _assert_eq(0, dx.efc_J[order][d.nefc :], 'efc_J') - _assert_eq(d.efc_aref, dx.efc_aref[order][: d.nefc], 'efc_aref') - _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref') - _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D') - _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos') - _assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos') + _assert_eq(d_efc_j, dx._impl.efc_J[order][: d.nefc], 'efc_J') + _assert_eq(0, dx._impl.efc_J[order][d.nefc :], 'efc_J') + _assert_eq(d.efc_aref, dx._impl.efc_aref[order][: d.nefc], 'efc_aref') + _assert_eq(0, dx._impl.efc_aref[order][d.nefc :], 'efc_aref') + _assert_eq(d.efc_D, dx._impl.efc_D[order][: d.nefc], 'efc_D') + _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos') + _assert_eq(dx._impl.efc_pos[order][d.nefc :], 0, 'efc_pos') _assert_eq( d.efc_frictionloss, - dx.efc_frictionloss[order][: d.nefc], + dx._impl.efc_frictionloss[order][: d.nefc], 'efc_frictionloss', ) @@ -104,7 +104,7 @@ class ConstraintTest(parameterized.TestCase): self.assertEqual(nl, 0) self.assertEqual(nc, 0) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 0) + self.assertEqual(dx._impl.efc_J.shape[0], 0) def test_disable_equality(self): m = test_util.load_test_file('constraints.xml') @@ -115,7 +115,9 @@ class ConstraintTest(parameterized.TestCase): self.assertEqual(nl, 5) self.assertEqual(nc, 180) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 187) # only joint/tendon limit, contact + self.assertEqual( + dx._impl.efc_J.shape[0], 187 + ) # only joint/tendon limit, contact def test_disable_contact(self): m = test_util.load_test_file('constraints.xml') @@ -126,7 +128,9 @@ class ConstraintTest(parameterized.TestCase): self.assertEqual(nl, 5) self.assertEqual(nc, 0) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 27) # only equality, joint/tendon limit + self.assertEqual( + dx._impl.efc_J.shape[0], 27 + ) # only equality, joint/tendon limit def test_disable_frictionloss(self): m = test_util.load_test_file('constraints.xml') @@ -137,7 +141,7 @@ class ConstraintTest(parameterized.TestCase): self.assertEqual(nl, 5) self.assertEqual(nc, 180) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 205) + self.assertEqual(dx._impl.efc_J.shape[0], 205) def test_margin(self): """Test margin.""" @@ -163,8 +167,8 @@ class ConstraintTest(parameterized.TestCase): dx = mjx.make_constraint(mx, dx) order = test_util.efc_order(m, d, dx) - _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos') - _assert_eq(d.efc_margin, dx.efc_margin[order][: d.nefc], 'efc_margin') + _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos') + _assert_eq(d.efc_margin, dx._impl.efc_margin[order][: d.nefc], 'efc_margin') if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index 85a18e60..b8bdb758 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -18,6 +18,7 @@ import copy import dataclasses import typing from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union +import warnings import jax import numpy as np diff --git a/mjx/mujoco/mjx/_src/derivative.py b/mjx/mujoco/mjx/_src/derivative.py new file mode 100644 index 00000000..5ca84a42 --- /dev/null +++ b/mjx/mujoco/mjx/_src/derivative.py @@ -0,0 +1,61 @@ +# Copyright 2025 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. +# ============================================================================== +"""Derivative functions.""" + +from typing import Optional + +import jax +from jax import numpy as jp +# 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 Model + + +def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]: + """Analytical derivative of smooth forces w.r.t velocities.""" + + qderiv = None + + # qDeriv += d qfrc_actuator / d qvel + if not m.opt.disableflags & DisableBit.ACTUATION: + affine_bias = m.actuator_biastype == BiasType.AFFINE + bias_vel = m.actuator_biasprm[:, 2] * affine_bias + affine_gain = m.actuator_gaintype == GainType.AFFINE + gain_vel = m.actuator_gainprm[:, 2] * affine_gain + ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act) + vel = bias_vel + gain_vel * ctrl + qderiv = d._impl.actuator_moment.T @ jax.vmap(jp.multiply)( + d._impl.actuator_moment, vel + ) + + # qDeriv += d qfrc_passive / d qvel + if not m.opt.disableflags & DisableBit.PASSIVE: + if qderiv is None: + qderiv = -jp.diag(m.dof_damping) + else: + qderiv -= jp.diag(m.dof_damping) + if m.ntendon: + qderiv -= d._impl.ten_J.T @ jp.diag(m.tendon_damping) @ d._impl.ten_J + # TODO(robotics-simulation): fluid drag model + if m.opt.has_fluid_params: # pytype: disable=attribute-error + raise NotImplementedError('fluid drag not supported for implicitfast') + + # TODO(team): rne derivative + + return qderiv diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 07f33a90..62535122 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -22,6 +22,7 @@ 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 derivative from mujoco.mjx._src import math from mujoco.mjx._src import passive from mujoco.mjx._src import scan @@ -32,12 +33,15 @@ 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 DataJAX 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 ModelJAX +from mujoco.mjx._src.types import TrnType # pylint: enable=g-importing-member import numpy as np @@ -79,10 +83,13 @@ def fwd_position(m: Model, d: Data) -> Data: @named_scope def fwd_velocity(m: Model, d: Data) -> Data: """Velocity-dependent computations.""" - d = d.replace( - actuator_velocity=d.actuator_moment @ d.qvel, - ten_velocity=d.ten_J @ d.qvel, - ) + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('fwd_velocity requires JAX backend implementation.') + + d = d.tree_replace({ + '_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel, + '_impl.ten_velocity': d._impl.ten_J @ d.qvel, + }) d = smooth.com_vel(m, d) d = passive.passive(m, d) d = smooth.rne(m, d) @@ -171,13 +178,41 @@ def fwd_actuation(m: Model, d: Data) -> Data: m.actuator_gainprm, m.actuator_biastype, m.actuator_biasprm, - d.actuator_length, - d.actuator_velocity, + d._impl.actuator_length, + d._impl.actuator_velocity, ctrl_act, jp.array(m.actuator_lengthrange), jp.array(m.actuator_acc0), group_by='u', ) + + # tendon total force clamping + if np.any(m.tendon_actfrclimited): + (tendon_actfrclimited_id,) = np.nonzero(m.tendon_actfrclimited) + actuator_tendon = m.actuator_trntype == TrnType.TENDON + + force_mask = [ + actuator_tendon & (m.actuator_trnid[:, 0] == tendon_id) + for tendon_id in tendon_actfrclimited_id + ] + force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask]) + force_mat = np.array(force_mask)[:, force_ids] + tendon_total_force = force_mat @ force[force_ids] + + force_scaling = jp.where( + tendon_total_force < m.tendon_actfrcrange[tendon_actfrclimited_id, 0], + m.tendon_actfrcrange[tendon_actfrclimited_id, 0] / tendon_total_force, + 1, + ) + force_scaling = jp.where( + tendon_total_force > m.tendon_actfrcrange[tendon_actfrclimited_id, 1], + m.tendon_actfrcrange[tendon_actfrclimited_id, 1] / tendon_total_force, + force_scaling, + ) + + tendon_forces = force[force_ids] * (force_mat.T @ force_scaling) + force = force.at[force_ids].set(tendon_forces) + forcerange = jp.where( m.actuator_forcelimited[:, None], m.actuator_forcerange, @@ -185,7 +220,7 @@ def fwd_actuation(m: Model, d: Data) -> Data: ) force = jp.clip(force, forcerange[:, 0], forcerange[:, 1]) - qfrc_actuator = d.actuator_moment.T @ force + qfrc_actuator = d._impl.actuator_moment.T @ force if m.ngravcomp: # actuator-level gravity compensation, skip if added as passive force @@ -301,13 +336,17 @@ def _advance( @named_scope def euler(m: Model, d: Data) -> Data: """Euler integrator, semi-implicit in velocity.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('euler requires JAX backend implementation.') + # integrate damping implicitly qacc = d.qacc if not m.opt.disableflags & DisableBit.EULERDAMP: if support.is_sparse(m): - dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)) + qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping) else: - dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping)) + qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping) + dh = d.tree_replace({'_impl.qM': qM}) dh = smooth.factor_m(m, dh) qfrc = d.qfrc_smooth + d.qfrc_constraint qacc = smooth.solve_m(m, dh, qfrc) @@ -317,7 +356,7 @@ def euler(m: Model, d: Data) -> Data: @named_scope def rungekutta4(m: Model, d: Data) -> Data: """Runge-Kutta explicit order 4 integrator.""" - d_t0 = d + d0 = 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) @@ -338,9 +377,9 @@ def rungekutta4(m: Model, d: Data) -> Data: 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 + kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d0.qpos, dqvel) + kact = d0.act + dact_dot * m.opt.timestep + kqvel = d0.qvel + dqacc * m.opt.timestep d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t) d = forward(m, d) @@ -352,44 +391,25 @@ def rungekutta4(m: Model, d: Data) -> Data: 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 + qvel, qacc, act_dot, _, d1 = out - d = _advance(m, d_t0, act_dot, qacc, qvel) + d = d1.replace(qpos=d0.qpos, qvel=d0.qvel, act=d0.act, time=d0.time) + d = _advance(m, d, act_dot, qacc, qvel) return d @named_scope def implicit(m: Model, d: Data) -> Data: """Integrates fully implicit in velocity.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('implicit requires JAX backend implementation.') - qderiv = None - - # qDeriv += d qfrc_actuator / d qvel - if not m.opt.disableflags & DisableBit.ACTUATION: - affine_bias = m.actuator_biastype == BiasType.AFFINE - bias_vel = m.actuator_biasprm[:, 2] * affine_bias - affine_gain = m.actuator_gaintype == GainType.AFFINE - gain_vel = m.actuator_gainprm[:, 2] * affine_gain - ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act) - vel = bias_vel + gain_vel * ctrl - qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment - - # qDeriv += d qfrc_passive / d qvel - if not m.opt.disableflags & DisableBit.PASSIVE: - if qderiv is None: - qderiv = -jp.diag(m.dof_damping) - else: - qderiv -= jp.diag(m.dof_damping) - if m.ntendon: - qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J - # TODO(robotics-simulation): fluid drag model - if m.opt.has_fluid_params: - raise NotImplementedError('fluid drag not supported for implicitfast') + qderiv = derivative.deriv_smooth_vel(m, d) qacc = d.qacc if qderiv is not None: # TODO(robotics-simulation): use smooth.factor_m / solve_m here: - qm = support.full_m(m, d) if support.is_sparse(m) else d.qM + qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM qm -= m.opt.timestep * qderiv qh, _ = jax.scipy.linalg.cho_factor(qm) qfrc = d.qfrc_smooth + d.qfrc_constraint @@ -401,6 +421,9 @@ def implicit(m: Model, d: Data) -> Data: @named_scope def forward(m: Model, d: Data) -> Data: """Forward dynamics.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('forward requires JAX backend implementation.') + d = fwd_position(m, d) d = sensor.sensor_pos(m, d) d = fwd_velocity(m, d) @@ -409,7 +432,7 @@ def forward(m: Model, d: Data) -> Data: d = fwd_acceleration(m, d) d = sensor.sensor_acc(m, d) - if d.efc_J.size == 0: + if d._impl.efc_J.size == 0: d = d.replace(qacc=d.qacc_smooth) return d diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index ebe94ebe..f90fae1a 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -17,6 +17,7 @@ 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 @@ -130,6 +131,7 @@ class ForwardTest(absltest.TestCase): _assert_attr_eq(d, dx, 'qpos') _assert_attr_eq(d, dx, 'act') _assert_attr_eq(d, dx, 'time') + _assert_attr_eq(d, dx, 'xpos') def test_eulerdamp(self): m = test_util.load_test_file('pendula.xml') @@ -195,6 +197,27 @@ class ActuatorTest(parameterized.TestCase): dx = jax.jit(mjx.euler)(mx, dx) _assert_attr_eq(d, dx, 'act') + def test_tendon_force_clamp(self): + m = test_util.load_test_file('actuator/tendon_force_clamp.xml') + d = mujoco.MjData(m) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + + dx = dx.replace(ctrl=jp.array([1.0, 1.0, 1.0, -4.0, 1.0, -20.0, 5.0, -5.0])) + dx = mjx.forward(mx, dx) + + _assert_eq( + dx.actuator_force, + jp.array([1.0, 1.0, 1.0, -4.0 / 3.0, 1.0 / 3.0, -10.0, 5.0, -5.0]), + 'actuator_force', + ) + + _assert_eq( + dx.sensordata, + jp.array([3.0, -1.0, -10.0, 0.0]), + 'sensordata', + ) + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/inverse.py b/mjx/mujoco/mjx/_src/inverse.py new file mode 100644 index 00000000..5df72de9 --- /dev/null +++ b/mjx/mujoco/mjx/_src/inverse.py @@ -0,0 +1,106 @@ +# Copyright 2025 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. +# ============================================================================== +"""Inverse dynamics functions.""" + +from jax import numpy as jp +from mujoco.mjx._src import derivative +from mujoco.mjx._src import forward +from mujoco.mjx._src import sensor +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 Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import EnableBit +from mujoco.mjx._src.types import IntegratorType +from mujoco.mjx._src.types import Model + + +def discrete_acc(m: Model, d: Data) -> Data: + """Convert discrete-time qacc to continuous-time qacc.""" + + if m.opt.integrator == IntegratorType.RK4: + raise RuntimeError( + 'discrete inverse dynamics is not supported by RK4 integrator' + ) + elif m.opt.integrator == IntegratorType.EULER: + dsbl_eulerdamp = m.opt.disableflags & DisableBit.EULERDAMP + no_dof_damping = (m.dof_damping == 0).all() + if dsbl_eulerdamp or no_dof_damping: + return d + + # set qfrc = (M + h*diag(B)) * qacc + qfrc = support.mul_m(m, d, d.qacc) + qfrc += m.opt.timestep * m.dof_damping * d.qacc + elif m.opt.integrator == IntegratorType.IMPLICITFAST: + qm = support.full_m(m, d) + + # compute analytical derivative qDeriv; skip rne derivative + qderiv = derivative.deriv_smooth_vel(m, d) + if qderiv is not None: + # M = M - dt*qDeriv + qm -= m.opt.timestep * qderiv + + # set qfrc = (M - dt*qDeriv) * qacc + qfrc = qm @ d.qacc + else: + raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') + + # solve for qacc: qfrc = M * qacc + qacc = smooth.solve_m(m, d, qfrc) + + return d.replace(qacc=qacc) + + +def inv_constraint(m: Model, d: Data) -> Data: + """Inverse constraint solver.""" + + # no constraints + if d._impl.efc_J.size == 0: # pytype: disable=attribute-error + return d.replace(qfrc_constraint=jp.zeros(m.nv)) + + # update + ctx = solver.Context.create(m, d, grad=False) + + return d.tree_replace({ + 'qfrc_constraint': ctx.qfrc_constraint, + '_impl.efc_force': ctx.efc_force, + }) + + +def inverse(m: Model, d: Data) -> Data: + """Inverse dynamics.""" + d = forward.fwd_position(m, d) + d = sensor.sensor_pos(m, d) + d = forward.fwd_velocity(m, d) + d = sensor.sensor_vel(m, d) + + qacc = d.qacc + if m.opt.enableflags & EnableBit.INVDISCRETE: + d = discrete_acc(m, d) + + d = inv_constraint(m, d) + d = smooth.rne(m, d, flg_acc=True) + d = sensor.sensor_acc(m, d) + + qfrc_inverse = ( + d.qfrc_bias + m.dof_armature * d.qacc - d.qfrc_passive - d.qfrc_constraint + ) + + if m.opt.enableflags & EnableBit.INVDISCRETE: + return d.replace(qfrc_inverse=qfrc_inverse, qacc=qacc) + else: + return d.replace(qfrc_inverse=qfrc_inverse) diff --git a/mjx/mujoco/mjx/_src/inverse_test.py b/mjx/mujoco/mjx/_src/inverse_test.py new file mode 100644 index 00000000..227406e3 --- /dev/null +++ b/mjx/mujoco/mjx/_src/inverse_test.py @@ -0,0 +1,114 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for inverse dynamics functions.""" +from absl.testing import absltest +from absl.testing import parameterized +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import support +import numpy as np + +# tolerance for difference between MuJoCo and MJX calculations - mostly +# due to float precision +_TOLERANCE = 1e-5 + + +def _assert_eq(a, b, name, tol=_TOLERANCE): + tol = tol * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) + + +class InverseTest(parameterized.TestCase): + + @parameterized.parameters( + (mujoco.mjtIntegrator.mjINT_EULER, False, False), + (mujoco.mjtIntegrator.mjINT_EULER, False, True), + (mujoco.mjtIntegrator.mjINT_EULER, True, False), + (mujoco.mjtIntegrator.mjINT_EULER, True, True), + (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, False, False), + (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, True, False), + ) + def test_forward_inverse_match(self, integrator, invdiscrete, eulerdamp): + m = mujoco.MjModel.from_xml_string(""" + + + """) + m.opt.integrator = integrator + if invdiscrete: + m.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_INVDISCRETE + if not eulerdamp: + m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_EULERDAMP + + d = mujoco.MjData(m) + d.qvel = np.random.uniform(low=-0.01, high=0.01, size=d.qvel.shape) + d.ctrl = np.random.uniform(low=-0.01, high=0.01, size=d.ctrl.shape) + d.qfrc_applied = np.random.uniform( + low=-0.01, high=0.01, size=d.qfrc_applied.shape + ) + d.xfrc_applied = np.random.uniform( + low=-0.01, high=0.01, size=d.xfrc_applied.shape + ) + mujoco.mj_step(m, d, 100) + + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + dx_next = mjx.step(mx, dx) + qacc_fd = (dx_next.qvel - dx.qvel) / mx.opt.timestep + + dx = mjx.forward(mx, dx) + + if invdiscrete: + dx = dx.replace(qacc=qacc_fd) + + dxinv = mjx.inverse(mx, dx) + + fwdinv0 = jp.linalg.norm( + dxinv.qfrc_constraint - dx.qfrc_constraint, ord=np.inf + ) + fwdinv1 = jp.linalg.norm( + dxinv.qfrc_inverse + - ( + dx.qfrc_applied + dx.qfrc_actuator + support.xfrc_accumulate(mx, dx) + ), + ord=np.inf, + ) + + self.assertLess(fwdinv0, 1.0e-3) + self.assertLess(fwdinv1, 1.0e-3) + _assert_eq(dxinv.qacc, dx.qacc, 'qacc') + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 1e8c2b5e..f8aaf565 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -15,7 +15,10 @@ """Functions to initialize, load, or save data.""" import copy -from typing import List, Tuple, Union +import logging +import os +from typing import Any, Dict, List, Optional, Tuple, Union +import warnings import jax from jax import numpy as jp @@ -29,6 +32,130 @@ import numpy as np import scipy +def _is_cuda_gpu_device(device: jax.Device) -> bool: + try: + cuda_devices = jax.devices('cuda') + except RuntimeError: + logging.info('No CUDA GPU devices found in jax.devices("cuda").') + return False + return device in cuda_devices + + +def _resolve_backend_impl( + device: jax.Device, +) -> types.BackendImpl: + """Pick a default backend impl based on the device specified.""" + if _is_cuda_gpu_device(device): + # TODO(btaba): Remove flag once Warp is ready to launch. + mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true' + if mjx_warp_enabled: + logging.debug('Picking default backend implementation: Warp.') + return types.BackendImpl.WARP + logging.info('MJX Warp is disabled via MJX_WARP_ENABLED=false.') + + if device.platform in ('gpu', 'tpu'): + logging.debug('Picking default backend implementation: JAX.') + return types.BackendImpl.JAX + + if device.platform == 'cpu': + mjx_c_default = ( + os.environ.get('MJX_C_DEFAULT_ENABLED', 'f').lower() == 'true' + ) + if mjx_c_default: + logging.debug('Picking default backend implementation: C.') + return types.BackendImpl.C + return types.BackendImpl.JAX + + raise ValueError(f'Unsupported device: {device}') + + +def _resolve_device( + backend_impl: types.BackendImpl, +) -> jax.Device: + """Resolves a device based on the backend implementation.""" + backend_impl = types.BackendImpl(backend_impl) + if backend_impl == types.BackendImpl.JAX: + device_0 = jax.devices()[0] + logging.debug('Picking default device: %s.', device_0) + return device_0 + + if backend_impl == types.BackendImpl.C: + cpu_0 = jax.devices('cpu')[0] + logging.debug('Picking default device: %s', cpu_0) + return cpu_0 + + if backend_impl == types.BackendImpl.WARP: + # WARP backend requires a CUDA GPU. + cuda_gpus = [d for d in jax.devices('cuda')] + if not cuda_gpus: + raise AssertionError( + 'No CUDA GPU devices found in' + f' jax.devices("cuda")={jax.devices("cuda")}.' + ) + + logging.debug('Picking default device: %s', cuda_gpus[0]) + return cuda_gpus[0] + + raise ValueError(f'Unsupported backend implementation: {backend_impl}') + + +def _check_backend_impl_device_compatibility( + backend_impl: Union[str, types.BackendImpl], + device: jax.Device, +) -> None: + """Checks that the backend implementation is compatible with the device.""" + if backend_impl is None: + raise ValueError('No backend implementation specified.') + + backend_impl = types.BackendImpl(backend_impl) + + if backend_impl == types.BackendImpl.WARP: + if not _is_cuda_gpu_device(device): + raise AssertionError( + 'Warp backend implementation requires a CUDA GPU device, got ' + f'{device}.' + ) + + mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true' + if not mjx_warp_enabled: + raise AssertionError( + 'Warp backend implementation is disabled via MJX_WARP_ENABLED=false.' + ) + + is_cpu_device = device.platform == 'cpu' + if backend_impl == types.BackendImpl.C: + if not is_cpu_device: + raise AssertionError( + f'C backend implementation requires a CPU device, got {device}.' + ) + + # NB: JAX backend works with any device. + + +def _resolve_backend_impl_and_device( + backend_impl: Optional[Union[str, types.BackendImpl]], + device: Optional[jax.Device] = None, +) -> Tuple[types.BackendImpl, jax.Device]: + """Resolves a backend implementation and device.""" + if backend_impl: + backend_impl = types.BackendImpl(backend_impl) + + has_backend_impl, has_device = backend_impl is not None, device is not None + if (has_backend_impl, has_device) == (True, True): + pass + elif (has_backend_impl, has_device) == (True, False): + device = _resolve_device(backend_impl) + elif (has_backend_impl, has_device) == (False, True): + backend_impl = _resolve_backend_impl(device) + else: + device = jax.devices(jax.default_backend())[0] + logging.info('Using JAX default device: %s.', device) + backend_impl = _resolve_backend_impl(device) + + _check_backend_impl_device_compatibility(backend_impl, device) + return backend_impl, device # pytype: disable=bad-return-type + + def _strip_weak_type(tree): def f(leaf): if isinstance(leaf, jax.Array): @@ -38,32 +165,27 @@ def _strip_weak_type(tree): return jax.tree_util.tree_map(f, tree) -def _make_option( - o: mujoco.MjOption, _full_compat: bool = False # pylint: disable=invalid-name +def _put_option( + o: mujoco.MjOption, + backend_impl: types.BackendImpl, + impl_fields: Optional[dict[str, Any]] = None, ) -> types.Option: """Returns mjx.Option given mujoco.MjOption.""" - if not _full_compat: - if o.integrator not in set(types.IntegratorType): - raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') + if o.integrator not in set(types.IntegratorType): + raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') - if o.cone not in set(types.ConeType): - raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') + if o.cone not in set(types.ConeType): + raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') - if o.jacobian not in set(types.JacobianType): - raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}') + if o.jacobian not in set(types.JacobianType): + raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}') - if o.solver not in set(types.SolverType): - raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') + if o.solver not in set(types.SolverType): + raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') - for i in range(mujoco.mjtEnableBit.mjNENABLE): - if o.enableflags & 2**i: - raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') - - has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() - implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST - if not _full_compat: - if implicitfast and has_fluid_params: - raise NotImplementedError('implicitfast not implemented for fluid drag.') + for i in range(mujoco.mjtEnableBit.mjNENABLE): + if o.enableflags & 2**i and 2 ** i not in set(types.EnableBit): + raise NotImplementedError(f'{mujoco.mjtEnableBit(2**i)}') fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()} fields['integrator'] = types.IntegratorType(o.integrator) @@ -71,12 +193,25 @@ def _make_option( fields['jacobian'] = types.JacobianType(o.jacobian) fields['solver'] = types.SolverType(o.solver) fields['disableflags'] = types.DisableBit(o.disableflags) - fields['has_fluid_params'] = has_fluid_params + fields['enableflags'] = types.EnableBit(o.enableflags) - return types.Option(**fields) + if backend_impl == types.BackendImpl.JAX: + has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() + implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST + if implicitfast and has_fluid_params: + raise NotImplementedError('implicitfast not implemented for fluid drag.') + fields['has_fluid_params'] = has_fluid_params + return types.OptionJAX(**fields, **(impl_fields or {})) + + if backend_impl == types.BackendImpl.C: + c_field_keys = types.OptionC.__annotations__.keys() - fields.keys() + c_fields = {k: getattr(o, k, None) for k in c_field_keys} + return types.OptionC(**fields, **c_fields, **(impl_fields or {})) + + raise NotImplementedError(f'Unsupported backend: {backend_impl}') -def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: +def _put_statistic(s: mujoco.MjStatistic) -> types.Statistic: """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" return types.Statistic( meaninertia=s.meaninertia, @@ -87,26 +222,16 @@ def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: ) -def put_model( - m: mujoco.MjModel, device=None, _full_compat: bool = False # pylint: disable=invalid-name +def _put_model_jax( + m: mujoco.MjModel, + device: Optional[jax.Device] = None, ) -> types.Model: - """Puts mujoco.MjModel onto a device, resulting in mjx.Model. - - Args: - m: the model to put onto device - device: which device to use - if unspecified picks the default device - _full_compat: put all MjModel fields onto device irrespective of MJX support - This is an experimental feature. Avoid using it for now. - - Returns: - an mjx.Model placed on device - """ - + """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" mesh_geomid = set() for g1, g2, ip in collision_driver.geom_pairs(m): t1, t2 = m.geom_type[[g1, g2]] # check collision function exists for type pair - if not collision_driver.has_collision_fn(t1, t2) and not _full_compat: + if not collision_driver.has_collision_fn(t1, t2): t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') # margin/gap not supported for meshes and height fields @@ -116,7 +241,7 @@ def put_model( margin = m.pair_margin[ip] else: margin = m.geom_margin[g1] + m.geom_margin[g2] - if margin.any() and not _full_compat: + if margin.any(): t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') for t, g in [(t1, g1), (t2, g2)]: @@ -128,6 +253,8 @@ def put_model( np.any(m.sensor_type == types.SensorType.ACCELEROMETER) | np.any(m.sensor_type == types.SensorType.FORCE) | np.any(m.sensor_type == types.SensorType.TORQUE) + | np.any(m.sensor_type == types.SensorType.FRAMELINACC) + | np.any(m.sensor_type == types.SensorType.FRAMEANGACC) ) eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any( m.eq_type == types.EqType.WELD @@ -148,38 +275,27 @@ def put_model( (m.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = set(enum_field) - set(enum_type) - if missing and not _full_compat: + if missing: raise NotImplementedError( f'{[mj_type(m) for m in missing]} not supported' ) - mj_field_names = { - f.name - for f in types.Model.fields() - if f.metadata.get('restricted_to') != 'mjx' - } + mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'} fields = {f: getattr(m, f) for f in mj_field_names} - - # zero out fields restricted to MuJoCo - if not _full_compat: - for f in types.Model.fields(): - if f.metadata.get('restricted_to') == 'mujoco' and isinstance( - fields[f.name], np.ndarray - ): - fields[f.name] = np.zeros((0,), dtype=fields[f.name].dtype) - - fields['dof_hasfrictionloss'] = fields['dof_frictionloss'] > 0 - fields['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0 - fields['geom_rbound_hfield'] = fields['geom_rbound'] fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) - fields['opt'] = _make_option(m.opt, _full_compat=_full_compat) - fields['stat'] = _make_statistic(m.stat) + fields['opt'] = _put_option(m.opt, types.BackendImpl.JAX) + fields['stat'] = _put_statistic(m.stat) + + fields_jax = {} + fields_jax['dof_hasfrictionloss'] = fields['dof_frictionloss'] > 0 + fields_jax['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0 + fields_jax['geom_rbound_hfield'] = fields['geom_rbound'] # spatial tendon wrap inside - fields['wrap_inside_maxiter'] = 5 - fields['wrap_inside_tolerance'] = 1.0e-4 - fields['wrap_inside_z_init'] = 1.0 - 1.0e-5 - fields['is_wrap_inside'] = np.zeros(0, dtype=bool) + fields_jax['wrap_inside_maxiter'] = 5 + fields_jax['wrap_inside_tolerance'] = 1.0e-4 + fields_jax['wrap_inside_z_init'] = 1.0 - 1.0e-5 + fields_jax['is_wrap_inside'] = np.zeros(0, dtype=bool) if m.nsite: # find sphere or cylinder geoms (if any exist) (wrap_id_geom,) = np.nonzero( @@ -195,28 +311,392 @@ def put_model( side = m.site_pos[side_id] # wrap inside flag - fields['is_wrap_inside'] = np.array( + fields_jax['is_wrap_inside'] = np.array( (np.linalg.norm(side - geom_pos, axis=1) < geom_size) & (side_id >= 0) ) # Pre-compile meshes for MJX collisions. - fields['mesh_convex'] = [None] * m.nmesh - if not _full_compat: - for i in mesh_geomid: - dataid = m.geom_dataid[i] - if fields['mesh_convex'][dataid] is None: - fields['mesh_convex'][dataid] = mesh.convex(m, dataid) # pytype: disable=unsupported-operands - fields['mesh_convex'] = tuple(fields['mesh_convex']) + fields_jax['mesh_convex'] = [None] * m.nmesh + for i in mesh_geomid: + dataid = m.geom_dataid[i] + if fields_jax['mesh_convex'][dataid] is None: + fields_jax['mesh_convex'][dataid] = mesh.convex(m, dataid) # pytype: disable=unsupported-operands + fields_jax['mesh_convex'] = tuple(fields_jax['mesh_convex']) - model = types.Model(**{k: copy.copy(v) for k, v in fields.items()}) + jax_impl = types.ModelJAX(**fields_jax) + model = types.Model( + **{k: copy.copy(v) for k, v in fields.items()}, _impl=jax_impl + ) model = jax.device_put(model, device=device) return _strip_weak_type(model) +def _put_model_c( + m: mujoco.MjModel, + device: Optional[jax.Device] = None, +) -> types.Model: + """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" + mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'} + fields = {f: getattr(m, f) for f in mj_field_names} + fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) + fields['opt'] = _put_option(m.opt, backend_impl=types.BackendImpl.C) + fields['stat'] = _put_statistic(m.stat) + + c_impl_keys = ( + types.ModelC.__annotations__.keys() - types.Model.__annotations__.keys() + ) + c_impl_dict = {k: getattr(m, k) for k in c_impl_keys} + c_impl_obj = types.ModelC(**{k: copy.copy(v) for k, v in c_impl_dict.items()}) + + model = types.Model( + **{k: copy.copy(v) for k, v in fields.items()}, _impl=c_impl_obj + ) + model = jax.device_put(model, device=device) + return _strip_weak_type(model) + + +def put_model( + m: mujoco.MjModel, + device: Optional[jax.Device] = None, + backend_impl: Optional[Union[str, types.BackendImpl]] = None, + _full_compat: bool = False, # pylint: disable=invalid-name +) -> types.Model: + """Puts mujoco.MjModel onto a device, resulting in mjx.Model. + + Args: + m: the model to put onto device + device: which device to use - if unspecified picks the default device + backend_impl: backend implementation to use + _full_compat: put all MjModel fields onto device irrespective of MJX support + This is an experimental feature. Avoid using it for now. + + Returns: + an mjx.Model placed on device + + Raises: + ValueError: if backend_impl is not supported + DeprecationWarning: if _full_compat is True + """ + + if _full_compat: + warnings.warn( + 'mjx.put_model(..., _full_compat=True) is deprecated. Use' + ' mjx.put_model(..., backend_impl=types.BackendImpl.C) instead.', + DeprecationWarning, + stacklevel=2, + ) + backend_impl = types.BackendImpl.C + + backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) + if backend_impl == types.BackendImpl.JAX: + return _put_model_jax(m, device) + elif backend_impl == types.BackendImpl.C: + return _put_model_c(m, device) + elif backend_impl == types.BackendImpl.WARP: + raise NotImplementedError('Warp backend not implemented yet.') + else: + raise ValueError(f'Unsupported backend implementation: {backend_impl}') + + +def _make_data_public_fields(m: types.Model) -> Dict[str, Any]: + """Create public fields for the Data object.""" + float_ = jp.zeros(1, float).dtype + zero_fields = { + 'time': (float_,), + 'qvel': (m.nv, float_), + 'act': (m.na, float_), + 'qacc_warmstart': (m.nv, float_), + 'ctrl': (m.nu, float_), + 'qfrc_applied': (m.nv, float_), + 'xfrc_applied': (m.nbody, 6, float_), + 'mocap_pos': (m.nmocap, 3, float_), + 'mocap_quat': (m.nmocap, 4, float_), + 'qacc': (m.nv, float_), + 'act_dot': (m.na, float_), + 'userdata': (m.nuserdata, float_), + 'sensordata': (m.nsensordata, float_), + 'xpos': (m.nbody, 3, float_), + 'xquat': (m.nbody, 4, float_), + 'xmat': (m.nbody, 3, 3, float_), + 'xipos': (m.nbody, 3, float_), + 'ximat': (m.nbody, 3, 3, float_), + 'xanchor': (m.njnt, 3, float_), + 'xaxis': (m.njnt, 3, float_), + 'geom_xpos': (m.ngeom, 3, float_), + 'geom_xmat': (m.ngeom, 3, 3, float_), + 'site_xpos': (m.nsite, 3, float_), + 'site_xmat': (m.nsite, 3, 3, float_), + 'cam_xpos': (m.ncam, 3, float_), + 'cam_xmat': (m.ncam, 3, 3, float_), + 'subtree_com': (m.nbody, 3, float_), + 'actuator_force': (m.nu, float_), + 'qfrc_bias': (m.nv, float_), + 'qfrc_gravcomp': (m.nv, float_), + 'qfrc_fluid': (m.nv, float_), + 'qfrc_passive': (m.nv, float_), + 'qfrc_actuator': (m.nv, float_), + 'qfrc_smooth': (m.nv, float_), + 'qacc_smooth': (m.nv, float_), + 'qfrc_constraint': (m.nv, float_), + 'qfrc_inverse': (m.nv, float_), + 'cvel': (m.nbody, 6, float_), + } + zero_fields = { + k: np.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items() + } + return zero_fields + + +def _make_data_contact_jax( + condim: np.ndarray, efc_address: np.ndarray +) -> types.Contact: + """Create contact for the Data object.""" + ncon = condim.size + float_ = jp.zeros(1, float).dtype + int_ = jp.zeros(1, int).dtype + contact = types.Contact( + dist=np.zeros((ncon,), dtype=float_), + pos=np.zeros((ncon, 3), dtype=float_), + frame=np.zeros((ncon, 3, 3), dtype=float_), + includemargin=np.zeros((ncon,), dtype=float_), + friction=np.zeros((ncon, 5), dtype=float_), + solref=np.zeros((ncon, mujoco.mjNREF), dtype=float_), + solreffriction=np.zeros((ncon, mujoco.mjNREF), dtype=float_), + solimp=np.zeros((ncon, mujoco.mjNIMP), dtype=float_), + dim=condim, + # let jax pick contact.geom int precision, for interop with + # jax_enable_x64 + geom1=np.full((ncon,), -1, dtype=int_), + geom2=np.full((ncon,), -1, dtype=int_), + geom=np.full((ncon, 2), -1, dtype=int_), + efc_address=efc_address, + ) + return contact + + +def _make_data_jax( + m: types.Model, + device: Optional[jax.Device] = None, +) -> types.Data: + """Allocate and initialize Data for the JAX implementation.""" + dim = collision_driver.make_condim(m) + efc_type = constraint.make_efc_type(m, dim) + ne, nf, nl, nc = constraint.counts(efc_type) + ncon, nefc = dim.size, ne + nf + nl + nc + efc_address = constraint.make_efc_address(m, dim, efc_type) + + float_ = jp.zeros(1, float).dtype + int_ = jp.zeros(1, int).dtype + contact = _make_data_contact_jax(dim, efc_address) + + if m.opt.cone == types.ConeType.ELLIPTIC and np.any(contact.dim == 1): + raise NotImplementedError( + 'condim=1 with ConeType.ELLIPTIC not implemented.' + ) + + zero_impl_fields = { + 'solver_niter': (int_,), + 'cdof': (m.nv, 6, float_), + 'cinert': (m.nbody, 10, float_), + 'ten_wrapadr': (m.ntendon, np.int32), + 'ten_wrapnum': (m.ntendon, np.int32), + 'ten_J': (m.ntendon, m.nv, float_), + 'ten_length': (m.ntendon, float_), + 'wrap_obj': (m.nwrap, 2, np.int32), + 'wrap_xpos': (m.nwrap, 6, float_), + 'actuator_length': (m.nu, float_), + 'actuator_moment': (m.nu, m.nv, float_), + 'crb': (m.nbody, 10, float_), + 'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_), + 'M': (m.nC, float_), + 'qLD': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_), + 'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_), + 'ten_velocity': (m.ntendon, float_), + 'actuator_velocity': (m.nu, float_), + 'cdof_dot': (m.nv, 6, float_), + 'cacc': (m.nbody, 6, float_), + 'cfrc_int': (m.nbody, 6, float_), + 'cfrc_ext': (m.nbody, 6, float_), + 'subtree_linvel': (m.nbody, 3, float_), + 'subtree_angmom': (m.nbody, 3, float_), + 'efc_J': (nefc, m.nv, float_), + 'efc_pos': (nefc, float_), + 'efc_margin': (nefc, float_), + 'efc_frictionloss': (nefc, float_), + 'efc_D': (nefc, float_), + 'efc_aref': (nefc, float_), + 'efc_force': (nefc, float_), + } + zero_impl_fields = { + k: np.zeros(v[:-1], dtype=v[-1]) for k, v in zero_impl_fields.items() + } + impl = types.DataJAX( + ne=ne, + nf=nf, + nl=nl, + nefc=nefc, + ncon=ncon, + contact=contact, + efc_type=efc_type, + **zero_impl_fields, + ) + + d = types.Data( + qpos=jp.array(m.qpos0, dtype=float_), + eq_active=m.eq_active0, + _impl=impl, + **_make_data_public_fields(m), + ) + + if m.nmocap: + # Set mocap_pos/quat = body_pos/quat for mocap bodies as done in C MuJoCo. + body_mask = m.body_mocapid >= 0 + body_pos = m.body_pos[body_mask] + body_quat = m.body_quat[body_mask] + d = d.replace( + mocap_pos=body_pos[m.body_mocapid[body_mask]], + mocap_quat=body_quat[m.body_mocapid[body_mask]], + ) + + d = jax.device_put(d, device=device) + return d + + +def _make_data_c( + m: types.Model, + device: Optional[jax.Device] = None, +) -> types.Data: + """Allocate and initialize Data for the C implementation.""" + # TODO(stunya): The C implementation should not use static dimensions, and + # the backend implementation details should be kept hidden from JAX + # altogether. + dim = collision_driver.make_condim(m) + efc_type = constraint.make_efc_type(m, dim) + efc_address = constraint.make_efc_address(m, dim, efc_type) + ne, nf, nl, nc = constraint.counts(efc_type) + ncon, nefc = dim.size, ne + nf + nl + nc + + float_ = jp.zeros(1, float).dtype + int_ = jp.zeros(1, int).dtype + # TODO(stunya): remove the JAX contact from C data. + contact = _make_data_contact_jax(dim, efc_address) + + nflexvert = m.nflexvert + nflexedge = m.nflexedge + nflexelem = m.nflexelem + nbvh = m.nbvh + nbvhdynamic = m.nbvhdynamic + zero_impl_fields = { + 'solver_niter': (int_,), + 'cdof': (m.nv, 6, float_), + 'cinert': (m.nbody, 10, float_), + 'light_xpos': (m.nlight, 3, float_), + 'light_xdir': (m.nlight, 3, float_), + 'flexvert_xpos': (nflexvert, 3, float_), + 'flexelem_aabb': (nflexelem, 6, float_), + 'flexedge_J_rownnz': (nflexedge, np.int32), + 'flexedge_J_rowadr': (nflexedge, np.int32), + 'flexedge_J_colind': (nflexedge, m.nv, np.int32), + 'flexedge_J': (nflexedge, m.nv, float_), + 'flexedge_length': (nflexedge, float_), + 'ten_J_rownnz': (m.ntendon, np.int32), + 'ten_J_rowadr': (m.ntendon, np.int32), + 'ten_J_colind': (m.ntendon, m.nv, np.int32), + 'ten_J': (m.ntendon, m.nv, float_), + 'ten_length': (m.ntendon, float_), + 'ten_wrapadr': (m.ntendon, np.int32), + 'ten_wrapnum': (m.ntendon, np.int32), + 'wrap_obj': (m.nwrap, 2, np.int32), + 'wrap_xpos': (m.nwrap, 6, float_), + 'actuator_length': (m.nu, float_), + 'moment_rownnz': (m.nu, np.int32), + 'moment_rowadr': (m.nu, np.int32), + 'moment_colind': (m.nJmom, np.int32), + 'actuator_moment': (m.nu, m.nv, float_), + 'bvh_aabb_dyn': (nbvhdynamic, 6, float_), + 'bvh_active': (nbvh, np.uint8), + 'flexedge_velocity': (nflexedge, float_), + 'crb': (m.nbody, 10, float_), + 'qM': (m.nM, float_), + 'M': (m.nC, float_), + 'qLD': (m.nC, float_), + 'qH': (m.nC, float_), + 'qHDiagInv': (m.nv, float_), + 'qLDiagInv': (m.nv, float_), + 'ten_velocity': (m.ntendon, float_), + 'actuator_velocity': (m.nu, float_), + 'plugin_data': (m.nplugin, np.uint64), + 'B_rownnz': (m.nbody, np.int32), + 'B_rowadr': (m.nbody, np.int32), + 'B_colind': (m.nB, np.int32), + 'M_rownnz': (m.nv, np.int32), + 'M_rowadr': (m.nv, np.int32), + 'M_colind': (m.nC, np.int32), + 'mapM2M': (m.nC, np.int32), + 'D_rownnz': (m.nv, np.int32), + 'D_rowadr': (m.nv, np.int32), + 'D_diag': (m.nv, np.int32), + 'D_colind': (m.nD, np.int32), + 'mapM2D': (m.nD, np.int32), + 'mapD2M': (m.nM, np.int32), + 'qDeriv': (m.nD, float_), + 'qLU': (m.nD, float_), + 'qfrc_spring': (m.nv, float_), + 'qfrc_damper': (m.nv, float_), + 'cdof_dot': (m.nv, 6, float_), + 'cacc': (m.nbody, 6, float_), + 'cfrc_int': (m.nbody, 6, float_), + 'cfrc_ext': (m.nbody, 6, float_), + 'subtree_linvel': (m.nbody, 3, float_), + 'subtree_angmom': (m.nbody, 3, float_), + 'efc_J': (nefc, m.nv, float_), + 'efc_pos': (nefc, float_), + 'efc_margin': (nefc, float_), + 'efc_frictionloss': (nefc, float_), + 'efc_D': (nefc, float_), + 'efc_aref': (nefc, float_), + 'efc_force': (nefc, float_), + } + zero_impl_fields = { + k: np.zeros(v[:-1], dtype=v[-1]) for k, v in zero_impl_fields.items() + } + impl = types.DataC( + ne=ne, + nf=nf, + nl=nl, + nefc=nefc, + ncon=ncon, + contact=contact, + efc_type=efc_type, + **zero_impl_fields, + ) + + d = types.Data( + qpos=jp.array(m.qpos0, dtype=float_), + eq_active=m.eq_active0, + _impl=impl, + **_make_data_public_fields(m), + ) + + if m.nmocap: + # Set mocap_pos/quat = body_pos/quat for mocap bodies as done in C MuJoCo. + body_mask = m.body_mocapid >= 0 + body_pos = m.body_pos[body_mask] + body_quat = m.body_quat[body_mask] + d = d.replace( + mocap_pos=body_pos[m.body_mocapid[body_mask]], + mocap_quat=body_quat[m.body_mocapid[body_mask]], + ) + + d = jax.device_put(d, device=device) + return d + + def make_data( m: Union[types.Model, mujoco.MjModel], - device=None, + device: Optional[jax.Device] = None, + backend_impl: Optional[Union[str, types.BackendImpl]] = None, _full_compat: bool = False, # pylint: disable=invalid-name ) -> types.Data: """Allocate and initialize Data. @@ -224,313 +704,48 @@ def make_data( Args: m: the model to use device: which device to use - if unspecified picks the default device - _full_compat: create all MjData fields on device irrespective of MJX support - This is an experimental feature. Avoid using it for now. If using this - flag, also use _full_compat for put_model. + backend_impl: backend implementation to use + _full_compat: put all fields onto device irrespective of MJX support This is + an experimental feature. Avoid using it for now. If using this flag, also + use _full_compat for put_model. Returns: an initialized mjx.Data placed on device - """ - dim = collision_driver.make_condim(m) - efc_type = constraint.make_efc_type(m, dim) - efc_address = constraint.make_efc_address(m, dim, efc_type) - ne, nf, nl, nc = constraint.counts(efc_type) - ncon, nefc = dim.size, ne + nf + nl + nc - with jax.default_device(device): - contact = types.Contact( - dist=jp.zeros((ncon,), dtype=float), - pos=jp.zeros((ncon, 3), dtype=float), - frame=jp.zeros((ncon, 3, 3), dtype=float), - includemargin=jp.zeros((ncon,), dtype=float), - friction=jp.zeros((ncon, 5), dtype=float), - solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float), - solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float), - solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float), - dim=dim, - # let jax pick contact.geom int precision, for interop with - # jax_enable_x64 - geom1=jp.full((ncon,), -1, dtype=int), - geom2=jp.full((ncon,), -1, dtype=int), - geom=jp.full((ncon, 2), -1, dtype=int), - efc_address=efc_address, + Raises: + ValueError: if the model's backend_impl does not match the make_data + backend_impl + NotImplementedError: if the backend_impl is not implemented yet + DeprecationWarning: if _full_compat is used + """ + if _full_compat: + warnings.warn( + 'mjx.make_data(..., _full_compat=True) is deprecated. Use' + ' mjx.make_data(..., backend_impl=types.BackendImpl.C) instead.', + DeprecationWarning, + stacklevel=2, + ) + backend_impl = types.BackendImpl.C + + backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) + + if isinstance(m, types.Model) and m.backend_impl != backend_impl: + raise ValueError( + f'Model backend_impl {m.backend_impl} does not match make_data ' + f'backend_impl {backend_impl}.' ) - if m.opt.cone == types.ConeType.ELLIPTIC and np.any(contact.dim == 1): - raise NotImplementedError( - 'condim=1 with ConeType.ELLIPTIC not implemented.' - ) + if backend_impl == types.BackendImpl.JAX: + return _make_data_jax(m, device) + elif backend_impl == types.BackendImpl.C: + return _make_data_c(m, device) - zero_fields = { - 'solver_niter': (int,), - 'time': (float,), - 'qvel': (m.nv, float), - 'act': (m.na, float), - 'qacc_warmstart': (m.nv, float), - 'ctrl': (m.nu, float), - 'qfrc_applied': (m.nv, float), - 'xfrc_applied': (m.nbody, 6, float), - 'mocap_pos': (m.nmocap, 3, float), - 'mocap_quat': (m.nmocap, 4, float), - 'qacc': (m.nv, float), - 'act_dot': (m.na, float), - 'userdata': (m.nuserdata, float), - 'sensordata': (m.nsensordata, float), - 'xpos': (m.nbody, 3, float), - 'xquat': (m.nbody, 4, float), - 'xmat': (m.nbody, 3, 3, float), - 'xipos': (m.nbody, 3, float), - 'ximat': (m.nbody, 3, 3, float), - 'xanchor': (m.njnt, 3, float), - 'xaxis': (m.njnt, 3, float), - 'geom_xpos': (m.ngeom, 3, float), - 'geom_xmat': (m.ngeom, 3, 3, float), - 'site_xpos': (m.nsite, 3, float), - 'site_xmat': (m.nsite, 3, 3, float), - 'cam_xpos': (m.ncam, 3, float), - 'cam_xmat': (m.ncam, 3, 3, float), - 'light_xpos': (m.nlight, 3, float), - 'light_xdir': (m.nlight, 3, float), - 'subtree_com': (m.nbody, 3, float), - 'cdof': (m.nv, 6, float), - 'cinert': (m.nbody, 10, float), - 'flexvert_xpos': (m.nflexvert, 3, float), - 'flexelem_aabb': (m.nflexelem, 6, float), - 'flexedge_J_rownnz': (m.nflexedge, jp.int32), - 'flexedge_J_rowadr': (m.nflexedge, jp.int32), - 'flexedge_J_colind': (m.nflexedge, m.nv, jp.int32), - 'flexedge_J': (m.nflexedge, m.nv, float), - 'flexedge_length': (m.nflexedge, float), - 'ten_wrapadr': (m.ntendon, jp.int32), - 'ten_wrapnum': (m.ntendon, jp.int32), - 'ten_J_rownnz': (m.ntendon, jp.int32), - 'ten_J_rowadr': (m.ntendon, jp.int32), - 'ten_J_colind': (m.ntendon, m.nv, jp.int32), - 'ten_J': (m.ntendon, m.nv, float), - 'ten_length': (m.ntendon, float), - 'wrap_obj': (m.nwrap, 2, jp.int32), - 'wrap_xpos': (m.nwrap, 6, float), - 'actuator_length': (m.nu, float), - 'moment_rownnz': (m.nu, jp.int32), - 'moment_rowadr': (m.nu, jp.int32), - 'moment_colind': (m.nJmom, jp.int32), - 'actuator_moment': (m.nu, m.nv, float), - 'crb': (m.nbody, 10, float), - 'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), - 'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), - 'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float), - 'bvh_aabb_dyn': (m.nbvhdynamic, 6, float), - 'bvh_active': (m.nbvh, jp.uint8), - 'flexedge_velocity': (m.nflexedge, float), - 'ten_velocity': (m.ntendon, float), - 'actuator_velocity': (m.nu, float), - 'cvel': (m.nbody, 6, float), - 'cdof_dot': (m.nv, 6, float), - 'qfrc_bias': (m.nv, float), - 'qfrc_spring': (m.nv, float), - 'qfrc_damper': (m.nv, float), - 'qfrc_gravcomp': (m.nv, float), - 'qfrc_fluid': (m.nv, float), - 'qfrc_passive': (m.nv, float), - 'subtree_linvel': (m.nbody, 3, float), - 'subtree_angmom': (m.nbody, 3, float), - 'qH': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), - 'qHDiagInv': (m.nv, float), - 'B_rownnz': (m.nbody, jp.int32), - 'B_rowadr': (m.nbody, jp.int32), - 'B_colind': (m.nB, jp.int32), - 'M_rownnz': (m.nv, jp.int32), - 'M_rowadr': (m.nv, jp.int32), - 'M_colind': (m.nM, jp.int32), - 'mapM2M': (m.nM, jp.int32), - 'C_rownnz': (m.nv, jp.int32), - 'C_rowadr': (m.nv, jp.int32), - 'C_colind': (m.nC, jp.int32), - 'mapM2C': (m.nC, jp.int32), - 'D_rownnz': (m.nv, jp.int32), - 'D_rowadr': (m.nv, jp.int32), - 'D_diag': (m.nv, jp.int32), - 'D_colind': (m.nD, jp.int32), - 'mapM2D': (m.nD, jp.int32), - 'mapD2M': (m.nM, jp.int32), - 'qDeriv': (m.nD, float), - 'qLU': (m.nD, float), - 'actuator_force': (m.nu, float), - 'qfrc_actuator': (m.nv, float), - 'qfrc_smooth': (m.nv, float), - 'qacc_smooth': (m.nv, float), - 'qfrc_constraint': (m.nv, float), - 'qfrc_inverse': (m.nv, float), - 'cacc': (m.nbody, 6, float), - 'cfrc_int': (m.nbody, 6, float), - 'cfrc_ext': (m.nbody, 6, float), - 'efc_J': (nefc, m.nv, float), - 'efc_pos': (nefc, float), - 'efc_margin': (nefc, float), - 'efc_frictionloss': (nefc, float), - 'efc_D': (nefc, float), - 'efc_aref': (nefc, float), - 'efc_force': (nefc, float), - '_qM_sparse': (m.nM, float), - '_qLD_sparse': (m.nM, float), - '_qLDiagInv_sparse': (m.nv, float), - } - - if not _full_compat: - for f in types.Data.fields(): - if f.metadata.get('restricted_to') in ('mujoco', 'mjx'): - zero_fields[f.name] = (0, zero_fields[f.name][-1]) - - zero_fields = { - k: jp.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items() - } - - d = types.Data( - ne=ne, - nf=nf, - nl=nl, - nefc=nefc, - ncon=ncon, - qpos=jp.array(m.qpos0), - contact=contact, - efc_type=efc_type, - eq_active=m.eq_active0, - **zero_fields, + raise NotImplementedError( + f'make_data for backend_impl "{backend_impl}" not implemented yet.' ) - return d - -def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact): - """Converts mjx.Contact to mujoco._structs._MjContactList.""" - con_id = np.nonzero(cx.dist <= 0)[0] - for field in types.Contact.fields(): - value = getattr(cx, field.name)[con_id] - if field.name == 'frame': - value = value.reshape((-1, 9)) - getattr(c, field.name)[:] = value - - -def get_data( - m: mujoco.MjModel, d: types.Data -) -> Union[mujoco.MjData, List[mujoco.MjData]]: - """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" - batched = len(d.qpos.shape) > 1 - batch_size = d.qpos.shape[0] if batched else 1 - - if batched: - result = [mujoco.MjData(m) for _ in range(batch_size)] - else: - result = mujoco.MjData(m) - - get_data_into(result, m, d) - - return result - - -def get_data_into( - result: Union[mujoco.MjData, List[mujoco.MjData]], - m: mujoco.MjModel, - d: types.Data, -): - """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" - batched = isinstance(result, list) - if batched and len(d.qpos.shape) < 2: - raise ValueError('dst is a list, but d is not batched.') - if not batched and len(d.qpos.shape) >= 2: - raise ValueError('dst is a an MjData, but d is batched.') - - d = jax.device_get(d) - - batch_size = d.qpos.shape[0] if batched else 1 - - dof_i, dof_j = [], [] - for i in range(m.nv): - j = i - while j > -1: - dof_i.append(i) - dof_j.append(j) - j = m.dof_parentid[j] - - for i in range(batch_size): - d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d - result_i = result[i] if batched else result - ncon = (d_i.contact.dist <= 0).sum() - efc_active = (d_i.efc_J != 0).any(axis=1) - nefc = int(efc_active.sum()) - result_i.nJ = nefc * m.nv - if ncon != result_i.ncon or nefc != result_i.nefc: - mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access - result_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc) - result_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv) - result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) - - for field in types.Data.fields(): - restricted_to = field.metadata.get('restricted_to') - if restricted_to == 'mjx': - continue - - if field.name == 'contact': - _get_contact(result_i.contact, d_i.contact) - # efc_address must be updated because rows were deleted above: - efc_map = np.cumsum(efc_active) - 1 - result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] - continue - - # MuJoCo actuator_moment is sparse, MJX uses a dense representation. - if field.name == 'actuator_moment': - moment_rownnz = np.zeros(m.nu, dtype=np.int32) - moment_rowadr = np.zeros(m.nu, dtype=np.int32) - moment_colind = np.zeros(m.nJmom, dtype=np.int32) - actuator_moment = np.zeros(m.nJmom) - if m.nu: - mujoco.mju_dense2sparse( - actuator_moment, - d_i.actuator_moment, - moment_rownnz, - moment_rowadr, - moment_colind, - ) - result_i.moment_rownnz[:] = moment_rownnz - result_i.moment_rowadr[:] = moment_rowadr - result_i.moment_colind[:] = moment_colind - result_i.actuator_moment[:] = actuator_moment - continue - - value = getattr(d_i, field.name) - - if field.name in ('nefc', 'ncon'): - value = {'nefc': nefc, 'ncon': ncon}[field.name] - elif field.name.endswith('xmat') or field.name == 'ximat': - value = value.reshape((-1, 9)) - elif field.name.startswith('efc_'): - value = value[efc_active] - if field.name == 'efc_J': - value = value.reshape(-1) - elif field.name == 'qM' and not support.is_sparse(m): - value = value[dof_i, dof_j] - elif field.name == 'qLD' and not support.is_sparse(m): - # TODO(erikfrey): provide correct qLDs - value = np.zeros(m.nM) - elif field.name == 'qLDiagInv' and not support.is_sparse(m): - value = np.ones(m.nv) - - if isinstance(value, np.ndarray) and value.shape: - if restricted_to in ('mujoco', 'mjx'): - continue # don't copy fields that are mujoco-only or MJX-only - else: - result_field = getattr(result_i, field.name) - if result_field.shape != value.shape: - raise ValueError( - f'Input field {field.name} has shape {value.shape}, but output' - f' has shape {result_field.shape}' - ) - result_field[:] = value - else: - setattr(result_i, field.name, value) - - -def _make_contact( +def _put_contact( c: mujoco._structs._MjContactList, dim: np.ndarray, efc_address: np.ndarray, @@ -565,22 +780,24 @@ def _make_contact( return types.Contact(**fields), contact_map -def put_data( - m: mujoco.MjModel, d: mujoco.MjData, device=None, _full_compat: bool = False # pylint: disable=invalid-name +def _put_data_public_fields(d: mujoco.MjData) -> Dict[str, Any]: + """Returns public fields from mujoco.MjData in a dictionary.""" + fields = { + f.name: getattr(d, f.name) + for f in types.Data.fields() + if f.name != '_impl' + } + # MJX uses square matrices for these fields: + for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): + fields[fname] = fields[fname].reshape((-1, 3, 3)) + + return fields + + +def _put_data_jax( + m: mujoco.MjModel, d: mujoco.MjData, device: Optional[jax.Device] = None ) -> types.Data: - """Puts mujoco.MjData onto a device, resulting in mjx.Data. - - Args: - m: the model to use - d: the data to put on device - device: which device to use - if unspecified picks the default device - _full_compat: put all MjModel fields onto device irrespective of MJX support - This is an experimental feature. Avoid using it for now. If using this - flag, also use _full_compat for put_model. - - Returns: - an mjx.Data placed on device - """ + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" dim = collision_driver.make_condim(m) efc_type = constraint.make_efc_type(m, dim) efc_address = constraint.make_efc_address(m, dim, efc_type) @@ -597,18 +814,16 @@ def put_data( if d_val > val: raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') - fields = { + fields = _put_data_public_fields(d) + + # Implementation specific fields. + impl_fields = { f.name: getattr(d, f.name) - for f in types.Data.fields() - if f.metadata.get('restricted_to') != 'mjx' + for f in types.DataJAX.fields() + if hasattr(d, f.name) } - - # MJX prefers square matrices for these fields: - for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): - fields[fname] = fields[fname].reshape((-1, 3, 3)) - # MJX does not support islanding, so only transfer the first solver_niter - fields['solver_niter'] = fields['solver_niter'][0] + impl_fields['solver_niter'] = impl_fields['solver_niter'][0] # convert sparse representation of actuator_moment to dense matrix moment = np.zeros((m.nu, m.nv)) @@ -619,9 +834,9 @@ def put_data( d.moment_rowadr, d.moment_colind, ) - fields['actuator_moment'] = moment + impl_fields['actuator_moment'] = moment - contact, contact_map = _make_contact(d.contact, dim, efc_address) + contact, contact_map = _put_contact(d.contact, dim, efc_address) # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints. # efc_J is also optionally column-sparse (typically for large nv). MJX is @@ -631,14 +846,16 @@ def put_data( efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv)) mujoco.mju_sparse2dense( efc_j, - fields['efc_J'], + impl_fields['efc_J'], d.efc_J_rownnz, d.efc_J_rowadr, d.efc_J_colind, ) - fields['efc_J'] = efc_j + impl_fields['efc_J'] = efc_j else: - fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv)) + impl_fields['efc_J'] = impl_fields['efc_J'].reshape( + (-1 if m.nv else 0, m.nv) + ) # move efc rows to their correct offsets for fname in ( @@ -655,7 +872,9 @@ def put_data( value_beg = sum([ne, nf][:i]) d_beg = sum([d.ne, d.nf][:i]) size = [d.ne, d.nf, d.nl][i] - value[value_beg : value_beg + size] = fields[fname][d_beg : d_beg + size] + value[value_beg : value_beg + size] = impl_fields[fname][ + d_beg : d_beg + size + ] # for nc, we may reorder contacts so they match MJX order: group by dim for id_to, id_from in enumerate(contact_map): @@ -667,43 +886,367 @@ def put_data( efc_i, efc_o = d.contact.efc_address[id_from], efc_address[id_to] if efc_i == -1: continue - value[efc_o : efc_o + num_rows] = fields[fname][efc_i : efc_i + num_rows] + value[efc_o : efc_o + num_rows] = impl_fields[fname][ + efc_i : efc_i + num_rows + ] - fields[fname] = value - - if _full_compat: - # full compatibility mode, we store sparse qM regardless of jacobian setting - fields['_qM_sparse'] = fields['qM'] - fields['_qLD_sparse'] = fields['qLD'] - fields['_qLDiagInv_sparse'] = fields['qLDiagInv'] - else: - fields['_qM_sparse'] = jp.zeros(0, dtype=float) - fields['_qLD_sparse'] = jp.zeros(0, dtype=float) - fields['_qLDiagInv_sparse'] = jp.zeros(0, dtype=float) - # otherwise clear out unused arrays - for f in types.Data.fields(): - if f.metadata.get('restricted_to') == 'mujoco' and isinstance( - fields[f.name], np.ndarray - ): - fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype) + impl_fields[fname] = value # convert qM and qLD if jacobian is dense if not support.is_sparse(m): - fields['qM'] = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, fields['qM'], d.qM) + impl_fields['qM'] = np.zeros((m.nv, m.nv)) + mujoco.mj_fullM(m, impl_fields['qM'], d.qM) # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing try: - fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM']) + impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['qM']) except scipy.linalg.LinAlgError: # this happens when qM is empty or unstable simulation - fields['qLD'] = np.zeros((m.nv, m.nv)) - fields['qLDiagInv'] = np.zeros(0) + impl_fields['qLD'] = np.zeros((m.nv, m.nv)) + impl_fields['qLDiagInv'] = np.zeros(0) - fields['contact'] = contact - fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type) + impl_fields['contact'] = contact + impl_fields.update( + ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type + ) # copy because device_put is async: - data = types.Data(**{k: copy.copy(v) for k, v in fields.items()}) + data_jax = types.DataJAX(**{k: copy.copy(v) for k, v in impl_fields.items()}) + data = types.Data( + **{k: copy.copy(v) for k, v in fields.items()}, _impl=data_jax + ) data = jax.device_put(data, device=device) return _strip_weak_type(data) + + +def _put_data_c( + m: mujoco.MjModel, d: mujoco.MjData, device: Optional[jax.Device] = None +) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" + # TODO(stunya): ncon, nefc should potentially be jax.Array, and contact/efc + # should not be materialized in JAX. + dim = collision_driver.make_condim(m) + efc_type = constraint.make_efc_type(m, dim) + efc_address = constraint.make_efc_address(m, dim, efc_type) + ne, nf, nl, nc = constraint.counts(efc_type) + ncon, nefc = dim.size, ne + nf + nl + nc + + # TODO(stunya): remove this check. + for d_val, val, name in ( + (d.ncon, ncon, 'ncon'), + (d.ne, ne, 'ne'), + (d.nf, nf, 'nf'), + (d.nl, nl, 'nl'), + (d.nefc, nefc, 'nefc'), + ): + if d_val > val: + raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') + + fields = _put_data_public_fields(d) + + # Implementation specific fields. + impl_fields = { + f.name: getattr(d, f.name) + for f in types.DataC.fields() + if hasattr(d, f.name) + } + + # TODO(stunya): support islanding via C backend impl. + impl_fields['solver_niter'] = impl_fields['solver_niter'][0] + + # TODO(btaba): remove dense actuator moment. + # convert sparse representation of actuator_moment to dense matrix + moment = np.zeros((m.nu, m.nv)) + mujoco.mju_sparse2dense( + moment, + d.actuator_moment, + d.moment_rownnz, + d.moment_rowadr, + d.moment_colind, + ) + impl_fields['actuator_moment'] = moment + + # TODO(btaba): remove reliance on JAX _put_contact. + contact, contact_map = _put_contact(d.contact, dim, efc_address) + + # TODO(btaba): remove reliance on dense efc_J. + if mujoco.mj_isSparse(m): + efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv)) + mujoco.mju_sparse2dense( + efc_j, + impl_fields['efc_J'], + d.efc_J_rownnz, + d.efc_J_rowadr, + d.efc_J_colind, + ) + impl_fields['efc_J'] = efc_j + else: + impl_fields['efc_J'] = impl_fields['efc_J'].reshape( + (-1 if m.nv else 0, m.nv) + ) + + # move efc rows to their correct offsets + for fname in ( + 'efc_J', + 'efc_pos', + 'efc_margin', + 'efc_frictionloss', + 'efc_D', + 'efc_aref', + 'efc_force', + ): + value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc) + for i in range(3): + value_beg = sum([ne, nf][:i]) + d_beg = sum([d.ne, d.nf][:i]) + size = [d.ne, d.nf, d.nl][i] + value[value_beg : value_beg + size] = impl_fields[fname][ + d_beg : d_beg + size + ] + + # for nc, we may reorder contacts so they match MJX order: group by dim + for id_to, id_from in enumerate(contact_map): + if id_from == -1: + continue + num_rows = dim[id_to] + if num_rows > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: + num_rows = (num_rows - 1) * 2 + efc_i, efc_o = d.contact.efc_address[id_from], efc_address[id_to] + if efc_i == -1: + continue + value[efc_o : efc_o + num_rows] = impl_fields[fname][ + efc_i : efc_i + num_rows + ] + + impl_fields[fname] = value + + impl_fields['contact'] = contact + impl_fields.update( + ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type + ) + + # copy because device_put is async: + data_jax = types.DataC(**{k: copy.copy(v) for k, v in impl_fields.items()}) + data = types.Data( + **{k: copy.copy(v) for k, v in fields.items()}, _impl=data_jax + ) + + data = jax.device_put(data, device=device) + return _strip_weak_type(data) + + +def put_data( + m: mujoco.MjModel, + d: mujoco.MjData, + device: Optional[jax.Device] = None, + backend_impl: Optional[Union[str, types.BackendImpl]] = None, + _full_compat: bool = False, # pylint: disable=invalid-name +) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data. + + Args: + m: the model to use + d: the data to put on device + device: which device to use - if unspecified picks the default device + backend_impl: backend implementation to use + _full_compat: put all MjModel fields onto device irrespective of MJX support + This is an experimental feature. Avoid using it for now. If using this + flag, also use _full_compat for put_model. + + Returns: + an mjx.Data placed on device + """ + if _full_compat: + warnings.warn( + 'mjx.put_data(..., _full_compat=True) is deprecated. Use' + ' mjx.put_data(..., backend_impl=types.BackendImpl.C) instead.', + DeprecationWarning, + stacklevel=2, + ) + backend_impl = types.BackendImpl.C + + backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) + if backend_impl == types.BackendImpl.JAX: + return _put_data_jax(m, d, device) + elif backend_impl == types.BackendImpl.C: + return _put_data_c(m, d, device) + + raise NotImplementedError( + f'put_data for backend_impl "{backend_impl}" not implemented yet.' + ) + + +def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact): + """Converts mjx.Contact to mujoco._structs._MjContactList.""" + con_id = np.nonzero(cx.dist <= 0)[0] + for field in types.Contact.fields(): + value = getattr(cx, field.name)[con_id] + if field.name == 'frame': + value = value.reshape((-1, 9)) + getattr(c, field.name)[:] = value + + +def _get_data_into( + result: Union[mujoco.MjData, List[mujoco.MjData]], + m: mujoco.MjModel, + d: types.Data, +): + """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" + batched = isinstance(result, list) + d = jax.device_get(d) + batch_size = d.qpos.shape[0] if batched else 1 + + dof_i, dof_j = [], [] + if d.backend_impl == types.BackendImpl.JAX: + for i in range(m.nv): + j = i + while j > -1: + dof_i.append(i) + dof_j.append(j) + j = m.dof_parentid[j] + + for i in range(batch_size): + d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d + result_i = result[i] if batched else result + ncon = (d_i._impl.contact.dist <= 0).sum() + efc_active = (d_i._impl.efc_J != 0).any(axis=1) + nefc = int(efc_active.sum()) + nj = (d_i._impl.efc_J != 0).sum() if support.is_sparse(m) else nefc * m.nv + + if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: + mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access + + if d.backend_impl == types.BackendImpl.JAX: + all_fields = types.Data.fields() + types.DataJAX.fields() + elif d.backend_impl == types.BackendImpl.C: + all_fields = types.Data.fields() + types.DataC.fields() + else: + raise NotImplementedError( + f'get_data_into for backend_impl "{d.backend_impl}" not implemented' + ' yet.' + ) + + for field in all_fields: + if field.name not in mujoco.MjData.__dict__.keys(): + continue + + if field.name == 'contact': + _get_contact(result_i.contact, d_i._impl.contact) + # efc_address must be updated because rows were deleted above: + efc_map = np.cumsum(efc_active) - 1 + result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] + continue + + # MuJoCo actuator_moment is sparse, MJX uses a dense representation. + if field.name == 'actuator_moment': + moment_rownnz = np.zeros(m.nu, dtype=np.int32) + moment_rowadr = np.zeros(m.nu, dtype=np.int32) + moment_colind = np.zeros(m.nJmom, dtype=np.int32) + actuator_moment = np.zeros(m.nJmom) + if m.nu: + mujoco.mju_dense2sparse( + actuator_moment, + d_i._impl.actuator_moment, + moment_rownnz, + moment_rowadr, + moment_colind, + ) + result_i.moment_rownnz[:] = moment_rownnz + result_i.moment_rowadr[:] = moment_rowadr + result_i.moment_colind[:] = moment_colind + result_i.actuator_moment[:] = actuator_moment + continue + + if hasattr(d_i._impl, field.name): + value = getattr(d_i._impl, field.name) + else: + value = getattr(d_i, field.name) + + if field.name in ('nefc', 'ncon'): + value = {'nefc': nefc, 'ncon': ncon}[field.name] + elif field.name.endswith('xmat') or field.name == 'ximat': + value = value.reshape((-1, 9)) + elif field.name == 'efc_J': + value = value[efc_active] + if support.is_sparse(m): + efc_J_rownnz = np.zeros(nefc, dtype=np.int32) + efc_J_rowadr = np.zeros(nefc, dtype=np.int32) + efc_J_colind = np.zeros(nj, dtype=np.int32) + efc_J = np.zeros(nj) + mujoco.mju_dense2sparse( + efc_J, + value, + efc_J_rownnz, + efc_J_rowadr, + efc_J_colind, + ) + result_i.efc_J_rownnz[:] = efc_J_rownnz + result_i.efc_J_rowadr[:] = efc_J_rowadr + result_i.efc_J_colind[:] = efc_J_colind + value = efc_J + else: + value = value.reshape(-1) + elif field.name.startswith('efc_'): + value = value[efc_active] + if d.backend_impl == types.BackendImpl.JAX: + if field.name == 'qM' and not support.is_sparse(m): + value = value[dof_i, dof_j] + elif field.name == 'qLD' and not support.is_sparse(m): + value = np.zeros(m.nC) + elif field.name == 'qLDiagInv' and not support.is_sparse(m): + value = np.ones(m.nv) + + if isinstance(value, np.ndarray) and value.shape: + result_field = getattr(result_i, field.name) + if result_field.shape != value.shape: + raise ValueError( + f'Input field {field.name} has shape {value.shape}, but output' + f' has shape {result_field.shape}' + ) + result_field[:] = value + else: + setattr(result_i, field.name, value) + + # recalculate qLD and qLDiagInv as MJX and MuJoCo have different + # representations of the Cholesky decomposition. + mujoco.mj_factorM(m, result_i) + + +def get_data_into( + result: Union[mujoco.MjData, List[mujoco.MjData]], + m: mujoco.MjModel, + d: types.Data, +): + """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" + is_batched = isinstance(result, list) + if is_batched and len(d.qpos.shape) < 2: + raise ValueError('destination is a list, but d is not batched.') + if not is_batched and len(d.qpos.shape) >= 2: + raise ValueError('destination is a an MjData, but d is batched.') + + d = jax.device_get(d) + + if d.backend_impl in (types.BackendImpl.JAX, types.BackendImpl.C): + # TODO(stunya): Split out _get_data_into once codepaths diverge enough. + return _get_data_into(result, m, d) + + raise NotImplementedError( + f'get_data_into for backend_impl "{d.backend_impl}" not implemented yet.' + ) + + +def get_data( + m: mujoco.MjModel, d: types.Data +) -> Union[mujoco.MjData, List[mujoco.MjData]]: + """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" + batched = len(d.qpos.shape) > 1 + batch_size = d.qpos.shape[0] if batched else 1 + + if batched: + result = [mujoco.MjData(m) for _ in range(batch_size)] + else: + result = mujoco.MjData(m) + + get_data_into(result, m, d) + + return result diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 494b2753..33565999 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -14,14 +14,19 @@ # ============================================================================== """Tests for io functions.""" +import os +from unittest import mock 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 io as mjx_io from mujoco.mjx._src import test_util + # pylint: disable=g-importing-member +from mujoco.mjx._src.types import BackendImpl from mujoco.mjx._src.types import ConeType # pylint: enable=g-importing-member import numpy as np @@ -91,14 +96,29 @@ _MULTIPLE_CONSTRAINTS = """ """ +_SIMPLE_BODY = """ + + + + + + + + + +""" + class ModelIOTest(parameterized.TestCase): """IO tests for mjx.Model.""" - @parameterized.parameters(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS) - def test_put_model(self, xml): + @parameterized.product( + xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS), + backend_impl=('jax', 'c'), + ) + def test_put_model(self, xml, backend_impl): m = mujoco.MjModel.from_xml_string(xml) - mx = mjx.put_model(m) + mx = mjx.put_model(m, backend_impl=backend_impl) def assert_not_weak_type(x): if isinstance(x, jax.Array): @@ -120,8 +140,14 @@ class ModelIOTest(parameterized.TestCase): self.assertEqual(mx.nM, m.nM) self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep) - # fields restricted to MuJoCo should not be populated - self.assertEqual(mx.bvh_aabb.shape, (0,)) + if backend_impl == 'jax': + # fields restricted to MuJoCo should not be populated + self.assertFalse(hasattr(mx, 'bvh_aabb')) + elif backend_impl == 'c': + # Options specific to C are populated. + self.assertEqual(mx.opt.apirate, m.opt.apirate) + # Fields private to C backend impl are populated. + self.assertTrue(hasattr(mx._impl, 'bvh_aabb')) np.testing.assert_allclose(mx.body_parentid, m.body_parentid) np.testing.assert_allclose(mx.geom_type, m.geom_type) @@ -150,7 +176,8 @@ class ModelIOTest(parameterized.TestCase): m = mjx.put_model( mujoco.MjModel.from_xml_string( '' - ) + ), + backend_impl='jax', ) self.assertTrue(m.opt.has_fluid_params) @@ -174,7 +201,8 @@ class ModelIOTest(parameterized.TestCase): def test_margin_gap_mesh_not_implemented(self): with self.assertRaises(NotImplementedError): - mjx.put_model(mujoco.MjModel.from_xml_string(""" + mjx.put_model( + mujoco.MjModel.from_xml_string(""" @@ -189,27 +217,32 @@ class ModelIOTest(parameterized.TestCase): - """)) + """), + backend_impl='jax', + ) def test_implicitfast_fluid_not_implemented(self): with self.assertRaises(NotImplementedError): - mjx.put_model(mujoco.MjModel.from_xml_string(""" + mjx.put_model( + mujoco.MjModel.from_xml_string(""" """)) + """), + backend_impl='jax', + ) def test_wrap_inside(self): m = test_util.load_test_file('tendon/wrap_sidesite.xml') - mx0 = mjx.put_model(m) + mx0 = mjx.put_model(m, backend_impl='jax') np.testing.assert_equal( - mx0.is_wrap_inside, + mx0._impl.is_wrap_inside, np.array([1, 0, 1, 0, 1, 1, 0]), ) m.site_pos[2] = m.site_pos[1] - mx1 = mjx.put_model(m) + mx1 = mjx.put_model(m, backend_impl='jax') np.testing.assert_equal( - mx1.is_wrap_inside, + mx1._impl.is_wrap_inside, np.array([0, 0, 1, 0, 1, 0, 0]), ) @@ -217,11 +250,11 @@ class ModelIOTest(parameterized.TestCase): class DataIOTest(parameterized.TestCase): """IO tests for mjx.Data.""" - def test_make_data(self): + @parameterized.parameters('jax', 'c') + def test_make_data(self, backend_impl: str): """Test that make_data returns the correct shapes.""" - m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) - d = mjx.make_data(m) + d = mjx.make_data(m, backend_impl=backend_impl) nq = 22 nbody = 5 @@ -230,7 +263,7 @@ class DataIOTest(parameterized.TestCase): nv = 19 nefc = 185 - self.assertEqual(d.nefc, nefc) + self.assertEqual(d._impl.nefc, nefc) self.assertEqual(d.qpos.shape, (nq,)) self.assertEqual(d.qvel.shape, (nv,)) self.assertEqual(d.act.shape, (0,)) @@ -251,57 +284,68 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d.geom_xpos.shape, (6, 3)) self.assertEqual(d.geom_xmat.shape, (6, 3, 3)) self.assertEqual(d.subtree_com.shape, (nbody, 3)) - self.assertEqual(d.cdof.shape, (nv, 6)) - self.assertEqual(d.cinert.shape, (nbody, 10)) - self.assertEqual(d.crb.shape, (nbody, 10)) - self.assertEqual(d.actuator_length.shape, (1,)) - self.assertEqual(d.actuator_moment.shape, (1, nv)) - self.assertEqual(d.qM.shape, (nv, nv)) - self.assertEqual(d.qLD.shape, (nv, nv)) - self.assertEqual(d.qLDiagInv.shape, (0,)) - self.assertEqual(d.contact.dist.shape, (ncon,)) - self.assertEqual(d.contact.pos.shape, (ncon, 3)) - self.assertEqual(d.contact.frame.shape, (ncon, 3, 3)) - self.assertEqual(d.contact.solref.shape, (ncon, 2)) - self.assertEqual(d.contact.solimp.shape, (ncon, 5)) - self.assertEqual(d.contact.geom1.shape, (ncon,)) - self.assertEqual(d.contact.geom2.shape, (ncon,)) - self.assertEqual(d.efc_J.shape, (nefc, nv)) - self.assertEqual(d.efc_frictionloss.shape, (nefc,)) - self.assertEqual(d.efc_D.shape, (nefc,)) - self.assertEqual(d.actuator_velocity.shape, (1,)) + self.assertEqual(d._impl.cdof.shape, (nv, 6)) + self.assertEqual(d._impl.cinert.shape, (nbody, 10)) + self.assertEqual(d._impl.crb.shape, (nbody, 10)) + self.assertEqual(d._impl.actuator_length.shape, (1,)) + self.assertEqual(d._impl.actuator_moment.shape, (1, nv)) + self.assertEqual(d._impl.contact.dist.shape, (ncon,)) + self.assertEqual(d._impl.contact.pos.shape, (ncon, 3)) + self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3)) + self.assertEqual(d._impl.contact.solref.shape, (ncon, 2)) + self.assertEqual(d._impl.contact.solimp.shape, (ncon, 5)) + self.assertEqual(d._impl.contact.geom1.shape, (ncon,)) + self.assertEqual(d._impl.contact.geom2.shape, (ncon,)) + self.assertEqual(d._impl.efc_J.shape, (nefc, nv)) + self.assertEqual(d._impl.efc_frictionloss.shape, (nefc,)) + self.assertEqual(d._impl.efc_D.shape, (nefc,)) + self.assertEqual(d._impl.actuator_velocity.shape, (1,)) self.assertEqual(d.cvel.shape, (nbody, 6)) - self.assertEqual(d.cdof_dot.shape, (nv, 6)) + self.assertEqual(d._impl.cdof_dot.shape, (nv, 6)) self.assertEqual(d.qfrc_bias.shape, (nv,)) self.assertEqual(d.qfrc_passive.shape, (nv,)) - self.assertEqual(d.efc_aref.shape, (nefc,)) + self.assertEqual(d._impl.efc_aref.shape, (nefc,)) self.assertEqual(d.qfrc_actuator.shape, (nv,)) self.assertEqual(d.qfrc_smooth.shape, (nv,)) self.assertEqual(d.qacc_smooth.shape, (nv,)) self.assertEqual(d.qfrc_constraint.shape, (nv,)) self.assertEqual(d.qfrc_inverse.shape, (nv,)) - self.assertEqual(d.efc_force.shape, (nefc,)) + self.assertEqual(d._impl.efc_force.shape, (nefc,)) + + if backend_impl == 'jax': + self.assertEqual(d._impl.qM.shape, (nv, nv)) + self.assertEqual(d._impl.qLD.shape, (nv, nv)) + self.assertEqual(d._impl.qLDiagInv.shape, (0,)) + elif backend_impl == 'c': + self.assertEqual(d._impl.qM.shape, (nm,)) + self.assertEqual(d._impl.qLD.shape, (nm,)) + self.assertEqual(d._impl.qLDiagInv.shape, (nv,)) # test sparse m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE - d = mjx.make_data(m) - self.assertEqual(d.qM.shape, (nm,)) - self.assertEqual(d.qLD.shape, (nm,)) - self.assertEqual(d.qLDiagInv.shape, (nv,)) + d = mjx.make_data(m, backend_impl=backend_impl) + self.assertEqual(d._impl.qM.shape, (nm,)) + self.assertEqual(d._impl.qLD.shape, (nm,)) + self.assertEqual(d._impl.qLDiagInv.shape, (nv,)) - def test_put_data(self): + if backend_impl == 'c': + # check C specific fields + self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3)) + self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,)) + + @parameterized.parameters('jax', 'c') + def test_put_data(self, backend_impl: str): """Test that put_data puts the correct data for dense and sparse.""" - m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) # check a few fields np.testing.assert_allclose(dx.qpos, d.qpos) np.testing.assert_allclose(dx.xpos, d.xpos) np.testing.assert_allclose(dx.cvel, d.cvel) - np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot) + np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot) # check that there are no weak types self.assertFalse( @@ -312,21 +356,26 @@ class DataIOTest(parameterized.TestCase): ) ) - # check that qM is transformed properly - qm = np.zeros((m.nv, m.nv), dtype=np.float64) - mujoco.mj_fullM(m, qm, d.qM) - np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx)) + if backend_impl == 'jax': + # check that qM is transformed properly + qm = np.zeros((m.nv, m.nv), dtype=np.float64) + mujoco.mj_fullM(m, qm, d.qM) + np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx)) + elif backend_impl == 'c': + np.testing.assert_allclose(dx._impl.qM, d.qM) + np.testing.assert_allclose(dx._impl.qLD, d.qLD) + np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv) # 4 contacts, 2 for each capsule against the plane - self.assertEqual(dx.contact.dist.shape, (4,)) + self.assertEqual(dx._impl.contact.dist.shape, (4,)) self.assertEqual(d.ncon, 1) # however only 1 contact in this step - np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0]) - self.assertTrue((dx.contact.dist[1:] > 0).all()) - self.assertEqual(dx.contact.frame.shape, (4, 3, 3)) + np.testing.assert_allclose(dx._impl.contact.dist[0], d.contact.dist[0]) + self.assertTrue((dx._impl.contact.dist[1:] > 0).all()) + self.assertEqual(dx._impl.contact.frame.shape, (4, 3, 3)) np.testing.assert_allclose( - dx.contact.frame[0].reshape(9), d.contact.frame[0] + dx._impl.contact.frame[0].reshape(9), d.contact.frame[0] ) - np.testing.assert_allclose(dx.contact.frame[1:], 0) + np.testing.assert_allclose(dx._impl.contact.frame[1:], 0) # xmat, ximat, geom_xmat are all shape transformed self.assertEqual(dx.xmat.shape, (3, 3, 3)) @@ -339,57 +388,68 @@ class DataIOTest(parameterized.TestCase): np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat) # tendon data is correct - np.testing.assert_allclose(dx.ten_length, d.ten_length) - np.testing.assert_equal(dx.ten_wrapadr, np.zeros((1,))) - np.testing.assert_equal(dx.ten_wrapnum, np.zeros((1,))) - np.testing.assert_equal(dx.wrap_obj, np.zeros((2, 2))) - np.testing.assert_equal(dx.wrap_xpos, np.zeros((2, 6))) + np.testing.assert_allclose(dx._impl.ten_length, d.ten_length) + np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,))) + np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,))) + np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2))) + np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6))) # efc_ are also shape transformed and padded - self.assertEqual(dx.efc_J.shape, (45, 8)) # nefc, nv + self.assertEqual(dx._impl.efc_J.shape, (45, 8)) # nefc, nv d_efc_j = d.efc_J.reshape((-1, 8)) - np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq - np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit - np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit - np.testing.assert_allclose(dx.efc_J[5:15], d_efc_j[4:14]) # contact - np.testing.assert_allclose(dx.efc_J[15:], 0) # no contact + np.testing.assert_allclose(dx._impl.efc_J[:3], d_efc_j[:3]) # connect eq + np.testing.assert_allclose( + dx._impl.efc_J[3], d_efc_j[3] + ) # one active limit + np.testing.assert_allclose(dx._impl.efc_J[4], 0) # one inactive limit + np.testing.assert_allclose(dx._impl.efc_J[5:15], d_efc_j[4:14]) # contact + np.testing.assert_allclose(dx._impl.efc_J[15:], 0) # no contact # check another efc_ too - self.assertEqual(dx.efc_aref.shape, (45,)) # nefc - np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3]) - np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3]) - np.testing.assert_allclose(dx.efc_aref[4], 0) - np.testing.assert_allclose(dx.efc_aref[5:15], d.efc_aref[4:14]) - np.testing.assert_allclose(dx.efc_aref[15:], 0) + self.assertEqual(dx._impl.efc_aref.shape, (45,)) # nefc + np.testing.assert_allclose(dx._impl.efc_aref[:3], d.efc_aref[:3]) + np.testing.assert_allclose(dx._impl.efc_aref[3], d.efc_aref[3]) + np.testing.assert_allclose(dx._impl.efc_aref[4], 0) + np.testing.assert_allclose(dx._impl.efc_aref[5:15], d.efc_aref[4:14]) + np.testing.assert_allclose(dx._impl.efc_aref[15:], 0) # check sparse transform is correct m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx_sparse = mjx.put_data(m, d) - np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8) + dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl) + np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8) # check sparse mass matrices are correct - np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8) - np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8) - np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8) + np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8) + np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8) + np.testing.assert_allclose( + dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8 + ) # check dense mass matrices are correct m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx_from_dense = mjx.put_data(m, d) - qm = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, qm, d.qM) - np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8) + dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl) + if backend_impl == 'jax': + qm = np.zeros((m.nv, m.nv)) + mujoco.mj_fullM(m, qm, d.qM) + np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8) + elif backend_impl == 'c': + np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8) - def test_get_data(self): + @parameterized.parameters( + ('jax', False), ('jax', True), ('c', False), ('c', True) + ) + def test_get_data(self, backend_impl: str, sparse: bool): """Test that get_data makes correct MjData.""" - m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + if sparse: + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) d_2: mujoco.MjData = mjx.get_data(m, dx) # check a few fields @@ -398,6 +458,8 @@ class DataIOTest(parameterized.TestCase): np.testing.assert_allclose(d_2.cvel, d.cvel) np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot) np.testing.assert_allclose(d_2.qM, d.qM) + np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6) + np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6) # only 1 contact active self.assertEqual(d_2.contact.dist.shape, (1,)) @@ -418,12 +480,45 @@ class DataIOTest(parameterized.TestCase): # efc_* are also shape transformed and filtered self.assertEqual(d_2.nefc, 14) - self.assertEqual(d_2.efc_J.shape, (112,)) # nefc * nv - np.testing.assert_allclose(d_2.efc_J, d.efc_J) + if sparse: + efc_j = np.zeros((d.nefc, m.nv)) + mujoco.mju_sparse2dense( + efc_j, + d.efc_J, + d.efc_J_rownnz, + d.efc_J_rowadr, + d.efc_J_colind, + ) + efc_j2 = np.zeros((d_2.nefc, m.nv)) + mujoco.mju_sparse2dense( + efc_j2, + d_2.efc_J, + d_2.efc_J_rownnz, + d_2.efc_J_rowadr, + d_2.efc_J_colind, + ) + np.testing.assert_allclose(efc_j, efc_j2) + else: + self.assertEqual(d_2.efc_J.shape, (112,)) # nefc * nv + np.testing.assert_allclose(d_2.efc_J, d.efc_J) self.assertEqual(d_2.efc_aref.shape, (14,)) # nefc np.testing.assert_allclose(d_2.efc_aref, d.efc_aref) np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address) + if backend_impl == 'c': + # check fields specific to the C implementation + np.testing.assert_allclose(d_2.bvh_active, d.bvh_active) + + def test_get_data_simplebody(self): + """Test that get_data works with simple bodies where nC < nM.""" + m = mujoco.MjModel.from_xml_string(_SIMPLE_BODY) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + d_2: mujoco.MjData = mjx.get_data(m, dx) + np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6) + np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6) + def test_get_data_runs(self): xml = """ @@ -444,13 +539,14 @@ class DataIOTest(parameterized.TestCase): dx = mjx.put_data(m, d) mjx.get_data(m, dx) - def test_get_data_batched(self): + @parameterized.parameters('jax', 'c') + def test_get_data_batched(self, backend_impl): """Test that get_data makes correct List[MjData] for batched Data.""" m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) # second data in batch has contact dist > 0, disables contact dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) ds = mjx.get_data(m, dx_b) @@ -460,13 +556,14 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(ds[0].ncon, 1) self.assertEqual(ds[1].ncon, 0) - def test_get_data_into(self): + @parameterized.parameters('jax', 'c') + def test_get_data_into(self, backend_impl): """Test that get_data_into correctly populates an MjData.""" m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) d_2 = mujoco.MjData(m) mjx.get_data_into(d_2, m, dx) @@ -482,32 +579,33 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d_2.contact.frame.shape, (1, 9)) np.testing.assert_allclose(d_2.contact.frame, d.contact.frame) - def test_get_data_into_wrong_shape(self): + @parameterized.parameters('jax', 'c') + def test_get_data_into_wrong_shape(self, backend_impl): """Tests that get_data_into throwsif input and output shapes don't match.""" m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) d_2 = mujoco.MjData(m_2) with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'): mjx.get_data_into(d_2, m, dx) - def test_make_matches_put(self): + @parameterized.parameters('jax', 'c') + def test_make_matches_put(self, backend_impl): """Test that make_data produces a pytree that matches put_data.""" - m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d) + dx = mjx.put_data(m, d, backend_impl=backend_impl) step_fn = lambda d: d.replace(time=d.time + 1) step_fn_jit = jax.jit(step_fn).lower(dx).compile() # placing an MjData onto device should yield the same treedef mjx.Data as # calling make_data. they should be interchangeable for jax functions: - step_fn_jit(mjx.make_data(m)) + step_fn_jit(mjx.make_data(m, backend_impl=backend_impl)) def test_contact_elliptic_condim1(self): """Test that condim=1 with ConeType.ELLIPTIC is not implemented.""" @@ -556,7 +654,327 @@ class DataIOTest(parameterized.TestCase): """) with self.assertRaises(NotImplementedError): - mjx.put_model(m) + mjx.put_model(m, backend_impl='jax') + + +class FullCompatTest(parameterized.TestCase): + """Tests for the _full_compat flag.""" + + def test_full_compat_deprecated(self): + """Tests that _full_compat is deprecated.""" + xml = """ + + + + + + + + + + + + """ + m = mujoco.MjModel.from_xml_string(xml) + with self.assertWarns(DeprecationWarning): + out = mjx_io.put_model(m, _full_compat=True) + self.assertEqual(out.backend_impl, BackendImpl.C) + with self.assertWarns(DeprecationWarning): + out = mjx_io.make_data(m, _full_compat=True) + self.assertEqual(out.backend_impl, BackendImpl.C) + + +# Test cases for `_resolve_backend_impl_and_device` where the device is +# specified by the user and the device is available. +_DEVICE_TEST_CASES = [ + # Arguments use the following format: + # (device_type_str, backend_impl_str, + # (expected_device, expected_backend_impl))) + # No backend specified. + ('cpu', None, ('cpu', BackendImpl.C)), + ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)), + ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)), + ('tpu', None, ('tpu', BackendImpl.JAX)), + # JAX backend specified. + ('cpu', 'jax', ('cpu', BackendImpl.JAX)), + ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)), + ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)), + ('tpu', 'jax', ('tpu', BackendImpl.JAX)), + # WARP backend specified. + ('cpu', 'warp', ('cpu', 'error')), + ('gpu-notnvidia', 'warp', ('cpu', 'error')), + ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)), + ('tpu', 'warp', ('tpu', 'error')), + # C backend specified. + ('cpu', 'c', ('cpu', BackendImpl.C)), + ('gpu-notnvidia', 'c', ('cpu', 'error')), + ('gpu-nvidia', 'c', ('cpu', 'error')), + ('tpu', 'c', ('tpu', 'error')), +] + +# Test cases for `_resolve_backend_impl_and_device` where the user does NOT +# specify a device. We mock the JAX default device. +_DEFAULT_DEVICE_TEST_CASES = [ + # Arguments use the following format: + # (jax.default_device, backend_impl_str, + # (expected_device, expected_backend_impl)) + # No backend impl specified. + ('cpu', None, ('cpu', BackendImpl.C)), + ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)), + ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)), + ('tpu', None, ('tpu', BackendImpl.JAX)), + # JAX backend impl specified. + ('cpu', 'jax', ('cpu', BackendImpl.JAX)), + ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)), + ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)), + ('tpu', 'jax', ('tpu', BackendImpl.JAX)), + # WARP backend impl specified. + ('cpu', 'warp', ('cpu', 'error')), + ('gpu-notnvidia', 'warp', ('cpu', 'error')), + ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)), + ('tpu', 'warp', ('tpu', 'error')), + # C backend impl specified, CPU should always be available. + ('cpu', 'c', ('cpu', BackendImpl.C)), + ('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)), + ('gpu-nvidia', 'c', ('cpu', BackendImpl.C)), + ('tpu', 'c', ('cpu', BackendImpl.C)), +] + + +class ResolveBackendImplAndDeviceTest(parameterized.TestCase): + """Tests for the _resolve_backend_impl_and_device function.""" + + def setUp(self): + super().setUp() + + # Create mock devices + self.mock_cpu = mock.Mock(spec=jax.Device) + self.mock_cpu.platform = 'cpu' + self.mock_cpu.device_kind = 'Mock CPU' + self.mock_cpu.id = 0 + + self.mock_nvidia_gpu = mock.Mock(spec=jax.Device) + self.mock_nvidia_gpu.platform = 'gpu' + self.mock_nvidia_gpu.device_kind = 'NVIDIA Mocked GPU' + self.mock_nvidia_gpu.id = 0 + + self.mock_other_gpu = mock.Mock(spec=jax.Device) + self.mock_other_gpu.platform = 'gpu' + self.mock_other_gpu.device_kind = 'Other Mocked GPU' + self.mock_other_gpu.id = 1 + + self.mock_tpu = mock.Mock(spec=jax.Device) + self.mock_tpu.platform = 'tpu' + self.mock_tpu.device_kind = 'Mock TPU' + self.mock_tpu.id = 0 + + # Patch jax.devices for the entire test class using enter_context + self.mock_jax_devices = self.enter_context(mock.patch('jax.devices')) + self.mock_default_backend = self.enter_context( + mock.patch('jax.default_backend') + ) + + @parameterized.named_parameters( + (f'{str(args[0])}_{str(args[1])}', *args) for args in _DEVICE_TEST_CASES + ) + @mock.patch.dict( + os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'} + ) + def test_resolve_with_device( + self, + device_type_str, + backend_impl_str, + expected, + ): + """Tests various combinations of device and backend impls.""" + input_device = { + 'cpu': self.mock_cpu, + 'gpu-nvidia': self.mock_nvidia_gpu, + 'gpu-notnvidia': self.mock_other_gpu, + 'tpu': self.mock_tpu, + }[device_type_str] + + def devices_side_effect(backend=None): + # assume the user-specified device is always available + if backend == 'cpu': + return [self.mock_cpu] + elif backend == 'gpu': + if 'nvidia' in device_type_str: + return [self.mock_nvidia_gpu] + return [self.mock_other_gpu] + elif backend == 'tpu': + return [self.mock_tpu] + elif backend == 'cuda': + return [self.mock_nvidia_gpu] + + raise AssertionError('Should not be called.') + + self.mock_jax_devices.side_effect = devices_side_effect + + expected_device, expected_backend_impl = expected + if expected_backend_impl == 'error': + with self.assertRaises(AssertionError): + mjx_io._resolve_backend_impl_and_device( + backend_impl=backend_impl_str, device=input_device + ) + return + + actual_backend_impl, actual_device = ( + mjx_io._resolve_backend_impl_and_device( + backend_impl=backend_impl_str, device=input_device + ) + ) + + self.assertEqual(actual_backend_impl, expected_backend_impl) + self.assertIsNotNone(actual_device) + self.assertEqual(actual_device.platform, expected_device) + + @parameterized.named_parameters( + (f'{str(args[0])}_{str(args[1])}', *args) + for args in _DEFAULT_DEVICE_TEST_CASES + ) + @mock.patch.dict( + os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'} + ) + def test_resolve_without_device( + self, + default_device_str, + backend_impl_str, + expected, + ): + """Tests various combinations of jax.default_device and backend impls.""" + default_devices = { + 'cpu': [self.mock_cpu], + 'gpu-nvidia': [self.mock_nvidia_gpu, self.mock_cpu], + 'gpu-notnvidia': [self.mock_other_gpu, self.mock_cpu], + 'tpu': [self.mock_tpu, self.mock_cpu], + }[default_device_str] + + def devices_side_effect(backend=None): + if backend == 'cpu': + return [self.mock_cpu] # CPU is always available + if backend == 'gpu' and default_device_str == 'gpu-notnvidia': + return [self.mock_other_gpu] + if backend == 'gpu' and default_device_str == 'gpu-nvidia': + return [self.mock_nvidia_gpu] + if backend == 'cuda' and default_device_str == 'gpu-nvidia': + return [self.mock_nvidia_gpu] + if backend == 'tpu' and default_device_str == 'tpu': + return [self.mock_tpu] + if backend is None: + return default_devices + if backend == 'cuda': + raise RuntimeError('cuda backend not supported') + raise AssertionError('jax.devices error') + + self.mock_jax_devices.side_effect = devices_side_effect + default_device_side_effect_str = { + 'cpu': 'cpu', + 'gpu-nvidia': 'gpu', + 'gpu-notnvidia': 'gpu', + 'tpu': 'tpu', + }[default_device_str] + self.mock_default_backend.side_effect = ( + lambda: default_device_side_effect_str + ) + + expected_device, expected_backend_impl = expected + if ( + expected_backend_impl == 'error' + and default_device_str != 'gpu-nvidia' + and backend_impl_str == 'warp' + ): + with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'): + mjx_io._resolve_backend_impl_and_device( + backend_impl=backend_impl_str, device=None + ) + return + + if expected_backend_impl == 'error': + with self.assertRaises(AssertionError): + mjx_io._resolve_backend_impl_and_device( + backend_impl=backend_impl_str, device=None + ) + return + + actual_backend_impl, actual_device = ( + mjx_io._resolve_backend_impl_and_device( + backend_impl=backend_impl_str, device=None + ) + ) + + self.assertEqual(actual_backend_impl, expected_backend_impl) + self.assertIsNotNone(actual_device) + self.assertEqual(actual_device.platform, expected_device) + + @mock.patch.dict(os.environ, {'MJX_WARP_ENABLED': 'false'}) + def test_resolve_warp_disabled(self): + """Tests behavior when MJX_WARP_ENABLED is false.""" + self.mock_jax_devices.side_effect = lambda backend=None: ( + [self.mock_nvidia_gpu, self.mock_cpu] + if backend is None + else ([self.mock_nvidia_gpu] if backend == 'gpu' else [self.mock_cpu]) + ) + self.mock_default_backend.side_effect = lambda: 'gpu' + + # Default to JAX instead of WARP on NVIDIA GPU. + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl=None, device=None + ) + self.assertEqual(backend_impl, BackendImpl.JAX) + self.assertEqual(device.platform, 'gpu') + + # Specifying an NVIDIA GPU should still choose JAX. + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl=None, device=self.mock_nvidia_gpu + ) + self.assertEqual(backend_impl, BackendImpl.JAX) + self.assertEqual(device.platform, 'gpu') + + # Requesting warp explicitly should fail since it is disabled. + with self.assertRaises(AssertionError): + mjx_io._resolve_backend_impl_and_device( + backend_impl='warp', device=self.mock_nvidia_gpu + ) + with self.assertRaises(AssertionError): + mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None) + + @mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'}) + def test_resolve_c_disabled(self): + """Tests behavior when MJX_C_DEFAULT_ENABLED is false.""" + # Users expect that CPU defaults to the JAX impl. But in the future, it will + # default to the C backend implementation. This test checks that + # MJX_C_DEFAULT_ENABLED=false defaults to the old behavior, until the + # migration to MJEP-15 is complete. + self.mock_jax_devices.side_effect = lambda backend=None: ([self.mock_cpu]) + self.mock_default_backend.side_effect = lambda: 'cpu' + + # Default to JAX instead of C on CPU. + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl=None, device=None + ) + self.assertEqual(backend_impl, BackendImpl.JAX) + self.assertEqual(device.platform, 'cpu') + + # Specifing CPU should still choose JAX. + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl=None, device=self.mock_cpu + ) + self.assertEqual(backend_impl, BackendImpl.JAX) + self.assertEqual(device.platform, 'cpu') + + # Specifying C should choose C! + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl='c', device=None + ) + self.assertEqual(backend_impl, BackendImpl.C) + self.assertEqual(device.platform, 'cpu') + + backend_impl, device = mjx_io._resolve_backend_impl_and_device( + backend_impl='c', device=self.mock_cpu + ) + self.assertEqual(backend_impl, BackendImpl.C) + self.assertEqual(device.platform, 'cpu') if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py index c289690f..b4a8b674 100644 --- a/mjx/mujoco/mjx/_src/passive.py +++ b/mjx/mujoco/mjx/_src/passive.py @@ -23,9 +23,12 @@ 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 DataJAX from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import JointType from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX +from mujoco.mjx._src.types import OptionJAX # pylint: enable=g-importing-member @@ -72,11 +75,11 @@ def _spring_damper(m: Model, d: Data) -> jax.Array: qfrc -= m.dof_damping * d.qvel # tendon-level spring-dampers - below, above = m.tendon_lengthspring.T - d.ten_length + below, above = m.tendon_lengthspring.T - d._impl.ten_length frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0) frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring) - frc_damper = -m.tendon_damping * d.ten_velocity - qfrc += d.ten_J.T @ (frc_spring + frc_damper) + frc_damper = -m.tendon_damping * d._impl.ten_velocity + qfrc += d._impl.ten_J.T @ (frc_spring + frc_damper) return qfrc @@ -113,6 +116,9 @@ def _fluid(m: Model, d: Data) -> jax.Array: def passive(m: Model, d: Data) -> Data: """Adds all passive forces.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('passive requires JAX backend implementation.') + if m.opt.disableflags & DisableBit.PASSIVE: return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv)) @@ -124,7 +130,7 @@ def passive(m: Model, d: Data) -> Data: # add gravcomp unless added via actuators qfrc_passive += qfrc_gravcomp * (1 - m.jnt_actgravcomp[m.dof_jntid]) - if m.opt.has_fluid_params: + if m.opt.has_fluid_params: # pytype: disable=attribute-error qfrc_passive += _fluid(m, d) d = d.replace(qfrc_passive=qfrc_passive, qfrc_gravcomp=qfrc_gravcomp) diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index 496ba4cd..0bcfc195 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -132,6 +132,8 @@ def _nvmap(f: Callable[..., Y], *args) -> Y: def _check_input(m: Model, args: Any, in_types: str) -> None: """Checks that scan input has the right shape.""" + if m.nv == 0: + raise ValueError('Scan across Model with zero DoFs unsupported.') size = { 'b': m.nbody, 'j': m.njnt, diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py index 74bf3224..758ebe89 100644 --- a/mjx/mujoco/mjx/_src/scan_test.py +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -61,10 +61,7 @@ class ScanTest(absltest.TestCase): 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)) + self.assertRaises(ValueError, scan.flat, m, fn, 'b', 'b', b_in) def test_flat_joints(self): """Tests scanning over bodies with joints of different types.""" diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 5d040052..41751527 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -22,16 +22,20 @@ from mujoco.mjx._src import math from mujoco.mjx._src import ray from mujoco.mjx._src import smooth from mujoco.mjx._src import support +from mujoco.mjx._src.types import BackendImpl from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX from mujoco.mjx._src.types import ObjType from mujoco.mjx._src.types import SensorType +from mujoco.mjx._src.types import TrnType # pylint: enable=g-importing-member import numpy as np -def apply_cutoff( +def _apply_cutoff( sensor: jax.Array, cutoff: jax.Array, data_type: int ) -> jax.Array: """Clip sensor to cutoff value.""" @@ -50,6 +54,8 @@ def apply_cutoff( def sensor_pos(m: Model, d: Data) -> Data: """Compute position-dependent sensors values.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('sensor_pos requires JAX backend implementation.') if m.opt.disableflags & DisableBit.SENSOR: return d @@ -162,15 +168,15 @@ def sensor_pos(m: Model, d: Data) -> Data: sensor, _ = jax.vmap( ray.ray, in_axes=(None, None, 0, 0, None, None, None) )(m, d, site_xpos, site_mat, (), True, sid) - sensors.append(apply_cutoff(sensor, cutoffs, data_type[0])) + sensors.append(_apply_cutoff(sensor, cutoffs, data_type[0])) adrs.append(adr[idxs]) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.JOINTPOS: sensor = d.qpos[m.jnt_qposadr[objid]] elif sensor_type == SensorType.TENDONPOS: - sensor = d.ten_length[objid] + sensor = d._impl.ten_length[objid] elif sensor_type == SensorType.ACTUATORPOS: - sensor = d.actuator_length[objid] + sensor = d._impl.actuator_length[objid] elif sensor_type == SensorType.BALLQUAT: jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None] quat = d.qpos[jnt_qposadr] @@ -193,7 +199,7 @@ def sensor_pos(m: Model, d: Data) -> Data: cutofft = cutoff[idxt] sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt) adrt = adr[idxt, None] + np.arange(3)[None] - sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type in frame_axis: @@ -213,7 +219,7 @@ def sensor_pos(m: Model, d: Data) -> Data: cutofft = cutoff[idxt] sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt) adrt = adr[idxt, None] + np.arange(3)[None] - sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.FRAMEQUAT: @@ -254,7 +260,7 @@ def sensor_pos(m: Model, d: Data) -> Data: ) )(quat, refquat, refidt) adrt = adr[idxt, None] + np.arange(4)[None] - sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.SUBTREECOM: @@ -266,7 +272,7 @@ def sensor_pos(m: Model, d: Data) -> Data: # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type - sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) adrs.append(adr) if not adrs: @@ -281,6 +287,8 @@ def sensor_pos(m: Model, d: Data) -> Data: def sensor_vel(m: Model, d: Data) -> Data: """Compute velocity-dependent sensors values.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('sensor_vel requires JAX backend implementation.') if m.opt.disableflags & DisableBit.SENSOR: return d @@ -332,9 +340,9 @@ def sensor_vel(m: Model, d: Data) -> Data: elif sensor_type == SensorType.JOINTVEL: sensor = d.qvel[m.jnt_dofadr[objid]] elif sensor_type == SensorType.TENDONVEL: - sensor = d.ten_velocity[objid] + sensor = d._impl.ten_velocity[objid] elif sensor_type == SensorType.ACTUATORVEL: - sensor = d.actuator_velocity[objid] + sensor = d._impl.actuator_velocity[objid] elif sensor_type == SensorType.BALLANGVEL: jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] sensor = d.qvel[jnt_dotadr] @@ -392,20 +400,20 @@ def sensor_vel(m: Model, d: Data) -> Data: adrt = adr[idxt, None] + np.arange(3)[None] - sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.SUBTREELINVEL: - sensor = d.subtree_linvel[objid] + sensor = d._impl.subtree_linvel[objid] adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) elif sensor_type == SensorType.SUBTREEANGMOM: - sensor = d.subtree_angmom[objid] + sensor = d._impl.subtree_angmom[objid] adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) else: # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type - sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) adrs.append(adr) if not adrs: @@ -420,6 +428,8 @@ def sensor_vel(m: Model, d: Data) -> Data: def sensor_acc(m: Model, d: Data) -> Data: """Compute acceleration/force-dependent sensors values.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('sensor_acc requires JAX backend implementation.') if m.opt.disableflags & DisableBit.SENSOR: return d @@ -441,6 +451,8 @@ def sensor_acc(m: Model, d: Data) -> Data: SensorType.ACCELEROMETER, SensorType.FORCE, SensorType.TORQUE, + SensorType.FRAMELINACC, + SensorType.FRAMEANGACC, }: d = smooth.rne_postconstraint(m, d) @@ -457,14 +469,14 @@ def sensor_acc(m: Model, d: Data) -> Data: # compute contact forces forces = [] condim_ids = [] - for dim in set(d.contact.dim): + for dim in set(d._impl.contact.dim): force, condim_id = support.contact_force_dim(m, d, dim) forces.append(force) condim_ids.append(condim_id) forces = jp.concatenate(forces)[np.argsort(np.concatenate(condim_ids))] # get bodies of contact geoms - conbody = jp.array(m.geom_bodyid)[d.contact.geom] + conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom] # get site information site_bodyid = m.site_bodyid[objid] @@ -474,12 +486,14 @@ def sensor_acc(m: Model, d: Data) -> Data: site_type = m.site_type[objid] conbody0 = site_bodyid[:, None] == conbody[:, 0] conbody1 = site_bodyid[:, None] == conbody[:, 1] - contacts = (d.contact.efc_address >= 0)[None] & (conbody0 | conbody1) + contacts = (d._impl.contact.efc_address >= 0)[None] & ( + conbody0 | conbody1 + ) # compute conray, flip if second body conray = jax.vmap( lambda frame, force: math.normalize(frame[0] * force[0]) - )(d.contact.frame, forces) + )(d._impl.contact.frame, forces) conray = jp.where(conbody1[..., None], -conray, conray) # compute distance, mapping over sites and contacts @@ -501,7 +515,7 @@ def sensor_acc(m: Model, d: Data) -> Data: site_xpos[dist_id_site], site_xmat[dist_id_site], st, - d.contact.pos, + d._impl.contact.pos, conray[dist_id_site], ) dist.append(jp.where(jp.isinf(dist_site), 0, dist_site)) @@ -523,14 +537,14 @@ def sensor_acc(m: Model, d: Data) -> Data: bodyid = m.site_bodyid[objid] rot = d.site_xmat[objid] cvel = d.cvel[bodyid] - cacc = d.cacc[bodyid] + cacc = d._impl.cacc[bodyid] dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]] sensor = _accelerometer(cvel, cacc, dif, rot) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) elif sensor_type == SensorType.FORCE: bodyid = m.site_bodyid[objid] - cfrc_int = d.cfrc_int[bodyid] + cfrc_int = d._impl.cfrc_int[bodyid] site_xmat = d.site_xmat[objid] sensor = jax.vmap(lambda mat, vec: mat.T @ vec)( site_xmat, cfrc_int[:, 3:] @@ -539,7 +553,7 @@ def sensor_acc(m: Model, d: Data) -> Data: elif sensor_type == SensorType.TORQUE: bodyid = m.site_bodyid[objid] rootid = m.body_rootid[bodyid] - cfrc_int = d.cfrc_int[bodyid] + cfrc_int = d._impl.cfrc_int[bodyid] site_xmat = d.site_xmat[objid] dif = d.site_xpos[objid] - d.subtree_com[rootid] sensor = jax.vmap( @@ -550,6 +564,15 @@ def sensor_acc(m: Model, d: Data) -> Data: sensor = d.actuator_force[objid] elif sensor_type == SensorType.JOINTACTFRC: sensor = d.qfrc_actuator[m.jnt_dofadr[objid]] + elif sensor_type == SensorType.TENDONACTFRC: + force_mask = [ + (m.actuator_trntype == TrnType.TENDON) + & (m.actuator_trnid[:, 0] == tendon_id) + for tendon_id in objid + ] + force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask]) + force_mat = np.array(force_mask)[:, force_ids] + sensor = force_mat @ d.actuator_force[force_ids] elif sensor_type in (SensorType.FRAMELINACC, SensorType.FRAMEANGACC): objtype = m.sensor_objtype[idx] @@ -559,7 +582,7 @@ def sensor_acc(m: Model, d: Data) -> Data: pos, bodyid = objtype_data[ot] pos = pos[objidt] bodyid = bodyid[objidt] - cacc = d.cacc[bodyid] + cacc = d._impl.cacc[bodyid] if sensor_type == SensorType.FRAMELINACC: @@ -589,7 +612,7 @@ def sensor_acc(m: Model, d: Data) -> Data: # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type - sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) + sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) adrs.append(adr) if not adrs: diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index fce6225c..cab50779 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/mjx/mujoco/mjx/_src/sensor_test.py @@ -64,14 +64,14 @@ class SensorTest(parameterized.TestCase): mujoco.mj_forward(m, d) mx = mjx.put_model(m) - dx = mjx.put_data(m, d).replace( - sensordata=jp.zeros_like(d.sensordata), - subtree_linvel=jp.zeros_like(d.subtree_linvel), - subtree_angmom=jp.zeros_like(d.subtree_angmom), - cacc=jp.zeros_like(d.cacc), - cfrc_int=jp.zeros_like(d.cfrc_int), - cfrc_ext=jp.zeros_like(d.cfrc_ext), - ) + dx = mjx.put_data(m, d).tree_replace({ + 'sensordata': jp.zeros_like(d.sensordata), + '_impl.subtree_linvel': jp.zeros_like(d.subtree_linvel), + '_impl.subtree_angmom': jp.zeros_like(d.subtree_angmom), + '_impl.cacc': jp.zeros_like(d.cacc), + '_impl.cfrc_int': jp.zeros_like(d.cfrc_int), + '_impl.cfrc_ext': jp.zeros_like(d.cfrc_ext), + }) dx = jax.jit(mjx.sensor_pos)(mx, dx) dx = jax.jit(mjx.sensor_vel)(mx, dx) dx = jax.jit(mjx.sensor_acc)(mx, dx) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 87750c04..f87c2ed3 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -23,10 +23,12 @@ from mujoco.mjx._src import support # pylint: disable=g-importing-member from mujoco.mjx._src.types import CamLightType from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX 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 +from mujoco.mjx._src.types import ModelJAX from mujoco.mjx._src.types import TrnType from mujoco.mjx._src.types import WrapType # pylint: enable=g-importing-member @@ -35,7 +37,6 @@ import numpy as np def kinematics(m: Model, d: Data) -> Data: """Converts position/velocity from generalized coordinates to maximal.""" - def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat): # calculate joint anchors, axes, body pos and quat in global frame # also normalize qpos while we're at it @@ -131,6 +132,8 @@ def kinematics(m: Model, d: Data) -> Data: def com_pos(m: Model, d: Data) -> Data: """Maps inertias and motion dofs to global frame centered at subtree-CoM.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('com_pos requires JAX backend implementation.') # calculate center of mass of each subtree def subtree_sum(carry, xipos, body_mass): @@ -162,7 +165,7 @@ def com_pos(m: Model, d: Data) -> Data: root_com = subtree_com[m.body_rootid] offset = d.xipos - root_com cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass) - d = d.replace(cinert=cinert) + d = d.tree_replace({'_impl.cinert': cinert}) # map motion dofs to global frame centered at subtree_com def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis): @@ -201,13 +204,16 @@ def com_pos(m: Model, d: Data) -> Data: d.xanchor, d.xaxis, ) - d = d.replace(cdof=cdof) + d = d.tree_replace({'_impl.cdof': cdof}) return d def camlight(m: Model, d: Data) -> Data: """Computes camera and light positions and orientations.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('camlight requires JAX backend implementation.') + if m.ncam == 0: return d.replace(cam_xpos=jp.zeros((0, 3)), cam_xmat=jp.zeros((0, 3, 3))) @@ -278,32 +284,33 @@ def camlight(m: Model, d: Data) -> Data: def crb(m: Model, d: Data) -> Data: """Runs composite rigid body inertia algorithm.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('crb requires JAX backend implementation.') 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 = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True) crb_body = crb_body.at[0].set(0.0) - d = d.replace(crb=crb_body) + d = d.tree_replace({'_impl.crb': crb_body}) crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0) - crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof) - qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature) - d = d.replace(qM=qm) - if support.is_sparse(m) and d._qM_sparse.size > 0: # pylint: disable=protected-access - d = d.replace(_qM_sparse=qm) - + crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof) + qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature) + d = d.tree_replace({'_impl.qM': qm}) return d def factor_m(m: Model, d: Data) -> Data: """Gets factorizaton of inertia-like matrix M, assumed spd.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('factor_m requires JAX backend implementation.') if not support.is_sparse(m): - qh, _ = jax.scipy.linalg.cho_factor(d.qM) - d = d.replace(qLD=qh) + qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM) + d = d.tree_replace({'_impl.qLD': qh}) return d # build up indices for where we will do backwards updates over qLD @@ -325,7 +332,7 @@ def factor_m(m: Model, d: Data) -> Data: (out_beg, out_end, madr_d, madr_ij) ) - qld = d.qM + qld = d._impl.qM for _, updates in sorted(updates.items(), reverse=True): # combine the updates into one update batch (per depth level) @@ -353,20 +360,17 @@ def factor_m(m: Model, d: Data) -> Data: qld_diag = qld[m.dof_Madr] qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag) - d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag) - if d._qLD_sparse.size > 0: # pylint: disable=protected-access - d = d.replace(_qLD_sparse=d.qLD) - if d._qLDiagInv_sparse.size > 0: # pylint: disable=protected-access - d = d.replace(_qLDiagInv_sparse=d.qLDiagInv) - + d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag}) return d def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array: """Computes sparse backsubstitution: x = inv(L'*D*L)*y .""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('solve_m requires JAX backend implementation.') if not support.is_sparse(m): - return jax.scipy.linalg.cho_solve((d.qLD, False), x) + return jax.scipy.linalg.cho_solve((d._impl.qLD, False), x) depth = [] for i in range(m.nv): @@ -385,21 +389,23 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array: # x <- inv(L') * x for _, vals in sorted(updates_j.items(), reverse=True): j, madr_ij, i = np.array(vals).T - x = x.at[j].add(-d.qLD[madr_ij] * x[i]) + x = x.at[j].add(-d._impl.qLD[madr_ij] * x[i]) # x <- inv(D) * x - x = x * d.qLDiagInv + x = x * d._impl.qLDiagInv # x <- inv(L) * x for _, vals in sorted(updates_i.items()): i, madr_ij, j = np.array(vals).T - x = x.at[i].add(-d.qLD[madr_ij] * x[j]) + x = x.at[i].add(-d._impl.qLD[madr_ij] * x[j]) return x def com_vel(m: Model, d: Data) -> Data: """Computes cvel, cdof_dot.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('com_vel requires JAX backend implementation.') # forward scan down tree: accumulate link center of mass velocity def fn(parent, jnt_typs, cdof, qvel): @@ -431,17 +437,19 @@ def com_vel(m: Model, d: Data) -> Data: 'jvv', 'bv', m.jnt_type, - d.cdof, + d._impl.cdof, d.qvel, ) - d = d.replace(cvel=cvel, cdof_dot=cdof_dot) + d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot}) return d def subtree_vel(m: Model, d: Data) -> Data: """Subtree linear velocity and angular momentum.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('subtree_vel requires JAX backend implementation.') # bodywise quantities def _forward(cvel, xipos, ximat, subtree_com_root, mass, inertia): @@ -529,14 +537,22 @@ def subtree_vel(m: Model, d: Data) -> Data: reverse=True, ) - return d.replace(subtree_linvel=subtree_linvel, subtree_angmom=subtree_angmom) + return d.tree_replace({ + '_impl.subtree_linvel': subtree_linvel, + '_impl.subtree_angmom': subtree_angmom, + }) -def rne(m: Model, d: Data) -> Data: - """Computes inverse dynamics using the recursive Newton-Euler algorithm.""" +def rne(m: Model, d: Data, flg_acc: bool = False) -> Data: + """Computes inverse dynamics using the recursive Newton-Euler algorithm. + + flg_acc=False removes inertial term. + """ + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('rne requires JAX backend implementation.') # forward scan over tree: accumulate link center of mass acceleration - def cacc_fn(cacc, cdof_dot, qvel): + def cacc_fn(cacc, cdof_dot, qvel, cdof, qacc): if cacc is None: if m.opt.disableflags & DisableBit.GRAVITY: cacc = jp.zeros((6,)) @@ -545,9 +561,15 @@ def rne(m: Model, d: Data) -> Data: cacc += jp.sum(jax.vmap(jp.multiply)(cdof_dot, qvel), axis=0) + # cacc += cdof * qacc + if flg_acc: + cacc += jp.sum(jax.vmap(jp.multiply)(cdof, qacc), axis=0) + return cacc - cacc = scan.body_tree(m, cacc_fn, 'vv', 'b', d.cdof_dot, d.qvel) + cacc = scan.body_tree( + m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc + ) def frc(cinert, cacc, cvel): frc = math.inert_mul(cinert, cacc) @@ -555,7 +577,7 @@ def rne(m: Model, d: Data) -> Data: return frc - loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel) + loc_cfrc = jax.vmap(frc)(d._impl.cinert, cacc, d.cvel) # backward scan up tree: accumulate body forces def cfrc_fn(cfrc_child, cfrc): @@ -564,7 +586,7 @@ def rne(m: Model, d: Data) -> Data: 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)]) + qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)]) d = d.replace(qfrc_bias=qfrc_bias) @@ -573,6 +595,8 @@ def rne(m: Model, d: Data) -> Data: def rne_postconstraint(m: Model, d: Data) -> Data: """RNE with complete data: compute cacc, cfrc_ext, cfrc_int.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('rne_postconstraint requires JAX backend implementation.') def _transform_force(frc, offset): force, torque = jp.split(frc, 2) @@ -593,7 +617,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data: # compute contact forces for each condim forces = [] condim_idx = [] - for dim in set(d.contact.dim): + for dim in set(d._impl.contact.dim): force, idx = support.contact_force_dim(m, d, dim) forces.append(force) condim_idx.append(idx) @@ -620,10 +644,10 @@ def rne_postconstraint(m: Model, d: Data) -> Data: ) condim_idx = jp.concatenate(condim_idx) - frame = d.contact.frame[condim_idx] - pos = d.contact.pos[condim_idx] - id1 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 0]] - id2 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 1]] + frame = d._impl.contact.frame[condim_idx] + pos = d._impl.contact.pos[condim_idx] + id1 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 0]] + id2 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 1]] com1 = d.subtree_com[jp.array(m.body_rootid)][id1] com2 = d.subtree_com[jp.array(m.body_rootid)][id2] @@ -659,8 +683,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data: ) # cacc = cacc_parent + cdofdot * qvel + cdof * qacc - cacc_vel = d.cdof_dot.T @ (mask * d.qvel) - cacc_acc = d.cdof.T @ (mask * d.qacc) + cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel) + cacc_acc = d._impl.cdof.T @ (mask * d.qacc) cacc = cacc_parent + cacc_vel + cacc_acc # cfrc_body = cinert * cacc + cvel x (cinert * cvel) @@ -678,7 +702,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data: 'bbbbb', 'bb', cfrc_ext, - d.cinert, + d._impl.cinert, d.cvel, jp.array(m.body_dofadr), jp.array(m.body_dofnum), @@ -695,11 +719,18 @@ def rne_postconstraint(m: Model, d: Data) -> Data: ) # update data - return d.replace(cacc=cacc, cfrc_int=cfrc_int, cfrc_ext=cfrc_ext) + return d.tree_replace({ + '_impl.cacc': cacc, + '_impl.cfrc_int': cfrc_int, + '_impl.cfrc_ext': cfrc_ext, + }) def tendon(m: Model, d: Data) -> Data: """Computes tendon lengths and moments.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('tendon requires JAX backend implementation.') + if not m.ntendon: return d @@ -827,8 +858,8 @@ def tendon(m: Model, d: Data) -> Data: # wrap inside # TODO(taylorhowell): check that is_wrap_inside is consistent with # site and geom relative positions - (wrap_inside_id,) = np.nonzero(m.is_wrap_inside) - (wrap_outside_id,) = np.nonzero(~m.is_wrap_inside) + (wrap_inside_id,) = np.nonzero(m._impl.is_wrap_inside) + (wrap_outside_id,) = np.nonzero(~m._impl.is_wrap_inside) # compute geom wrap length and connect points (if wrap occurs) v_wrap = jax.vmap( @@ -844,9 +875,9 @@ def tendon(m: Model, d: Data) -> Data: has_sidesite[wrap_inside_id], is_sphere[wrap_inside_id], True, - m.wrap_inside_maxiter, - m.wrap_inside_tolerance, - m.wrap_inside_z_init, + m._impl.wrap_inside_maxiter, + m._impl.wrap_inside_tolerance, + m._impl.wrap_inside_z_init, ) lengths_outside, pnt0_outside, pnt1_outside = v_wrap( @@ -859,9 +890,9 @@ def tendon(m: Model, d: Data) -> Data: has_sidesite[wrap_outside_id], is_sphere[wrap_outside_id], False, - m.wrap_inside_maxiter, - m.wrap_inside_tolerance, - m.wrap_inside_z_init, + m._impl.wrap_inside_maxiter, + m._impl.wrap_inside_tolerance, + m._impl.wrap_inside_z_init, ) wrap_id = np.argsort(np.concatenate([wrap_inside_id, wrap_outside_id])) @@ -943,12 +974,14 @@ def tendon(m: Model, d: Data) -> Data: ) # assemble length and moment - ten_length = jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt) + ten_length = ( + jp.zeros_like(d._impl.ten_length).at[tendon_id_jnt].set(length_jnt) + ) ten_length = ten_length.at[tendon_id_site].add(length_site) ten_length = ten_length.at[tendon_id_geom].add(length_geom) ten_moment = ( - jp.zeros_like(d.ten_J) + jp.zeros_like(d._impl.ten_J) .at[adr_moment_jnt, dofadr_moment_jnt] .set(moment_jnt) ) @@ -1011,14 +1044,14 @@ def tendon(m: Model, d: Data) -> Data: [wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)] ).reshape((m.nwrap, 2)) - return d.replace( - ten_length=ten_length, - ten_J=ten_moment, - ten_wrapadr=jp.array(ten_wrapadr, dtype=int), - ten_wrapnum=jp.array(ten_wrapnum, dtype=int), - wrap_xpos=wrap_xpos, - wrap_obj=jp.array(wrap_obj, dtype=int), - ) + return d.tree_replace({ + '_impl.ten_length': ten_length, + '_impl.ten_J': ten_moment, + '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int), + '_impl.ten_wrapnum': jp.array(ten_wrapnum, dtype=int), + '_impl.wrap_xpos': wrap_xpos, + '_impl.wrap_obj': jp.array(wrap_obj, dtype=int), + }) def _site_dof_mask(m: Model) -> np.ndarray: @@ -1052,6 +1085,9 @@ def _site_dof_mask(m: Model) -> np.ndarray: def transmission(m: Model, d: Data) -> Data: """Computes actuator/transmission lengths and moments.""" + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('transmission requires JAX backend implementation.') + if not m.nu: return d @@ -1111,8 +1147,8 @@ def transmission(m: Model, d: Data) -> Data: wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:])) moment = jac @ wrench elif trntype == TrnType.TENDON: - length = d.ten_length[trnid[0]] * gear[:1] - moment = d.ten_J[trnid[0]] * gear[0] + length = d._impl.ten_length[trnid[0]] * gear[:1] + moment = d._impl.ten_J[trnid[0]] * gear[0] else: raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}') @@ -1144,5 +1180,7 @@ def transmission(m: Model, d: Data) -> Data: length = length.reshape((m.nu,)) moment = moment.reshape((m.nu, m.nv)) - d = d.replace(actuator_length=length, actuator_moment=moment) + d = d.tree_replace( + {'_impl.actuator_length': length, '_impl.actuator_moment': moment} + ) return d diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 01ce8abb..96470a21 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -20,7 +20,7 @@ import jax import mujoco from mujoco import mjx from mujoco.mjx._src import test_util -from mujoco.mjx._src.types import ConeType +from mujoco.mjx._src.types import ConeType # pylint: disable=g-importing-member import numpy as np # tolerance for difference between MuJoCo and MJX smooth calculations - mostly @@ -78,33 +78,37 @@ class SmoothTest(absltest.TestCase): # com_pos dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'subtree_com') - _assert_attr_eq(d, dx, 'cinert') - _assert_attr_eq(d, dx, 'cdof') + _assert_attr_eq(d, dx._impl, 'cinert') + _assert_attr_eq(d, dx._impl, 'cdof') # camlight dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'cam_xpos') _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat') # crb dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d)) - _assert_attr_eq(d, dx, 'crb') - _assert_attr_eq(d, dx, 'qM') - _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse') + _assert_attr_eq(d, dx._impl, 'crb') + _assert_attr_eq(d, dx._impl, 'qM') # factor_m dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d)) qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name - for i in range(m.nM): + for i in range(m.nC): qLDLegacy[d.mapM2M[i]] = d.qLD[i] - _assert_eq(qLDLegacy, dx.qLD, 'qLD') - _assert_attr_eq(d, dx, 'qLDiagInv') - _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse') - _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse') + _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD') + _assert_attr_eq(d, dx._impl, 'qLDiagInv') # com_vel dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'cvel') - _assert_attr_eq(d, dx, 'cdof_dot') + _assert_attr_eq(d, dx._impl, 'cdof_dot') # rne dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'qfrc_bias') + # rne (flg_acc=True) + qfrc_bias = np.zeros(m.nv) + mujoco.mj_rne(m, d, 1, qfrc_bias) + dx = jax.jit(mjx.rne, static_argnums=(2,))( + mx, mjx.put_data(m, d), flg_acc=True + ) + _assert_eq(dx.qfrc_bias, qfrc_bias, 'qfrc_bias') # set dense jacobian for tendon: m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE @@ -115,11 +119,11 @@ class SmoothTest(absltest.TestCase): mujoco.mj_forward(m, d) # tendon dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d)) - _assert_attr_eq(d, dx, 'ten_J') - _assert_attr_eq(d, dx, 'ten_length') + _assert_attr_eq(d, dx._impl, 'ten_J') + _assert_attr_eq(d, dx._impl, 'ten_length') # transmission dx = jax.jit(mjx.transmission)(mx, dx) - _assert_attr_eq(d, dx, 'actuator_length') + _assert_attr_eq(d, dx._impl, 'actuator_length') # convert sparse actuator_moment to dense representation moment = np.zeros((m.nu, m.nv)) @@ -130,7 +134,7 @@ class SmoothTest(absltest.TestCase): d.moment_rowadr, d.moment_colind, ) - _assert_eq(moment, dx.actuator_moment, 'actuator_moment') + _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment') def test_disable_gravity(self): m = mujoco.MjModel.from_xml_string(""" @@ -190,7 +194,7 @@ class SmoothTest(absltest.TestCase): mujoco.mj_transmission(m, d) dx = jax.jit(mjx.transmission)(mx, dx) - _assert_attr_eq(d, dx, 'actuator_length') + _assert_attr_eq(d, dx._impl, 'actuator_length') # convert sparse actuator_moment to dense representation moment = np.zeros((m.nu, m.nv)) @@ -201,7 +205,7 @@ class SmoothTest(absltest.TestCase): d.moment_rowadr, d.moment_colind, ) - _assert_eq(moment, dx.actuator_moment, 'actuator_moment') + _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment') def test_subtree_vel(self): """Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel.""" @@ -219,8 +223,8 @@ class SmoothTest(absltest.TestCase): mujoco.mj_subtreeVel(m, d) dx = jax.jit(mjx.subtree_vel)(mx, dx) - _assert_attr_eq(d, dx, 'subtree_linvel') - _assert_attr_eq(d, dx, 'subtree_angmom') + _assert_attr_eq(d, dx._impl, 'subtree_linvel') + _assert_attr_eq(d, dx._impl, 'subtree_angmom') class RnePostConstraintTest(parameterized.TestCase): @@ -271,9 +275,9 @@ class RnePostConstraintTest(parameterized.TestCase): mujoco.mj_rnePostConstraint(m, d) dx = jax.jit(mjx.rne_postconstraint)(mx, dx) - _assert_eq(d.cacc, dx.cacc, 'cacc') - _assert_eq(d.cfrc_ext, dx.cfrc_ext, 'cfrc_ext') - _assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int') + _assert_eq(d.cacc, dx._impl.cacc, 'cacc') + _assert_eq(d.cfrc_ext, dx._impl.cfrc_ext, 'cfrc_ext') + _assert_eq(d.cfrc_int, dx._impl.cfrc_int, 'cfrc_int') class TendonTest(parameterized.TestCase): @@ -305,12 +309,12 @@ class TendonTest(parameterized.TestCase): mujoco.mj_forward(m, d) dx = jax.jit(mjx.forward)(mx, dx) - _assert_eq(d.ten_length, dx.ten_length, 'ten_length') - _assert_eq(d.ten_J, dx.ten_J, 'ten_J') - _assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum') - _assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr') - _assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj') - _assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos') + _assert_eq(d.ten_length, dx._impl.ten_length, 'ten_length') + _assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J') + _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum') + _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr') + _assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj') + _assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos') if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index 6d3d4f21..98ec3d27 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -24,13 +24,15 @@ from mujoco.mjx._src import support from mujoco.mjx._src.dataclasses import PyTreeNode from mujoco.mjx._src.types import ConeType from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ModelJAX from mujoco.mjx._src.types import SolverType # pylint: enable=g-importing-member -class _Context(PyTreeNode): +class Context(PyTreeNode): """Data updated during each solver iteration. Attributes: @@ -72,25 +74,30 @@ class _Context(PyTreeNode): h: jax.Array @classmethod - def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context': - jaref = d.efc_J @ d.qacc - d.efc_aref + def create(cls, m: Model, d: Data, grad: bool = True) -> 'Context': + if not isinstance(d._impl, DataJAX): + raise ValueError( + 'Constraint context requires JAX backend implementation.' + ) + + jaref = d._impl.efc_J @ d.qacc - d._impl.efc_aref # TODO(robotics-team): determine nv at which sparse mul is faster ma = support.mul_m(m, d, d.qacc) nv_0 = jp.zeros(m.nv) fri = 0.0 if m.opt.cone == ConeType.ELLIPTIC: - friction = d.contact.friction[d.contact.dim > 1] - dim = d.contact.dim[d.contact.dim > 1] + friction = d._impl.contact.friction[d._impl.contact.dim > 1] + dim = d._impl.contact.dim[d._impl.contact.dim > 1] mu = friction[:, 0] / jp.sqrt(m.opt.impratio) fri = jp.concatenate((mu[:, None], friction), axis=1) for condim in (3, 4, 6): fri = fri.at[dim == condim, condim:].set(0) - ctx = _Context( + ctx = Context( qacc=d.qacc, qfrc_constraint=d.qfrc_constraint, Jaref=jaref, - efc_force=d.efc_force, + efc_force=d._impl.efc_force, Ma=ma, grad=nv_0, Mgrad=nv_0, @@ -133,7 +140,7 @@ class _LSPoint(PyTreeNode): cls, m: Model, d: Data, - ctx: _Context, + ctx: Context, alpha: jax.Array, jv: jax.Array, quad: jax.Array, @@ -145,18 +152,20 @@ class _LSPoint(PyTreeNode): ) -> '_LSPoint': """Creates a linesearch point with first and second derivatives.""" # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('LSPoint requires JAX backend implementation.') cost, deriv_0, deriv_1 = 0.0, 0.0, 0.0 quad_total = quad_gauss x = ctx.Jaref + alpha * jv - active = (x < 0).at[: d.ne + d.nf].set(True) + active = (x < 0).at[: d._impl.ne + d._impl.nf].set(True) - dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss + dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss if (dof_fl.any() or ten_fl.any()) and not ( m.opt.disableflags & DisableBit.FRICTIONLOSS ): - f = d.efc_frictionloss - r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL) + f = d._impl.efc_frictionloss + r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL) rf, z = r * f, jp.zeros_like(f) linear_neg = (x <= -rf)[:, None] linear_pos = (x >= rf)[:, None] @@ -174,13 +183,13 @@ class _LSPoint(PyTreeNode): middle_zone = (tsqr > 0) & (n < (mu * t)) & ((mu * n + t) > 0) # quadratic cost for equality, friction, limits, frictionless contacts - dim1 = d.contact.efc_address[d.contact.dim == 1] - nefl = d.ne + d.nf + d.nl + dim1 = d._impl.contact.efc_address[d._impl.contact.dim == 1] + nefl = d._impl.ne + d._impl.nf + d._impl.nl active = active.at[nefl:].set(False).at[dim1].set(active[dim1]) quad_efld = jax.vmap(jp.multiply)(quad, active) quad_total += jp.sum(quad_efld, axis=0) # elliptic bottom zone: quadratic cost - efc_elliptic = d.contact.efc_address[d.contact.dim > 1] + efc_elliptic = d._impl.contact.efc_address[d._impl.contact.dim > 1] quad_c = jax.vmap(jp.multiply)(quad[efc_elliptic], bottom_zone) quad_total += jp.sum(quad_c, axis=0) # elliptic middle zone @@ -241,7 +250,7 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter): return jax.lax.scan(_fun, init, None, length=max_iter)[0][0] -def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: +def _update_constraint(m: Model, d: Data, ctx: Context) -> Context: """Updates constraint force and resulting cost given last solver iteration. Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c @@ -254,17 +263,20 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: Returns: context with new constraint force and costs """ + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('_update_constraint requires JAX backend implementation.') + # ne constraints are always active, nf are conditionally active, others are # non-negative constraints. - active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True) + active = (ctx.Jaref < 0).at[: d._impl.ne + d._impl.nf].set(True) - floss_force, floss_cost = jp.zeros(d.nefc), 0.0 - dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss + floss_force, floss_cost = jp.zeros(d._impl.nefc), 0.0 + dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss if (dof_fl.any() or ten_fl.any()) and not ( m.opt.disableflags & DisableBit.FRICTIONLOSS ): - f = d.efc_frictionloss - r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL) + f = d._impl.efc_frictionloss + r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL) linear_neg = (ctx.Jaref <= -r * f) * (f > 0) linear_pos = (ctx.Jaref >= r * f) * (f > 0) active = active & ~linear_neg & ~linear_pos @@ -274,13 +286,13 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: floss_cost = floss_cost.sum() if m.opt.cone == ConeType.PYRAMIDAL: - efc_force = d.efc_D * -ctx.Jaref * active + floss_force - cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active) + efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force + cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active) dm, u, h = 0.0, 0.0, 0.0 elif m.opt.cone == ConeType.ELLIPTIC: - friction = d.contact.friction[d.contact.dim > 1] - efc_address = d.contact.efc_address[d.contact.dim > 1] - dim = d.contact.dim[d.contact.dim > 1] + friction = d._impl.contact.friction[d._impl.contact.dim > 1] + efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1] + dim = d._impl.contact.dim[d._impl.contact.dim > 1] # to prevent out of range append zeros to ctx.Jaref slice_fn = jax.vmap( lambda x: jax.lax.dynamic_slice( @@ -297,12 +309,12 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: adr_i.extend(range(addr, addr + condim)) adr_j.extend([i] * condim) active = active.at[jp.array(adr_i)].set(bottom_zone[jp.array(adr_j)]) - efc_force = d.efc_D * -ctx.Jaref * active + floss_force - cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active) + efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force + cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active) # middle zone: cone middle_zone = (t > 0) & (n < (mu * t)) & ((mu * n + t) > 0) - dm = d.efc_D[efc_address] / jp.maximum( + dm = d._impl.efc_D[efc_address] / jp.maximum( mu * mu * (1 + mu * mu), mujoco.mjMINVAL ) nmt = n - mu * t @@ -339,7 +351,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: else: raise NotImplementedError(f'unsupported cone type: {m.opt.cone}') - qfrc_constraint = d.efc_J.T @ efc_force + qfrc_constraint = d._impl.efc_J.T @ efc_force gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth) ctx = ctx.replace( qfrc_constraint=qfrc_constraint, @@ -356,7 +368,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: return ctx -def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: +def _update_gradient(m: Model, d: Data, ctx: Context) -> Context: """Updates grad and M / grad given latest solver iteration. Corresponds to CGupdateGradient in mujoco/src/engine/engine_solver.c @@ -371,6 +383,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: Raises: NotImplementedError: for unsupported solver type """ + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('_update_gradient requires JAX backend implementation.') grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint @@ -378,16 +392,16 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: mgrad = smooth.solve_m(m, d, grad) elif m.opt.solver == SolverType.NEWTON: if m.opt.cone == ConeType.ELLIPTIC: - cm = jp.diag(d.efc_D * ctx.active) - efc_address = d.contact.efc_address[d.contact.dim > 1] - dim = d.contact.dim[d.contact.dim > 1] + cm = jp.diag(d._impl.efc_D * ctx.active) + efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1] + dim = d._impl.contact.dim[d._impl.contact.dim > 1] # set efc of cone H along diagonal for i, (condim, addr) in enumerate(zip(dim, efc_address)): h_cone = ctx.h[i, :condim, :condim] cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone) - h = d.efc_J.T @ cm @ d.efc_J + h = d._impl.efc_J.T @ cm @ d._impl.efc_J else: - h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J + h = (d._impl.efc_J.T * d._impl.efc_D * ctx.active) @ d._impl.efc_J h = support.full_m(m, d) + h h_ = jax.scipy.linalg.cho_factor(h) mgrad = jax.scipy.linalg.cho_solve(h_, grad) @@ -403,7 +417,7 @@ def _rescale(m: Model, value: jax.Array) -> jax.Array: return value / (m.stat.meaninertia * max(1, m.nv)) -def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: +def _linesearch(m: Model, d: Data, ctx: Context) -> Context: """Performs a zoom linesearch to find optimal search step size. Args: @@ -414,12 +428,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: Returns: updated context with new qacc, Ma, Jaref """ + if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): + raise ValueError('_lineasearch requires JAX backend implementation.') + smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv) gtol = m.opt.tolerance * m.opt.ls_tolerance * smag # compute Mv, Jv mv = support.mul_m(m, d, ctx.search) - jv = d.efc_J @ ctx.search + jv = d._impl.efc_J @ ctx.search # prepare quadratics quad_gauss = jp.stack(( @@ -428,13 +445,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: 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 + quad = (quad * d._impl.efc_D).T uu, v0, uv, vv = 0.0, 0.0, 0.0, 0.0 if m.opt.cone == ConeType.ELLIPTIC: - mask = d.contact.dim > 1 + mask = d._impl.contact.dim > 1 # complete vector quadratic (for bottom zone) efc_con, efc_fri = [], [] - for condim, addr in zip(d.contact.dim[mask], d.contact.efc_address[mask]): + for condim, addr in zip( + d._impl.contact.dim[mask], d._impl.contact.efc_address[mask] + ): efc_con.extend([addr] * (condim - 1)) efc_fri.extend(range(addr + 1, addr + condim)) quad = quad.at[jp.array(efc_con)].add(quad[jp.array(efc_fri)]) @@ -446,7 +465,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: jp.concatenate((jv, jp.zeros(3))), (x,), (6,) ) ) - efc_elliptic = d.contact.efc_address[mask] + efc_elliptic = d._impl.contact.efc_address[mask] v = jv_fn(efc_elliptic) * ctx.fri uu = jp.sum(ctx.u[:, 1:] * ctx.u[:, 1:], axis=1) v0 = v[:, 0] @@ -529,7 +548,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: def solve(m: Model, d: Data) -> Data: """Finds forces that satisfy constraints using conjugate gradient descent.""" - def cond(ctx: _Context) -> jax.Array: + def cond(ctx: Context) -> jax.Array: improvement = _rescale(m, ctx.prev_cost - ctx.cost) gradient = _rescale(m, math.norm(ctx.grad)) @@ -539,7 +558,7 @@ def solve(m: Model, d: Data) -> Data: return ~done - def body(ctx: _Context) -> _Context: + def body(ctx: Context) -> Context: ctx = _linesearch(m, d, ctx) prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad # pylint: disable=invalid-name ctx = _update_constraint(m, d, ctx) @@ -560,22 +579,22 @@ def solve(m: Model, d: Data) -> Data: # warmstart: qacc = d.qacc_smooth if not m.opt.disableflags & DisableBit.WARMSTART: - warm = _Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False) - smth = _Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False) + warm = Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False) + smth = Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False) qacc = jp.where(warm.cost < smth.cost, d.qacc_warmstart, d.qacc_smooth) d = d.replace(qacc=qacc) - ctx = _Context.create(m, d) + ctx = Context.create(m, d) if m.opt.iterations == 1: ctx = body(ctx) else: ctx = jax.lax.while_loop(cond, body, ctx) - d = d.replace( - qacc_warmstart=ctx.qacc, - qacc=ctx.qacc, - qfrc_constraint=ctx.qfrc_constraint, - efc_force=ctx.efc_force, - ) + d = d.tree_replace({ + 'qacc_warmstart': ctx.qacc, + 'qfrc_constraint': ctx.qfrc_constraint, + 'qacc': ctx.qacc, + '_impl.efc_force': ctx.efc_force, + }) return d diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py index 57d4a116..0b1b3b5a 100644 --- a/mjx/mujoco/mjx/_src/solver_test.py +++ b/mjx/mujoco/mjx/_src/solver_test.py @@ -74,7 +74,7 @@ class SolverTest(parameterized.TestCase): # compare costs mj_cost = cost(d.qacc) - ctx = solver._Context.create(mjx.put_model(m), mjx.put_data(m, d)) + ctx = solver.Context.create(mjx.put_model(m), mjx.put_data(m, d)) mjx_cost = ctx.cost - ctx.gauss _assert_eq(mj_cost, mjx_cost, 'cost') @@ -87,8 +87,8 @@ class SolverTest(parameterized.TestCase): # MJX finds very similar solutions with the newton solver if solver_ == mujoco.mjtSolver.mjSOL_NEWTON: - nnz = dx.efc_J.any(axis=1) - _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + nnz = dx._impl.efc_J.any(axis=1) + _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force') _assert_attr_eq(d, dx, 'qfrc_constraint') _assert_attr_eq(d, dx, 'qacc') @@ -108,9 +108,9 @@ class SolverTest(parameterized.TestCase): mujoco.mj_forward(m, d) mx = mjx.put_model(m) dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d)) - nnz = dx.efc_J.any(axis=1) + nnz = dx._impl.efc_J.any(axis=1) # even without warmstart, newton converges quickly - _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-4) + _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force', tol=2e-4) def test_sparse(self): """Test solver works with sparse mass matrices.""" @@ -130,8 +130,8 @@ class SolverTest(parameterized.TestCase): _assert_attr_eq(d, dx, 'qacc') _assert_attr_eq(d, dx, 'qfrc_constraint') - nnz = dx.efc_J.any(axis=1) - _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + nnz = dx._impl.efc_J.any(axis=1) + _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force') def test_quad_frictionloss(self): """Test a case with quadratic frictionloss constraints.""" @@ -144,8 +144,8 @@ class SolverTest(parameterized.TestCase): _assert_attr_eq(d, dx, 'qacc') _assert_attr_eq(d, dx, 'qfrc_constraint') - nnz = dx.efc_J.any(axis=1) - _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + nnz = dx._impl.efc_J.any(axis=1) + _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force') # TODO(taylorhowell): condim=1 with ConeType.ELLIPTIC @parameterized.product(condim=(3, 4, 6), cone=tuple(ConeType)) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 591ee0ee..15715021 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================== """Engine support functions.""" + from collections.abc import Iterable, Sequence from typing import Optional, Tuple, Union @@ -92,7 +93,7 @@ def full_m(m: Model, d: Data) -> jax.Array: """Reconstitute dense mass matrix from qM.""" if not is_sparse(m): - return d.qM + return d._impl.qM # pytype: disable=attribute-error ij = [] for i in range(m.nv): @@ -103,7 +104,7 @@ def full_m(m: Model, d: Data) -> jax.Array: i, j = (jp.array(x) for x in zip(*ij)) - mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM) + mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM) # pytype: disable=attribute-error # also set upper triangular mat = mat + jp.tril(mat, -1).T @@ -115,9 +116,9 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array: """Multiply vector by inertia matrix.""" if not is_sparse(m): - return d.qM @ vec + return d._impl.qM @ vec # pytype: disable=attribute-error - diag_mul = d.qM[jp.array(m.dof_Madr)] * vec + diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec # pytype: disable=attribute-error is_, js, madr_ijs = [], [], [] for i in range(m.nv): @@ -131,8 +132,8 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array: 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]) + out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j]) # pytype: disable=attribute-error + out = out.at[j].add(d._impl.qM[madr_ij] * vec[i]) # pytype: disable=attribute-error return out @@ -147,9 +148,9 @@ def jac( 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(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof) # pytype: disable=attribute-error jacp = jax.vmap(jp.multiply)(jacp, mask) - jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask) + jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask) # pytype: disable=attribute-error return jacp, jacr @@ -295,73 +296,56 @@ class BindModel(object): self.prefix = '' ids = [] for spec in specs: - if not spec.name: - raise KeyError(f'cannot bind spec with empty name') + if model.signature != spec.signature: + raise ValueError( + 'mjSpec signature does not match mjx.Model signature:' + f' {spec.signature} != {model.signature}' + ) + elif spec.id < 0: + raise KeyError(f'invalid id: {spec.id}') elif isinstance(spec, mujoco.MjsBody): self.prefix = 'body_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name) elif isinstance(spec, mujoco.MjsJoint): self.prefix = 'jnt_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name) elif isinstance(spec, mujoco.MjsGeom): self.prefix = 'geom_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name) elif isinstance(spec, mujoco.MjsSite): self.prefix = 'site_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name) elif isinstance(spec, mujoco.MjsLight): self.prefix = 'light_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name) elif isinstance(spec, mujoco.MjsCamera): self.prefix = 'cam_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name) elif isinstance(spec, mujoco.MjsMesh): self.prefix = 'mesh_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_MESH, spec.name) elif isinstance(spec, mujoco.MjsHField): self.prefix = 'hfield_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, spec.name) elif isinstance(spec, mujoco.MjsPair): self.prefix = 'pair_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_PAIR, spec.name) elif isinstance(spec, mujoco.MjsTendon): self.prefix = 'tendon_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name) elif isinstance(spec, mujoco.MjsActuator): self.prefix = 'actuator_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name) elif isinstance(spec, mujoco.MjsSensor): self.prefix = 'sensor_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name) elif isinstance(spec, mujoco.MjsNumeric): self.prefix = 'numeric_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_NUMERIC, spec.name) elif isinstance(spec, mujoco.MjsText): self.prefix = 'text_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_TEXT, spec.name) elif isinstance(spec, mujoco.MjsTuple): self.prefix = 'tuple_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_TUPLE, spec.name) elif isinstance(spec, mujoco.MjsKey): self.prefix = 'key_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_KEY, spec.name) elif isinstance(spec, mujoco.MjsEquality): self.prefix = 'eq_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name) elif isinstance(spec, mujoco.MjsExclude): self.prefix = 'exclude_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_EXCLUDE, spec.name) elif isinstance(spec, mujoco.MjsSkin): self.prefix = 'skin_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_SKIN, spec.name) elif isinstance(spec, mujoco.MjsMaterial): self.prefix = 'material_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, spec.name) else: raise ValueError('invalid spec type') - if idx < 0: - raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error - ids.append(idx) + ids.append(spec.id) if len(ids) == 1: self.id = ids[0] else: @@ -402,42 +386,36 @@ class BindData(object): self.prefix = '' ids = [] for spec in specs: - if not spec.name: - raise KeyError(f'cannot bind spec with empty name') + if model.signature != spec.signature: + raise ValueError( + 'mjSpec signature does not match mjx.Model signature:' + f' {spec.signature} != {model.signature}' + ) + if spec.id < 0: + raise KeyError(f'invalid id: {spec.id}') elif isinstance(spec, mujoco.MjsBody): - idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name) + pass elif isinstance(spec, mujoco.MjsJoint): self.prefix = 'jnt_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name) elif isinstance(spec, mujoco.MjsGeom): self.prefix = 'geom_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name) elif isinstance(spec, mujoco.MjsSite): self.prefix = 'site_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name) elif isinstance(spec, mujoco.MjsLight): self.prefix = 'light_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name) elif isinstance(spec, mujoco.MjsCamera): self.prefix = 'cam_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name) elif isinstance(spec, mujoco.MjsTendon): self.prefix = 'ten_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name) elif isinstance(spec, mujoco.MjsActuator): self.prefix = 'actuator_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name) elif isinstance(spec, mujoco.MjsSensor): self.prefix = 'sensor_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name) elif isinstance(spec, mujoco.MjsEquality): self.prefix = 'eq_' - idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name) else: raise ValueError('invalid spec type') - if idx < 0: - raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error - ids.append(idx) + ids.append(spec.id) if len(ids) == 1: self.id = ids[0] else: @@ -455,11 +433,13 @@ class BindData(object): return name else: raise AttributeError('ctrl is not available for this type') - if name == 'qpos' or name == 'qvel' or name == 'qacc': + if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'): if self.prefix == 'jnt_': return name else: - raise AttributeError('qpos, qvel, qacc are not available for this type') + raise AttributeError( + 'qpos, qvel, qacc, qfrc are not available for this type' + ) else: return self.prefix + name @@ -473,7 +453,9 @@ class BindData(object): return var[..., idx, :] def __getattr__(self, name: str): - if name in ('sensordata', 'qpos', 'qvel', 'qacc'): + if name in ('sensordata', 'qpos', 'qvel', 'qacc') or ( + name.startswith('qfrc_') + ): adr = num = 0 if name == 'sensordata': adr = self.model.sensor_adr[self.id] @@ -482,7 +464,7 @@ class BindData(object): adr = self.model.jnt_qposadr[self.id] typ = self.model.jnt_type[self.id] num = sum((typ == jt) * jt.qpos_width() for jt in JointType) - elif name == 'qvel' or name == 'qacc': + elif name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'): adr = self.model.jnt_dofadr[self.id] typ = self.model.jnt_type[self.id] num = sum((typ == jt) * jt.dof_width() for jt in JointType) @@ -495,6 +477,8 @@ class BindData(object): return self._slice(self.__getname(name), slice(adr, adr + num)) else: return self._slice(self.__getname(name), adr) + elif name in ('mocap_pos', 'mocap_quat'): + return self._slice(self.__getname(name), self.model.body_mocapid[self.id]) return self._slice(self.__getname(name), self.id) def set(self, name: str, value: jax.Array) -> Data: @@ -507,7 +491,7 @@ class BindData(object): iter(value) except TypeError: value = [value] - if name in ('qpos', 'qvel', 'qacc'): + if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'): adr = num = 0 if name == 'qpos': adr = self.model.jnt_qposadr[self.id] @@ -517,16 +501,23 @@ class BindData(object): adr = self.model.jnt_dofadr[self.id] typ = self.model.jnt_type[self.id] num = sum((typ == jt) * jt.dof_width() for jt in JointType) + elif name == 'mocap_pos': + adr = self.model.body_mocapid[self.id] * 3 + num = np.ones_like(self.id, dtype=int) * 3 + elif name == 'mocap_quat': + adr = self.model.body_mocapid[self.id] * 4 + num = np.ones_like(self.id, dtype=int) * 4 if not isinstance(self.id, list): adr = [adr] num = [num] elif isinstance(self.id, list): - adr = self.id * dim + adr = (np.array(self.id) * dim).tolist() num = [dim for _ in range(len(self.id))] else: adr = [self.id * dim] num = [dim] i = 0 + value = jax.numpy.array(value).flatten() for a, n in zip(adr, num): shape = array.shape array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape) @@ -571,20 +562,20 @@ def contact_force( m: Model, d: Data, contact_id: int, to_world_frame: bool = False ) -> jax.Array: """Extract 6D force:torque for one contact, in contact frame by default.""" - efc_address = d.contact.efc_address[contact_id] - condim = d.contact.dim[contact_id] + efc_address = d._impl.contact.efc_address[contact_id] # pytype: disable=attribute-error + condim = d._impl.contact.dim[contact_id] # pytype: disable=attribute-error if m.opt.cone == ConeType.PYRAMIDAL: force = _decode_pyramid( - d.efc_force[efc_address:], d.contact.friction[contact_id], condim + d._impl.efc_force[efc_address:], d._impl.contact.friction[contact_id], condim # pytype: disable=attribute-error ) elif m.opt.cone == ConeType.ELLIPTIC: - force = d.efc_force[efc_address : efc_address + condim] + force = d._impl.efc_force[efc_address : efc_address + condim] # pytype: disable=attribute-error force = jp.concatenate([force, jp.zeros((6 - condim))]) else: raise ValueError(f'Unknown cone type: {m.opt.cone}') if to_world_frame: - force = force.reshape((-1, 3)) @ d.contact.frame[contact_id] + force = force.reshape((-1, 3)) @ d._impl.contact.frame[contact_id] # pytype: disable=attribute-error force = force.reshape(-1) return force * (efc_address >= 0) @@ -595,21 +586,21 @@ def contact_force_dim( ) -> Tuple[jax.Array, np.ndarray]: """Extract 6D force:torque for contacts with dimension dim.""" # valid contact and condim indices - idx_dim = (d.contact.efc_address >= 0) & (d.contact.dim == dim) + idx_dim = (d._impl.contact.efc_address >= 0) & (d._impl.contact.dim == dim) # pytype: disable=attribute-error # contact force from efc if m.opt.cone == ConeType.PYRAMIDAL: efc_address = ( - d.contact.efc_address[idx_dim, None] + d._impl.contact.efc_address[idx_dim, None] # pytype: disable=attribute-error + np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None] ) - efc_force = d.efc_force[efc_address] + efc_force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))( - efc_force, d.contact.friction[idx_dim], dim + efc_force, d._impl.contact.friction[idx_dim], dim # pytype: disable=attribute-error ) elif m.opt.cone == ConeType.ELLIPTIC: - efc_address = d.contact.efc_address[idx_dim, None] + np.arange(dim)[None] - force = d.efc_force[efc_address] + efc_address = d._impl.contact.efc_address[idx_dim, None] + np.arange(dim)[None] # pytype: disable=attribute-error + force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error force = jp.hstack([force, jp.zeros((force.shape[0], 6 - dim))]) else: raise ValueError(f'Unknown cone type: {m.opt.cone}.') diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 5948f9e1..3303aa62 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -173,18 +173,20 @@ class SupportTest(parameterized.TestCase): + + - - - + + + - - - + + + """ @@ -247,6 +249,11 @@ class SupportTest(parameterized.TestCase): dx.bind(mx, s.joints[i]).qacc, d.qacc[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6 ) + np.testing.assert_array_almost_equal( + dx.bind(mx, s.joints[i]).qfrc_actuator, + d.qfrc_actuator[m.jnt_dofadr[i] : m.jnt_dofadr[i] + dofnum[i]], + decimal=6, + ) np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl) for i in range(m.nu): @@ -308,30 +315,49 @@ class SupportTest(parameterized.TestCase): np.testing.assert_array_equal( dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0] ) + dx10 = dx.bind(mx, s.bodies[0:2]).set( + 'xfrc_applied', + np.array([np.array([1, 1, 1, 1, 1, 1]), np.array([2, 2, 2, 2, 2, 2])]), + ) + np.testing.assert_array_equal( + dx10.bind(mx, s.bodies[0]).xfrc_applied, [1, 1, 1, 1, 1, 1] + ) + np.testing.assert_array_equal( + dx10.bind(mx, s.bodies[1]).xfrc_applied, [2, 2, 2, 2, 2, 2] + ) + dx11 = dx.bind(mx, s.bodies[-1]).set( + 'mocap_pos', + [1, 2, 3], + ) + np.testing.assert_array_equal( + dx11.bind(mx, s.bodies[-2]).mocap_pos, [100, 110, 120] + ) + np.testing.assert_array_equal( + dx11.bind(mx, s.bodies[-1]).mocap_pos, [1, 2, 3] + ) - # test invalid name - with self.assertRaises( - AttributeError, msg='ctrl is not available for this type' + # test attribute and type mismatches + with self.assertRaisesRegex( + AttributeError, 'ctrl is not available for this type' ): print(dx.bind(mx, s.geoms).ctrl) - with self.assertRaises( - KeyError, msg='actuator_actuator_ctrl' - ): + with self.assertRaises(KeyError): print(dx.bind(mx, s.actuators).actuator_ctrl) - with self.assertRaises( - AttributeError, msg='actuator_actuator_ctrl' + with self.assertRaisesRegex( + AttributeError, + "'Data' object has no attribute 'actuator_actuator_ctrl'", ): print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3])) - with self.assertRaises( - AttributeError, msg='qpos, qvel, qacc are not available for this type' + with self.assertRaisesRegex( + AttributeError, 'qpos, qvel, qacc, qfrc are not available for this type' ): print(dx.bind(mx, s.geoms).qpos) - with self.assertRaises(KeyError, msg='invalid name: invalid_actuator_name'): - s.actuators[0].name = 'invalid_actuator_name' - print(dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3])) - with self.assertRaises(KeyError, msg='invalid name: invalid_geom_name'): - s.geoms[0].name = 'invalid_geom_name' - print(mx.bind(s.geoms).pos) + + # test that modified names do not raise an error + s.actuators[0].name = 'modified_actuator_name' + np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl) + s.geoms[0].name = 'modified_geom_name' + np.testing.assert_array_equal(mx.bind(s.geoms[0]).pos, m.geom_pos[0, :]) # test batched data batch_size = 16 @@ -343,12 +369,15 @@ class SupportTest(parameterized.TestCase): vdx.bind(mx, s.bodies[i]).xpos, [d.xpos[i, :]] * batch_size ) - # test emtpy name + # test that adding a body requires recompilation s.worldbody.add_body() - m = s.compile() - mx = mjx.put_model(m) - with self.assertRaises(KeyError, msg='cannot bind spec with empty name'): + with self.assertRaises(ValueError) as e: mx.bind(s.bodies) + self.assertEqual( + str(e.exception), + 'mjSpec signature does not match mjx.Model signature:' + ' 15297169659434471387 != 2785811613804955188', + ) _CONTACTS = """ @@ -383,7 +412,7 @@ class SupportTest(parameterized.TestCase): # map MJX contacts to MJ ones def _find(g): - val = (g == dx.contact.geom).sum(axis=1) + val = (g == dx._impl.contact.geom).sum(axis=1) return np.where(val == 2)[0][0] contact_id_map = {i: _find(d.contact.geom[i]) for i in range(d.ncon)} @@ -397,7 +426,7 @@ class SupportTest(parameterized.TestCase): np.testing.assert_allclose(result, force, rtol=1e-5, atol=2) # check for zeros after first condim elements - condim = dx.contact.dim[j] + condim = dx._impl.contact.dim[j] if condim < 6: np.testing.assert_allclose(force[condim:], 0, rtol=1e-5, atol=1e-5) @@ -410,8 +439,8 @@ class SupportTest(parameterized.TestCase): ), )(mx, dx, j, True) # back to contact frame - force = force.at[:3].set(dx.contact.frame[j] @ force[:3]) - force = force.at[3:].set(dx.contact.frame[j] @ force[3:]) + force = force.at[:3].set(dx._impl.contact.frame[j] @ force[:3]) + force = force.at[3:].set(dx._impl.contact.frame[j] @ force[3:]) np.testing.assert_allclose(result, force, rtol=1e-5, atol=2) def test_wrap_inside(self): diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index f91ee9af..142a2eaa 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -106,20 +106,20 @@ def benchmark( def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray: - """Returns a sort order such that dx.efc_*[order][:d.nefc] == d.efc_*.""" + """Returns a sort order such that dx.efc_*[order][:d._impl.nefc] == d.efc_*.""" # pytype: disable=attribute-error # reorder efc rows to skip inactive constraints and match contact order - efl = dx.ne + dx.nf + dx.nl + efl = dx._impl.ne + dx._impl.nf + dx._impl.nl # pytype: disable=attribute-error order = np.arange(efl) - order[(dx.efc_J[:efl] == 0).all(axis=1)] = 2**16 # move empty rows to end - for i in range(dx.ncon): - num_rows = dx.contact.dim[i] - if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: - num_rows = (dx.contact.dim[i] - 1) * 2 - if dx.contact.dist[i] > 0: # move empty contacts to end + order[(dx._impl.efc_J[:efl] == 0).all(axis=1)] = 2**16 # move empty rows to end # pytype: disable=attribute-error + for i in range(dx._impl.ncon): # pytype: disable=attribute-error + num_rows = dx._impl.contact.dim[i] # pytype: disable=attribute-error + if dx._impl.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: # pytype: disable=attribute-error + num_rows = (dx._impl.contact.dim[i] - 1) * 2 # pytype: disable=attribute-error + if dx._impl.contact.dist[i] > 0: # move empty contacts to end # pytype: disable=attribute-error order = np.append(order, np.repeat(2**16, num_rows)) continue - contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1) - contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1) + contact_match = (d.contact.geom == dx._impl.contact.geom[i]).all(axis=-1) # pytype: disable=attribute-error + contact_match &= (d.contact.pos == dx._impl.contact.pos[i]).all(axis=-1) # pytype: disable=attribute-error assert contact_match.any(), f'contact {i} not found' contact_id = np.nonzero(contact_match)[0][0] order = np.append(order, np.repeat(efl + contact_id, num_rows)) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 228f3f3d..a60a1dd6 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -14,20 +14,33 @@ # ============================================================================== """Base types used in MJX.""" -import dataclasses import enum -from typing import Tuple +from typing import Tuple, Union +import warnings + import jax import mujoco from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member import numpy as np -def _restricted_to(platform: str): - """Specifies whether a field exists in only MuJoCo or MJX.""" - if platform not in ('mujoco', 'mjx'): - raise ValueError(f'unknown platform: {platform}') - return dataclasses.field(metadata={'restricted_to': platform}) +class BackendImpl(enum.Enum): + """Backend implementation to use.""" + + C = 'c' + JAX = 'jax' + WARP = 'warp' + + @classmethod + def _missing_(cls, value): + # This method is called only when lookup by value fails + # (e.g., BackendImpl('JAX') fails initially because 'JAX' != 'jax') + if not isinstance(value, str): + return None + for member in cls: + if member.value == value.lower(): + return member + return None class DisableBit(enum.IntFlag): @@ -65,6 +78,17 @@ class DisableBit(enum.IntFlag): # unsupported: MIDPHASE +class EnableBit(enum.IntFlag): + """Enable optional feature bitflags. + + Members: + INVDISCRETE: discrete-time inverse dynamics + """ + + INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE + # unsupported: OVERRIDE, ENERGY, FWDINV, MULTICCD, ISLAND + + class JointType(enum.IntEnum): """Type of degree of freedom. @@ -359,6 +383,7 @@ class SensorType(enum.IntEnum): TORQUE: torque ACTUATORFRC: scalar actuator force JOINTACTFRC: scalar actuator force, measured at the joint + TENDONACTFRC: scalar actuator force, measured at the tendon FRAMELINACC: 3D linear acceleration FRAMEANGACC: 3D angular acceleration """ @@ -393,6 +418,7 @@ class SensorType(enum.IntEnum): TORQUE = mujoco.mjtSensor.mjSENS_TORQUE ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC + TENDONACTFRC = mujoco.mjtSensor.mjSENS_TENDONACTFRC FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC @@ -417,77 +443,6 @@ class ObjType(PyTreeNode): CAMERA = mujoco.mjtObj.mjOBJ_CAMERA -class Option(PyTreeNode): - """Physics options. - - Attributes: - timestep: timestep - apirate: update rate for remote API (Hz) (not used) - impratio: ratio of friction-to-normal contact impedance - tolerance: main solver tolerance - ls_tolerance: CG/Newton linesearch tolerance - noslip_tolerance: noslip solver tolerance (not used) - ccd_tolerance: CCD solver tolerance (not used) - gravity: gravitational acceleration (3,) - wind: wind (for lift, drag and viscosity) - magnetic: global magnetic flux (not used) - density: density of medium - viscosity: viscosity of medium - o_margin: contact solver override: margin (not used) - o_solref: contact solver override: solref (not used) - o_solimp: contact solver override: solimp (not used) - o_friction[5]: contact solver override: friction (not used) - has_fluid_params: automatically set by mjx if wind/density/viscosity are - nonzero. Not used by mj - integrator: integration mode - cone: type of friction cone - jacobian: matrix layout for mass matrices (dense or sparse) - (note that this is different from MuJoCo, where jacobian - specifies whether efc_J and its accompanying matrices - are dense or sparse. - solver: solver algorithm - iterations: number of main solver iterations - ls_iterations: maximum number of CG/Newton linesearch iterations - noslip_iterations: maximum number of noslip solver iterations (not used) - ccd_iterations: maximum number of CCD solver iterations (not used) - disableflags: bit flags for disabling standard features - enableflags: bit flags for enabling optional features (not used) - disableactuator: bit flags for disabling actuators by group id (not used) - sdf_initpoints: number of starting points for gradient descent (not used) - sdf_iterations: max number of iterations for gradient descent (not used) - """ # fmt: skip - timestep: jax.Array - apirate: jax.Array = _restricted_to('mujoco') - impratio: jax.Array - tolerance: jax.Array - ls_tolerance: jax.Array - noslip_tolerance: jax.Array = _restricted_to('mujoco') - ccd_tolerance: jax.Array = _restricted_to('mujoco') - gravity: jax.Array - wind: jax.Array - magnetic: jax.Array - density: jax.Array - viscosity: jax.Array - o_margin: jax.Array - o_solref: jax.Array - o_solimp: jax.Array - o_friction: jax.Array - has_fluid_params: bool = _restricted_to('mjx') - integrator: IntegratorType - cone: ConeType - jacobian: JacobianType - solver: SolverType - iterations: int - ls_iterations: int - noslip_iterations: int = _restricted_to('mujoco') - ccd_iterations: int = _restricted_to('mujoco') - disableflags: DisableBit - enableflags: int - disableactuator: int - sdf_initpoints: int = _restricted_to('mujoco') - sdf_iterations: int = _restricted_to('mujoco') - - class Statistic(PyTreeNode): """Model statistics (in qpos0). @@ -506,367 +461,146 @@ class Statistic(PyTreeNode): center: jax.Array -class Model(PyTreeNode): - """Static model of the scene that remains unchanged with each physics step. +class Option(PyTreeNode): + """Physics options.""" # fmt: skip + timestep: jax.Array + impratio: jax.Array + tolerance: jax.Array + ls_tolerance: jax.Array + gravity: jax.Array + wind: jax.Array + magnetic: jax.Array + density: jax.Array + viscosity: jax.Array + o_margin: jax.Array + o_solref: jax.Array + o_solimp: jax.Array + o_friction: jax.Array + integrator: IntegratorType + cone: ConeType + jacobian: JacobianType + solver: SolverType + iterations: int + ls_iterations: int + disableflags: DisableBit + enableflags: int + disableactuator: int + sdf_initpoints: int - 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 - nbvh: number of total bounding volumes in all bodies - nbvhstatic: number of static bounding volumes (aabb stored in mjModel) - nbvhdynamic: number of dynamic bounding volumes (aabb stored in mjData) - njnt: number of joints - ngeom: number of geoms - nsite: number of sites - ncam: number of cameras - nlight: number of lights - nflex: number of flexes - nflexvert: number of vertices in all flexes - nflexedge: number of edges in all flexes - nflexelem: number of elements in all flexes - nflexelemdata: number of element vertex ids in all flexes - nflexshelldata: number of shell fragment vertex ids in all flexes - nflexevpair: number of element-vertex pairs in all flexes - nflextexcoord: number of vertices with texture coordinates - nmesh: number of meshes - nmeshvert: number of vertices in all meshes - nmeshnormal: number of normals in all meshes - nmeshtexcoord: number of texcoords in all meshes - nmeshface: number of triangular faces in all meshes - nmeshgraph: number of ints in mesh auxiliary data - nhfield: number of heightfields - nhfielddata: number of data points in all heightfields - ntex: number of textures - ntexdata: number of bytes in texture rgb data - nmat: number of materials - npair: number of predefined geom pairs - nexclude: number of excluded geom pairs - neq: number of equality constraints - ntendon: number of tendons - nwrap: number of wrap objects in all tendon paths - nsensor: number of sensors - nnumeric: number of numeric custom fields - ntuple: number of tuple custom fields - nkey: number of keyframes - nmocap: number of mocap bodies - nM: number of non-zeros in sparse inertia matrix - nD: number of non-zeros in sparse dof-dof matrix - nB: number of non-zeros in sparse body-dof matrix - nC: number of non-zeros in sparse reduced dof-dof matrix - nD: number of non-zeros in sparse dof-dof matrix - nJmom: number of non-zeros in sparse actuator_moment matrix - ntree: number of kinematic trees under world body - ngravcomp: number of bodies with nonzero gravcomp - nuserdata: size of userdata array - nsensordata: number of mjtNums in sensor data vector - narena: number of bytes in the mjData arena (inclusive of stack) - 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_treeid: id of body's kinematic tree; -1: static (nbody,) - body_geomnum: number of geoms (nbody,) - body_geomadr: start addr of geoms; -1: no geoms (nbody,) - body_simple: 1: diag M; 2: diag M, sliders only (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_gravcomp: antigravity force, units of body weight (nbody,) - body_margin: MAX over all geom margins (nbody,) - body_contype: OR over all geom contypes (nbody,) - body_conaffinity: OR over all geom conaffinities (nbody,) - body_bvhadr: address of bvh root (nbody,) - body_bvhnum: number of bounding volumes (nbody,) - bvh_child: left and right children in tree (nbvh, 2) - bvh_nodeid: geom or elem id of node; -1: non-leaf (nbvh,) - bvh_aabb: local bounding box (center, size) (nbvhstatic, 6) - 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_actfrclimited: does joint have actuator force limits (njnt,) - jnt_actgravcomp: is gravcomp force applied via actuators (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_treeid: id of dof's kinematic tree (nv,) - dof_Madr: dof address in M-diagonal (nv,) - dof_simplenum: number of consecutive simple dofs (nv,) - dof_solref: constraint solver reference:frictionloss (nv, mjNREF) - dof_solimp: constraint solver impedance:frictionloss (nv, mjNIMP) - dof_frictionloss: dof friction loss (nv,) - dof_hasfrictionloss: dof has >0 frictionloss (MJX) (nv,) - dof_armature: dof armature inertia/mass (nv,) - dof_damping: damping coefficient (nv,) - dof_invweight0: diag. inverse inertia in qpos0 (nv,) - dof_M0: diag. inertia in qpos0 (nv,) - geom_type: geometric type (mjtGeom) (ngeom,) - geom_contype: geom contact type (ngeom,) - geom_conaffinity: geom contact affinity (ngeom,) - geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) - geom_bodyid: id of geom's body (ngeom,) - geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) - geom_group: group for visibility (ngeom,) - geom_matid: material id for rendering (ngeom,) - geom_priority: geom contact priority (ngeom,) - geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) - geom_solref: constraint solver reference: contact (ngeom, mjNREF) - geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) - geom_size: geom-specific size parameters (ngeom, 3) - geom_aabb: bounding box, (center, size) (ngeom, 6) - geom_rbound: radius of bounding sphere (ngeom,) - geom_rbound_hfield: static rbound for hfield grid bounds (ngeom,) - 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 dist0 frictionloss (MJX) (ntendon,) - wrap_type: wrap object type (mjtWrap) (nwrap,) - wrap_objid: object id: geom, site, joint (nwrap,) - wrap_prm: divisor, joint coef, or site id (nwrap,) - wrap_inside_maxiter: maximum iterations for wrap_inside - wrap_inside_tolerance: tolerance for wrap_inside - wrap_inside_z_init: initialization for wrap_inside - is_wrap_inside: spatial tendon sidesite inside geom (nwrapinside,) - 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_group: group for visibility (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_actearly: step activation before force (nu,) - 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) - actuator_cranklength: crank length for slider-crank (nu,) - actuator_acc0: acceleration from unit force in qpos0 (nu,) - actuator_lengthrange: feasible actuator length range (nu, 2) - sensor_type: sensor type (mjtSensor) (nsensor,) - sensor_datatype: numeric data type (mjtDataType) (nsensor,) - sensor_needstage: required compute stage (mjtStage) (nsensor,) - sensor_objtype: type of sensorized object (mjtObj) (nsensor,) - sensor_objid: id of sensorized object (nsensor,) - sensor_reftype: type of reference frame (mjtObj) (nsensor,) - sensor_refid: id of reference frame; -1: global frame (nsensor,) - sensor_dim: number of scalar outputs (nsensor,) - sensor_adr: address in sensor array (nsensor,) - sensor_cutoff: cutoff for real and positive; 0: ignore (nsensor,) - numeric_adr: address of field in numeric_data (nnumeric,) - numeric_data: array of all numeric fields (nnumericdata,) - tuple_adr: address of text in text_data (ntuple,) - tuple_size: number of objects in tuple (ntuple,) - tuple_objtype: array of object types in all tuples (ntupledata,) - tuple_objid: array of object ids in all tuples (ntupledata,) - tuple_objprm: array of object params in all tuples (ntupledata,) - key_time: key time (nkey,) - key_qpos: key position (nkey, nq) - key_qvel: key velocity (nkey, nv) - key_act: key activation (nkey, na) - key_mpos: key mocap position (nkey, nmocap, 3) - key_mquat: key mocap quaternion (nkey, nmocap, 4) - key_ctrl: key control (nkey, nu) - name_bodyadr: body name pointers (nbody,) - name_jntadr: joint name pointers (njnt,) - name_geomadr: geom name pointers (ngeom,) - name_siteadr: site name pointers (nsite,) - name_camadr: camera name pointers (ncam,) - name_meshadr: mesh name pointers (nmesh,) - name_pairadr: geom pair name pointers (npair,) - name_eqadr: equality constraint name pointers (neq,) - name_tendonadr: tendon name pointers (ntendon,) - name_actuatoradr: actuator name pointers (nu,) - name_sensoradr: sensor name pointers (nsensor,) - name_numericadr: numeric name pointers (nnumeric,) - name_tupleadr: tuple name pointers (ntuple,) - name_keyadr: keyframe name pointers (nkey,) - names: names of all objects, 0-terminated (nnames,) - """ + +class OptionC(Option): + """C-specific option.""" + + apirate: jax.Array + noslip_tolerance: jax.Array + ccd_tolerance: jax.Array + noslip_iterations: int + ccd_iterations: int + sdf_iterations: int + + +class OptionJAX(Option): + """JAX-specific option.""" + + has_fluid_params: bool + + +class ModelC(PyTreeNode): + """CPU-specific model data.""" + + nbvh: jax.Array + nbvhstatic: jax.Array + nbvhdynamic: jax.Array + nflex: jax.Array + nflexvert: jax.Array + nflexedge: jax.Array + nflexelem: jax.Array + nflexelemdata: jax.Array + nflexshelldata: jax.Array + nflexevpair: jax.Array + nflextexcoord: jax.Array + nplugin: jax.Array + ntree: jax.Array + narena: jax.Array + body_bvhadr: jax.Array + body_bvhnum: jax.Array + bvh_child: jax.Array + bvh_nodeid: jax.Array + bvh_aabb: jax.Array + geom_plugin: jax.Array + light_bodyid: jax.Array + light_targetbodyid: jax.Array + flex_contype: jax.Array + flex_conaffinity: jax.Array + flex_condim: jax.Array + flex_priority: jax.Array + flex_solmix: jax.Array + flex_solref: jax.Array + flex_solimp: jax.Array + flex_friction: jax.Array + flex_margin: jax.Array + flex_gap: jax.Array + flex_internal: jax.Array + flex_selfcollide: jax.Array + flex_activelayers: jax.Array + flex_dim: jax.Array + flex_vertadr: jax.Array + flex_vertnum: jax.Array + flex_edgeadr: jax.Array + flex_edgenum: jax.Array + flex_elemadr: jax.Array + flex_elemnum: jax.Array + flex_elemdataadr: jax.Array + flex_evpairadr: jax.Array + flex_evpairnum: jax.Array + flex_vertbodyid: jax.Array + flex_edge: jax.Array + flex_elem: jax.Array + flex_elemlayer: jax.Array + flex_evpair: jax.Array + flex_vert: jax.Array + flexedge_length0: jax.Array + flexedge_invweight0: jax.Array + flex_radius: jax.Array + flex_edgestiffness: jax.Array + flex_edgedamping: jax.Array + flex_edgeequality: jax.Array + flex_rigid: jax.Array + flexedge_rigid: jax.Array + flex_centered: jax.Array + flex_bvhadr: jax.Array + flex_bvhnum: jax.Array + actuator_plugin: jax.Array + sensor_plugin: jax.Array + plugin: jax.Array + + +class ModelJAX(PyTreeNode): + """JAX-specific model data.""" + + dof_hasfrictionloss: np.ndarray + geom_rbound_hfield: np.ndarray + mesh_convex: Tuple[ConvexMesh, ...] + tendon_hasfrictionloss: np.ndarray + wrap_inside_maxiter: int + wrap_inside_tolerance: float + wrap_inside_z_init: float + is_wrap_inside: np.ndarray + + +class Model(PyTreeNode): + """Static model of the scene that remains unchanged with each physics step.""" nq: int nv: int nu: int na: int nbody: int - nbvh: int = _restricted_to('mujoco') - nbvhstatic: int = _restricted_to('mujoco') - nbvhdynamic: int = _restricted_to('mujoco') njnt: int ngeom: int nsite: int ncam: int nlight: int - nflex: int = _restricted_to('mujoco') - nflexvert: int = _restricted_to('mujoco') - nflexedge: int = _restricted_to('mujoco') - nflexelem: int = _restricted_to('mujoco') - nflexelemdata: int = _restricted_to('mujoco') - nflexshelldata: int = _restricted_to('mujoco') - nflexevpair: int = _restricted_to('mujoco') - nflextexcoord: int = _restricted_to('mujoco') nmesh: int nmeshvert: int nmeshnormal: int @@ -893,11 +627,9 @@ class Model(PyTreeNode): nC: int # pylint:disable=invalid-name nD: int # pylint:disable=invalid-name nJmom: int # pylint:disable=invalid-name - ntree: int = _restricted_to('mujoco') ngravcomp: int nuserdata: int nsensordata: int - narena: int = _restricted_to('mujoco') opt: Option stat: Statistic qpos0: jax.Array @@ -926,11 +658,6 @@ class Model(PyTreeNode): body_margin: np.ndarray body_contype: np.ndarray body_conaffinity: np.ndarray - body_bvhadr: np.ndarray = _restricted_to('mujoco') - body_bvhnum: np.ndarray = _restricted_to('mujoco') - bvh_child: np.ndarray = _restricted_to('mujoco') - bvh_nodeid: np.ndarray = _restricted_to('mujoco') - bvh_aabb: np.ndarray = _restricted_to('mujoco') body_invweight0: jax.Array jnt_type: np.ndarray jnt_qposadr: np.ndarray @@ -956,7 +683,6 @@ class Model(PyTreeNode): dof_solref: jax.Array dof_solimp: jax.Array dof_frictionloss: jax.Array - dof_hasfrictionloss: np.ndarray = _restricted_to('mjx') dof_armature: jax.Array dof_damping: jax.Array dof_invweight0: jax.Array @@ -977,7 +703,6 @@ class Model(PyTreeNode): geom_size: jax.Array geom_aabb: np.ndarray geom_rbound: jax.Array - geom_rbound_hfield: np.ndarray = _restricted_to('mjx') geom_pos: jax.Array geom_quat: jax.Array geom_friction: jax.Array @@ -1004,9 +729,7 @@ class Model(PyTreeNode): cam_sensorsize: np.ndarray cam_intrinsic: np.ndarray light_mode: np.ndarray - light_bodyid: np.ndarray = _restricted_to('mujoco') - light_targetbodyid: np.ndarray = _restricted_to('mujoco') - light_directional: jax.Array + light_type: jax.Array light_castshadow: jax.Array light_pos: jax.Array light_dir: jax.Array @@ -1014,46 +737,6 @@ class Model(PyTreeNode): light_pos0: np.ndarray light_dir0: np.ndarray light_cutoff: jax.Array - flex_contype: np.ndarray = _restricted_to('mujoco') - flex_conaffinity: np.ndarray = _restricted_to('mujoco') - flex_condim: np.ndarray = _restricted_to('mujoco') - flex_priority: np.ndarray = _restricted_to('mujoco') - flex_solmix: np.ndarray = _restricted_to('mujoco') - flex_solref: np.ndarray = _restricted_to('mujoco') - flex_solimp: np.ndarray = _restricted_to('mujoco') - flex_friction: np.ndarray = _restricted_to('mujoco') - flex_margin: np.ndarray = _restricted_to('mujoco') - flex_gap: np.ndarray = _restricted_to('mujoco') - flex_internal: np.ndarray = _restricted_to('mujoco') - flex_selfcollide: np.ndarray = _restricted_to('mujoco') - flex_activelayers: np.ndarray = _restricted_to('mujoco') - flex_dim: np.ndarray = _restricted_to('mujoco') - flex_vertadr: np.ndarray = _restricted_to('mujoco') - flex_vertnum: np.ndarray = _restricted_to('mujoco') - flex_edgeadr: np.ndarray = _restricted_to('mujoco') - flex_edgenum: np.ndarray = _restricted_to('mujoco') - flex_elemadr: np.ndarray = _restricted_to('mujoco') - flex_elemnum: np.ndarray = _restricted_to('mujoco') - flex_elemdataadr: np.ndarray = _restricted_to('mujoco') - flex_evpairadr: np.ndarray = _restricted_to('mujoco') - flex_evpairnum: np.ndarray = _restricted_to('mujoco') - flex_vertbodyid: np.ndarray = _restricted_to('mujoco') - flex_edge: np.ndarray = _restricted_to('mujoco') - flex_elem: np.ndarray = _restricted_to('mujoco') - flex_elemlayer: np.ndarray = _restricted_to('mujoco') - flex_evpair: np.ndarray = _restricted_to('mujoco') - flex_vert: np.ndarray = _restricted_to('mujoco') - flexedge_length0: np.ndarray = _restricted_to('mujoco') - flexedge_invweight0: np.ndarray = _restricted_to('mujoco') - flex_radius: np.ndarray = _restricted_to('mujoco') - flex_edgestiffness: np.ndarray = _restricted_to('mujoco') - flex_edgedamping: np.ndarray = _restricted_to('mujoco') - flex_edgeequality: np.ndarray = _restricted_to('mujoco') - flex_rigid: np.ndarray = _restricted_to('mujoco') - flexedge_rigid: np.ndarray = _restricted_to('mujoco') - flex_centered: np.ndarray = _restricted_to('mujoco') - flex_bvhadr: np.ndarray = _restricted_to('mujoco') - flex_bvhnum: np.ndarray = _restricted_to('mujoco') mesh_vertadr: np.ndarray mesh_vertnum: np.ndarray mesh_faceadr: np.ndarray @@ -1065,7 +748,6 @@ class Model(PyTreeNode): mesh_graph: np.ndarray mesh_pos: np.ndarray mesh_quat: np.ndarray - mesh_convex: Tuple[ConvexMesh, ...] = _restricted_to('mjx') mesh_texcoordadr: np.ndarray mesh_texcoordnum: np.ndarray mesh_texcoord: np.ndarray @@ -1104,26 +786,24 @@ class Model(PyTreeNode): tendon_adr: np.ndarray tendon_num: np.ndarray tendon_limited: np.ndarray + tendon_actfrclimited: np.ndarray tendon_solref_lim: jax.Array tendon_solimp_lim: jax.Array tendon_solref_fri: jax.Array tendon_solimp_fri: jax.Array tendon_range: jax.Array + tendon_actfrcrange: jax.Array tendon_margin: jax.Array tendon_stiffness: jax.Array tendon_damping: jax.Array + tendon_armature: jax.Array tendon_frictionloss: jax.Array tendon_lengthspring: jax.Array tendon_length0: jax.Array tendon_invweight0: jax.Array - tendon_hasfrictionloss: np.ndarray = _restricted_to('mjx') wrap_type: np.ndarray wrap_objid: np.ndarray wrap_prm: np.ndarray - wrap_inside_maxiter: int = _restricted_to('mjx') - wrap_inside_tolerance: float = _restricted_to('mjx') - wrap_inside_z_init: float = _restricted_to('mjx') - is_wrap_inside: np.ndarray = _restricted_to('mjx') actuator_trntype: np.ndarray actuator_dyntype: np.ndarray actuator_gaintype: np.ndarray @@ -1146,7 +826,6 @@ class Model(PyTreeNode): actuator_cranklength: np.ndarray actuator_acc0: jax.Array actuator_lengthrange: np.ndarray - actuator_plugin: np.ndarray = _restricted_to('mujoco') sensor_type: np.ndarray sensor_datatype: np.ndarray sensor_needstage: np.ndarray @@ -1187,7 +866,38 @@ class Model(PyTreeNode): name_tupleadr: np.ndarray name_keyadr: np.ndarray names: bytes + signature: np.uint64 _sizes: jax.Array + _impl: Union[ModelC, ModelJAX] + + @property + def backend_impl(self) -> BackendImpl: + return { + ModelC: BackendImpl.C, + ModelJAX: BackendImpl.JAX, + }[type(self._impl)] + + def __getattr__(self, name: str): + if name == 'value': + # Special case for NNX, the value attribute may not exist on the parent + # PyTreeNode, before it exists on the child PyTreeNode. Thanks NNX. + return object.__getattribute__(self, 'value') + + try: + impl_instsance = object.__getattribute__(self, '_impl') + val = getattr(impl_instsance, name) + warnings.warn( + f'Accessing `{name}` directly from `Model` is deprecated. ' + f'Access it via `model._impl.{name}` instead.', + DeprecationWarning, + stacklevel=2, + ) + except AttributeError: + # raise the standard exception + raise AttributeError( # pylint: disable=raise-missing-from + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + return val class Contact(PyTreeNode): @@ -1225,144 +935,140 @@ class Contact(PyTreeNode): efc_address: np.ndarray -class Data(PyTreeNode): - r"""\Dynamic state that updates each step. +class DataC(PyTreeNode): + """C-specific data.""" - Attributes: - ne: number of equality constraints - nf: number of friction constraints - nl: number of limit constraints - nefc: number of constraints - ncon: number of contacts - solver_niter: number of solver iterations - 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,) - mocap_pos: positions of mocap bodies (nmocap x 3) - mocap_quat: orientations of mocap bodies (nmocap x 4) - qacc: acceleration (nv,) - act_dot: time-derivative of actuator activation (na,) - userdata: user data, not touched by engine (nuserdata,) - sensordata: sensor data array (nsensordata,) - xpos: Cartesian position of body frame (nbody, 3) - xquat: Cartesian orientation of body frame (nbody, 4) - xmat: Cartesian orientation of body frame (nbody, 3, 3) - xipos: Cartesian position of body com (nbody, 3) - ximat: Cartesian orientation of body inertia (nbody, 3, 3) - xanchor: Cartesian position of joint anchor (njnt, 3) - xaxis: Cartesian joint axis (njnt, 3) - geom_xpos: Cartesian geom position (ngeom, 3) - geom_xmat: Cartesian geom orientation (ngeom, 3, 3) - site_xpos: Cartesian site position (nsite, 3) - site_xmat: Cartesian site orientation (nsite, 3, 3) - cam_xpos: Cartesian camera position (ncam, 3) - cam_xmat: Cartesian camera orientation (ncam, 3, 3) - light_xpos: Cartesian light position (nlight, 3) - light_xdir: Cartesian light direction (nlight, 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) - flexvert_xpos: Cartesian flex vertex positions (nflexvert, 3) - flexelem_aabb: flex element bounding boxes (center, size) (nflexelem, 6) - flexedge_J_rownnz: number of non-zeros in Jacobian row (nflexedge,) - flexedge_J_rowadr: row start address in colind array (nflexedge,) - flexedge_J_colind: column indices in sparse Jacobian (nflexedge, nv) - flexedge_J: flex edge Jacobian (nflexedge, nv) - flexedge_length: flex edge lengths (nflexedge,) - ten_wrapadr: start address of tendon's path (ntendon,) - ten_wrapnum: number of wrap points in path (ntendon,) - ten_J_rownnz: number of non-zeros in Jacobian row (ntendon,) - ten_J_rowadr: row start address in colind array (ntendon,) - ten_J_colind: column indices in sparse Jacobian (ntendon, nv) - ten_J: tendon Jacobian (ntendon, nv) - ten_length: tendon lengths (ntendon,) - wrap_obj: geom id; -1: site; -2: pulley (nwrap*2,) - wrap_xpos: Cartesian 3D points in all path (nwrap*2, 3) - actuator_length: actuator lengths (nu,) - moment_rownnz: number of non-zeros in actuator_moment row (nu,) - moment_rowadr: row start address in colind array (nu,) - moment_colind: column indices in sparse Jacobian (nJmom,) - actuator_moment: actuator moments (nJmom,) - crb: com-based composite inertia and mass (nbody, 10) - qM: total inertia if sparse: (nM,) - if dense: (nv, nv) - qLD: L'*D*L (or Cholesky) factorization of M. if sparse: (nM,) - if dense: (nv, nv) - qLDiagInv: 1/diag(D) if sparse: (nv,) - if dense: (0,) - bvh_aabb_dyn: global bounding box (center, size) (nbvhdynamic, 6) - bvh_active: volume has been added to collisions (nbvh,) - flexedge_velocity: flex edge velocities (nflexedge,) - ten_velocity: tendon velocities (ntendon,) - 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_spring: passive spring force (nv,) - qfrc_damper: passive damper force (nv,) - qfrc_gravcomp: passive gravity compensation force (nv,) - qfrc_fluid: passive fluid force (nv,) - qfrc_passive: total passive force (nv,) - subtree_linvel: linear velocity of subtree com (nbody, 3) - subtree_angmom: angular momentum about subtree com (nbody, 3) - qH: L'*D*L factorization of modified M (nM,) - qHDiagInv: 1/diag(D) of modified M (nv,) - B_rownnz: body-dof: non-zeros in each row (nbody,) - B_rowadr: body-dof: address of each row in B_colind (nbody,) - B_colind: body-dof: column indices of non-zeros (nB,) - M_rownnz: inertia: non-zeros in each row (nv,) - M_rowadr: inertia: address of each row in M_colind (nv,) - M_colind: inertia: column indices of non-zeros (nM,) - mapM2M: index mapping from M (legacy) to M (CSR) (nM,) - C_rownnz: reduced dof-dof: non-zeros in each row (nv,) - C_rowadr: reduced dof-dof: address of each row in C_colind (nv,) - C_colind: reduced dof-dof: column indices of non-zeros (nC,) - mapM2C: index mapping from M to C (nC,) - D_rownnz: dof-dof: non-zeros in each row (nv,) - D_rowadr: dof-dof: address of each row in D_colind (nv,) - D_diag: dof-dof: index of diagonal element (nv,) - D_colind: dof-dof: column indices of non-zeros (nD,) - mapM2D: index mapping from M to D (nD,) - mapD2M: index mapping from D to M (nM,) - qDeriv: d (passive + actuator - bias) / d qvel (nD,) - qLU: sparse LU of (qM - dt*qDeriv) (nD,) - 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 - cacc: com-based acceleration (nbody, 6) - cfrc_int: com-based interaction force with parent (nbody, 6) - cfrc_ext: com-based external force on body (nbody, 6) - contact: all detected contacts (ncon,) - efc_type: constraint type (nefc,) - efc_J: constraint Jacobian (nefc, nv) - efc_pos: constraint position (equality, contact) (nefc,) - efc_margin: inclusion margin (contact) (nefc,) - efc_frictionloss: frictionloss (friction) (nefc,) - efc_D: constraint mass (nefc,) - efc_aref: reference pseudo-acceleration (nefc,) - efc_force: constraint force in constraint space (nefc,) - _qM_sparse: qM in sparse representation (nM,) - _qLD_sparse: qLD in sparse representation (nM,) - _qLDiagInv_sparse: qLDiagInv in sparse representation (nv,) - """ # fmt: skip # constant sizes: + # TODO(stunya): make these sizes jax.Array? ne: int nf: int nl: int nefc: int ncon: int - # solver statistics: + # TODO(stunya): remove most of these fields solver_niter: jax.Array + cdof: jax.Array + cinert: jax.Array + light_xpos: jax.Array + light_xdir: jax.Array + flexvert_xpos: jax.Array + flexelem_aabb: jax.Array + flexedge_J_rownnz: jax.Array # pylint:disable=invalid-name + flexedge_J_rowadr: jax.Array # pylint:disable=invalid-name + flexedge_J_colind: jax.Array # pylint:disable=invalid-name + flexedge_J: jax.Array # pylint:disable=invalid-name + flexedge_length: jax.Array + ten_wrapadr: jax.Array + ten_wrapnum: jax.Array + ten_J_rownnz: jax.Array # pylint:disable=invalid-name + ten_J_rowadr: jax.Array # pylint:disable=invalid-name + ten_J_colind: jax.Array # pylint:disable=invalid-name + ten_J: jax.Array # pylint:disable=invalid-name + ten_length: jax.Array + wrap_obj: jax.Array + wrap_xpos: jax.Array + actuator_length: jax.Array + moment_rownnz: jax.Array # pylint:disable=invalid-name + moment_rowadr: jax.Array # pylint:disable=invalid-name + moment_colind: jax.Array # pylint:disable=invalid-name + actuator_moment: jax.Array + crb: jax.Array + qM: jax.Array # pylint:disable=invalid-name + M: jax.Array # pylint:disable=invalid-name + qLD: jax.Array # pylint:disable=invalid-name + qLDiagInv: jax.Array # pylint:disable=invalid-name + bvh_aabb_dyn: jax.Array + bvh_active: jax.Array + # position, velocity dependent: + flexedge_velocity: jax.Array + ten_velocity: jax.Array + actuator_velocity: jax.Array + cdof_dot: jax.Array + plugin_data: jax.Array + qH: jax.Array # pylint:disable=invalid-name + qHDiagInv: jax.Array # pylint:disable=invalid-name + B_rownnz: jax.Array # pylint:disable=invalid-name + B_rowadr: jax.Array # pylint:disable=invalid-name + B_colind: jax.Array # pylint:disable=invalid-name + M_rownnz: jax.Array # pylint:disable=invalid-name + M_rowadr: jax.Array # pylint:disable=invalid-name + M_colind: jax.Array # pylint:disable=invalid-name + mapM2M: jax.Array # pylint:disable=invalid-name + D_rownnz: jax.Array # pylint:disable=invalid-name + D_rowadr: jax.Array # pylint:disable=invalid-name + D_diag: jax.Array # pylint:disable=invalid-name + D_colind: jax.Array # pylint:disable=invalid-name + mapM2D: jax.Array # pylint:disable=invalid-name + mapD2M: jax.Array # pylint:disable=invalid-name + qDeriv: jax.Array # pylint:disable=invalid-name + qLU: jax.Array # pylint:disable=invalid-name + qfrc_spring: jax.Array + qfrc_damper: jax.Array + cacc: jax.Array + cfrc_int: jax.Array + cfrc_ext: jax.Array + subtree_linvel: jax.Array + subtree_angmom: jax.Array + # dynamically sized arrays which are made static for the frontend JAX API + # TODO(stunya): remove these dynamic fields entirely + contact: Contact + efc_type: jax.Array + efc_J: jax.Array # pylint:disable=invalid-name + efc_pos: jax.Array + efc_margin: jax.Array + efc_frictionloss: jax.Array + efc_D: jax.Array # pylint:disable=invalid-name + efc_aref: jax.Array + efc_force: jax.Array + + +class DataJAX(PyTreeNode): + """JAX-specific data.""" + + ne: int + nf: int + nl: int + nefc: int + ncon: int + solver_niter: jax.Array + cdof: jax.Array + cinert: jax.Array + ten_wrapadr: jax.Array + ten_wrapnum: jax.Array + ten_J: jax.Array # pylint:disable=invalid-name + ten_length: jax.Array + wrap_obj: jax.Array + wrap_xpos: jax.Array + actuator_length: jax.Array + actuator_moment: jax.Array + crb: jax.Array + qM: jax.Array # pylint:disable=invalid-name + M: jax.Array # pylint:disable=invalid-name + qLD: jax.Array # pylint:disable=invalid-name + qLDiagInv: jax.Array # pylint:disable=invalid-name + ten_velocity: jax.Array + actuator_velocity: jax.Array + cdof_dot: jax.Array + cacc: jax.Array + cfrc_int: jax.Array + cfrc_ext: jax.Array + subtree_linvel: jax.Array + subtree_angmom: jax.Array + # dynamically sized data which are made static due to JAX limitations + contact: Contact + efc_type: jax.Array + efc_J: jax.Array # pylint:disable=invalid-name + efc_pos: jax.Array + efc_margin: jax.Array + efc_frictionloss: jax.Array + efc_D: jax.Array # pylint:disable=invalid-name + efc_aref: jax.Array + efc_force: jax.Array + + +class Data(PyTreeNode): + """Dynamic state that updates each step.""" + # global properties: time: jax.Array # state: @@ -1398,98 +1104,45 @@ class Data(PyTreeNode): site_xmat: jax.Array cam_xpos: jax.Array cam_xmat: jax.Array - light_xpos: jax.Array = _restricted_to('mujoco') - light_xdir: jax.Array = _restricted_to('mujoco') subtree_com: jax.Array - cdof: jax.Array - cinert: jax.Array - flexvert_xpos: jax.Array = _restricted_to('mujoco') - flexelem_aabb: jax.Array - flexedge_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - flexedge_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - flexedge_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - flexedge_J: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - flexedge_length: jax.Array = _restricted_to('mujoco') - ten_wrapadr: jax.Array - ten_wrapnum: jax.Array - ten_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - ten_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - ten_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - ten_J: jax.Array # pylint:disable=invalid-name - ten_length: jax.Array - wrap_obj: jax.Array - wrap_xpos: jax.Array - actuator_length: jax.Array - moment_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - moment_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - moment_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - actuator_moment: jax.Array - crb: jax.Array - qM: jax.Array # pylint:disable=invalid-name - qLD: jax.Array # pylint:disable=invalid-name - qLDiagInv: jax.Array # pylint:disable=invalid-name - bvh_aabb_dyn: jax.Array = _restricted_to('mujoco') - bvh_active: jax.Array = _restricted_to('mujoco') - # position, velocity dependent: - flexedge_velocity: jax.Array = _restricted_to('mujoco') - ten_velocity: jax.Array - actuator_velocity: jax.Array cvel: jax.Array - cdof_dot: jax.Array qfrc_bias: jax.Array - qfrc_spring: jax.Array = _restricted_to('mujoco') - qfrc_damper: jax.Array = _restricted_to('mujoco') qfrc_gravcomp: jax.Array qfrc_fluid: jax.Array qfrc_passive: jax.Array - subtree_linvel: jax.Array - subtree_angmom: jax.Array - qH: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - qHDiagInv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - M_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - M_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - M_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - mapM2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - C_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - C_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - mapM2C: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - D_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - D_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - D_diag: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - D_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - mapM2D: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - mapD2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - qDeriv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - qLU: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name - # position, velocity, control & acceleration dependent: qfrc_actuator: jax.Array actuator_force: jax.Array qfrc_smooth: jax.Array qacc_smooth: jax.Array qfrc_constraint: jax.Array qfrc_inverse: jax.Array - cacc: jax.Array - cfrc_int: jax.Array - cfrc_ext: jax.Array - # dynamically sized - contact: Contact - # dynamically sized - position dependent: - efc_type: jax.Array - efc_J: jax.Array # pylint:disable=invalid-name - efc_pos: jax.Array - efc_margin: jax.Array - efc_frictionloss: jax.Array - efc_D: jax.Array # pylint:disable=invalid-name - # dynamically sized - position & velocity dependent: - efc_aref: jax.Array - # dynamically sized - position, velocity, control & acceleration dependent: - efc_force: jax.Array - # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo - # when in dense mode - _qM_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name - _qLD_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name - _qLDiagInv_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name + _impl: Union[DataC, DataJAX] + + @property + def backend_impl(self) -> BackendImpl: + return { + DataC: BackendImpl.C, + DataJAX: BackendImpl.JAX, + }[type(self._impl)] + + def __getattr__(self, name: str): + if name == 'value': + # Special case for NNX, the value attribute may not exist on the parent + # PyTreeNode, before it exists on the child PyTreeNode. Thanks NNX. + return object.__getattribute__(self, 'value') + + try: + impl_instsance = object.__getattribute__(self, '_impl') + val = getattr(impl_instsance, name) + warnings.warn( + f'Accessing `{name}` directly from `Data` is deprecated. ' + f'Access it via `data._impl.{name}` instead.', + DeprecationWarning, + stacklevel=2, + ) + except AttributeError: + # raise the standard exception + raise AttributeError( # pylint: disable=raise-missing-from + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + return val diff --git a/mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml b/mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml new file mode 100644 index 00000000..51952b3a --- /dev/null +++ b/mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml @@ -0,0 +1,55 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index 1d102175..ee03aa0a 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name="mujoco-mjx" -version = "3.3.1" +version = "3.3.3" authors = [ {name = "Google DeepMind", email = "mujoco@deepmind.com"}, ] @@ -30,7 +30,7 @@ dependencies = [ "etils[epath]", "jax", "jaxlib", - "mujoco>=3.3.1.dev0", + "mujoco>=3.3.3.dev0", "scipy", "trimesh", ] @@ -41,9 +41,9 @@ mjx-viewer = "mujoco.mjx.viewer:main" [project.urls] Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx" -Documentation = "https://mujoco.readthedocs.io/en/3.3.1" +Documentation = "https://mujoco.readthedocs.io/en/3.3.3" Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx" -Changelog = "https://mujoco.readthedocs.io/en/3.3.1/changelog.html" +Changelog = "https://mujoco.readthedocs.io/en/3.3.3/changelog.html" [tool.isort] force_single_line = true diff --git a/mjx/requirements.txt b/mjx/requirements.txt index 3747653a..9b8cf8c1 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -76,8 +76,9 @@ scipy==1.13.1; python_version == '3.9' \ --hash=sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004 \ --hash=sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24 \ --hash=sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5 -setuptools==75.5.0 \ - --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829 +setuptools==78.1.1 \ + --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ + --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d trimesh==4.5.2 \ --hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73 wheel==0.45.0 \ diff --git a/model/flex/bunny.xml b/model/flex/bunny.xml index 64e081ef..ec31684e 100644 --- a/model/flex/bunny.xml +++ b/model/flex/bunny.xml @@ -31,7 +31,7 @@ - + diff --git a/model/flex/bunny_with_uv.xml b/model/flex/bunny_with_uv.xml index 5b9e42e4..c3cde418 100644 --- a/model/flex/bunny_with_uv.xml +++ b/model/flex/bunny_with_uv.xml @@ -37,7 +37,7 @@ - + diff --git a/model/plugin/elasticity/flag_flex.xml b/model/flex/flag.xml similarity index 82% rename from model/plugin/elasticity/flag_flex.xml rename to model/flex/flag.xml index 667a8c36..c7bfa371 100644 --- a/model/plugin/elasticity/flag_flex.xml +++ b/model/flex/flag.xml @@ -26,10 +26,6 @@ - - - - @@ -37,12 +33,7 @@ - - - - - - + diff --git a/model/plugin/elasticity/mannequin.xml b/model/flex/mannequin.xml similarity index 100% rename from model/plugin/elasticity/mannequin.xml rename to model/flex/mannequin.xml diff --git a/model/plugin/elasticity/pancake_flex.xml b/model/flex/pancake.xml similarity index 80% rename from model/plugin/elasticity/pancake_flex.xml rename to model/flex/pancake.xml index 38e839f9..8e3d1ec1 100644 --- a/model/plugin/elasticity/pancake_flex.xml +++ b/model/flex/pancake.xml @@ -16,10 +16,6 @@ - - - - diff --git a/model/plugin/elasticity/plate_flex.xml b/model/flex/plate.xml similarity index 81% rename from model/plugin/elasticity/plate_flex.xml rename to model/flex/plate.xml index b490b3ae..9278c7fe 100644 --- a/model/plugin/elasticity/plate_flex.xml +++ b/model/flex/plate.xml @@ -16,10 +16,6 @@ - - - - diff --git a/model/plugin/elasticity/poncho_flex.xml b/model/flex/poncho.xml similarity index 99% rename from model/plugin/elasticity/poncho_flex.xml rename to model/flex/poncho.xml index 72783d7e..c2462e56 100644 --- a/model/plugin/elasticity/poncho_flex.xml +++ b/model/flex/poncho.xml @@ -19,10 +19,6 @@ - - - - @@ -1418,13 +1414,8 @@ 398 399 418 398 376 378"> + - - - - - - diff --git a/model/flex/poncho_vertcollide.xml b/model/flex/poncho_vertcollide.xml new file mode 100644 index 00000000..cb75a7cd --- /dev/null +++ b/model/flex/poncho_vertcollide.xml @@ -0,0 +1,1423 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/model/flex/trampoline.xml b/model/flex/trampoline.xml index 8d436e5d..bb538201 100644 --- a/model/flex/trampoline.xml +++ b/model/flex/trampoline.xml @@ -39,7 +39,7 @@ radius=".001" mass="10" name="plate" dim="2"> - + diff --git a/model/plugin/sensor/touch_stress.xml b/model/plugin/sensor/touch_stress.xml new file mode 100644 index 00000000..025777a4 --- /dev/null +++ b/model/plugin/sensor/touch_stress.xml @@ -0,0 +1,64 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/plugin/elasticity/CMakeLists.txt b/plugin/elasticity/CMakeLists.txt index f6186261..97b2f412 100644 --- a/plugin/elasticity/CMakeLists.txt +++ b/plugin/elasticity/CMakeLists.txt @@ -22,8 +22,6 @@ set(MUJOCO_ELASTICITY_SRCS elasticity.cc elasticity.h register.cc - shell.cc - shell.h ) add_library(elasticity SHARED) diff --git a/plugin/elasticity/README.md b/plugin/elasticity/README.md index 2e86c1a7..7ccf5712 100644 --- a/plugin/elasticity/README.md +++ b/plugin/elasticity/README.md @@ -18,15 +18,3 @@ Parameters: - `bend` [Pa]: bending stiffness. - `flat` [bool]: if true, the stress-equilibrium configuration is that of a straight cable; if false or unspecified, it is the configuration defined in the XML. - `vmax` [N/m^2]: If greater than zero, the cable is colored using mechanical stresses; the value represent the maximum stress in the color scale. - -### Shell - -Implemented in [shell.cc](shell.cc). - -The shell plugin discretizes an inextensible 2D continuum. It is intended to simulate the bending of plates where the stretching is negligible compared to other deformation modes. - -Parameters: - - - `young` [Pa]: Young's modulus. - - `poisson` [Pa]: Poisson's ratio; if 0, then the material only opposed shear deformations; if near 0.5, then the material is nearly incompressible (rubber-like). - - `thickness` [m]: shell thickness, used to scale the bending stiffness. diff --git a/plugin/elasticity/register.cc b/plugin/elasticity/register.cc index edde5252..ab8d283a 100644 --- a/plugin/elasticity/register.cc +++ b/plugin/elasticity/register.cc @@ -14,13 +14,11 @@ #include #include "cable.h" -#include "shell.h" namespace mujoco::plugin::elasticity { mjPLUGIN_LIB_INIT { Cable::RegisterPlugin(); - Shell::RegisterPlugin(); } } // namespace mujoco::plugin::elasticity diff --git a/plugin/elasticity/shell.cc b/plugin/elasticity/shell.cc deleted file mode 100644 index 30b02051..00000000 --- a/plugin/elasticity/shell.cc +++ /dev/null @@ -1,271 +0,0 @@ -// 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. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include "elasticity.h" -#include "shell.h" - - -namespace mujoco::plugin::elasticity { -namespace { - -// local tetrahedron numbering -constexpr int kNumEdges = Stencil2D::kNumEdges; -constexpr int kNumVerts = Stencil2D::kNumVerts; -constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}}; - -// cotangent between two edges -mjtNum cot(mjtNum* x, int v0, int v1, int v2) { - mjtNum normal[3]; - mjtNum edge1[3]; - mjtNum edge2[3]; - - mju_sub3(edge1, x+3*v1, x+3*v0); - mju_sub3(edge2, x+3*v2, x+3*v0); - mju_cross(normal, edge1, edge2); - - return mju_dot3(edge1, edge2) / mju_norm3(normal); -} - -// area of a triangle -mjtNum ComputeVolume(const mjtNum* x, const int v[kNumVerts]) { - mjtNum normal[3]; - mjtNum edge1[3]; - mjtNum edge2[3]; - - mju_sub3(edge1, x+3*v[1], x+3*v[0]); - mju_sub3(edge2, x+3*v[2], x+3*v[0]); - mju_cross(normal, edge1, edge2); - - return mju_norm3(normal) / 2; -} - -} // namespace - -// factory function -std::optional Shell::Create(const mjModel* m, mjData* d, int instance) { - if (CheckAttr("face", m, instance) && - CheckAttr("edge", m, instance) && - CheckAttr("poisson", m, instance) && - CheckAttr("young", m, instance) && - CheckAttr("thickness", m, instance)) { - mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr); - mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr); - mjtNum thick = - strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr); - std::vector face, edge; - String2Vector(mj_getPluginConfig(m, instance, "face"), face); - String2Vector(mj_getPluginConfig(m, instance, "edge"), edge); - return Shell(m, d, instance, nu, E, thick, face, edge); - } else { - mju_warning("Invalid parameter specification in shell plugin"); - return std::nullopt; - } -} - -// create map from triangles to vertices and edges and from edges to vertices -void Shell::CreateStencils(const std::vector& simplex, - const std::vector& edgeidx) { - // populate stencil - nt = simplex.size() / kNumVerts; - elements.resize(nt); - for (int t = 0; t < nt; t++) { - for (int v = 0; v < kNumVerts; v++) { - elements[t].vertices[v] = simplex[kNumVerts*t+v]; - } - } - - // map from edge vertices to their index in `edges` vector - std::unordered_map, int, PairHash> edge_indices; - - // loop over all triangles - for (int t = 0; t < nt; t++) { - int* v = elements[t].vertices; - - // compute edges to vertices map for fast computations - for (int e = 0; e < kNumEdges; e++) { - auto pair = std::pair( - std::min(v[edge[e][0]], v[edge[e][1]]), - std::max(v[edge[e][0]], v[edge[e][1]]) - ); - - // if edge is already present in the vector only store its index - auto [it, inserted] = edge_indices.insert({pair, ne}); - - if (inserted) { - StencilFlap flap; - flap.vertices[0] = v[edge[e][0]]; - flap.vertices[1] = v[edge[e][1]]; - flap.vertices[2] = v[(edge[e][1]+1) % 3]; - flap.vertices[3] = -1; - flaps.push_back(flap); - elements[t].edges[e] = ne++; - } else { - elements[t].edges[e] = it->second; - flaps[it->second].vertices[3] = v[(edge[e][1]+1) % 3]; - } - - if (!edgeidx.empty()) { - assert(elements[t].edges[e] == edgeidx[kNumEdges*t+e]); - } - } - } -} - -// plugin constructor -Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, - mjtNum thick, const std::vector& face, - const std::vector& edgeidx) - : thickness(thick) { - // count plugin bodies - nv = ne = 0; - for (int i = 1; i < m->nbody; i++) { - if (m->body_plugin[i] == instance) { - if (!nv++) { - i0 = i; - } - } - } - - // generate triangles from the vertices - CreateStencils(face, edgeidx); - - // material parameters - mjtNum mu = E / (2*(1+nu)); - - // loop over all triangles - for (int t = 0; t < nt; t++) { - int* v = elements[t].vertices; - for (int i = 0; i < kNumVerts; i++) { - if (m->body_plugin[i0+v[i]] != instance) { - mju_error("This body does not have the requested plugin instance"); - } - } - } - - // allocate array - position.assign(nv*3, 0); - bending.assign(ne*16, 0); - - // store previous positions - mju_copy(position.data(), m->body_pos+3*i0, 3*nv); - - // assemble bending Hessian - for (int e = 0; e < ne; e++) { - int* v = flaps[e].vertices; - int vadj[3] = {v[1], v[0], v[3]}; - - if (v[3]== -1) { - // skip boundary edges - continue; - } - - // cotangent operator from Wardetzky at al., "Discrete Quadratic Curvature - // Energies", https://cims.nyu.edu/gcl/papers/wardetzky2007dqb.pdf - - mjtNum a01 = cot(m->body_pos+3*i0, v[0], v[1], v[2]); - mjtNum a02 = cot(m->body_pos+3*i0, v[0], v[3], v[1]); - mjtNum a03 = cot(m->body_pos+3*i0, v[1], v[2], v[0]); - mjtNum a04 = cot(m->body_pos+3*i0, v[1], v[0], v[3]); - mjtNum c[4] = {a03 + a04, a01 + a02, -(a01 + a03), -(a02 + a04)}; - mjtNum volume = ComputeVolume(m->body_pos+3*i0, v) + - ComputeVolume(m->body_pos+3*i0, vadj); - - for (int v1 = 0; v1 < StencilFlap::kNumVerts; v1++) { - for (int v2 = 0; v2 < StencilFlap::kNumVerts; v2++) { - bending[16 * e + 4 * v1 + v2] += - 1.5 * c[v1] * c[v2] / volume * mu * pow(thickness, 3) / 12; - } - } - } -} - -void Shell::Compute(const mjModel* m, mjData* d, int instance) { - for (int e = 0; e < ne; e++) { - int* v = flaps[e].vertices; - mjtNum force[12] = {0}; - if (v[3] == -1) { - // skip boundary edges - continue; - } - for (int i = 0; i < StencilFlap::kNumVerts; i++) { - for (int j = 0; j < StencilFlap::kNumVerts; j++) { - for (int x = 0; x < 3; x++) { - force[3*i+x] += bending[16*e+4*i+j] * d->xpos[3*(i0+v[j])+x]; - } - } - } - - // update stored positions - mju_copy(position.data(), d->xpos+3*i0, 3*nv); - - // insert into global force - for (int i = 0; i < StencilFlap::kNumVerts; i++) { - for (int x = 0; x < 3; x++) { - d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x]; - } - } - } -} - - - -void Shell::RegisterPlugin() { - mjpPlugin plugin; - mjp_defaultPlugin(&plugin); - - plugin.name = "mujoco.elasticity.shell"; - plugin.capabilityflags |= mjPLUGIN_PASSIVE; - - const char* attributes[] = {"face", "edge", "young", - "poisson", "thickness", "damping"}; - plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]); - plugin.attributes = attributes; - plugin.nstate = +[](const mjModel* m, int instance) { return 0; }; - - plugin.init = +[](const mjModel* m, mjData* d, int instance) { - auto elasticity_or_null = Shell::Create(m, d, instance); - if (!elasticity_or_null.has_value()) { - return -1; - } - d->plugin_data[instance] = reinterpret_cast( - new Shell(std::move(*elasticity_or_null))); - return 0; - }; - plugin.destroy = +[](mjData* d, int instance) { - delete reinterpret_cast(d->plugin_data[instance]); - d->plugin_data[instance] = 0; - }; - plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) { - auto* elasticity = reinterpret_cast(d->plugin_data[instance]); - elasticity->Compute(m, d, instance); - }; - - mjp_registerPlugin(&plugin); -} - -} // namespace mujoco::plugin::elasticity diff --git a/plugin/elasticity/shell.h b/plugin/elasticity/shell.h deleted file mode 100644 index f92fffbb..00000000 --- a/plugin/elasticity/shell.h +++ /dev/null @@ -1,74 +0,0 @@ -// 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. - -#ifndef MUJOCO_PLUGIN_ELASTICITY_SHELL_H_ -#define MUJOCO_PLUGIN_ELASTICITY_SHELL_H_ - -#include -#include - -#include -#include -#include -#include "elasticity.h" - - -namespace mujoco::plugin::elasticity { - -struct StencilFlap { - static constexpr int kNumVerts = 4; - int vertices[kNumVerts]; -}; - -class Shell { - public: - // Returns a new Shell instance or nullopt on failure. - static std::optional Create(const mjModel* m, mjData* d, - int instance); - Shell(Shell&&) = default; - - Shell& operator=(Shell&& other) = default; - - void Compute(const mjModel* m, mjData* d, int instance); - - static void RegisterPlugin(); - - int i0; // index of first body - int nc; // number of quads in the grid - int nv; // number of vertices (bodies) in the Shell - int nt; // number of area elements (triangles) - int ne; // number of edges in the Shell - - // connectivity info for mapping tetrahedra to edges and vertices - std::vector elements; // triangles (nt x 6) - std::vector flaps; // adjacent triangles (ne x 4) - - // precomputed quantities - std::vector position; // previous-step positions (nv x 3) - std::vector bending; // bending Hessian (ne x 16) - - mjtNum thickness; - - private: - Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E, - mjtNum thick, const std::vector& face, - const std::vector& edgeidx); - - void CreateStencils(const std::vector& simplex, - const std::vector& edgeidx); -}; - -} // namespace mujoco::plugin::elasticity - -#endif // MUJOCO_PLUGIN_ELASTICITY_SHELL_H_ diff --git a/plugin/sdf/CMakeLists.txt b/plugin/sdf/CMakeLists.txt index 3e216fc4..8b834971 100644 --- a/plugin/sdf/CMakeLists.txt +++ b/plugin/sdf/CMakeLists.txt @@ -34,19 +34,19 @@ set(MUJOCO_SDF_SRCS torus.h ) -add_library(sdf SHARED) -target_sources(sdf PRIVATE ${MUJOCO_SDF_SRCS}) -target_include_directories(sdf PRIVATE ${MUJOCO_SDF_INCLUDE}) -target_link_libraries(sdf PRIVATE mujoco SdfLib) +add_library(sdf_plugin SHARED) +target_sources(sdf_plugin PRIVATE ${MUJOCO_SDF_SRCS}) +target_include_directories(sdf_plugin PRIVATE ${MUJOCO_SDF_INCLUDE}) +target_link_libraries(sdf_plugin PRIVATE mujoco SdfLib) target_compile_options( - sdf + sdf_plugin PRIVATE ${AVX_COMPILE_OPTIONS} ${MUJOCO_MACOS_COMPILE_OPTIONS} ${EXTRA_COMPILE_OPTIONS} ${MUJOCO_CXX_FLAGS} ) target_link_options( - sdf + sdf_plugin PRIVATE ${MUJOCO_MACOS_LINK_OPTIONS} ${EXTRA_LINK_OPTIONS} diff --git a/plugin/sensor/CMakeLists.txt b/plugin/sensor/CMakeLists.txt index 4dac6d0a..d5fde827 100644 --- a/plugin/sensor/CMakeLists.txt +++ b/plugin/sensor/CMakeLists.txt @@ -18,9 +18,11 @@ set(MUJOCO_SENSOR_INCLUDE ) set(MUJOCO_SENSOR_SRCS - sensor.cc + register.cc touch_grid.h touch_grid.cc + touch_stress.h + touch_stress.cc ) add_library(sensor SHARED) diff --git a/plugin/sensor/README.md b/plugin/sensor/README.md index e6b4e9d3..1b72db7d 100644 --- a/plugin/sensor/README.md +++ b/plugin/sensor/README.md @@ -9,6 +9,9 @@ plugins](https://mujoco.readthedocs.io/en/latest/programming/extension.html#engi - [Illustration of foveal deformation](#illustration-of-foveal-deformation) - [Illustration combining resolution, fields-of-view and foveal deformation](#illustration-combining-resolution-fields-of-view-and-foveal-deformation) +- [Touch Stress](#touch-stress) + - [Example model with analytical SDF](#example-model-with-analytical-sdf) + ## [Touch Grid](touch_grid.h) This sensor aggregates contact forces into "taxels": a rectangular array of pixel-like elements. @@ -23,7 +26,7 @@ The output of the sensor is a stack of 1 to 6 "touch images" corresponding to fo and torques in the frame of the sensor. Forces and torques are in the in [z, x, y] order, corresponding to the ordering in contact frames: [normal, tangent, tangent] and [torsional, rolling, rolling]. Each "taxel" corresponds to an angular bin -in spherical coordinates, and aggregates all the forces occuring inside this bin, which occur +in spherical coordinates, and aggregates all the forces occurring inside this bin, which occur between the body in which the sensor's site is defined and any other body. The sensor is parametrized by 6 numbers: @@ -84,3 +87,64 @@ See [touch_grid.xml](../../model/plugin/sensor/touch_grid.xml) to play with the ### Illustration combining resolution, fields-of-view and foveal deformation [![touch grid illustration](https://img.youtube.com/vi/YScjmR8LwQI/0.jpg)](https://www.youtube.com/watch?v=YScjmR8LwQI) + +## [Touch Stress](touch_stress.h) + +This sensor is based on similar concepts and parametrization as the `touch_grid`, +while overcoming some of its limitations. In particular, the `touch_grid` can +only provide sparse information, depending on the number of contact points +generated. The `touch_stress` sensor can instead generate a high-resolution +touch image. In order to do this, it requires a signed distance function (SDF) +of the object that is in contact with the sensor. This is handled internally for +primitives or it must be declared explicitly in the model using SDF plugins. + +There is one important difference with respect to the `touch_grid`: in this case, +the force is computed in the local taxel frame and not in the frame of the sensor. +This allows for a more intuitive interpretation of normal and tangential stresses, +as shown in the images below. + +Note that in this case, the absolute values of the stresses reported by the +sensor are unrelated to the contact forces. They are purely based on geometric +and kinematic considerations, i.e. the SDF for the normal stress and the sliding +velocity for the tangential contributions. + +### Example model with analytical SDF + +```xml + + + + + + + + + + + + + + + ... + + + + + + ... + + + + + ... + + + + +``` +The images below show a static sphere over a gear described by an analytic SDF +and the same sphere dragged along the x and y axes. + + + + diff --git a/plugin/sensor/images/normal.png b/plugin/sensor/images/normal.png new file mode 100644 index 00000000..f40dd0c3 Binary files /dev/null and b/plugin/sensor/images/normal.png differ diff --git a/plugin/sensor/images/tangential1.png b/plugin/sensor/images/tangential1.png new file mode 100644 index 00000000..50d3c962 Binary files /dev/null and b/plugin/sensor/images/tangential1.png differ diff --git a/plugin/sensor/images/tangential2.png b/plugin/sensor/images/tangential2.png new file mode 100644 index 00000000..68018d8e Binary files /dev/null and b/plugin/sensor/images/tangential2.png differ diff --git a/plugin/sensor/sensor.cc b/plugin/sensor/register.cc similarity index 86% rename from plugin/sensor/sensor.cc rename to plugin/sensor/register.cc index f120665e..a1561ec4 100644 --- a/plugin/sensor/sensor.cc +++ b/plugin/sensor/register.cc @@ -14,9 +14,13 @@ #include #include "touch_grid.h" +#include "touch_stress.h" namespace mujoco::plugin::sensor { -mjPLUGIN_LIB_INIT { TouchGrid::RegisterPlugin(); } +mjPLUGIN_LIB_INIT { + TouchGrid::RegisterPlugin(); + TouchStress::RegisterPlugin(); +} } // namespace mujoco::plugin::sensor diff --git a/plugin/sensor/touch_stress.cc b/plugin/sensor/touch_stress.cc new file mode 100644 index 00000000..ba7b1d4c --- /dev/null +++ b/plugin/sensor/touch_stress.cc @@ -0,0 +1,557 @@ +// 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. + +#include "touch_stress.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace mujoco::plugin::sensor { + +namespace { + +// Checks that a plugin config attribute exists. +bool CheckAttr(const std::string& input) { + char* end; + std::string value = input; + value.erase(std::remove_if(value.begin(), value.end(), isspace), value.end()); + strtod(value.c_str(), &end); + return end == value.data() + value.size(); +} + +// Converts a string into a numeric vector +template +void ReadVector(std::vector& output, const std::string& input) { + std::stringstream ss(input); + std::string item; + char delim = ' '; + while (getline(ss, item, delim)) { + CheckAttr(item); + output.push_back(strtod(item.c_str(), nullptr)); + } +} + +// Evenly spaced numbers over a specified interval. +void LinSpace(mjtNum lower, mjtNum upper, int n, mjtNum array[]) { + mjtNum increment = n > 1 ? (upper - lower) / (n - 1) : 0; + for (int i = 0; i < n; ++i) { + *array = lower; + ++array; + lower += increment; + } +} + +// Parametrized linear/quintic interpolated nonlinearity. +mjtNum Fovea(mjtNum x, mjtNum gamma) { + // Quick return. + if (!gamma) return x; + + // Foveal deformation. + mjtNum g = mjMAX(0, mjMIN(1, gamma)); + return g*mju_pow(x, 5) + (1 - g)*x; +} + +// Make bin edges. +void BinEdges(mjtNum* x_edges, mjtNum* y_edges, int size[2], mjtNum fov[2], + mjtNum gamma) { + // Make unit bin edges. + LinSpace(-1, 1, size[0] + 1, x_edges); + LinSpace(-1, 1, size[1] + 1, y_edges); + + // Apply foveal deformation. + for (int i = 0; i < size[0] + 1; i++) { + x_edges[i] = Fovea(x_edges[i], gamma); + } + for (int i = 0; i < size[1] + 1; i++) { + y_edges[i] = Fovea(y_edges[i], gamma); + } + + // Scale by field-of-view. + mju_scl(x_edges, x_edges, fov[0]*mjPI / 180, size[0] + 1); + mju_scl(y_edges, y_edges, fov[1]*mjPI / 180, size[1] + 1); +} + +// Permute 3-vector from 0,1,2 to 2,0,1. +static void xyz2zxy(mjtNum* x) { + mjtNum z = x[2]; + x[2] = x[1]; + x[1] = x[0]; + x[0] = z; +} + +// Transform spherical (azimuth, elevation, radius) to Cartesian (x,y,z). +void SphericalToCartesian(const mjtNum aer[3], mjtNum xyz[3]) { + mjtNum a = aer[0], e = aer[1], r = aer[2]; + xyz[0] = r * mju_cos(e) * mju_sin(a); + xyz[1] = r * mju_sin(e); + xyz[2] = -r * mju_cos(e) * mju_cos(a); +} + +// Tangent frame in Cartesian coordinates. +void TangentFrame(const mjtNum aer[3], mjtNum mat[9]) { + mjtNum a = aer[0], e = aer[1], r = aer[2]; + mjtNum ta[3] = {r * mju_cos(e) * mju_cos(a), 0, + r * mju_cos(e) * mju_sin(a)}; + mjtNum te[3] = {-r * mju_sin(e) * mju_sin(a), r * mju_cos(e), + r * mju_sin(e) * mju_cos(a)}; + mju_normalize3(ta); + mju_normalize3(te); + mju_copy3(mat, ta); + mju_copy3(mat+3, te); + mju_cross(mat+6, te, ta); +} + +} // namespace + +// Creates a TouchStress instance if all config attributes are defined and +// within their allowed bounds. +TouchStress* TouchStress::Create(const mjModel* m, mjData* d, + int instance) { + if (CheckAttr(std::string(mj_getPluginConfig(m, instance, "gamma"))) && + CheckAttr(std::string(mj_getPluginConfig(m, instance, "nchannel")))) { + // nchannel + int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr); + if (!nchannel) nchannel = 1; + if (nchannel < 1 || nchannel > 3) { + mju_error("nchannel must be between 1 and 3"); + return nullptr; + } + + // size + std::vector size; + std::string size_str = std::string(mj_getPluginConfig(m, instance, "size")); + ReadVector(size, size_str.c_str()); + if (size.size()!= 2) { + mju_error("Both horizontal and vertical resolutions must be specified"); + return nullptr; + } + if (size[0] <= 0 || size[1] <= 0) { + mju_error("Horizontal and vertical resolutions must be positive"); + return nullptr; + } + + // field of view + std::vector fov; + std::string fov_str = std::string(mj_getPluginConfig(m, instance, "fov")); + ReadVector(fov, fov_str.c_str()); + if (fov.size()!= 2) { + mju_error( + "Both horizontal and vertical fields of view must be specified"); + return nullptr; + } + if (fov[0] <= 0 || fov[0] > 180) { + mju_error("`fov[0]` must be a float between (0, 180] degrees"); + return nullptr; + } + if (fov[1] <= 0 || fov[1] > 90) { + mju_error("`fov[1]` must be a float between (0, 90] degrees"); + return nullptr; + } + + // gamma + mjtNum gamma = strtod(mj_getPluginConfig(m, instance, "gamma"), nullptr); + if (gamma < 0 || gamma > 1) { + mju_error("`gamma` must be a nonnegative float between [0, 1]"); + return nullptr; + } + + return new TouchStress(m, d, instance, nchannel, size.data(), fov.data(), + gamma); + } else { + mju_error("Invalid or missing parameters in touch_grid sensor plugin"); + return nullptr; + } +} + +TouchStress::TouchStress(const mjModel* m, mjData* d, int instance, + int nchannel, int size[2], mjtNum fov[2], mjtNum gamma) + : nchannel_(nchannel), + size_{size[0], size[1]}, + fov_{fov[0], fov[1]}, + gamma_(gamma) { + // Make sure sensor is attached to a site. + for (int i = 0; i < m->nsensor; ++i) { + if (m->sensor_type[i] == mjSENS_PLUGIN && m->sensor_plugin[i] == instance) { + if (m->sensor_objtype[i] != mjOBJ_SITE) { + mju_error("Touch Grid sensor must be attached to a site"); + } + } + } + + // Get sensor id. + for (id_ = 0; id_ < m->nsensor; ++id_) { + if (m->sensor_type[id_] == mjSENS_PLUGIN && + m->sensor_plugin[id_] == instance) { + break; + } + } + + // Get parent weld id. + int site_id = m->sensor_objid[id_]; + int parent_body = m->body_weldid[m->site_bodyid[site_id]]; + parent_weld_ = m->body_weldid[parent_body]; + + // Get geom id. + if (m->body_geomnum[parent_body] != 1) { + mju_error("Touch sensor must be attached to a body with exactly one geom"); + } + geom_id_ = m->body_geomadr[parent_body]; + + // Create bin edges. + x_edges_.assign(size[0] + 1, 0); + y_edges_.assign(size[1] + 1, 0); + BinEdges(x_edges_.data(), y_edges_.data(), size_, fov_, gamma_); + dist_.resize(size[0]*size[1], 0); + pos_.resize(3*size[0]*size[1], 0); + mat_.resize(9*size[0]*size[1], 0); + + // Precompute spherical coordinates. + for (int i = 0; i < size[0]; i++) { + for (int j = 0; j < size[1]; j++) { + mjtNum aer[3]; + aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]); + aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]); + aer[2] = m->geom_size[3*geom_id_]; + SphericalToCartesian(aer, pos_.data() + 3 * (i * size[1] + j)); + dist_[i*size[1]+j] = mju_abs(aer[2]); + TangentFrame(aer, mat_.data() + 9 * (i * size[1] + j)); + } + } +} + +void TouchStress::Reset(const mjModel* m, int instance) {} + +void TouchStress::Compute(const mjModel* m, mjData* d, int instance) { + mj_markStack(d); + + // Clear sensordata and distance matrix. + mjtNum* sensordata = d->sensordata + m->sensor_adr[id_]; + mju_zero(sensordata, m->sensor_dim[id_]); + + // Get site id. + int site_id = m->sensor_objid[id_]; + + // Count contacts and get contact geom ids. + std::unordered_set contact_geom_ids; + for (int i = 0; i < d->ncon; i++) { + int body1 = m->body_weldid[m->geom_bodyid[d->contact[i].geom1]]; + int body2 = m->body_weldid[m->geom_bodyid[d->contact[i].geom2]]; + if (body1 == parent_weld_) { + contact_geom_ids.insert(d->contact[i].geom2); + } + if (body2 == parent_weld_) { + contact_geom_ids.insert(d->contact[i].geom1); + } + } + + // No contacts, return. + if (contact_geom_ids.empty()) { + mj_freeStack(d); + return; + } + + // All of the quadrature points are contact points. + int ncon = size_[0]*size_[1]; + + // Get site frame. + mjtNum* site_pos = d->site_xpos + 3*site_id; + mjtNum* site_mat = d->site_xmat + 9*site_id; + + // Allocate contact forces and positions. + mjtNum* forces = mj_stackAllocNum(d, ncon*3); + mjtNum* forcesT = mj_stackAllocNum(d, ncon*3); + + // Iterate over colliding geoms. + for (auto geom : contact_geom_ids) { + int body = m->geom_bodyid[geom]; + + // Get sdf plugin of the geoms. + int sdf_instance[2] = {-1, geom_id_}; + mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE}; + const mjpPlugin* sdf_ptr[2] = {NULL, NULL}; + if (m->geom_type[geom] == mjGEOM_SDF) { + sdf_instance[0] = m->geom_plugin[geom]; + sdf_ptr[0] = mjc_getSDF(m, geom); + } else { + sdf_instance[0] = geom; + geomtype[0] = (mjtGeom)m->geom_type[geom]; + } + + // Set SDF parameters. + mjSDF geom_sdf; + geom_sdf.id = &sdf_instance[0]; + geom_sdf.type = mjSDFTYPE_SINGLE; + geom_sdf.plugin = &sdf_ptr[0]; + geom_sdf.geomtype = &geomtype[0]; + + mjSDF sensor_sdf; + sensor_sdf.id = &sdf_instance[1]; + sensor_sdf.type = mjSDFTYPE_SINGLE; + sensor_sdf.plugin = &sdf_ptr[1]; + sensor_sdf.geomtype = &geomtype[1]; + + // Get forces and positions in spherical coordinates. + int node = 0; + for (int j = 0; j < size_[1]; j++) { + for (int i = 0; i < size_[0]; i++) { + // Position in site frame. + mjtNum* pos = pos_.data() + 3*(i*size_[1] + j); + mjtNum* mat = mat_.data() + 9*(i*size_[1] + j); + + // Position in global frame. + mjtNum xpos[3]; + mju_mulMatVec3(xpos, site_mat, pos); + mju_addTo3(xpos, site_pos); + + // Position in other geom frame. + mjtNum lpos[3], tmp[3]; + mju_sub3(tmp, xpos, d->geom_xpos + 3*geom); + mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp); + + // Add mesh position if needed. + if (m->geom_type[geom] == mjGEOM_MESH || + m->geom_type[geom] == mjGEOM_SDF) { + mjtNum mesh_mat[9]; + mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]); + mju_mulMatVec3(lpos, mesh_mat, lpos); + mju_addTo3(lpos, m->mesh_pos + 3 * m->geom_dataid[geom]); + } + + // Compute distance. + mjtNum depth = mju_min(mjc_distance(m, d, &geom_sdf, lpos), 0); + if (depth == 0) { + mju_zero3(forces + 3*node); + node++; + continue; + } + + // Get velocity in global frame. + mjtNum vel_sensor[6], vel_other[6], vel_rel[3]; + mju_transformSpatial( + vel_sensor, d->cvel + 6 * parent_weld_, 0, xpos, + d->subtree_com + 3 * m->body_rootid[parent_weld_], NULL); + mju_transformSpatial( + vel_other, d->cvel + 6 * body, 0, d->geom_xpos + 3 * geom, + d->subtree_com + 3 * m->body_rootid[body], NULL); + mju_sub3(vel_rel, vel_sensor+3, vel_other+3); + + // Get contact force/torque, rotate into node frame. + mjtNum tmp_force[3], normal[3]; + mjtNum kMaxDepth = 0.05; + mjtNum pressure = 1 / (kMaxDepth - depth) - 1 / kMaxDepth; + mjc_gradient(m, d, &sensor_sdf, normal, pos); + mju_scl3(tmp_force, normal, pressure); + mju_mulMatTVec3(forces + 3*node, mat, tmp_force); + forces[3*node+0] = mju_abs(mju_dot3(vel_rel, mat + 0)); + forces[3*node+1] = mju_abs(mju_dot3(vel_rel, mat + 3)); + + // Permute forces from x,y,z to z,x,y (normal, tangent, tangent) + xyz2zxy(forces + 3*node); + node++; + } + } + + // Transpose forces. + mju_transpose(forcesT, forces, ncon, 3); + + // Compute sensor output. + for (int c = 0; c < nchannel_; c++) { + if (!mju_isZero(forcesT + c*ncon, ncon)) { + mju_addTo(sensordata + c*ncon, forcesT + c*ncon, size_[0]*size_[1]); + } + } + } + + mj_freeStack(d); +} + +// Thickness of taxel-visualization boxes relative to contact distance. +static const mjtNum kRelativeThickness = 0.02; + +void TouchStress::Visualize(const mjModel* m, mjData* d, const mjvOption* opt, + mjvScene* scn, int instance) { + mj_markStack(d); + + // Get sensor data. + mjtNum* sensordata = d->sensordata + m->sensor_adr[id_]; + + // Get maximum absolute normal force. + mjtNum maxval = 0; + int frame = size_[0]*size_[1]; + for (int j=0; j < frame; j++) { + maxval = mju_max(maxval, mju_abs(sensordata[j])); + } + + // If no normal force readings, quick return. + if (!maxval) { + mj_freeStack(d); + return; + } + + // Get site id and frame. + int site_id = m->sensor_objid[id_]; + mjtNum* site_pos = d->site_xpos + 3*site_id; + mjtNum* site_mat = d->site_xmat + 9*site_id; + mjtNum site_quat[4]; + mju_mat2Quat(site_quat, site_mat); + + // Draw geoms. + for (int i=0; i < size_[0]; i++) { + for (int j=0; j < size_[1]; j++) { + mjtNum dist = dist_[i*size_[1]+j]; + if (!dist) { + continue; + } + if (scn->ngeom >= scn->maxgeom) { + mj_warning(d, mjWARN_VGEOMFULL, scn->maxgeom); + mj_freeStack(d); + return; + } else { + // size + mjtNum size[3]; + size[0] = dist*0.5*(x_edges_[i+1]-x_edges_[i]); + size[1] = dist*0.5*(y_edges_[j+1]-y_edges_[j]); + size[2] = dist*kRelativeThickness; + + // position + mjtNum pos[3]; + mjtNum aer[3]; + aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]); + aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]); + aer[2] = dist*(1-kRelativeThickness); + SphericalToCartesian(aer, pos); + mju_mulMatVec3(pos, site_mat, pos); + mju_addTo3(pos, site_pos); + + // orientation + mjtNum a_quat[4]; + mjtNum site_y[3] = {-site_mat[1], -site_mat[4], -site_mat[7]}; + mju_axisAngle2Quat(a_quat, site_y, aer[0]); + mjtNum e_quat[4]; + mjtNum site_x[3] = {site_mat[0], site_mat[3], site_mat[6]}; + mju_axisAngle2Quat(e_quat, site_x, aer[1]); + mjtNum quat[4]; + mju_mulQuat(quat, e_quat, site_quat); + mju_mulQuat(quat, a_quat, quat); + mjtNum mat[9]; + mju_quat2Mat(mat, quat); + + // color + float rgba[4] = {1, 1, 1, 1.0}; + for (int k=0; k < mjMIN(nchannel_, 3); k++) { + rgba[k] = mju_abs(sensordata[k*frame + j*size_[0] + i]) / maxval; + } + + // draw box geom + mjvGeom* thisgeom = scn->geoms + scn->ngeom; + mjv_initGeom(thisgeom, mjGEOM_BOX, size, pos, mat, rgba); + thisgeom->objtype = mjOBJ_UNKNOWN; + thisgeom->objid = id_; + thisgeom->category = mjCAT_DECOR; + thisgeom->segid = scn->ngeom; + scn->ngeom++; + } + } + } + + mj_freeStack(d); +} + + +void TouchStress::RegisterPlugin() { + mjpPlugin plugin; + mjp_defaultPlugin(&plugin); + + plugin.name = "mujoco.sensor.touch_stress"; + plugin.capabilityflags |= mjPLUGIN_SENSOR; + + // Parameterized by 4 attributes. + const char* attributes[] = {"nchannel", "size", "fov", "gamma"}; + plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]); + plugin.attributes = attributes; + + // Stateless. + plugin.nstate = +[](const mjModel* m, int instance) { return 0; }; + + // Sensor dimension = nchannel * size[0] * size[1] + plugin.nsensordata = +[](const mjModel* m, int instance, int sensor_id) { + int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr); + if (!nchannel) nchannel = 1; + std::vector size; + std::string size_str = std::string(mj_getPluginConfig(m, instance, "size")); + ReadVector(size, size_str.c_str()); + return nchannel * size[0] * size[1]; + }; + + // Can only run after forces have been computed. + plugin.needstage = mjSTAGE_ACC; + + // Initialization callback. + plugin.init = +[](const mjModel* m, mjData* d, int instance) { + auto* TouchStress = TouchStress::Create(m, d, instance); + if (!TouchStress) { + return -1; + } + d->plugin_data[instance] = reinterpret_cast(TouchStress); + return 0; + }; + + // Destruction callback. + plugin.destroy = +[](mjData* d, int instance) { + delete reinterpret_cast(d->plugin_data[instance]); + d->plugin_data[instance] = 0; + }; + + // Reset callback. + plugin.reset = +[](const mjModel* m, mjtNum* plugin_state, void* plugin_data, + int instance) { + auto* TouchStress = reinterpret_cast(plugin_data); + TouchStress->Reset(m, instance); + }; + + // Compute callback. + plugin.compute = + +[](const mjModel* m, mjData* d, int instance, int capability_bit) { + auto* TouchStress = + reinterpret_cast(d->plugin_data[instance]); + TouchStress->Compute(m, d, instance); + }; + + // Visualization callback. + plugin.visualize = +[](const mjModel* m, mjData* d, const mjvOption* opt, + mjvScene* scn, int instance) { + auto* TouchStress = + reinterpret_cast(d->plugin_data[instance]); + TouchStress->Visualize(m, d, opt, scn, instance); + }; + + // Register the plugin. + mjp_registerPlugin(&plugin); +} + +} // namespace mujoco::plugin::sensor diff --git a/plugin/sensor/touch_stress.h b/plugin/sensor/touch_stress.h new file mode 100644 index 00000000..7453d87a --- /dev/null +++ b/plugin/sensor/touch_stress.h @@ -0,0 +1,80 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_ +#define MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_ + +#include + +#include +#include +#include +#include + +namespace mujoco::plugin::sensor { + +// A touch grid sensor is associated with a site and senses contact stresses +// between the site's parent body and all other bodies. The site's +// frame determines the orientation of the sensor with the same convention used +// for cameras and lights: the sensor points in the frame's negative-Z +// direction, so the X and Y axes correspond to horizontal and vertical +// directions respectively. +// +// The output of the sensor is a stack of 3 "touch images" corresponding to +// forces in the local frame of the taxels. Forces are in the [z, x, y] order, +// corresponding to the ordering in contact frames: [normal, tangent, tangent]. +// +// The sensor has 6 parameters: +// 1. (int) Number of channels [1-3]. Defaults to 1. +// 2. (int) Horizontal resolution. +// 3. (int) Vertical resolution. +// 4. (float) Horizontal field-of-view (fov_x), in degrees. +// 5. (float) Vertical field-of-view (fov_y), in degrees. +// 6. (float) Foveal deformation. Defaults to 0. +class TouchStress { + public: + static TouchStress* Create(const mjModel* m, mjData* d, int instance); + TouchStress(TouchStress&&) = default; + ~TouchStress() = default; + + void Reset(const mjModel* m, int instance); + void Compute(const mjModel* m, mjData* d, int instance); + void Visualize(const mjModel* m, mjData* d, const mjvOption* opt, + mjvScene* scn, int instance); + + static void RegisterPlugin(); + + int nchannel_; // number of channels (1-3) + int size_[2]; // horizontal and vertical resolution + mjtNum fov_[2]; // horizontal and vertical field of view, in degrees + mjtNum gamma_; // foveal deformation + + private: + TouchStress(const mjModel* m, mjData* d, int instance, int nchannel, + int* size, mjtNum* fov_x, mjtNum gamma); + + std::vector x_edges_; + std::vector y_edges_; + std::vector dist_; + std::vector pos_; + std::vector mat_; + + int id_; + int parent_weld_; + int geom_id_; +}; + +} // namespace mujoco::plugin::sensor + +#endif // MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_ diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 53475f10..0c1173df 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -57,8 +57,9 @@ PyOpenGL==3.1.7 \ --hash=sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a pytest==8.3.3 \ --hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2 -setuptools==75.5.0 \ - --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829 +setuptools==78.1.1 \ + --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ + --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d wheel==0.45.0 \ --hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7 diff --git a/python/make_sdist_requirements.txt b/python/make_sdist_requirements.txt index 10f531c1..f961524c 100644 --- a/python/make_sdist_requirements.txt +++ b/python/make_sdist_requirements.txt @@ -4,8 +4,9 @@ build==1.2.2.post1 \ --hash=sha256:1d61c0887fa860c01971625baae8bdd338e517b836a2f70dd1f7aa3a6b2fc5b5 pip==24.3.1 \ --hash=sha256:3790624780082365f47549d032f3770eeb2b1e8bd1f7b2e02dace1afa361b4ed -setuptools==75.5.0 \ - --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829 +setuptools==78.1.1 \ + --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ + --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d # Transitive dependencies of build colorama==0.4.6; platform_system == 'Windows' \ diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb index 925da6bc..ff4b997e 100644 --- a/python/mjspec.ipynb +++ b/python/mjspec.ipynb @@ -10,7 +10,8 @@ "\n", "#

Model Editing

\n", "\n", - "This notebook provides an introductory tutorial for model editing in MuJoCo using the `mjSpec` API. This notebook assumes that the reader is already familiar with MuJoCo basic concepts, as demostrated in the [introductory tutorial](https://github.com/google-deepmind/mujoco?tab=readme-ov-file#getting-started). Documentation for this API can be found in the [Model Editing](https://mujoco.readthedocs.io/en/latest/programming/modeledit.html) chapter in the documentation (C API) and in the [Python chapter](https://mujoco.readthedocs.io/en/latest/python.html#model-editing). Here we use the Python API.\n", + "This notebook provides an introductory tutorial for model editing in MuJoCo using the `mjSpec` API. This notebook assumes that the reader is already familiar with MuJoCo basic concepts, as\n", + "demonstrated in the [introductory tutorial](https://github.com/google-deepmind/mujoco?tab=readme-ov-file#getting-started). Documentation for this API can be found in the [Model Editing](https://mujoco.readthedocs.io/en/latest/programming/modeledit.html) chapter in the documentation (C API) and in the [Python chapter](https://mujoco.readthedocs.io/en/latest/python.html#model-editing). Here we use the Python API.\n", "\n", "The goal of the API is to allow users to easily interact with and modify MuJoCo\n", "models in Python, similarly to what the JavaScript DOM does for HTML.\n", @@ -80,7 +81,7 @@ "print('Setting environment variable to use GPU rendering:')\n", "%env MUJOCO_GL=egl\n", "\n", - "# Check if installation was succesful.\n", + "# Check if installation was successful.\n", "try:\n", " print('Checking that the installation succeeded:')\n", " import mujoco as mj\n", @@ -96,6 +97,7 @@ "\n", "# Other imports and helper functions\n", "import numpy as np\n", + "import random\n", "from scipy.signal import convolve2d\n", "\n", "# Graphics and plotting.\n", @@ -124,12 +126,12 @@ " highlighted = pygments.highlight(xml_string, lexer, formatter)\n", " display(HTML(f\"{highlighted}\"))\n", "\n", - "def render(model, data=None, height=300):\n", + "def render(model, data=None, height=300, camera=-1):\n", " if data is None:\n", " data = mj.MjData(model)\n", " with mj.Renderer(model, 480, 640) as renderer:\n", " mj.mj_forward(model, data)\n", - " renderer.update_scene(data)\n", + " renderer.update_scene(data, camera)\n", " media.show_image(renderer.render(), height=height)" ] }, @@ -148,6 +150,7 @@ "cell_type": "code", "execution_count": 0, "metadata": { + "cellView": "form", "id": "oummB7I7EfSq" }, "outputs": [], @@ -537,7 +540,7 @@ "\n", " Args:\n", " shape: The shape of the generated array (tuple of two ints).\n", - " This must be a multple of res.\n", + " This must be a multiple of res.\n", " res: The number of periods of noise to generate along each\n", " axis (tuple of two ints). Note shape must be a multiple of\n", " res.\n", @@ -635,7 +638,7 @@ "outputs": [], "source": [ "def add_hfield(spec=None, hsize=10, vsize=4):\n", - " \"\"\" Function that adds a heighfield with countours\"\"\"\n", + " \"\"\" Function that adds a height field with contours\"\"\"\n", "\n", " # Initialize spec\n", " if spec is None:\n", @@ -852,16 +855,1524 @@ { "cell_type": "markdown", "metadata": { - "id": "3N4YEIVt75_T" + "id": "bZ-mpAKhSBHw" }, "source": [ - "# `dm_control` example" + "## Terrain Generation\n", + "Here we will create a terrain out of different tiles. We start by creating each sinlge tile. Then we move on to putting tiles side by side to create a complete terrain." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "cellView": "form", + "id": "d9oHHP6ZSZeK" + }, + "outputs": [], + "source": [ + "#@title Utilities\n", + "def render_tile(tile_func, direction=None, cam_distance=6, cam_elevation=-30):\n", + " arena_xml = \"\"\"\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \"\"\"\n", + "\n", + " spec = mj.MjSpec.from_string(arena_xml)\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " name = 'base_tile'\n", + "\n", + " spec.worldbody.add_body(pos=[-3, 0, 0], name=name)\n", + " if direction:\n", + " tile_func(spec, direction=direction)\n", + " else:\n", + " tile_func(spec)\n", + "\n", + " model = spec.compile()\n", + " data = mj.MjData(model)\n", + "\n", + " cam = mj.MjvCamera()\n", + " mj.mjv_defaultCamera(cam)\n", + " cam.lookat = [0, 0, 0]\n", + " cam.distance = cam_distance\n", + " cam.elevation = cam_elevation\n", + "\n", + " height = 300\n", + "\n", + " with mj.Renderer(model, 480, 640) as renderer:\n", + " mj.mj_forward(model, data)\n", + " renderer.update_scene(data,cam)\n", + " media.show_image(renderer.render(), height=height)\n", + "\n", + "\n", + "def interpolant(t):\n", + " return t*t*t*(t*(t*6 - 15) + 10)\n", + "\n", + "def perlin(shape, res, tileable=(False, False), interpolant=interpolant):\n", + " \"\"\"Generate a 2D numpy array of perlin noise.\n", + "\n", + " Args:\n", + " shape: The shape of the generated array (tuple of two ints).\n", + " This must be a multiple of res.\n", + " res: The number of periods of noise to generate along each\n", + " axis (tuple of two ints). Note shape must be a multiple of\n", + " res.\n", + " tileable: If the noise should be tileable along each axis\n", + " (tuple of two bools). Defaults to (False, False).\n", + " interpolant: The interpolation function, defaults to\n", + " t*t*t*(t*(t*6 - 15) + 10).\n", + "\n", + " Returns:\n", + " A numpy array of shape shape with the generated noise.\n", + "\n", + " Raises:\n", + " ValueError: If shape is not a multiple of res.\n", + " \"\"\"\n", + " delta = (res[0] / shape[0], res[1] / shape[1])\n", + " d = (shape[0] // res[0], shape[1] // res[1])\n", + " grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]]\\\n", + " .transpose(1, 2, 0) % 1\n", + " # Gradients\n", + " angles = 2*np.pi*np.random.rand(res[0]+1, res[1]+1)\n", + " gradients = np.dstack((np.cos(angles), np.sin(angles)))\n", + " if tileable[0]:\n", + " gradients[-1,:] = gradients[0,:]\n", + " if tileable[1]:\n", + " gradients[:,-1] = gradients[:,0]\n", + " gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)\n", + " g00 = gradients[ :-d[0], :-d[1]]\n", + " g10 = gradients[d[0]: , :-d[1]]\n", + " g01 = gradients[ :-d[0],d[1]: ]\n", + " g11 = gradients[d[0]: ,d[1]: ]\n", + " # Ramps\n", + " n00 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1] )) * g00, 2)\n", + " n10 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1] )) * g10, 2)\n", + " n01 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1]-1)) * g01, 2)\n", + " n11 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1]-1)) * g11, 2)\n", + " # Interpolation\n", + " t = interpolant(grid)\n", + " n0 = n00*(1-t[:,:,0]) + t[:,:,0]*n10\n", + " n1 = n01*(1-t[:,:,0]) + t[:,:,0]*n11\n", + " return np.sqrt(2)*((1-t[:,:,1])*n0 + t[:,:,1]*n1)\n", + "\n", + "def edge_slope(size, border_width=5, blur_iterations=20):\n", + " \"\"\"Creates a grayscale image with a white center and fading black edges using convolution.\"\"\"\n", + " img = np.ones((size, size), dtype=np.float32)\n", + " img[:border_width, :] = 0\n", + " img[-border_width:, :] = 0\n", + " img[:, :border_width] = 0\n", + " img[:, -border_width:] = 0\n", + "\n", + " kernel = np.array([[1, 1, 1],\n", + " [1, 1, 1],\n", + " [1, 1, 1]]) / 9.0\n", + "\n", + " for _ in range(blur_iterations):\n", + " img = convolve2d(img, kernel, mode='same', boundary='symm')\n", + "\n", + " return img" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "dtFWSSwWSgeD" + }, + "outputs": [], + "source": [ + "# @title Stairs\n", + "def stairs(spec=None, grid_loc=[0, 0] , num_stairs=4, direction=1, name='stair'):\n", + " SQUARE_LENGTH = 2\n", + " V_SIZE = 0.076\n", + " H_SIZE = 0.12\n", + " H_STEP = H_SIZE * 2\n", + " V_STEP = V_SIZE * 2\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " # Offset\n", + " x_beginning, y_end = [-SQUARE_LENGTH + H_SIZE] * 2\n", + " x_end, y_beginning = [SQUARE_LENGTH - H_SIZE] * 2\n", + " # Dimension\n", + " size_one = [H_SIZE, SQUARE_LENGTH, V_SIZE]\n", + " size_two = [SQUARE_LENGTH, H_SIZE, V_SIZE]\n", + " # Geoms positions\n", + " x_pos_l = [x_beginning, 0, direction * V_SIZE]\n", + " x_pos_r = [x_end, 0, direction * V_SIZE]\n", + " y_pos_up = [0, y_beginning, direction * V_SIZE]\n", + " y_pos_down = [0, y_end, direction * V_SIZE]\n", + "\n", + " for i in range(num_stairs):\n", + " size_one[1] = SQUARE_LENGTH - H_STEP * i\n", + " size_two[0] = SQUARE_LENGTH - H_STEP * i\n", + "\n", + " x_pos_l[2], x_pos_r[2], y_pos_up[2], y_pos_down[2] = [\n", + " direction * ( V_SIZE + V_STEP * i)] * 4\n", + "\n", + " # Left side\n", + " x_pos_l[0] = x_beginning + H_STEP * i\n", + " body.add_geom(pos=x_pos_l, size=size_one, rgba=BROWN)\n", + " # Right side\n", + " x_pos_r[0] = x_end - H_STEP * i\n", + " body.add_geom(pos=x_pos_r, size=size_one, rgba=BROWN)\n", + " # Top\n", + " y_pos_up[1] = y_beginning - H_STEP * i\n", + " body.add_geom(pos=y_pos_up, size=size_two, rgba=BROWN)\n", + " # Bottom\n", + " y_pos_down[1] = y_end + H_STEP * i\n", + " body.add_geom(pos=y_pos_down, size=size_two, rgba=BROWN)\n", + "\n", + " # Closing\n", + " size = [SQUARE_LENGTH - H_STEP * num_stairs,\n", + " SQUARE_LENGTH - H_STEP * num_stairs,\n", + " V_SIZE]\n", + " pos = [0, 0,\n", + " direction * (V_SIZE + V_STEP * num_stairs)]\n", + " body.add_geom(pos=pos, size=size, rgba=BROWN)\n", + "\n", + "render_tile(stairs, direction=random.choice([-1, 1]))" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "YwX50AtzSmnT" + }, + "outputs": [], + "source": [ + "# @title Debris (Geoms)\n", + "def debris_with_simple_geoms(spec=None, grid_loc=[0, 0], name='plane'):\n", + " SQUARE_LENGTH = 2\n", + " THICKNESS = 0.05\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + " RED = [0.6, 0.12, 0.15, 1.0]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Create tile\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN )\n", + "\n", + " # Simple Geoms\n", + " x_beginning, y_end = [-SQUARE_LENGTH + THICKNESS] * 2\n", + " x_end, y_beginning = [SQUARE_LENGTH - THICKNESS] * 2\n", + "\n", + " x_grid = np.linspace(x_beginning, x_end, 10)\n", + " y_grid = np.linspace(y_beginning, y_end, 10)\n", + "\n", + " for i in range(10):\n", + " x = np.random.choice(x_grid)\n", + " y = np.random.choice(y_grid)\n", + "\n", + " pos=[grid_loc[0] + x, grid_loc[1] + y, 0.2]\n", + "\n", + " g_type = None\n", + " size = None\n", + " if random.randint(0, 1):\n", + " g_type = mj.mjtGeom.mjGEOM_BOX\n", + " size = [0.1, 0.1, 0.02]\n", + " else:\n", + " g_type = mj.mjtGeom.mjGEOM_CYLINDER\n", + " size = [0.1, 0.02, 0]\n", + "\n", + " body = spec.worldbody.add_body(pos=pos, name=f'g{i}_{name}', mass=1)\n", + " body.add_geom(type=g_type, size=size, rgba=RED)\n", + " body.add_freejoint()\n", + "\n", + "render_tile(debris_with_simple_geoms)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Wm_0DHK8S2kQ" + }, + "outputs": [], + "source": [ + "# @title Debris (Mesh)\n", + "def debris(spec=None, grid_loc=[0, 0] , name='debris'):\n", + " SQUARE_LENGTH = 2\n", + " THICKNESS = 0.05\n", + " STEP = THICKNESS * 8\n", + " SCALE = 0.1\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + " RED = [0.6, 0.12, 0.15, 1.0]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + " main.mesh.scale = np.array([SCALE]*3, dtype=np.float64)\n", + "\n", + " x_beginning = -SQUARE_LENGTH + THICKNESS\n", + " y_beginning = SQUARE_LENGTH - THICKNESS\n", + "\n", + " # Create tile\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n", + "\n", + " # Place debris on the tile\n", + " for i in range(10):\n", + " for j in range(10):\n", + " # draw on xy plane\n", + " drawing = np.random.normal(size=(4, 2))\n", + " drawing /= np.linalg.norm(drawing, axis=1, keepdims=True)\n", + " z = np.zeros((drawing.shape[0], 1))\n", + " # Add z value to drawing\n", + " base = np.concatenate((drawing, z), axis=1)\n", + " # Extrude drawing\n", + " z_extrusion = np.full((drawing.shape[0], 1), THICKNESS * 4)\n", + " top = np.concatenate((drawing, z_extrusion), axis=1)\n", + " # Combine to get a mesh\n", + " mesh = np.vstack((base, top))\n", + "\n", + " # Create body and add the mesh to the geom of the body\n", + " spec.add_mesh(name=f'd{i}_{j}_{name}', uservert=mesh.flatten())\n", + " pos=[grid_loc[0] + x_beginning + i * STEP,\n", + " grid_loc[1] + y_beginning - j * STEP,\n", + " 0.2]\n", + "\n", + " body = spec.worldbody.add_body(pos=pos, name=f'd{i}_{j}_{name}', mass=1)\n", + " body.add_geom(type=mj.mjtGeom.mjGEOM_MESH, meshname=f'd{i}_{j}_{name}',\n", + " rgba=RED)\n", + " body.add_freejoint()\n", + "\n", + "render_tile(debris)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "BqTG_K8uS3v_" + }, + "outputs": [], + "source": [ + "# @title Boxy Terrain\n", + "def boxy_terrain(spec=None, grid_loc=[0, 0], name='boxy_terrain'):\n", + " SQUARE_LENGTH = 2\n", + " CUBE_LENGTH = 0.05\n", + " GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n", + " STEP = CUBE_LENGTH * 2\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + "\n", + " if spec == None:\n", + " spec=mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Create tile\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + "\n", + " x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n", + " y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n", + " for i in range(GRID_SIZE):\n", + " for j in range(GRID_SIZE):\n", + " body.add_geom(\n", + " pos=[x_beginning + i * STEP ,\n", + " y_beginning - j * STEP ,\n", + " random.randint(-1, 1) * CUBE_LENGTH\n", + " ],\n", + " size=[CUBE_LENGTH] * 3,\n", + " rgba=BROWN\n", + " )\n", + "\n", + "render_tile(boxy_terrain)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "gfQ1CafnS7hs" + }, + "outputs": [], + "source": [ + "# @title Box (Extrusion | Cut)\n", + "def box_extrusions(spec=None, grid_loc=[0, 0], complex=False,\n", + " name='box_extrusions'):\n", + " # Warning! complex sometimes leads to creation of holes\n", + " SQUARE_LENGTH = 2\n", + " CUBE_LENGTH = 0.05\n", + " GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n", + " STEP = CUBE_LENGTH * 2\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Create tile\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + "\n", + " x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n", + " y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n", + "\n", + " # Create initial grid and store geoms ref\n", + " grid = [[ 0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]\n", + " for i in range(GRID_SIZE):\n", + " for j in range(GRID_SIZE):\n", + " ref = body.add_geom(\n", + " pos=[x_beginning + i * STEP, y_beginning - j * STEP, 0],\n", + " size=[CUBE_LENGTH] * 3,\n", + " rgba = BROWN\n", + " )\n", + " grid[i][j] = ref\n", + "\n", + " # Extrude or Cut operation using the boxes\n", + " for _ in range(random.randint(4, 50)):\n", + " box = None\n", + " while box == None:\n", + " # Create a box\n", + " start = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE - 2))\n", + " dim = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE-2))\n", + " # Make suer box is valid\n", + " if start[0] + dim [0] < len(grid) and start[1] + dim [1] < len(grid):\n", + " box = {\"start\":start, \"dim\":dim}\n", + "\n", + " # Use the box to Cut or Extrude\n", + " operation = random.choice([1, -1])\n", + " start = box[\"start\"]\n", + " dim = box[\"dim\"]\n", + " for i in range(start[0], dim[0]):\n", + " for j in range(start[1], dim[1]):\n", + " tile = grid[i][j]\n", + " if complex:\n", + " tile.pos[2] += operation * CUBE_LENGTH\n", + " else:\n", + " tile.pos[2] = operation * CUBE_LENGTH\n", + "\n", + "render_tile(box_extrusions)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "e1PEGQ37S_ee" + }, + "outputs": [], + "source": [ + "# @title Heightfield\n", + "def h_field(spec=None, grid_loc=[0, 0], name='h_field'):\n", + " SQUARE_LENGTH = 2\n", + " HEIGHT = 0.1\n", + " BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n", + "\n", + " if spec is None:\n", + " spec = mj.MjSpec()\n", + "\n", + " size = 128\n", + " noise = perlin((size, size), (8, 8))\n", + "\n", + " # Remap noise to 0 to 1\n", + " noise = (noise + 1)/2\n", + " noise -= np.min(noise)\n", + " noise /= np.max(noise)\n", + "\n", + " # Makes the edges slope down to avoid sharp boundary\n", + " noise *= edge_slope(size)\n", + "\n", + " # Create height field\n", + " hfield = spec.add_hfield(name=name,\n", + " size=[SQUARE_LENGTH, SQUARE_LENGTH,\n", + " HEIGHT, HEIGHT/10],\n", + " nrow=noise.shape[0],\n", + " ncol=noise.shape[1],\n", + " userdata=noise.flatten())\n", + "\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(type=mj.mjtGeom.mjGEOM_HFIELD, hfieldname=name,\n", + " rgba=BROWN_RGBA)\n", + "\n", + "render_tile(h_field)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "UwDVd6aUTRKF" + }, + "outputs": [], + "source": [ + "# @title Floating platform\n", + "def floating_platform(spec=None, gird_loc=[0, 0, 0], name='platform'):\n", + " PLATFORM_LENGTH = 0.5\n", + " WIDTH = 0.12\n", + " INWARD_OFFSET = 0.008\n", + " THICKNESS = 0.005\n", + " SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n", + " TENDON_LENGTH = 0.5\n", + " Z_OFFSET = 0.1\n", + "\n", + " GOLD = [0.850, 0.838, 0.119, 1]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Platform with sites\n", + " gird_loc[2] += Z_OFFSET\n", + " platform = spec.worldbody.add_body(pos=gird_loc, name=name)\n", + " platform.add_geom(size=SIZE, rgba=GOLD)\n", + " platform.add_freejoint()\n", + "\n", + " for x_dir in [-1, 1]:\n", + " for y_dir in [-1, 1]:\n", + " # Add site to world\n", + " vector = np.array([x_dir * PLATFORM_LENGTH,\n", + " y_dir * (WIDTH - INWARD_OFFSET)])\n", + " x_w = gird_loc[0] + vector[0]\n", + " y_w = gird_loc[1] + vector[1]\n", + " z_w = gird_loc[2] + TENDON_LENGTH\n", + " # Rotate sites by theta\n", + " spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n", + " pos=[ x_w, y_w, z_w],\n", + " size=[0.01, 0, 0])\n", + " # Add site to platform\n", + " x_p = x_dir * PLATFORM_LENGTH\n", + " y_p = y_dir * (WIDTH - INWARD_OFFSET)\n", + " platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n", + " pos=[ x_p, y_p, THICKNESS * 2],\n", + " size=[0.01, 0, 0])\n", + "\n", + " # Connect tendon to sites\n", + " thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}',\n", + " limited=True,\n", + " range=[0, TENDON_LENGTH], width=0.01 )\n", + " thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n", + " thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n", + "\n", + "render_tile(floating_platform, cam_distance=2, cam_elevation=-20)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Wb4GmZewTSVI" + }, + "outputs": [], + "source": [ + "# @title Simple stairs\n", + "def simple_suspended_stair(spec=None, grid_loc=[0, 0], num_stair=20,\n", + " name=\"simple_suspended_stair\"):\n", + " BROWN= [0.460, 0.362, 0.216, 1.0]\n", + " SQUARE_LENGTH = 2\n", + " THICKNESS = 0.05\n", + " OFFSET_Y = -4/5 * SQUARE_LENGTH\n", + "\n", + " V_STEP = 0.076\n", + " H_STEP = 0.12\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Create tile\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n", + "\n", + " # Create Stairs\n", + " for i in range(num_stair):\n", + " floating_platform(spec,[grid_loc[0],\n", + " OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n", + " i * V_STEP],\n", + " name =f'{name}_p_{i}')\n", + "\n", + "render_tile(simple_suspended_stair,cam_distance=7, cam_elevation=-30)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "AaoLjrKeTXu-" + }, + "outputs": [], + "source": [ + "# @title Sinusoidal stairs\n", + "def sin_suspended_stair(spec, grid_loc=[0, 0], num_stair=40,\n", + " name=\"sin_suspended_stair\"):\n", + " BROWN = [0.460, 0.362, 0.216, 1.0]\n", + " SQUARE_LENGTH = 2\n", + " THICKNESS = 0.05\n", + " OFFSET_Y = -4/5 * SQUARE_LENGTH\n", + "\n", + " V_STEP = 0.076\n", + " H_STEP = 0.12\n", + " AMPLITUDE = 0.2\n", + " FREQUENCY = 0.5\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + "\n", + " # Plane\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n", + "\n", + " for i in range(num_stair):\n", + " x_step = AMPLITUDE * np.sin(2 * np.pi * FREQUENCY * (i * H_STEP))\n", + " floating_platform(spec, [grid_loc[0] + x_step,\n", + " OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n", + " i * V_STEP],\n", + " name=f'{name}_p_{i}')\n", + "\n", + "render_tile(sin_suspended_stair,cam_distance=7, cam_elevation=-30)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "2_Ux8h-PTqbW" + }, + "outputs": [], + "source": [ + "# @title Floating platform for circular stair\n", + "def floating_platform_for_circular_stair(spec=None, gird_loc=[0, 0, 0] ,theta=0,\n", + " name='platform'):\n", + " PLATFORM_LENGTH = 0.5\n", + " TENDON_LENGTH = 0.5\n", + " WIDTH = 0.12/4 # Platform (body) is made of 4 separate geoms\n", + " THICKNESS = 0.005\n", + " SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n", + " Z_OFFSET = 0.1\n", + "\n", + " GOLD = [0.850, 0.838, 0.119, 1]\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + " spec.compiler.degree = False\n", + "\n", + " # Platform with sites\n", + " gird_loc[2] += Z_OFFSET\n", + " platform = spec.worldbody.add_body(pos=gird_loc, name=name, euler=[0, 0, theta])\n", + " platform.add_geom(pos=[0, 0, 0] , size=SIZE, euler=[0, 0, 0],rgba=GOLD)\n", + " platform.add_geom(pos=[0, 0.02, 0], size=SIZE, euler=[0, 0, 0.05],rgba=GOLD)\n", + " platform.add_geom(pos=[0, 0.05, 0], size=SIZE, euler=[0, 0, 0.1],rgba=GOLD)\n", + " platform.add_geom(pos=[0, 0.08, 0], size=SIZE, euler=[0, 0, 0.15],rgba=GOLD)\n", + " platform.add_freejoint()\n", + "\n", + " for i, x_dir in enumerate([-1, 1]):\n", + " for j, y_dir in enumerate([-1, 1]):\n", + " # Rotate sites by theta\n", + " rotation_matrix = np.array([[np.cos(-theta), -np.sin(-theta)],\n", + " [np.sin(-theta), np.cos(-theta)]])\n", + " vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * WIDTH ])\n", + " if i + j == 2:\n", + " vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * 6 * WIDTH ])\n", + " vector = np.dot(vector , rotation_matrix)\n", + " x_w = gird_loc[0] + vector[0]\n", + " y_w = gird_loc[1] + vector[1]\n", + " z_w = gird_loc[2] + TENDON_LENGTH\n", + "\n", + " # Add site to world\n", + " spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n", + " pos=[ x_w, y_w, z_w],\n", + " size=[0.01, 0, 0])\n", + " # Add site to platform\n", + " x_p = x_dir * PLATFORM_LENGTH\n", + " y_p = y_dir * WIDTH\n", + " if i + j == 2:\n", + " y_p = y_dir * 6 * WIDTH\n", + " platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n", + " pos=[x_p, y_p, THICKNESS * 2],\n", + " size=[0.01, 0, 0])\n", + "\n", + " # Connect tendon to sites\n", + " thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}', limited=True,\n", + " range=[0, TENDON_LENGTH], width=0.01 )\n", + " thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n", + " thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n", + "\n", + "render_tile(floating_platform_for_circular_stair,cam_distance=2, cam_elevation=-40)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "UXSDdVF8TwmQ" + }, + "outputs": [], + "source": [ + "# @title Circular stairs\n", + "def circular_stairs(spec, grid_loc=[0, 0], num_stair=60, name=\"circular_stairs\"):\n", + " BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n", + " SQUARE_LENGTH = 2\n", + " THICKNESS = 0.05\n", + "\n", + " RADIUS = 1.5\n", + " V_STEP = 0.076\n", + "\n", + " if spec == None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Defaults\n", + " main = spec.default\n", + " main.geom.type = mj.mjtGeom.mjGEOM_BOX\n", + " spec.compiler.degree = False\n", + "\n", + " # Plane\n", + " body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n", + " body.add_geom(size = [SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba = BROWN_RGBA )\n", + "\n", + " theta_step = 2 * np.pi / num_stair\n", + " for i in range(num_stair):\n", + " theta = i * theta_step\n", + " x = grid_loc[0] + RADIUS * np.cos(theta)\n", + " y = grid_loc[1] + RADIUS * np.sin(theta)\n", + " z = i * V_STEP\n", + "\n", + " floating_platform_for_circular_stair(spec, [x, y, z], theta=theta, name=f'{name}_p_{i}')\n", + "\n", + "render_tile(circular_stairs,cam_distance=12, cam_elevation=-30)" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "lsDky1KjT30W" + }, + "outputs": [], + "source": [ + "# @title Tile Generator\n", + "def add_tile(spec=None, grid_loc=[0, 0]):\n", + " if spec is None:\n", + " spec = mj.MjSpec()\n", + "\n", + " tile_type = random.randint(0, 9)\n", + "\n", + " if tile_type == 0:\n", + " debris_with_simple_geoms(spec, grid_loc, name=f\"plane_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 1:\n", + " stairs(spec, grid_loc, name=f\"stairs_up_{grid_loc[0]}_{grid_loc[1]}\",direction=1)\n", + " elif tile_type == 2:\n", + " stairs(spec, grid_loc, name=f\"stairs_down_{grid_loc[0]}_{grid_loc[1]}\",direction=-1)\n", + " elif tile_type == 3:\n", + " debris(spec, grid_loc, name=f\"debris_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 4:\n", + " box_extrusions(spec, grid_loc, name=f\"box_extrusions_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 5:\n", + " boxy_terrain(spec, grid_loc, name=f\"boxy_terrain_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 6:\n", + " h_field(spec, grid_loc, name=f\"h_field_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 7:\n", + " simple_suspended_stair(spec, grid_loc, name=f\"sss_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 8:\n", + " sin_suspended_stair(spec, grid_loc, name=f\"sinss_{grid_loc[0]}_{grid_loc[1]}\")\n", + " elif tile_type == 9:\n", + " circular_stairs(spec, grid_loc, name=f\"circular_s_{grid_loc[0]}_{grid_loc[1]}\")\n", + " return spec" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "P2F0ObC2T8w8" + }, + "outputs": [], + "source": [ + "# @title Generate Terrain\n", + "arena_xml = \"\"\"\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + "\n", + "\"\"\"\n", + "\n", + "spec = mj.MjSpec.from_string(arena_xml)\n", + "\n", + "spec.option.enableflags |= mj.mjtEnableBit.mjENBL_OVERRIDE\n", + "spec.option.enableflags |= mj.mjtEnableBit.mjENBL_MULTICCD\n", + "spec.option.timestep = 0.0001\n", + "spec.compiler.degree = False\n", + "\n", + "main = spec.default\n", + "main.geom.solref = [0.001, 1]\n", + "\n", + "# Add lights\n", + "for x in [-1, 1]:\n", + " for y in [-1, 1]:\n", + " spec.worldbody.add_light(pos=[x, y, 40], dir=[-x, -y, -15])\n", + "\n", + "SQUARE_LENGTH = 2\n", + "for i in range(-2, 2):\n", + " for j in range(-2, 2):\n", + " add_tile(spec=spec, grid_loc=[i * 2 * SQUARE_LENGTH, j * 2 * SQUARE_LENGTH])\n", + "\n", + "model = spec.compile()\n", + "data = mj.MjData(model)\n", + "\n", + "cam = mj.MjvCamera()\n", + "mj.mjv_defaultCamera(cam)\n", + "cam.lookat = [-2, 0, -2]\n", + "cam.distance = 18\n", + "cam.elevation = -30\n", + "\n", + "with mj.Renderer(model, 720, 1280) as renderer:\n", + " mj.mj_forward(model, data)\n", + " renderer.update_scene(data,cam)\n", + " media.show_image(renderer.render())" ] }, { "cell_type": "markdown", "metadata": { - "id": "TcQuv56BwaJf" + "id": "IGd0uD64LdEJ" + }, + "source": [ + "# Model editing" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "cellView": "form", + "id": "223KzKAzLdEJ" + }, + "outputs": [], + "source": [ + "# @title Get resources\n", + "\n", + "# Get Models\n", + "print('Getting MuJoCo humanoid XML description from GitHub:')\n", + "!git clone https://github.com/google-deepmind/mujoco\n", + "humanoid_file = 'mujoco/model/humanoid/humanoid.xml'\n", + "humanoid100_file = 'mujoco/model/humanoid/humanoid100.xml'\n", + "print('Getting MuJoCo Fly and Franka XML description from GitHub:')\n", + "!git clone https://github.com/google-deepmind/mujoco_menagerie\n", + "fly_file = 'mujoco_menagerie/flybody/fruitfly.xml'\n", + "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n", + "\n", + "# Camera options\n", + "cam = mj.MjvCamera()\n", + "mj.mjv_defaultCamera(cam)\n", + "cam.elevation = -10\n", + "cam.lookat = [0, 0, 1]\n", + "cam.distance = 4\n", + "cam.azimuth = 135\n", + "\n", + "# Arena\n", + "arena_xml = \"\"\"\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + "\n", + "\"\"\"\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "eGgXNjQ8LdEK" + }, + "source": [ + "`mjSpec` elements can be traversed in two ways:\n", + "- For elements inside the kinematic tree, the tree can be traversed using the `first` and `next` functions.\n", + "- For all other elements, we provide a list.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Len0o_idLdEK" + }, + "outputs": [], + "source": [ + "#@title Traversing the spec {vertical-output: true}\n", + "\n", + "spec = mj.MjSpec.from_file(humanoid_file)\n", + "\n", + "# Function that recursively prints all body names\n", + "def print_bodies(parent, level=0):\n", + " body = parent.first_body()\n", + " while body:\n", + " print(''.join(['-' for i in range(level)]) + body.name)\n", + " print_bodies(body, level + 1)\n", + " body = parent.next_body(body)\n", + "\n", + "print(\"The spec has the following actuators:\")\n", + "for actuator in spec.actuators:\n", + " print(actuator.name)\n", + "\n", + "print(\"\\nThe spec has the following bodies:\")\n", + "print_bodies(spec.worldbody)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "GeiFFBYxLdEK" + }, + "source": [ + "An `mjSpec` can be compiled multiple times. If the state has to be preserved between different compilations, then the function `recompile()` must be used, which returns a new `mjData` that contains the mapped state, possibly having a different dimension from the origin." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "eiRXgh9OLdEK" + }, + "outputs": [], + "source": [ + "#@title Model re-compilation with state preservation {vertical-output: true}\n", + "\n", + "spec = mj.MjSpec.from_file(humanoid100_file)\n", + "model = spec.compile()\n", + "data = mj.MjData(model)\n", + "\n", + "# Run for 5 seconds\n", + "for i in range(1000):\n", + " mj.mj_step(model, data)\n", + "\n", + "# Show result\n", + "render(model, data)\n", + "\n", + "# Create list of all bodies we want to delete\n", + "body = spec.worldbody.first_body()\n", + "delete_list = []\n", + "while body:\n", + " geom_type = body.first_geom().type\n", + " if (geom_type == mj.mjtGeom.mjGEOM_BOX or\n", + " geom_type == mj.mjtGeom.mjGEOM_ELLIPSOID):\n", + " delete_list.append(body)\n", + " body = spec.worldbody.next_body(body)\n", + "\n", + "# Remove all bodies in the list from the spec\n", + "for body in delete_list:\n", + " spec.detach_body(body)\n", + "\n", + "# # Add another humanoid\n", + "spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n", + "attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n", + "attachment_frame.attach_body(spec_humanoid.body('torso'), 'a', 'b')\n", + "\n", + "# Recompile preserving the state\n", + "new_model, new_data = spec.recompile(model, data)\n", + "\n", + "# Show result\n", + "render(new_model, new_data)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "kuTWD415LdEK" + }, + "source": [ + "Let us load the humanoid model and inspect it." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "5d1wmQM2LdEK" + }, + "outputs": [], + "source": [ + "#@title Humanoid model {vertical-output: true}\n", + "\n", + "spec = mj.MjSpec.from_file(humanoid_file)\n", + "\n", + "model = spec.compile()\n", + "render(model)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "38PXB1rWLdEK" + }, + "source": [ + "We wish to remove the arms and replace them with the legs. This can be done by first storing the arm positions into frames attached to the torso. Then we can detach the arms and self-attach the legs into the frames." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "0eaNq0Q7LdEK" + }, + "outputs": [], + "source": [ + "#@title Humanoid with arms replaced by legs {vertical-output: true}\n", + "\n", + "spec = mj.MjSpec.from_file(humanoid_file)\n", + "spec.copy_during_attach = True\n", + "\n", + "# Get the torso, arm, and leg bodies\n", + "arm_left = spec.body('upper_arm_left')\n", + "arm_right = spec.body('upper_arm_right')\n", + "leg_left = spec.body('thigh_left')\n", + "leg_right = spec.body('thigh_right')\n", + "torso = spec.body('torso')\n", + "\n", + "# Attach frames at the arm positions\n", + "shoulder_left = torso.add_frame(pos=arm_left.pos)\n", + "shoulder_right = torso.add_frame(pos=arm_right.pos)\n", + "\n", + "# Remove the arms\n", + "spec.detach_body(arm_left)\n", + "spec.detach_body(arm_right)\n", + "\n", + "# Add new legs\n", + "shoulder_left.attach_body(leg_left, 'shoulder', 'left')\n", + "shoulder_right.attach_body(leg_right, 'shoulder', 'right')\n", + "\n", + "model = spec.compile()\n", + "render(model, height=400)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "HfhxL2EqLdEK" + }, + "source": [ + "Similarly, different models can be attach together. Here, the right arm is detached and a robot arm from a different model is attached in its place." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "uS4LGbI7LdEK" + }, + "outputs": [], + "source": [ + "#@title Humanoid with Franka arm {vertical-output: true}\n", + "\n", + "spec = mj.MjSpec.from_file(humanoid_file)\n", + "franka = mj.MjSpec.from_file(franka_file)\n", + "\n", + "if hasattr(spec, 'compiler'):\n", + " spec.compiler.degree = False # MuJoCo dev (next release).\n", + "else:\n", + " spec.degree = False # MuJoCo release\n", + "\n", + "# Replace right arm with frame\n", + "arm_right = spec.body('upper_arm_right')\n", + "torso = spec.body('torso')\n", + "shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n", + "spec.detach_body(arm_right)\n", + "\n", + "# Attach Franka arm to humanoid\n", + "franka_arm = franka.body('fr3_link2')\n", + "shoulder_right.attach_body(franka_arm, 'franka', '')\n", + "\n", + "model = spec.compile()\n", + "render(model, height=400)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "CWXYy_1uLdEK" + }, + "source": [ + "When doing this, the actuators and all other objects referenced by the attached sub-tree are imported in the new model. All assets are currently imported, referenced or not." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "UwWDD-NHLdEK" + }, + "outputs": [], + "source": [ + "#@title Imported actuators {vertical-output: true}\n", + "\n", + "for actuator in spec.actuators:\n", + " print(actuator.name)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hDvt3vcxLdEK" + }, + "source": [ + "Domain randomization can be performed by attaching multiple times the same spec, edited each time with a new instance of randomized parameters." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "oPPFbWawLdEK" + }, + "outputs": [], + "source": [ + "#@title Humanoid with randomized heads and arm poses {vertical-output: true}\n", + "\n", + "humanoid = mj.MjSpec.from_file(humanoid_file)\n", + "spec = mj.MjSpec()\n", + "spec.copy_during_attach = True\n", + "\n", + "# Delete all key frames to avoid name conflicts\n", + "while humanoid.keys:\n", + " humanoid.keys[-1].delete()\n", + "\n", + "# Create a grid of humanoids by attaching humanoid to spec multiple times\n", + "for i in range(4):\n", + " for j in range(4):\n", + " humanoid.materials[0].rgba = [\n", + " np.random.uniform(), np.random.uniform(),\n", + " np.random.uniform(), 1] # Randomize color\n", + " humanoid.body('head').first_geom().size = [\n", + " .18*np.random.uniform(), 0, 0] # Randomize head size\n", + " humanoid.body('upper_arm_left').quat = [\n", + " np.random.uniform(), np.random.uniform(),\n", + " np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n", + " humanoid.body('upper_arm_right').quat = [\n", + " np.random.uniform(), np.random.uniform(),\n", + " np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n", + "\n", + " # attach randomized humanoid to parent spec\n", + " frame = spec.worldbody.add_frame(pos=[i, j, 0])\n", + " frame.attach_body(humanoid.body('torso'), str(i), str(j))\n", + "\n", + "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", + " targetbody='3torso3', diffuse=[.8, .8, .8],\n", + " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", + "model = spec.compile()\n", + "render(model, height=400)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PML2pxYgLdEK" + }, + "source": [ + "## Model scaling" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "pcUNLmQBLdEK" + }, + "outputs": [], + "source": [ + "#@title Uniformly scale humanoid {vertical-output: true}\n", + "\n", + "def scale_spec(spec, scale):\n", + " scaled_spec = spec.copy()\n", + " # Traverse the kinematic tree, scaling all geoms\n", + " def scale_bodies(parent, scale=1.0):\n", + " body = parent.first_body()\n", + " while body:\n", + " if body.pos is not None:\n", + " body.pos = body.pos * scale\n", + " for geom in body.geoms:\n", + " geom.fromto = geom.fromto * scale\n", + " geom.size = geom.size * scale\n", + " if geom.pos is not None:\n", + " geom.pos = geom.pos * scale\n", + " scale_bodies(body, scale)\n", + " body = parent.next_body(body)\n", + "\n", + " scale_bodies(scaled_spec.body('world'), scale)\n", + " return scaled_spec\n", + "\n", + "spec = mj.MjSpec.from_string(arena_xml)\n", + "humanoid = mj.MjSpec.from_file(humanoid_file)\n", + "small_humanoid = scale_spec(humanoid, 0.75)\n", + "large_humanoid = scale_spec(humanoid, 1.25)\n", + "\n", + "# Create a line-up of humanoids\n", + "frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(humanoid.body('torso'), str(0))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(small_humanoid.body('torso'), str(1))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n", + "frame.attach_body(large_humanoid.body('torso'), str(2))\n", + "\n", + "\n", + "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", + " targetbody='1torso', diffuse=[.8, .8, .8],\n", + " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", + "model = spec.compile()\n", + "render(model, height=400, camera=cam)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "RYbaTPNmLdEK" + }, + "source": [ + "We can scale the size of a model by traversing the kinematic tree and applying the scale to the relevant geoms. Above we can see humanoids of three different sizes." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "u-ejx8lKLdEK" + }, + "outputs": [], + "source": [ + "# @title Scaling actuator forces {vertical-output: true}\n", + "\n", + "def scale_spec(spec, scale, scale_actuators=False):\n", + " scaled_spec = spec.copy()\n", + " # Traverse the kinematic tree, scaling all geoms\n", + " def scale_bodies(parent, scale=1.0):\n", + " body = parent.first_body()\n", + " while body:\n", + " if body.pos is not None:\n", + " body.pos = body.pos * scale\n", + " for geom in body.geoms:\n", + " geom.fromto = geom.fromto * scale\n", + " geom.size = geom.size * scale\n", + " if geom.pos is not None:\n", + " geom.pos = geom.pos * scale\n", + " scale_bodies(body, scale)\n", + " body = parent.next_body(body)\n", + "\n", + " if scale_actuators:\n", + " # scale gear\n", + " for actuator in scaled_spec.actuators:\n", + " # scale the actuator gear by (scale ** 2),\n", + " # this is because muscle force-generating capacity\n", + " # scales with the cross-sectional area of the muscle\n", + " actuator.gear = actuator.gear * scale * scale\n", + "\n", + " # scale the z-position of the humanoid for all keypoints\n", + " for keypoint in scaled_spec.keys:\n", + " qpos = keypoint.qpos\n", + " qpos[2] = qpos[2] * scale\n", + " keypoint.qpos = qpos\n", + " keypoint.qpos[2] = keypoint.qpos[2] * scale\n", + "\n", + " scale_bodies(scaled_spec.body('world'), scale)\n", + " return scaled_spec\n", + "\n", + "# Create specs\n", + "scale = 0.6\n", + "spec = mj.MjSpec.from_string(arena_xml)\n", + "humanoid = mj.MjSpec.from_file(humanoid_file)\n", + "small_humanoid = scale_spec(humanoid, scale)\n", + "small_humanoid_actuators_scaled = scale_spec(humanoid, scale, True)\n", + "\n", + "# Create a line-up of humanoids\n", + "squat_qpos = []\n", + "\n", + "# Add unscaled humanoid\n", + "frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(humanoid.body('torso'), str(0))\n", + "# Record squat pose\n", + "humanoid_squat = humanoid.key('squat').qpos\n", + "humanoid_squat[:2] = frame.pos[:2]\n", + "humanoid_squat[3:7] = frame.quat\n", + "squat_qpos.append(humanoid_squat)\n", + "\n", + "# Add small humanoid\n", + "frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(small_humanoid.body('torso'), str(1))\n", + "# Record squat pose\n", + "humanoid_squat = small_humanoid.key('squat').qpos\n", + "humanoid_squat[:2] = frame.pos[:2]\n", + "humanoid_squat[3:7] = frame.quat\n", + "squat_qpos.append(humanoid_squat)\n", + "\n", + "# Add small humanoid with scaled actuators\n", + "frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n", + "frame.attach_body(small_humanoid_actuators_scaled.body('torso'), str(2))\n", + "# Record squat pose\n", + "humanoid_squat = small_humanoid_actuators_scaled.key('squat').qpos\n", + "humanoid_squat[:2] = frame.pos[:2]\n", + "humanoid_squat[3:7] = frame.quat\n", + "squat_qpos.append(humanoid_squat)\n", + "squat_qpos = np.concatenate(squat_qpos)\n", + "\n", + "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", + " targetbody='1torso', diffuse=[.8, .8, .8],\n", + " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", + "model = spec.compile()\n", + "\n", + "# Initialize to squat position\n", + "data = mj.MjData(model)\n", + "data.qpos = squat_qpos\n", + "\n", + "# jumping motion\n", + "u_t = lambda t: 10.0 * t / duration\n", + "\n", + "# Simulate and display video.\n", + "duration = 2 # (seconds)\n", + "framerate = 30 # (Hz)\n", + "frames = []\n", + "\n", + "with mj.Renderer(model, 480, 640) as renderer:\n", + " while data.time < duration:\n", + " data.ctrl = u_t(data.time)\n", + " mj.mj_step(model, data)\n", + " if len(frames) < data.time * framerate:\n", + " renderer.update_scene(data, camera=cam)\n", + " pixels = renderer.render()\n", + " frames.append(pixels)\n", + "\n", + "media.show_video(frames, fps=framerate, height=400)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "uKhDI_IfLdEK" + }, + "source": [ + "We can also apply scaling to the actuators. In the humanoid case, scaling the geoms without scaling the `gear` parameter for the actuators results in a humanoid that can jump higher proportional to its size." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "ovQIAxn7LdEK" + }, + "outputs": [], + "source": [ + "# @title Long-limbed humanoid {vertical-output: true}\n", + "\n", + "def scale_spec(spec, scale):\n", + " scaled_spec = spec.copy()\n", + " # Traverse the kinematic tree, scaling all geoms\n", + " def scale_bodies(parent, scale=1.0):\n", + " if parent is not None:\n", + " for geom in parent.geoms:\n", + " # Only scale fromto, not size to scale length of capsules\n", + " geom.fromto = geom.fromto * scale\n", + " if geom.pos is not None:\n", + " geom.pos = geom.pos * scale\n", + " body = parent.first_body()\n", + " while body:\n", + " if body.pos is not None:\n", + " body.pos = body.pos * scale\n", + " scale_bodies(body, scale)\n", + " body = parent.next_body(body)\n", + "\n", + " # Scale all the limbs\n", + " scale_bodies(scaled_spec.body('upper_arm_right'), scale)\n", + " scale_bodies(scaled_spec.body('upper_arm_left'), scale)\n", + " scale_bodies(scaled_spec.body('thigh_right'), scale)\n", + " scale_bodies(scaled_spec.body('thigh_left'), scale)\n", + " return scaled_spec\n", + "\n", + "spec = mj.MjSpec.from_string(arena_xml)\n", + "humanoid = mj.MjSpec.from_file(humanoid_file)\n", + "small_humanoid = scale_spec(humanoid, 1.25)\n", + "large_humanoid = scale_spec(humanoid, 2)\n", + "\n", + "# Create a line-up of humanoids by attaching\n", + "frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(humanoid.body('torso'), str(0), str(0))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[0, 0, 0.2],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(small_humanoid.body('torso'), str(0), str(1))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[1, 0, 0.8],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n", + "frame.attach_body(large_humanoid.body('torso'), str(0), str(2))\n", + "\n", + "\n", + "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", + " targetbody='0torso1', diffuse=[.8, .8, .8],\n", + " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", + "model = spec.compile()\n", + "\n", + "# camera options\n", + "render(model, height=400, camera=cam)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8o1daXIOLdEK" + }, + "source": [ + "We can also apply scaling to the model non-uniformly. In this instance we scale the humanoid to have long limbs, by only applying the scale to the length of the capsule geoms for the arms, legs and feet." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "UBSo2nfQLdEK" + }, + "outputs": [], + "source": [ + "# @title Meshes {vertical-output: true}\n", + "\n", + "def scale_spec(spec, scale):\n", + " scaled_spec = spec.copy()\n", + " # scale all meshes\n", + " for mesh in scaled_spec.meshes:\n", + " if mesh.scale is None:\n", + " mesh.scale = np.ones(3)\n", + " mesh.scale = mesh.scale * scale\n", + "\n", + " # Traverse the kinematic tree\n", + " def scale_bodies(parent, scale=1.0):\n", + " if parent is not None:\n", + " for geom in parent.geoms:\n", + " if geom.pos is not None:\n", + " geom.pos = geom.pos * scale\n", + " body = parent.first_body()\n", + " while body:\n", + " if body.pos is not None:\n", + " body.pos = body.pos * scale\n", + " scale_bodies(body, scale)\n", + " body = parent.next_body(body)\n", + "\n", + " # Scale all the limbs\n", + " scale_bodies(scaled_spec.body('world'), scale)\n", + "\n", + " return scaled_spec\n", + "\n", + "spec = mj.MjSpec.from_string(arena_xml)\n", + "fly = mj.MjSpec.from_file(fly_file)\n", + "# Remove lights from fly so they are not duplicated in line-up\n", + "for light in fly.lights:\n", + " light.delete()\n", + "\n", + "small_fly = scale_spec(fly, 1.25)\n", + "large_fly = scale_spec(fly, 2)\n", + "\n", + "# Create a line-up of flys by attaching\n", + "frame = spec.worldbody.add_frame(pos=[-1, 0, 0.25],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(fly.body('thorax'), str(0), str(0))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[0, 0, 0.25],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n", + "frame.attach_body(small_fly.body('thorax'), str(0), str(1))\n", + "\n", + "frame = spec.worldbody.add_frame(pos=[1, 0, 0.25],\n", + " quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n", + "frame.attach_body(large_fly.body('thorax'), str(0), str(2))\n", + "\n", + "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", + " targetbody='0thorax1', diffuse=[.8, .8, .8],\n", + " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", + "model = spec.compile()\n", + "render(model, height=400, camera=cam)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "qhXwxLe3LdEK" + }, + "source": [ + "# dm_control example" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "enuJ_YIqLdEK" }, "source": [ "A key feature is the ability to easily attach multiple models into a larger one. Disambiguation of duplicated names from different\n", @@ -874,7 +2385,7 @@ "cell_type": "code", "execution_count": 0, "metadata": { - "id": "7C-hfbtj8nRV" + "id": "4p3P_dP8LdEK" }, "outputs": [], "source": [ @@ -921,7 +2432,7 @@ { "cell_type": "markdown", "metadata": { - "id": "MQGsxnIB_RLO" + "id": "Mqr8rXLILdEK" }, "source": [ "The `Leg` class describes an abstract articulated leg, with two joints and corresponding proportional-derivative actuators.\n", @@ -929,14 +2440,14 @@ "Note that:\n", "\n", "- MJCF attributes correspond directly to arguments of the `add_()` methods.\n", - "- When referencing elements, e.g when specifying the joint to which an actuator is attached, the name string of the MJCF elements is used." + "- When referencing elements, e.g. when specifying the joint to which an actuator is attached, the name string of the MJCF elements is used." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { - "id": "kMiuMyZW_XoB" + "id": "1z2NBpAPLdEK" }, "outputs": [], "source": [ @@ -973,7 +2484,7 @@ { "cell_type": "markdown", "metadata": { - "id": "QMQ3jc6-_toj" + "id": "865FGuntLdEL" }, "source": [ "The `make_creature` function uses the `attach()` method to procedurally attach legs to the torso. Note that at this stage both the torso and hip attachment sites are children of the `worldbody`, since their parent body has yet to be instantiated. We'll now make an arena with a chequered floor and two lights, and place our creatures in a grid." @@ -983,11 +2494,11 @@ "cell_type": "code", "execution_count": 0, "metadata": { - "id": "vt2JwXd__1cT" + "id": "2fPaSkgfLdEL" }, "outputs": [], "source": [ - "#@title Six Creatures on a floor.{vertical-output: true}\n", + "#@title Six Creatures on a floor {vertical-output: true}\n", "\n", "arena = mj.MjSpec()\n", "\n", @@ -1032,7 +2543,7 @@ { "cell_type": "markdown", "metadata": { - "id": "mPUGkrCzAFMg" + "id": "tq5mKlc_LdEL" }, "source": [ "Multi-legged creatures, ready to roam! Let's inject some controls and watch them move. We'll generate a sinusoidal open-loop control signal of fixed frequency and random phase, recording both video frames and the horizontal positions of the torso geoms, in order to plot the movement trajectories." @@ -1042,11 +2553,11 @@ "cell_type": "code", "execution_count": 0, "metadata": { - "id": "7gz9FfNzGxPO" + "id": "i37FpwCeLdEL" }, "outputs": [], "source": [ - "#@title Video of the movement{vertical-output: true}\n", + "#@title Video of the movement {vertical-output: true}\n", "\n", "data = mj.MjData(model)\n", "duration = 10 # (Seconds)\n", @@ -1090,11 +2601,11 @@ "cell_type": "code", "execution_count": 0, "metadata": { - "id": "qt2L52e_Tcgt" + "id": "uFrvaih4LdEL" }, "outputs": [], "source": [ - "#@title Movement trajectories{vertical-output: true}\n", + "#@title Movement trajectories {vertical-output: true}\n", "\n", "creature_colors = [torso.rgba[:3] for torso in torsos_model]\n", "fig, ax = plt.subplots(figsize=(4, 4))\n", @@ -1105,321 +2616,18 @@ { "cell_type": "markdown", "metadata": { - "id": "kSEUoxifxYJ4" + "id": "FMW4l-fSLdEL" }, "source": [ "The plot above shows the corresponding movement trajectories of creature positions. Note how `mjSpec` attribute `id` were used to access both `xpos` and `rgba` values. This attribute is valid only after a model is compiled." ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "QZ8alJZz8cB1" - }, - "source": [ - "# Model editing" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "JN3Z4v0PyXKa" - }, - "source": [ - "`mjSpec` elements can be traversed in two ways:\n", - "- For elements inside the kinematic tree, the tree can be traversed using the `first` and `next` functions.\n", - "- For all other elements, we provide a list.\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "8IcB7nezblyT" - }, - "outputs": [], - "source": [ - "#@title Traversing the spec.{vertical-output: true}\n", - "\n", - "# Get MuJoCo's humanoid model.\n", - "print('Getting MuJoCo humanoid XML description from GitHub:')\n", - "!git clone https://github.com/google-deepmind/mujoco\n", - "humanoid_file = 'mujoco/model/humanoid/humanoid.xml'\n", - "humanoid100_file = 'mujoco/model/humanoid/humanoid100.xml'\n", - "\n", - "spec = mj.MjSpec.from_file(humanoid_file)\n", - "\n", - "# Function that recursively prints all body names\n", - "def print_bodies(parent, level=0):\n", - " body = parent.first_body()\n", - " while body:\n", - " print(''.join(['-' for i in range(level)]) + body.name)\n", - " print_bodies(body, level + 1)\n", - " body = parent.next_body(body)\n", - "\n", - "print(\"The spec has the following actuators:\")\n", - "for actuator in spec.actuators:\n", - " print(actuator.name)\n", - "\n", - "print(\"\\nThe spec has the following bodies:\")\n", - "print_bodies(spec.worldbody)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "hcGI4orhyzvc" - }, - "source": [ - "An `mjSpec` can be compiled multiple times. If the state has to be preserved between different compilations, then the function `recompile()` must be used, which returns a new `mjData` that contains the mapped state, possibly having a different dimension from the origin." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "uh_N1Fkqk-Mi" - }, - "outputs": [], - "source": [ - "#@title Model re-compilation with state preservation.{vertical-output: true}\n", - "\n", - "spec = mj.MjSpec.from_file(humanoid100_file)\n", - "model = spec.compile()\n", - "data = mj.MjData(model)\n", - "\n", - "# Run for 5 seconds\n", - "for i in range(1000):\n", - " mj.mj_step(model, data)\n", - "\n", - "# Show result\n", - "render(model, data)\n", - "\n", - "# Create list of all bodies we want to delete\n", - "body = spec.worldbody.first_body()\n", - "delete_list = []\n", - "while body:\n", - " geom_type = body.first_geom().type\n", - " if (geom_type == mj.mjtGeom.mjGEOM_BOX or\n", - " geom_type == mj.mjtGeom.mjGEOM_ELLIPSOID):\n", - " delete_list.append(body)\n", - " body = spec.worldbody.next_body(body)\n", - "\n", - "# Remove all bodies in the list from the spec\n", - "for body in delete_list:\n", - " spec.detach_body(body)\n", - "\n", - "# # Add another humanoid\n", - "spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n", - "attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n", - "attachment_frame.attach_body(spec_humanoid.body('torso'), 'a', 'b')\n", - "\n", - "# Recompile preserving the state\n", - "new_model, new_data = spec.recompile(model, data)\n", - "\n", - "# Show result\n", - "render(new_model, new_data)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "XmSlXirVzLqt" - }, - "source": [ - "Let us load the humanoid model and inspect it." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "UywMzsp5Hnk2" - }, - "outputs": [], - "source": [ - "#@title Humanoid model.{vertical-output: true}\n", - "\n", - "spec = mj.MjSpec.from_file(humanoid_file)\n", - "\n", - "model = spec.compile()\n", - "render(model)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "owcmKeuSzQRy" - }, - "source": [ - "We wish to remove the arms and replace them with the legs. This can be done by first storing the arm positions into frames attached to the torso. Then we can detach the arms and self-attach the legs into the frames." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "qZCyv-B0IGiG" - }, - "outputs": [], - "source": [ - "#@title Humanoid with arms replaced by legs.{vertical-output: true}\n", - "\n", - "spec = mj.MjSpec.from_file(humanoid_file)\n", - "spec.copy_during_attach = True\n", - "\n", - "# Get the torso, arm, and leg bodies\n", - "arm_left = spec.body('upper_arm_left')\n", - "arm_right = spec.body('upper_arm_right')\n", - "leg_left = spec.body('thigh_left')\n", - "leg_right = spec.body('thigh_right')\n", - "torso = spec.body('torso')\n", - "\n", - "# Attach frames at the arm positions\n", - "shoulder_left = torso.add_frame(pos=arm_left.pos)\n", - "shoulder_right = torso.add_frame(pos=arm_right.pos)\n", - "\n", - "# Remove the arms\n", - "spec.detach_body(arm_left)\n", - "spec.detach_body(arm_right)\n", - "\n", - "# Add new legs\n", - "shoulder_left.attach_body(leg_left, 'shoulder', 'left')\n", - "shoulder_right.attach_body(leg_right, 'shoulder', 'right')\n", - "\n", - "model = spec.compile()\n", - "render(model, height=400)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "LnEwEjW3zdua" - }, - "source": [ - "Similarly, different models can be attach together. Here, the right arm is detached and a robot arm from a different model is attached in its place." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "w-NdFhSIIrLL" - }, - "outputs": [], - "source": [ - "#@title Humanoid with Franka arm.{vertical-output: true}\n", - "\n", - "# Get Franka arm from the MuJoCo Menagerie.\n", - "!git clone https://github.com/google-deepmind/mujoco_menagerie\n", - "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n", - "\n", - "spec = mj.MjSpec.from_file(humanoid_file)\n", - "franka = mj.MjSpec.from_file(franka_file)\n", - "\n", - "if hasattr(spec, 'compiler'):\n", - " spec.compiler.degree = False # MuJoCo dev (next release).\n", - "else:\n", - " spec.degree = False # MuJoCo release\n", - "\n", - "# Replace right arm with frame\n", - "arm_right = spec.body('upper_arm_right')\n", - "torso = spec.body('torso')\n", - "shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n", - "spec.detach_body(arm_right)\n", - "\n", - "# Attach Franka arm to humanoid\n", - "franka_arm = franka.body('fr3_link2')\n", - "shoulder_right.attach_body(franka_arm, 'franka', '')\n", - "\n", - "model = spec.compile()\n", - "render(model, height=400)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "e_idaggAznXu" - }, - "source": [ - "When doing this, the actuators and all other objects referenced by the attached sub-tree are imported in the new model. All assets are currently imported, referenced or not." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "50lOgJ7mQ2bV" - }, - "outputs": [], - "source": [ - "#@title Imported actuators.{vertical-output: true}\n", - "\n", - "for actuator in spec.actuators:\n", - " print(actuator.name)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "APDoWK4mz0aJ" - }, - "source": [ - "Domain randomization can be performed by attaching multiple times the same spec, edited each time with a new instance of randomized parameters." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "oHjdgkISNLKy" - }, - "outputs": [], - "source": [ - "#@title Humanoid with randomized heads and arm poses.{vertical-output: true}\n", - "\n", - "humanoid = mj.MjSpec.from_file(humanoid_file)\n", - "spec = mj.MjSpec()\n", - "spec.copy_during_attach = True\n", - "\n", - "# Delete all key frames to avoid name conflicts\n", - "while humanoid.keys:\n", - " humanoid.keys[-1].delete()\n", - "\n", - "# Create a grid of humanoids by attaching humanoid to spec multiple times\n", - "for i in range(4):\n", - " for j in range(4):\n", - " humanoid.materials[0].rgba = [\n", - " np.random.uniform(), np.random.uniform(),\n", - " np.random.uniform(), 1] # Randomize color\n", - " humanoid.body('head').first_geom().size = [\n", - " .18*np.random.uniform(), 0, 0] # Randomize head size\n", - " humanoid.body('upper_arm_left').quat = [\n", - " np.random.uniform(), np.random.uniform(),\n", - " np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n", - " humanoid.body('upper_arm_right').quat = [\n", - " np.random.uniform(), np.random.uniform(),\n", - " np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n", - "\n", - " # attach randomized humanoid to parent spec\n", - " frame = spec.worldbody.add_frame(pos=[i, j, 0])\n", - " frame.attach_body(humanoid.body('torso'), str(i), str(j))\n", - "\n", - "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", - " targetbody='3torso3', diffuse=[.8, .8, .8],\n", - " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", - "model = spec.compile()\n", - "render(model, height=400)" - ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ - "sJFuNetilv4m", - "yXY7HGfVsVlo" + "sJFuNetilv4m" ], "gpuClass": "premium", "private_outputs": true, diff --git a/python/mujoco/CMakeLists.txt b/python/mujoco/CMakeLists.txt index e6c067f7..b1aec9ed 100644 --- a/python/mujoco/CMakeLists.txt +++ b/python/mujoco/CMakeLists.txt @@ -84,7 +84,7 @@ if(NOT TARGET mujoco) if(MUJOCO_FRAMEWORK) message("MuJoCo framework is at ${MUJOCO_FRAMEWORK}/mujoco.framework") set(MUJOCO_LIBRARY - ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.1.dylib + ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.3.dylib ) target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK}) endif() @@ -92,7 +92,7 @@ if(NOT TARGET mujoco) if(NOT MUJOCO_FRAMEWORK) find_library( - MUJOCO_LIBRARY mujoco mujoco.3.3.1 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED + MUJOCO_LIBRARY mujoco mujoco.3.3.3 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED ) find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED) message("MuJoCo is at ${MUJOCO_LIBRARY}") @@ -140,7 +140,7 @@ findorfetch( GIT_REPO https://github.com/abseil/abseil-cpp GIT_TAG - 9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0 + d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1 TARGETS ${MUJOCO_PYTHON_ABSL_TARGETS} EXCLUDE_FROM_ALL @@ -173,7 +173,7 @@ findorfetch( GIT_REPO https://gitlab.com/libeigen/eigen GIT_TAG - 66f7f51b7e069d0a03a21157fa60b24aece69aeb + 464c1d097891a1462ab28bf8bb763c1683883892 TARGETS Eigen3::Eigen EXCLUDE_FROM_ALL @@ -278,6 +278,25 @@ target_link_libraries( raw ) +add_library(structs_wrappers STATIC + structs_wrappers.cc + serialization.h + indexers.cc +) +target_include_directories(structs_wrappers PRIVATE ${Python3_INCLUDE_DIRS}) +target_link_libraries( + structs_wrappers + PRIVATE absl::flat_hash_map + absl::span + crossplatform + mujoco + raw + structs_header + pybind11::headers + Eigen3::Eigen +) + + add_library(functions_header INTERFACE) target_sources(functions_header INTERFACE functions.h) set_target_properties(functions_header PROPERTIES PUBLIC_HEADER functions.h) @@ -403,6 +422,7 @@ target_link_libraries( func_wrap function_traits structs_header + structs_wrappers ) if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/specs.cc.inc) @@ -417,6 +437,8 @@ endif() mujoco_pybind11_module( _specs specs.cc + specs_wrapper.h + specs_wrapper.cc specs.cc.inc ) target_link_libraries( diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index fb2ed004..99ca73ab 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -623,10 +623,12 @@ class MuJoCoBindingsTest(parameterized.TestCase): ncon = 13 nefc = 17 - mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc) + nj = 21 + mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc, nJ=nj) self.assertLen(self.data.contact, ncon) self.assertEqual(self.data.efc_id.shape, (nefc,)) + self.assertEqual(self.data.efc_J.shape, (nj,)) self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4)) expected_error = 'insufficient arena memory available' @@ -1083,6 +1085,20 @@ Euler integrator, semi-implicit in velocity. ): mujoco.mj_forward(self.model, self.data) + def test_timer_installed_by_default(self): + timer_step = mujoco.mjtTimer.mjTIMER_STEP + self.assertEqual(self.data.timer[timer_step].number, 0) + self.assertEqual(self.data.timer[timer_step].duration, 0.0) + + mujoco.mj_step(self.model, self.data) + self.assertEqual(self.data.timer[timer_step].number, 1) + duration_1 = self.data.timer[timer_step].duration + self.assertGreater(duration_1, 0.0) + + mujoco.mj_step(self.model, self.data, 5) + self.assertEqual(self.data.timer[timer_step].number, 6) + self.assertGreater(self.data.timer[timer_step].duration, duration_1) + def test_mjcb_time(self): class CallCounter: diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index ba412275..1db57b3c 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -84,6 +84,7 @@ def _value_binding_code( fulltype = fulltype.replace('mjOption', 'raw::MjOption') fulltype = fulltype.replace('mjVisual', 'raw::MjVisual') fulltype = fulltype.replace('mjStatistic', 'raw::MjStatistic') + element = '.element' if fullvarname == 'plugin' else '' def_property_args = ( f'"{varname}"', @@ -91,7 +92,7 @@ def _value_binding_code( return self.{fullvarname}; }}""", f"""[]({rawclassname}& self, {fulltype} {varname}) {{ - self.{fullvarname} = {varname}; + self.{fullvarname}{element} = {varname}{element}; }}""", ) @@ -613,12 +614,44 @@ def generate_find() -> None: print(code) +def generate_signature() -> None: + """Generate signature functions.""" + for key, _, _, _, _ in SPECS: + elem = key.removeprefix('mjs') + titlecase = 'Mjs' + elem + code = f"""\n + {key}.def_property_readonly("signature", + [](raw::{titlecase}& self) -> uint64_t {{ + return mjs_getSpec(self.element)->element->signature; + }}); + """ + print(code) + + +def generate_id() -> None: + """Generate id functions.""" + for key, _, _, _, _ in SPECS: + if key == 'mjsPlugin': + continue + elem = key.removeprefix('mjs') + titlecase = 'Mjs' + elem + code = f"""\n + {key}.def_property_readonly("id", + [](raw::{titlecase}& self) -> int {{ + return mjs_getId(self.element); + }}); + """ + print(code) + + def main(argv: Sequence[str]) -> None: if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') generate() generate_add() generate_find() + generate_signature() + generate_id() if __name__ == '__main__': diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 7fed04f7..b23c6d1e 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -211,6 +211,7 @@ PYBIND11_MODULE(_functions, pymodule) { Def(pymodule); Def(pymodule); Def(pymodule); + Def(pymodule); Def(pymodule); DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM, "n")( pymodule, @@ -1467,10 +1468,10 @@ PYBIND11_MODULE(_functions, pymodule) { pymodule.def( "_realloc_con_efc", - [](MjDataWrapper& d, int ncon, int nefc) { + [](MjDataWrapper& d, int ncon, int nefc, int nJ) { raw::MjData* data = d.get(); - auto cleanup = [](raw::MjData* data) { + auto cleanup = [](raw::MjData* data, int nJ) { #ifdef ADDRESS_SANITIZER ASAN_POISON_MEMORY_REGION( static_cast(data->arena), @@ -1479,6 +1480,7 @@ PYBIND11_MODULE(_functions, pymodule) { data->parena = 0; data->ncon = 0; data->nefc = 0; + if (nJ > -1) data->nJ = 0; data->contact = static_cast(data->arena); #define X(type, name, nr, nc) data->name = nullptr; MJDATA_ARENA_POINTERS_SOLVER @@ -1486,14 +1488,15 @@ PYBIND11_MODULE(_functions, pymodule) { #undef X }; - cleanup(data); + cleanup(data, nJ); data->ncon = ncon; data->nefc = nefc; + if (nJ > -1) data->nJ = nJ; data->contact = static_cast(InterceptMjErrors(::mj_arenaAllocByte)( data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact))); if (!data->contact) { - cleanup(data); + cleanup(data, nJ); throw FatalError("insufficient arena memory available"); } @@ -1505,7 +1508,7 @@ PYBIND11_MODULE(_functions, pymodule) { data->name = static_cast(InterceptMjErrors(::mj_arenaAllocByte)( \ data, sizeof(type) * (nr) * (nc), alignof(type))); \ if (!data->name) { \ - cleanup(data); \ + cleanup(data, nJ); \ throw FatalError("insufficient arena memory available"); \ } @@ -1519,7 +1522,7 @@ PYBIND11_MODULE(_functions, pymodule) { #undef MJ_M #define MJ_M(x) x }, - py::arg("d"), py::arg("ncon"), py::arg("nefc"), + py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1, py::call_guard()); } // PYBIND11_MODULE NOLINT(readability/fn_size) } // namespace diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h index f39d2da2..9f26000d 100644 --- a/python/mujoco/indexer_xmacro.h +++ b/python/mujoco/indexer_xmacro.h @@ -152,7 +152,7 @@ X( int, light_, mode, nlight, 1 ) \ X( int, light_, bodyid, nlight, 1 ) \ X( int, light_, targetbodyid, nlight, 1 ) \ - X( mjtByte, light_, directional, nlight, 1 ) \ + X( int, light_, type, nlight, 1 ) \ X( mjtByte, light_, castshadow, nlight, 1 ) \ X( mjtByte, light_, active, nlight, 1 ) \ X( mjtNum, light_, pos, nlight, 3 ) \ diff --git a/python/mujoco/introspect/codegen/generate_functions.py b/python/mujoco/introspect/codegen/generate_functions.py index 1966e328..06beeaf2 100644 --- a/python/mujoco/introspect/codegen/generate_functions.py +++ b/python/mujoco/introspect/codegen/generate_functions.py @@ -99,8 +99,13 @@ class MjFunctionVisitor: return ''.join(strings) def visit(self, node: ClangJsonNode) -> None: - if (node.get('kind') == 'FunctionDecl' and - node.get('name', '').startswith('mj')): + # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated + # objects and doesn't need a python wrapper. + if ( + node.get('kind') == 'FunctionDecl' + and node.get('name', '').startswith('mj') + and node.get('name', '') != 'mjs_setUserValueWithCleanup' + ): func_decl = self._make_function(node) self._functions[func_decl.name] = func_decl diff --git a/python/mujoco/introspect/codegen/generate_structs.py b/python/mujoco/introspect/codegen/generate_structs.py index 1c50db22..24da6e2e 100644 --- a/python/mujoco/introspect/codegen/generate_structs.py +++ b/python/mujoco/introspect/codegen/generate_structs.py @@ -47,7 +47,7 @@ _EXCLUDED = ( 'mjResource_', ) -_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s\s+\((.+) x (.+)\)\Z') +_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s+\(([^\(\)]+) x ([^\(\)]+)\)\Z') def traverse(node, visitor): @@ -93,16 +93,19 @@ class MjStructVisitor: # No valid normalization, just parse the declname. return type_parsing.parse_type(declname) - def _make_comment(self, node: ClangJsonNode) -> str: + def _make_comment(self, node: ClangJsonNode, strip: bool = True) -> str: """Makes a comment string from a Clang AST FullComment node.""" kind = node.get('kind') if kind == 'TextComment': - return node['text'].replace('\N{NO-BREAK SPACE}', ' ').strip() + retval = node['text'].replace('\N{NO-BREAK SPACE}', ' ') else: strings = [] for child in node['inner']: - strings.append(self._make_comment(child)) - return ''.join(strings).strip() + strings.append(self._make_comment(child, strip=False)) + retval = ''.join(strings) + if strip: + retval = retval.strip() + return retval def _make_field( self, node: ClangJsonNode diff --git a/python/mujoco/introspect/enums.py b/python/mujoco/introspect/enums.py index 66b937a2..7cbe18bc 100644 --- a/python/mujoco/introspect/enums.py +++ b/python/mujoco/introspect/enums.py @@ -111,6 +111,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjCAMLIGHT_TARGETBODYCOM', 4), ]), )), + ('mjtLightType', + EnumDecl( + name='mjtLightType', + declname='enum mjtLightType_', + values=dict([ + ('mjLIGHT_SPOT', 0), + ('mjLIGHT_DIRECTIONAL', 1), + ('mjLIGHT_POINT', 2), + ('mjLIGHT_IMAGE', 3), + ]), + )), ('mjtTexture', EnumDecl( name='mjtTexture', @@ -139,6 +150,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjNTEXROLE', 10), ]), )), + ('mjtColorSpace', + EnumDecl( + name='mjtColorSpace', + declname='enum mjtColorSpace_', + values=dict([ + ('mjCOLORSPACE_AUTO', 0), + ('mjCOLORSPACE_LINEAR', 1), + ('mjCOLORSPACE_SRGB', 2), + ]), + )), ('mjtIntegrator', EnumDecl( name='mjtIntegrator', @@ -288,6 +309,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjNOBJECT', 26), ('mjOBJ_FRAME', 100), ('mjOBJ_DEFAULT', 101), + ('mjOBJ_MODEL', 102), ]), )), ('mjtConstraint', @@ -339,34 +361,35 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjSENS_ACTUATORVEL', 14), ('mjSENS_ACTUATORFRC', 15), ('mjSENS_JOINTACTFRC', 16), - ('mjSENS_BALLQUAT', 17), - ('mjSENS_BALLANGVEL', 18), - ('mjSENS_JOINTLIMITPOS', 19), - ('mjSENS_JOINTLIMITVEL', 20), - ('mjSENS_JOINTLIMITFRC', 21), - ('mjSENS_TENDONLIMITPOS', 22), - ('mjSENS_TENDONLIMITVEL', 23), - ('mjSENS_TENDONLIMITFRC', 24), - ('mjSENS_FRAMEPOS', 25), - ('mjSENS_FRAMEQUAT', 26), - ('mjSENS_FRAMEXAXIS', 27), - ('mjSENS_FRAMEYAXIS', 28), - ('mjSENS_FRAMEZAXIS', 29), - ('mjSENS_FRAMELINVEL', 30), - ('mjSENS_FRAMEANGVEL', 31), - ('mjSENS_FRAMELINACC', 32), - ('mjSENS_FRAMEANGACC', 33), - ('mjSENS_SUBTREECOM', 34), - ('mjSENS_SUBTREELINVEL', 35), - ('mjSENS_SUBTREEANGMOM', 36), - ('mjSENS_GEOMDIST', 37), - ('mjSENS_GEOMNORMAL', 38), - ('mjSENS_GEOMFROMTO', 39), - ('mjSENS_E_POTENTIAL', 40), - ('mjSENS_E_KINETIC', 41), - ('mjSENS_CLOCK', 42), - ('mjSENS_PLUGIN', 43), - ('mjSENS_USER', 44), + ('mjSENS_TENDONACTFRC', 17), + ('mjSENS_BALLQUAT', 18), + ('mjSENS_BALLANGVEL', 19), + ('mjSENS_JOINTLIMITPOS', 20), + ('mjSENS_JOINTLIMITVEL', 21), + ('mjSENS_JOINTLIMITFRC', 22), + ('mjSENS_TENDONLIMITPOS', 23), + ('mjSENS_TENDONLIMITVEL', 24), + ('mjSENS_TENDONLIMITFRC', 25), + ('mjSENS_FRAMEPOS', 26), + ('mjSENS_FRAMEQUAT', 27), + ('mjSENS_FRAMEXAXIS', 28), + ('mjSENS_FRAMEYAXIS', 29), + ('mjSENS_FRAMEZAXIS', 30), + ('mjSENS_FRAMELINVEL', 31), + ('mjSENS_FRAMEANGVEL', 32), + ('mjSENS_FRAMELINACC', 33), + ('mjSENS_FRAMEANGACC', 34), + ('mjSENS_SUBTREECOM', 35), + ('mjSENS_SUBTREELINVEL', 36), + ('mjSENS_SUBTREEANGMOM', 37), + ('mjSENS_GEOMDIST', 38), + ('mjSENS_GEOMNORMAL', 39), + ('mjSENS_GEOMFROMTO', 40), + ('mjSENS_E_POTENTIAL', 41), + ('mjSENS_E_KINETIC', 42), + ('mjSENS_CLOCK', 43), + ('mjSENS_PLUGIN', 44), + ('mjSENS_USER', 45), ]), )), ('mjtStage', @@ -426,6 +449,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjFLEXSELF_AUTO', 4), ]), )), + ('mjtSDFType', + EnumDecl( + name='mjtSDFType', + declname='enum mjtSDFType_', + values=dict([ + ('mjSDFTYPE_SINGLE', 0), + ('mjSDFTYPE_INTERSECTION', 1), + ('mjSDFTYPE_MIDSURFACE', 2), + ('mjSDFTYPE_COLLISION', 3), + ]), + )), ('mjtTaskStatus', EnumDecl( name='mjtTaskStatus', diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py index 737de191..1c4877b9 100644 --- a/python/mujoco/introspect/functions.py +++ b/python/mujoco/introspect/functions.py @@ -248,6 +248,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Compile spec to model.', )), + ('mj_copyBack', + FunctionDecl( + name='mj_copyBack', + return_type=ValueType(name='int'), + parameters=( + FunctionParameterDecl( + name='s', + type=PointerType( + inner_type=ValueType(name='mjSpec'), + ), + ), + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + ), + doc='Copy real-valued arrays from model to spec, returns 1 on success.', # pylint: disable=line-too-long + )), ('mj_recompile', FunctionDecl( name='mj_recompile', @@ -745,6 +765,34 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Copy mjData. m is only required to contain the size fields from MJMODEL_INTS.', # pylint: disable=line-too-long )), + ('mjv_copyData', + FunctionDecl( + name='mjv_copyData', + return_type=PointerType( + inner_type=ValueType(name='mjData'), + ), + parameters=( + FunctionParameterDecl( + name='dest', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='src', + type=PointerType( + inner_type=ValueType(name='mjData', is_const=True), + ), + ), + ), + doc='Copy mjData, skip large arrays not required for visualization.', + )), ('mj_resetData', FunctionDecl( name='mj_resetData', @@ -1801,6 +1849,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Run composite rigid body inertia algorithm (CRB).', )), + ('mj_makeM', + FunctionDecl( + name='mj_makeM', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + ), + doc='Make inertia matrix.', + )), ('mj_factorM', FunctionDecl( name='mj_factorM', @@ -2900,7 +2968,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc='Add inertia matrix to destination matrix. Destination can be sparse uncompressed, or dense when all int* are NULL', # pylint: disable=line-too-long + doc='Add inertia matrix to destination matrix. Destination can be sparse or dense when all int* are NULL.', # pylint: disable=line-too-long )), ('mj_applyFT', FunctionDecl( @@ -3983,44 +4051,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Move camera with mouse; action is mjtMouse.', )), - ('mjv_moveCameraFromState', - FunctionDecl( - name='mjv_moveCameraFromState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState', is_const=True), - ), - ), - FunctionParameterDecl( - name='action', - type=ValueType(name='int'), - ), - FunctionParameterDecl( - name='reldx', - type=ValueType(name='mjtNum'), - ), - FunctionParameterDecl( - name='reldy', - type=ValueType(name='mjtNum'), - ), - FunctionParameterDecl( - name='scn', - type=PointerType( - inner_type=ValueType(name='mjvScene', is_const=True), - ), - ), - FunctionParameterDecl( - name='cam', - type=PointerType( - inner_type=ValueType(name='mjvCamera'), - ), - ), - ), - doc='Move camera with mouse given a scene state; action is mjtMouse.', - )), ('mjv_movePerturb', FunctionDecl( name='mjv_movePerturb', @@ -4065,44 +4095,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Move perturb object with mouse; action is mjtMouse.', )), - ('mjv_movePerturbFromState', - FunctionDecl( - name='mjv_movePerturbFromState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState', is_const=True), - ), - ), - FunctionParameterDecl( - name='action', - type=ValueType(name='int'), - ), - FunctionParameterDecl( - name='reldx', - type=ValueType(name='mjtNum'), - ), - FunctionParameterDecl( - name='reldy', - type=ValueType(name='mjtNum'), - ), - FunctionParameterDecl( - name='scn', - type=PointerType( - inner_type=ValueType(name='mjvScene', is_const=True), - ), - ), - FunctionParameterDecl( - name='pert', - type=PointerType( - inner_type=ValueType(name='mjvPerturb'), - ), - ), - ), - doc='Move perturb object with mouse given a scene state; action is mjtMouse.', # pylint: disable=line-too-long - )), ('mjv_moveModel', FunctionDecl( name='mjv_moveModel', @@ -4532,48 +4524,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Update entire scene given model state.', )), - ('mjv_updateSceneFromState', - FunctionDecl( - name='mjv_updateSceneFromState', - return_type=ValueType(name='int'), - parameters=( - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState', is_const=True), - ), - ), - FunctionParameterDecl( - name='opt', - type=PointerType( - inner_type=ValueType(name='mjvOption', is_const=True), - ), - ), - FunctionParameterDecl( - name='pert', - type=PointerType( - inner_type=ValueType(name='mjvPerturb', is_const=True), - ), - ), - FunctionParameterDecl( - name='cam', - type=PointerType( - inner_type=ValueType(name='mjvCamera'), - ), - ), - FunctionParameterDecl( - name='catmask', - type=ValueType(name='int'), - ), - FunctionParameterDecl( - name='scn', - type=PointerType( - inner_type=ValueType(name='mjvScene'), - ), - ), - ), - doc='Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings.', # pylint: disable=line-too-long - )), ('mjv_copyModel', FunctionDecl( name='mjv_copyModel', @@ -4594,96 +4544,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Copy mjModel, skip large arrays not required for abstract visualization.', # pylint: disable=line-too-long )), - ('mjv_defaultSceneState', - FunctionDecl( - name='mjv_defaultSceneState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState'), - ), - ), - ), - doc='Set default scene state.', - )), - ('mjv_makeSceneState', - FunctionDecl( - name='mjv_makeSceneState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='m', - type=PointerType( - inner_type=ValueType(name='mjModel', is_const=True), - ), - ), - FunctionParameterDecl( - name='d', - type=PointerType( - inner_type=ValueType(name='mjData', is_const=True), - ), - ), - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState'), - ), - ), - FunctionParameterDecl( - name='maxgeom', - type=ValueType(name='int'), - ), - ), - doc='Allocate resources and initialize a scene state object.', - )), - ('mjv_freeSceneState', - FunctionDecl( - name='mjv_freeSceneState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState'), - ), - ), - ), - doc='Free scene state.', - )), - ('mjv_updateSceneState', - FunctionDecl( - name='mjv_updateSceneState', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='m', - type=PointerType( - inner_type=ValueType(name='mjModel', is_const=True), - ), - ), - FunctionParameterDecl( - name='d', - type=PointerType( - inner_type=ValueType(name='mjData'), - ), - ), - FunctionParameterDecl( - name='opt', - type=PointerType( - inner_type=ValueType(name='mjvOption', is_const=True), - ), - ), - FunctionParameterDecl( - name='scnstate', - type=PointerType( - inner_type=ValueType(name='mjvSceneState'), - ), - ), - ), - doc='Update a scene state from model and data.', - )), ('mjv_addGeoms', FunctionDecl( name='mjv_addGeoms', @@ -8570,6 +8430,99 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Sigmoid function over 0<=x<=1 using quintic polynomial.', )), + ('mjc_getSDF', + FunctionDecl( + name='mjc_getSDF', + return_type=PointerType( + inner_type=ValueType(name='mjpPlugin', is_const=True), + ), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='id', + type=ValueType(name='int'), + ), + ), + doc='get sdf from geom id', + )), + ('mjc_distance', + FunctionDecl( + name='mjc_distance', + return_type=ValueType(name='mjtNum'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData', is_const=True), + ), + ), + FunctionParameterDecl( + name='s', + type=PointerType( + inner_type=ValueType(name='mjSDF', is_const=True), + ), + ), + FunctionParameterDecl( + name='x', + type=ArrayType( + inner_type=ValueType(name='mjtNum', is_const=True), + extents=(3,), + ), + ), + ), + doc='signed distance function', + )), + ('mjc_gradient', + FunctionDecl( + name='mjc_gradient', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData', is_const=True), + ), + ), + FunctionParameterDecl( + name='s', + type=PointerType( + inner_type=ValueType(name='mjSDF', is_const=True), + ), + ), + FunctionParameterDecl( + name='gradient', + type=ArrayType( + inner_type=ValueType(name='mjtNum'), + extents=(3,), + ), + ), + FunctionParameterDecl( + name='x', + type=ArrayType( + inner_type=ValueType(name='mjtNum', is_const=True), + extents=(3,), + ), + ), + ), + doc='gradient of sdf', + )), ('mjd_transitionFD', FunctionDecl( name='mjd_transitionFD', @@ -9000,23 +8953,23 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Wait for a task to complete.', )), - ('mjs_attachBody', + ('mjs_attach', FunctionDecl( - name='mjs_attachBody', + name='mjs_attach', return_type=PointerType( - inner_type=ValueType(name='mjsBody'), + inner_type=ValueType(name='mjsElement'), ), parameters=( FunctionParameterDecl( name='parent', type=PointerType( - inner_type=ValueType(name='mjsFrame'), + inner_type=ValueType(name='mjsElement'), ), ), FunctionParameterDecl( name='child', type=PointerType( - inner_type=ValueType(name='mjsBody', is_const=True), + inner_type=ValueType(name='mjsElement', is_const=True), ), ), FunctionParameterDecl( @@ -9032,109 +8985,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc='Attach child body to a parent frame, return the attached body if success or NULL otherwise.', # pylint: disable=line-too-long - )), - ('mjs_attachFrame', - FunctionDecl( - name='mjs_attachFrame', - return_type=PointerType( - inner_type=ValueType(name='mjsFrame'), - ), - parameters=( - FunctionParameterDecl( - name='parent', - type=PointerType( - inner_type=ValueType(name='mjsBody'), - ), - ), - FunctionParameterDecl( - name='child', - type=PointerType( - inner_type=ValueType(name='mjsFrame', is_const=True), - ), - ), - FunctionParameterDecl( - name='prefix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - FunctionParameterDecl( - name='suffix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - ), - doc='Attach child frame to a parent body, return the attached frame if success or NULL otherwise.', # pylint: disable=line-too-long - )), - ('mjs_attachToSite', - FunctionDecl( - name='mjs_attachToSite', - return_type=PointerType( - inner_type=ValueType(name='mjsBody'), - ), - parameters=( - FunctionParameterDecl( - name='parent', - type=PointerType( - inner_type=ValueType(name='mjsSite'), - ), - ), - FunctionParameterDecl( - name='child', - type=PointerType( - inner_type=ValueType(name='mjsBody', is_const=True), - ), - ), - FunctionParameterDecl( - name='prefix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - FunctionParameterDecl( - name='suffix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - ), - doc='Attach child body to a parent site, return the attached body if success or NULL otherwise.', # pylint: disable=line-too-long - )), - ('mjs_attachFrameToSite', - FunctionDecl( - name='mjs_attachFrameToSite', - return_type=PointerType( - inner_type=ValueType(name='mjsFrame'), - ), - parameters=( - FunctionParameterDecl( - name='parent', - type=PointerType( - inner_type=ValueType(name='mjsSite'), - ), - ), - FunctionParameterDecl( - name='child', - type=PointerType( - inner_type=ValueType(name='mjsFrame', is_const=True), - ), - ), - FunctionParameterDecl( - name='prefix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - FunctionParameterDecl( - name='suffix', - type=PointerType( - inner_type=ValueType(name='char', is_const=True), - ), - ), - ), - doc='Attach child frame to a parent site, return the attached frame if success or NULL otherwise.', # pylint: disable=line-too-long + doc='Attach child to a parent, return the attached element if success or NULL otherwise.', # pylint: disable=line-too-long )), ('mjs_detachBody', FunctionDecl( @@ -10400,6 +10251,22 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Get double array contents and optionally its size.', )), + ('mjs_getPluginAttributes', + FunctionDecl( + name='mjs_getPluginAttributes', + return_type=PointerType( + inner_type=ValueType(name='void', is_const=True), + ), + parameters=( + FunctionParameterDecl( + name='plugin', + type=PointerType( + inner_type=ValueType(name='mjsPlugin', is_const=True), + ), + ), + ), + doc='Get plugin attributes.', + )), ('mjs_setDefault', FunctionDecl( name='mjs_setDefault', @@ -10423,7 +10290,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ('mjs_setFrame', FunctionDecl( name='mjs_setFrame', - return_type=ValueType(name='void'), + return_type=ValueType(name='int'), parameters=( FunctionParameterDecl( name='dest', @@ -10438,7 +10305,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc="Set element's enclosing frame.", + doc="Set element's enclosing frame, return 0 on success.", )), ('mjs_resolveOrientation', FunctionDecl( diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index fc75c798..49b8752b 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -1744,14 +1744,16 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='constraint solver reference:frictionloss (nv x mjNREF)', + doc='constraint solver reference:frictionloss', + array_extent=('nv', 'mjNREF'), ), StructFieldDecl( name='dof_solimp', type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='constraint solver impedance:frictionloss (nv x mjNIMP)', + doc='constraint solver impedance:frictionloss', + array_extent=('nv', 'mjNIMP'), ), StructFieldDecl( name='dof_frictionloss', @@ -1958,7 +1960,8 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='detect contact if dist 0) + clamped_upper = (x == bounds[1]) & (grad < 0) + clamped = clamped_lower | clamped_upper + grad_free = grad[~clamped] + + # Check termination condition on gradient norm. + g_norm = np.linalg.norm(grad_free) + if g_norm <= gtol: + status = Status.G_TOL + if g_norm == 0: + print('Zero gradient norm: exact minimum found?', file=output) + break + # Bounds relative to x dlower = None if bounds is None else bounds[0] - x dupper = None if bounds is None else bounds[1] - x @@ -353,13 +376,15 @@ def least_squares( # Append log to trace, call iter_callback. log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu) if verbose >= Verbosity.FULLITER.value: - log = dataclasses.replace(log, residual=r, jacobian=jac, step=dx) + log = dataclasses.replace( + log, residual=r, jacobian=jac, grad=grad, step=dx + ) trace.append(log) if iter_callback is not None: iter_callback(trace) - # Check for success. - if dx_norm < tol: + # Check termination condition on step norm. + if dx_norm < xtol * (xtol + np.linalg.norm(x)): status = Status.DX_TOL break @@ -376,8 +401,25 @@ def least_squares( # Append final log to trace, call iter_callback. # Note: unlike other iter logs, values are computed at the end point. yfinal = norm.value(r) - red = np.float64(0.0) # No reduction sice we didn't take a step. + red = np.float64(0.0) # No reduction since we didn't take a step. log = IterLog(candidate=x, objective=yfinal, reduction=red, regularizer=mu) + + # If full verbosity requested, compute values at the final point. + if verbose >= Verbosity.FULLITER.value: + # Get Jacobian jac. + t_start = time.time() + if jacobian is None: + jac, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds) + t_res += time.time() - t_start + else: + jac = jacobian(x, r) + t_jac += time.time() - t_start + n_jac += 1 + + # Get gradient, add to log. + grad, _ = norm.grad_hess(r, jac) + log = dataclasses.replace(log, residual=r, jacobian=jac, grad=grad) + trace.append(log) if iter_callback is not None: iter_callback(trace) @@ -430,13 +472,15 @@ def jacobian_fd( """ n = x.size if bounds is None: - eps_vec = eps * np.ones(n) + eps_vec = eps * np.ones((n, 1)) else: mid = 0.5 * (bounds[1] - bounds[0]) - eps_vec = np.where(x > mid, -eps, eps).flatten() - xh = x + np.diag(eps_vec) + eps_vec = np.where(x > mid, -eps, eps) + eps_vec *= np.maximum(1.0, np.abs(x)) + eps_vec = (eps_vec + x) - x + xh = x + np.diag(eps_vec.flatten()) rh = residual(xh) - jac = (rh - r) / eps_vec + jac = (rh - r) / eps_vec.T return jac, n_res + n diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index e642a8e1..1792643d 100644 --- a/python/mujoco/minimize_test.py +++ b/python/mujoco/minimize_test.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for minimize.py.""" import io @@ -32,7 +31,7 @@ class MinimizeTest(absltest.TestCase): x, _ = minimize.least_squares(x0, residual, output=out) expected_x = np.array((1.0, 1.0)) np.testing.assert_array_almost_equal(x, expected_x) - self.assertIn('norm(dx) < tol', out.getvalue()) + self.assertIn('norm(gradient) < tol', out.getvalue()) def test_start_at_minimum(self) -> None: def residual(x): @@ -43,7 +42,7 @@ class MinimizeTest(absltest.TestCase): x, _ = minimize.least_squares(x0, residual, output=out) expected_x = np.array((1.0, 1.0)) np.testing.assert_array_almost_equal(x, expected_x) - self.assertIn('norm(dx) < tol', out.getvalue()) + self.assertIn('norm(gradient) < tol', out.getvalue()) self.assertIn('exact minimum found', out.getvalue()) def test_jac_callback(self) -> None: @@ -61,7 +60,7 @@ class MinimizeTest(absltest.TestCase): ) expected_x = np.array((1.0, 1.0)) np.testing.assert_array_almost_equal(x, expected_x) - self.assertIn('norm(dx) < tol', out.getvalue()) + self.assertIn('norm(gradient) < tol', out.getvalue()) self.assertIn('Jacobian matches', out.getvalue()) # Try with bad Jacobian, ask least_squares to check it. @@ -116,7 +115,7 @@ class MinimizeTest(absltest.TestCase): x0, residual, bounds=bounds_types['inbounds'], output=out ) np.testing.assert_array_almost_equal(x, expected_x) - self.assertIn('norm(dx) < tol', out.getvalue()) + self.assertIn('norm(gradient) < tol', out.getvalue()) # Test different bounds conditions. for bounds in bounds_types.values(): @@ -128,8 +127,8 @@ class MinimizeTest(absltest.TestCase): output=out, verbose=minimize.Verbosity.FULLITER, ) - self.assertIn(' < tol', out.getvalue()) - grad = trace[-2].jacobian.T @ trace[-2].residual + self.assertIn('norm(gradient) < tol', out.getvalue()) + grad = trace[-1].grad # If x_i is on the boundary, gradient points out, otherwise it is 0. for i, xi in enumerate(x): if xi == bounds[0][i]: diff --git a/python/mujoco/mjpython/Info.plist b/python/mujoco/mjpython/Info.plist index 4006fd6d..475027c5 100644 --- a/python/mujoco/mjpython/Info.plist +++ b/python/mujoco/mjpython/Info.plist @@ -7,13 +7,13 @@ CFBundleIdentifier org.mujoco.mjpython CFBundleVersion - 3.3.1 + 3.3.3 CFBundleGetInfoString - 3.3.1 + 3.3.3 CFBundleLongVersionString - 3.3.1 + 3.3.3 CFBundleShortVersionString - 3.3.1 + 3.3.3 CFBundleExecutable mjpython CFBundleIconFile diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 56fe4437..a62d1b6b 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -13,17 +13,22 @@ // limitations under the License. #include -#include +#include // NOLINT(build/c++11) #include #include +#include #include -#include +#include // NOLINT(build/c++11) +#include #include +#include #include #include #include +#include #include "errors.h" +#include "indexers.h" #include "structs.h" #include #include @@ -134,19 +139,112 @@ class SimulateWrapper { void SetFigures( const std::vector>& viewports_figures) { - // Pairs of [viewport, figure], where viewport corresponds to the location - // of the figure on the viewer window. - std::vector> user_figures; - for (const auto& [viewport, figure] : viewports_figures) { - mjvFigure casted_figure = *figure.cast().get(); - user_figures.push_back(std::make_pair(viewport, casted_figure)); + + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newfigurerequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); } - // Set them all at once to prevent figure flickering. - simulate_->user_figures_ = user_figures; + // Pairs of [viewport, figure], where viewport corresponds to the location + // of the figure on the viewer window. + for (const auto& [viewport, figure] : viewports_figures) { + mjvFigure casted_figure = *figure.cast().get(); + simulate_->user_figures_new_.push_back(std::make_pair(viewport, casted_figure)); + } + + int value = 0; + simulate_->newfigurerequest.compare_exchange_strong(value, 1); } - void ClearFigures() { simulate_->user_figures_.clear(); } + void ClearFigures() { + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newfigurerequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + simulate_->user_figures_new_.clear(); + + int value = 0; + simulate_->newfigurerequest.compare_exchange_strong(value, 1); + } + + void SetTexts( + const std::vector>& + texts) { + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newtextrequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + // Collection of [font, gridpos, text1, text2] tuples for overlay text + for (const auto& [font, gridpos, text1, text2] : texts) { + simulate_->user_texts_new_.push_back(std::make_tuple(font, gridpos, text1, text2)); + } + + int value = 0; + simulate_->newtextrequest.compare_exchange_strong(value, 1); + } + + void ClearTexts() { + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newtextrequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + simulate_->user_texts_new_.clear(); + + int value = 0; + simulate_->newtextrequest.compare_exchange_strong(value, 1); + } + + void SetImages( + const std::vector> viewports_images + ) { + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newimagerequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + for (const auto& [viewport, image] : viewports_images) { + auto buf = image.request(); + if (buf.ndim != 3) { + throw std::invalid_argument("image must have 3 dimensions (H, W, C)"); + } + if (static_cast(buf.shape[2]) != 3) { + throw std::invalid_argument("image must have 3 channels"); + } + if (buf.itemsize != sizeof(unsigned char)) { + throw std::invalid_argument("image must be uint8 format"); + } + + // Calculate size of the image data + size_t height = buf.shape[0]; + size_t width = buf.shape[1]; + size_t size = height * width * 3; + + // Make a copy of the image data since Python is + // not required to keep it + std::unique_ptr image_copy(new unsigned char[size]()); + std::memcpy(image_copy.get(), buf.ptr, size); + + simulate_->user_images_new_.push_back(std::make_tuple(viewport, std::move(image_copy))); + } + + int value = 0; + simulate_->newimagerequest.compare_exchange_strong(value, 1); + } + + void ClearImages() { + // TODO: replace with atomic wait when we migrate to C++20 + while (simulate_ && simulate_->newimagerequest.load() != 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + simulate_->user_images_new_.clear(); + + int value = 0; + simulate_->newimagerequest.compare_exchange_strong(value, 1); + } private: mujoco::Simulate* simulate_; @@ -249,6 +347,11 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("set_figures", &SimulateWrapper::SetFigures, py::arg("viewports_figures")) .def("clear_figures", &SimulateWrapper::ClearFigures) + .def("set_texts", &SimulateWrapper::SetTexts, py::arg("overlay_texts")) + .def("clear_texts", &SimulateWrapper::ClearTexts) + .def("set_images", &SimulateWrapper::SetImages, + py::arg("viewports_images")) + .def("clear_images", &SimulateWrapper::ClearImages) .def_property_readonly("m", &SimulateWrapper::GetModel) .def_property_readonly("d", &SimulateWrapper::GetData) .def_property_readonly("viewport", &SimulateWrapper::GetViewport) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index bc9e5299..51846d2a 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -30,6 +30,7 @@ #include #include "errors.h" #include "indexers.h" // IWYU pragma: keep +#include "specs_wrapper.h" // IWYU pragma: keep #include "raw.h" #include "structs.h" // IWYU pragma: keep #include @@ -70,100 +71,6 @@ using MjDoubleRef10 = Eigen::Ref>; using MjDoubleRef11 = Eigen::Ref>; using MjDoubleRefVec = Eigen::Ref; -struct MjSpec { - MjSpec() : ptr(mj_makeSpec()) {} - MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {}) : ptr(ptr) { - for (const auto [key, value] : assets_) { - assets[key] = value; - } - } - - // copy constructor and assignment - MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) { - override_assets = other.override_assets; - for (const auto [key, value] : other.assets) { - assets[key] = value; - } - parent = other.parent; - } - MjSpec& operator=(const MjSpec& other) { - override_assets = other.override_assets; - ptr = mj_copySpec(other.ptr); - for (const auto [key, value] : other.assets) { - assets[key] = value; - } - parent = other.parent; - return *this; - } - - // move constructor and move assignment - MjSpec(MjSpec&& other) : ptr(other.ptr) { - override_assets = other.override_assets; - other.ptr = nullptr; - for (const auto [key, value] : other.assets) { - assets[key] = value; - } - other.assets.clear(); - parent = other.parent; - other.parent = nullptr; - } - MjSpec& operator=(MjSpec&& other) { - override_assets = other.override_assets; - ptr = other.ptr; - other.ptr = nullptr; - for (const auto [key, value] : other.assets) { - assets[key] = value; - } - other.assets.clear(); - parent = other.parent; - other.parent = nullptr; - return *this; - } - - ~MjSpec() { - mj_deleteSpec(ptr); - } - - raw::MjModel* Compile() { - if (assets.empty()) { - auto m = mj_compile(ptr, 0); - if (!m || mjs_isWarning(ptr)) { - throw py::value_error(mjs_getError(ptr)); - } - return m; - } - mjVFS vfs; - mj_defaultVFS(&vfs); - for (const auto& asset : assets) { - std::string buffer_name = - _impl::StripPath(py::cast(asset.first).c_str()); - std::string buffer = py::cast(asset.second); - const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( - &vfs, buffer_name.c_str(), buffer.c_str(), buffer.size()); - if (vfs_error) { - mj_deleteVFS(&vfs); - if (vfs_error == 2) { - throw py::value_error("Repeated file name in assets dict: " + - buffer_name); - } else { - throw py::value_error("Asset failed to load: " + buffer_name); - } - } - } - auto m = mj_compile(ptr, &vfs); - if (!m || mjs_isWarning(ptr)) { - throw py::value_error(mjs_getError(ptr)); - } - mj_deleteVFS(&vfs); - return m; - } - - raw::MjSpec* ptr; - py::dict assets; - bool override_assets = true; - MjSpec* parent = nullptr; -}; - template static raw::MjSpec* LoadSpecFileImpl( const std::string& filename, @@ -261,7 +168,7 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) { // this should never happen throw pybind11::value_error( "body.find_all supports the types: body, frame, geom, site, " - "light, camera."); + "joint, light, camera."); break; } el = mjs_nextChild(&body, el, recursive); @@ -269,16 +176,6 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) { return list; // list of pointers, so they can be copied } -void SetFrame(raw::MjsBody* body, mjtObj objtype, raw::MjsFrame* frame) { - mjsElement* el = mjs_firstChild(body, objtype, 0); - while (el) { - if (frame->element != el && mjs_getFrame(el) == nullptr) { - mjs_setFrame(el, frame); - } - el = mjs_nextChild(body, el, 0); - } -} - PYBIND11_MODULE(_specs, m) { auto structs_m = py::module::import("mujoco._structs"); py::function mjmodel_from_raw_ptr = @@ -526,21 +423,9 @@ PYBIND11_MODULE(_specs, m) { throw pybind11::value_error( "Only one of frame or site can be specified."); } - auto worldbody = mjs_findBody(child.ptr, "world"); - if (!worldbody) { - throw pybind11::value_error("Child does not have a world body."); - } - auto worldframe = mjs_addFrame(worldbody, nullptr); - SetFrame(worldbody, mjOBJ_BODY, worldframe); - SetFrame(worldbody, mjOBJ_SITE, worldframe); - SetFrame(worldbody, mjOBJ_FRAME, worldframe); - SetFrame(worldbody, mjOBJ_JOINT, worldframe); - SetFrame(worldbody, mjOBJ_GEOM, worldframe); - SetFrame(worldbody, mjOBJ_LIGHT, worldframe); - SetFrame(worldbody, mjOBJ_CAMERA, worldframe); const char* p = prefix.has_value() ? prefix.value().c_str() : ""; const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - raw::MjsFrame* attached_frame = nullptr; + raw::MjsElement* attached_frame = nullptr; if (frame.has_value()) { raw::MjsFrame* frame_ptr = nullptr; try { @@ -556,15 +441,11 @@ PYBIND11_MODULE(_specs, m) { throw pybind11::value_error( "Frame spec does not match parent spec."); } - raw::MjsBody* parent_body = mjs_getParent(frame_ptr->element); - if (!parent_body) { - throw pybind11::value_error("Frame does not have a parent body."); - } - attached_frame = mjs_attachFrame(parent_body, worldframe, p, s); + attached_frame = + mjs_attach(frame_ptr->element, child.ptr->element, p, s); if (!attached_frame) { throw pybind11::value_error(mjs_getError(self.ptr)); } - mjs_setFrame(attached_frame->element, frame_ptr); } if (site.has_value()) { raw::MjsSite* site_ptr = nullptr; @@ -581,7 +462,8 @@ PYBIND11_MODULE(_specs, m) { throw pybind11::value_error( "Site spec does not match parent spec."); } - attached_frame = mjs_attachFrameToSite(site_ptr, worldframe, p, s); + attached_frame = + mjs_attach(site_ptr->element, child.ptr->element, p, s); if (!attached_frame) { throw pybind11::value_error(mjs_getError(self.ptr)); } @@ -595,12 +477,33 @@ PYBIND11_MODULE(_specs, m) { self.assets[asset.first] = asset.second; } child.parent = &self; - return attached_frame; + return mjs_asFrame(attached_frame); }, py::arg("child"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), py::arg("site") = py::none(), py::arg("frame") = py::none(), py::return_value_policy::reference_internal); + mjSpec.def( + "activate_plugin", + [](MjSpec& self, std::string& name) { + mjs_activatePlugin(self.ptr, name.c_str()); + }, + py::arg("name"), + py::return_value_policy::reference_internal); + mjSpec.def_static( + "resolve_orientation", + [](bool degree, const MjTypeVec& sequence, + const raw::MjsOrientation* orientation) -> std::array { + std::array quat = {0, 0, 0, 0}; + const char* err = mjs_resolveOrientation(quat.data(), degree, + sequence.ptr, orientation); + if (err) { + throw pybind11::value_error(err); + } + return quat; + }, + py::arg("degree"), py::arg("sequence") = py::none(), + py::arg("orientation"), py::return_value_policy::copy); // ============================= MJSBODY ===================================== mjsBody.def( @@ -637,10 +540,11 @@ PYBIND11_MODULE(_specs, m) { return out; }, py::return_value_policy::reference_internal); - mjsBody.def("set_frame", - [](raw::MjsBody& self, raw::MjsFrame& frame) -> void { - mjs_setFrame(self.element, &frame); - }); + mjsBody.def("set_frame", [](raw::MjsBody& self, raw::MjsFrame& frame) { + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } + }); mjsBody.def_property( "classname", [](raw::MjsBody& self) -> raw::MjsDefault* { @@ -826,11 +730,11 @@ PYBIND11_MODULE(_specs, m) { std::optional& suffix) -> raw::MjsFrame* { const char* p = prefix.has_value() ? prefix.value().c_str() : ""; const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - auto new_frame = mjs_attachFrame(&self, &frame, p, s); + auto new_frame = mjs_attach(self.element, frame.element, p, s); if (!new_frame) { throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); } - return new_frame; + return mjs_asFrame(new_frame); }, py::arg("frame"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), @@ -846,11 +750,20 @@ PYBIND11_MODULE(_specs, m) { return frame; }, py::return_value_policy::reference_internal); + mjsBody.def_property_readonly( + "frame", + [](raw::MjsBody* self) -> raw::MjsFrame* { + return mjs_getFrame(self->element); + }, + py::return_value_policy::reference_internal); + // ============================= MJSFRAME ==================================== mjsFrame.def("delete", [](raw::MjsFrame& self) { mjs_delete(self.element); }); mjsFrame.def("set_frame", [](raw::MjsFrame& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsFrame.def_property_readonly( "parent", @@ -865,21 +778,29 @@ PYBIND11_MODULE(_specs, m) { std::optional& suffix) -> raw::MjsBody* { const char* p = prefix.has_value() ? prefix.value().c_str() : ""; const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - auto new_body = mjs_attachBody(&self, &body, p, s); + auto new_body = mjs_attach(self.element, body.element, p, s); if (!new_body) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } - return new_body; + return mjs_asBody(new_body); }, py::arg("body"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); + mjsFrame.def_property_readonly( + "frame", + [](raw::MjsFrame& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSGEOM ===================================== mjsGeom.def("delete", [](raw::MjsGeom& self) { mjs_delete(self.element); }); mjsGeom.def("set_frame", [](raw::MjsGeom& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsGeom.def_property_readonly( "parent", @@ -895,11 +816,19 @@ PYBIND11_MODULE(_specs, m) { [](raw::MjsGeom& self, raw::MjsDefault& default_) -> void { mjs_setDefault(self.element, &default_); }); + mjsGeom.def_property_readonly( + "frame", + [](raw::MjsGeom& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSJOINT ==================================== mjsJoint.def("delete", [](raw::MjsJoint& self) { mjs_delete(self.element); }); mjsJoint.def("set_frame", [](raw::MjsJoint& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsJoint.def_property_readonly( "parent", @@ -915,11 +844,19 @@ PYBIND11_MODULE(_specs, m) { [](raw::MjsJoint& self, raw::MjsDefault& default_) -> void { mjs_setDefault(self.element, &default_); }); + mjsJoint.def_property_readonly( + "frame", + [](raw::MjsJoint& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSSITE ===================================== mjsSite.def("delete", [](raw::MjsSite& self) { mjs_delete(self.element); }); mjsSite.def("set_frame", [](raw::MjsSite& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsSite.def_property_readonly( "parent", @@ -942,22 +879,30 @@ PYBIND11_MODULE(_specs, m) { std::optional& suffix) -> raw::MjsBody* { const char* p = prefix.has_value() ? prefix.value().c_str() : ""; const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - auto new_body = mjs_attachToSite(&self, &body, p, s); + auto new_body = mjs_attach(self.element, body.element, p, s); if (!new_body) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } - return new_body; + return mjs_asBody(new_body); }, py::arg("body"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); + mjsSite.def_property_readonly( + "frame", + [](raw::MjsSite& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSCAMERA =================================== mjsCamera.def("delete", [](raw::MjsCamera& self) { mjs_delete(self.element); }); mjsCamera.def("set_frame", [](raw::MjsCamera& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsCamera.def_property_readonly( "parent", @@ -973,11 +918,19 @@ PYBIND11_MODULE(_specs, m) { [](raw::MjsCamera& self, raw::MjsDefault& default_) -> void { mjs_setDefault(self.element, &default_); }); + mjsCamera.def_property_readonly( + "frame", + [](raw::MjsCamera& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSLIGHT ==================================== mjsLight.def("delete", [](raw::MjsLight& self) { mjs_delete(self.element); }); mjsLight.def("set_frame", [](raw::MjsLight& self, raw::MjsFrame& frame) { - mjs_setFrame(self.element, &frame); + if (mjs_setFrame(self.element, &frame) != 0) { + throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); + } }); mjsLight.def_property_readonly( "parent", @@ -993,6 +946,12 @@ PYBIND11_MODULE(_specs, m) { [](raw::MjsLight& self, raw::MjsDefault& default_) -> void { mjs_setDefault(self.element, &default_); }); + mjsLight.def_property_readonly( + "frame", + [](raw::MjsLight& self) -> raw::MjsFrame* { + return mjs_getFrame(self.element); + }, + py::return_value_policy::reference_internal); // ============================= MJSMATERIAL ================================= mjsMaterial.def("delete", @@ -1130,6 +1089,30 @@ PYBIND11_MODULE(_specs, m) { }); mjsPlugin.def("delete", [](raw::MjsPlugin& self) { mjs_delete(self.element); }); + mjsPlugin.def_property( + "config", + [](raw::MjsPlugin& self) -> py::dict { + const std::map>* config_attribs = + static_cast>*>( + mjs_getPluginAttributes(&self)); + py::dict config; + for (const auto& [key, value] : *config_attribs) { + config[py::str(key)] = value; + } + return config; + }, + [](raw::MjsPlugin& self, py::dict& config) { + std::map> config_attribs; + for (const auto& [key, value] : config) { + std::string key_str = key.cast(); + if (config_attribs.find(key_str) != config_attribs.end()) { + throw pybind11::value_error("Duplicate config key: " + key_str); + } + config_attribs[key_str] = value.cast(); + } + mjs_setPluginAttributes(&self, &config_attribs); + }, + py::return_value_policy::reference_internal); // ============================= MJVISUAL ==================================== mjVisual.def_property( "global_", diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index c946f420..41a2043c 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -16,6 +16,7 @@ import gc import inspect +import math import os import textwrap import typing @@ -129,6 +130,18 @@ class SpecsTest(absltest.TestCase): """) self.assertEqual(spec.to_xml(), xml) + def test_resolve_orientation(self): + spec = mujoco.MjSpec() + body = spec.worldbody.add_body(euler=[0, 0, 90]) + quat = mujoco.MjSpec.resolve_orientation( + degree=spec.compiler.degree, + sequence=spec.compiler.eulerseq, + orientation=body.alt, + ) + np.testing.assert_array_almost_equal( + quat, [math.sqrt(2) / 2, 0, 0, math.sqrt(2) / 2] + ) + def test_kwarg(self): # Create a spec. spec = mujoco.MjSpec() @@ -636,6 +649,8 @@ class SpecsTest(absltest.TestCase): + + @@ -649,6 +664,8 @@ class SpecsTest(absltest.TestCase): self.assertLen(spec.sites, 5) self.assertLen(spec.worldbody.find_all('body'), 4) self.assertLen(spec.worldbody.find_all('site'), 5) + self.assertLen(spec.worldbody.find_all('joint'), 1) + self.assertLen(spec.worldbody.find_all('geom'), 1) self.assertEqual(spec.bodies[1].name, 'body1') self.assertEqual(spec.bodies[2].name, 'body2') self.assertEqual(spec.bodies[3].name, 'body3') @@ -831,22 +848,21 @@ class SpecsTest(absltest.TestCase): self.assertEqual(model.nsensor, 9) def test_plugin(self): - xml = """ - - - - - - """ - - spec = mujoco.MjSpec.from_string(xml) - self.assertIsNotNone(spec.worldbody) + spec = mujoco.MjSpec() + spec.activate_plugin('mujoco.elasticity.cable') + plugin = spec.add_plugin( + name='instance_name', + plugin_name='mujoco.elasticity.cable', + active=True, + info='info' + ) + plugin.config = {'twist': '10', 'bend': '1'} + self.assertEqual(plugin.config, {'twist': '10', 'bend': '1'}) body = spec.worldbody.add_body() + body.plugin = plugin body.plugin.plugin_name = 'mujoco.elasticity.cable' - body.plugin.id = spec.add_plugin() body.plugin.active = True - self.assertEqual(body.plugin.id, 0) geom = body.add_geom() geom.type = mujoco.mjtGeom.mjGEOM_BOX @@ -857,7 +873,19 @@ class SpecsTest(absltest.TestCase): model = spec.compile() self.assertIsNotNone(model) self.assertEqual(model.nplugin, 1) + self.assertEqual(model.npluginattr, 7) self.assertEqual(model.body_plugin[1], 0) + attributes = (''.join([chr(i) for i in model.plugin_attr]).split(chr(0))) + self.assertEqual(attributes[:2], ['10', '1']) + + copy = spec.copy() # before assigning the new config + wrong_config = {'wrong': '10', 'bend': '1'} + for s in [spec, copy]: + s.plugins[0].config = wrong_config + with self.assertRaisesRegex( + ValueError, "Error: unrecognized attribute 'plugin:wrong'" + ): + s.compile() def test_recompile_error(self): main_xml = """ @@ -1054,6 +1082,15 @@ class SpecsTest(absltest.TestCase): frame = body.to_frame() np.testing.assert_array_equal(frame.pos, [1, 2, 3]) + def test_get_frame(self): + spec = mujoco.MjSpec() + body = spec.worldbody.add_body() + frame = body.add_frame() + geom = body.add_geom() + geom.set_frame(frame) + self.assertIsNotNone(frame) + self.assertIs(geom.frame, frame) + def test_attach_to_frame(self): parent = mujoco.MjSpec() parent.assets = {'cube.obj': 'cube_content'} @@ -1153,6 +1190,13 @@ class SpecsTest(absltest.TestCase): AttributeError, "object has no attribute 'invalid'" ): print(mj_model.bind(joints).invalid) + invalid_spec = mujoco.MjSpec() + invalid_spec.worldbody.add_body(name='main') + with self.assertRaisesRegex( + ValueError, + 'The mjSpec does not match mjModel. Please recompile the mjSpec.', + ): + print(mj_model.bind(invalid_spec.body('main'))) def test_incorrect_hfield_size(self): nrow = 300 diff --git a/python/mujoco/specs_wrapper.cc b/python/mujoco/specs_wrapper.cc new file mode 100644 index 00000000..18ffd3b5 --- /dev/null +++ b/python/mujoco/specs_wrapper.cc @@ -0,0 +1,124 @@ +// Copyright 2024 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. + +#include "specs_wrapper.h" + +#include // IWYU pragma: keep +#include +#include // IWYU pragma: keep +#include // IWYU pragma: keep + +#include // IWYU pragma: keep +#include +#include "errors.h" +#include "indexers.h" // IWYU pragma: keep +#include "raw.h" +#include "structs.h" // IWYU pragma: keep +#include +#include +#include +#include + +namespace py = ::pybind11; + +namespace mujoco::python { + +MjSpec::MjSpec() : ptr(mj_makeSpec()) {} +MjSpec::MjSpec(raw::MjSpec* ptr, const py::dict& assets_) : ptr(ptr) { + for (const auto [key, value] : assets_) { + assets[key] = value; + } +} + +// copy constructor and assignment +MjSpec::MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) { + override_assets = other.override_assets; + for (const auto [key, value] : other.assets) { + assets[key] = value; + } + parent = other.parent; +} + +MjSpec& MjSpec::operator=(const MjSpec& other) { + override_assets = other.override_assets; + ptr = mj_copySpec(other.ptr); + for (const auto [key, value] : other.assets) { + assets[key] = value; + } + parent = other.parent; + return *this; +} + +// move constructor and move assignment +MjSpec::MjSpec(MjSpec&& other) : ptr(other.ptr) { + override_assets = other.override_assets; + other.ptr = nullptr; + for (const auto [key, value] : other.assets) { + assets[key] = value; + } + other.assets.clear(); + parent = other.parent; + other.parent = nullptr; +} + +MjSpec& MjSpec::operator=(MjSpec&& other) { + override_assets = other.override_assets; + ptr = other.ptr; + other.ptr = nullptr; + for (const auto [key, value] : other.assets) { + assets[key] = value; + } + other.assets.clear(); + parent = other.parent; + other.parent = nullptr; + return *this; +} + +MjSpec::~MjSpec() { mj_deleteSpec(ptr); } + +raw::MjModel* MjSpec::Compile() { + if (assets.empty()) { + auto m = mj_compile(ptr, 0); + if (!m || mjs_isWarning(ptr)) { + throw py::value_error(mjs_getError(ptr)); + } + return m; + } + mjVFS vfs; + mj_defaultVFS(&vfs); + for (const auto& asset : assets) { + std::string buffer_name = + _impl::StripPath(py::cast(asset.first).c_str()); + std::string buffer = py::cast(asset.second); + const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( + &vfs, buffer_name.c_str(), buffer.c_str(), buffer.size()); + if (vfs_error) { + mj_deleteVFS(&vfs); + if (vfs_error == 2) { + throw py::value_error("Repeated file name in assets dict: " + + buffer_name); + } else { + throw py::value_error("Asset failed to load: " + buffer_name); + } + } + } + auto m = mj_compile(ptr, &vfs); + if (!m || mjs_isWarning(ptr)) { + throw py::value_error(mjs_getError(ptr)); + } + mj_deleteVFS(&vfs); + return m; +} + +} // namespace mujoco::python diff --git a/python/mujoco/specs_wrapper.h b/python/mujoco/specs_wrapper.h new file mode 100644 index 00000000..2ebaac60 --- /dev/null +++ b/python/mujoco/specs_wrapper.h @@ -0,0 +1,45 @@ +// Copyright 2024 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. + +#include "raw.h" +#include +#include +#include +#include + +namespace py = ::pybind11; + +namespace mujoco::python { + +struct MjSpec { + MjSpec(); + MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {}); + + // copy constructor and assignment + MjSpec(const MjSpec& other); + MjSpec& operator=(const MjSpec& other); + + // move constructor and move assignment + MjSpec(MjSpec&& other); + MjSpec& operator=(MjSpec&& other); + ~MjSpec(); + + raw::MjModel* Compile(); + + raw::MjSpec* ptr; + py::dict assets; + bool override_assets = true; + MjSpec* parent = nullptr; +}; +} // namespace mujoco::python diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index f9272721..1b92c61a 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -16,36 +16,23 @@ #include -#include -#include -#include -#include #include -#include -#include #include #include #include -#include #include #include #include -#include -#include -#include #include #include -#include #include #include #include "errors.h" #include "function_traits.h" #include "indexer_xmacro.h" #include "indexers.h" -#include "private.h" #include "raw.h" -#include "serialization.h" #include #include #include @@ -68,13 +55,9 @@ namespace { // (dim0, dim1). #define X_ARRAY_SHAPE(dim0, dim1) XArrayShapeImpl(#dim1)((dim0), (dim1)) -std::vector XArrayShapeImpl1D(int dim0, int dim1) { - return {dim0}; -} +std::vector XArrayShapeImpl1D(int dim0, int dim1) { return {dim0}; } -std::vector XArrayShapeImpl2D(int dim0, int dim1) { - return {dim0, dim1}; -} +std::vector XArrayShapeImpl2D(int dim0, int dim1) { return {dim0, dim1}; } constexpr auto XArrayShapeImpl(const std::string_view dim1_str) { if (dim1_str == "1") { @@ -84,334 +67,6 @@ constexpr auto XArrayShapeImpl(const std::string_view dim1_str) { } } -inline std::size_t NConMax(const mjData* d) { - return d->narena / sizeof(mjContact); -} - -} // namespace - -// ==================== MJOPTION =============================================== -#define X(var, dim) , var(InitPyArray(std::array{dim}, ptr_->var, owner_)) -MjOptionWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjOption* const opt = new raw::MjOption; - mj_defaultOption(opt); - return opt; - }()) - MJOPTION_VECTORS {} - -MjOptionWrapper::MjWrapper(raw::MjOption* ptr, py::handle owner) - : WrapperBase(ptr, owner) - MJOPTION_VECTORS {} -#undef X - -MjOptionWrapper::MjWrapper(const MjOptionWrapper& other) - : MjOptionWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVISUAL =============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjVisualHeadlightWrapper::MjWrapper() - : WrapperBase(new raw::MjVisualHeadlight{}), - X(ambient), - X(diffuse), - X(specular) {} - -MjVisualHeadlightWrapper::MjWrapper( - raw::MjVisualHeadlight* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(ambient), - X(diffuse), - X(specular) {} -#undef X - -MjVisualHeadlightWrapper::MjWrapper(const MjVisualHeadlightWrapper& other) - : MjVisualHeadlightWrapper() { - *this->ptr_ = *other.ptr_; -} - -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjVisualRgbaWrapper::MjWrapper() - : WrapperBase(new raw::MjVisualRgba{}), - X(fog), - X(haze), - X(force), - X(inertia), - X(joint), - X(actuator), - X(actuatornegative), - X(actuatorpositive), - X(com), - X(camera), - X(light), - X(selectpoint), - X(connect), - X(contactpoint), - X(contactforce), - X(contactfriction), - X(contacttorque), - X(contactgap), - X(rangefinder), - X(constraint), - X(slidercrank), - X(crankbroken), - X(frustum) {} - -MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(fog), - X(haze), - X(force), - X(inertia), - X(joint), - X(actuator), - X(actuatornegative), - X(actuatorpositive), - X(com), - X(camera), - X(light), - X(selectpoint), - X(connect), - X(contactpoint), - X(contactforce), - X(contactfriction), - X(contacttorque), - X(contactgap), - X(rangefinder), - X(constraint), - X(slidercrank), - X(crankbroken), - X(frustum) {} -#undef X - -MjVisualRgbaWrapper::MjWrapper(const MjVisualRgbaWrapper& other) - : MjVisualRgbaWrapper() { - *this->ptr_ = *other.ptr_; -} - -MjVisualWrapper::MjWrapper() - : WrapperBase(new raw::MjVisual{}), - headlight(&ptr_->headlight, owner_), - rgba(&ptr_->rgba, owner_) {} - -MjVisualWrapper::MjWrapper(raw::MjVisual* ptr, py::handle owner) - : WrapperBase(ptr, owner), - headlight(&ptr_->headlight, owner_), - rgba(&ptr_->rgba, owner_) {} - - -MjVisualWrapper::MjWrapper(const MjVisualWrapper& other) - : MjVisualWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJMODEL ================================================ -static void MjModelCapsuleDestructor(PyObject* pyobj) { - mj_deleteModel( - static_cast(PyCapsule_GetPointer(pyobj, nullptr))); -} - -static absl::flat_hash_map& -MjModelRawPointerMap() { - static auto* hash_map = - new absl::flat_hash_map(); - return *hash_map; -} - -MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept { - try { - auto& map = MjModelRawPointerMap(); - { - py::gil_scoped_acquire gil; - auto found = map.find(m); - return found != map.end() ? found->second : nullptr; - } - } catch (...) { - return nullptr; - } -} - -#undef MJ_M -#define MJ_M(x) ptr_->x -#define X(dtype, var, dim0, dim1) \ - , var (InitPyArray(X_ARRAY_SHAPE(ptr_->dim0, dim1), ptr_->var, owner_)) -MjModelWrapper::MjWrapper(raw::MjModel* ptr) - : WrapperBase(ptr, &MjModelCapsuleDestructor), - opt(&ptr->opt, owner_), - vis(&ptr->vis, owner_), - stat(&ptr->stat, owner_) - MJMODEL_POINTERS, - text_data_bytes(ptr->text_data, ptr->ntextdata), - names_bytes(ptr->names, ptr->nnames), - paths_bytes(ptr->paths, ptr->npaths), - indexer_(ptr, owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second; - } - if (!is_newly_inserted) { - throw UnexpectedError( - "MjModelWrapper(mjModel*): MjModelRawPointerMap already contains this " - "raw mjModel*"); - } -} - -MjModelWrapper::MjWrapper(MjModelWrapper&& other) - : WrapperBase(other.ptr_, other.owner_), - opt(&ptr_->opt, owner_), - vis(&ptr_->vis, owner_), - stat(&ptr_->stat, owner_) - MJMODEL_POINTERS, - text_data_bytes(ptr_->text_data, ptr_->ntextdata), - names_bytes(ptr_->names, ptr_->nnames), - paths_bytes(ptr_->paths, ptr_->npaths), - indexer_(ptr_, owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = - MjModelRawPointerMap().insert_or_assign(ptr_, this).second; - } - if (is_newly_inserted) { - throw UnexpectedError( - "MjModelRawPointerMap does not contains the moved-from mjModel*"); - } - other.ptr_ = nullptr; -} -#undef X -#undef MJ_M -#define MJ_M(x) x - -// Delegating to the MjModelWrapper::MjWrapper(raw::MjModel*) constructor, -// no need to modify MjModelRawPointerMap here. -MjModelWrapper::MjWrapper(const MjModelWrapper& other) - : MjModelWrapper(InterceptMjErrors(mj_copyModel)(NULL, other.get())) {} - -MjModelWrapper::~MjWrapper() { - if (ptr_) { - bool erased = false; - { - py::gil_scoped_acquire gil; - erased = MjModelRawPointerMap().erase(ptr_); - } - if (!erased) { - std::cerr << "MjModelRawPointerMap does not contain this raw mjModel*" - << std::endl; - std::terminate(); - } - } -} - -// Helper function for both LoadXMLFile and LoadBinaryFile. -// Creates a temporary MJB from the assets dictionary if one is supplied. -template -static raw::MjModel* LoadModelFileImpl( - const std::string& filename, - const std::vector& assets, - LoadFunc&& loadfunc) { - mjVFS vfs; - mjVFS* vfs_ptr = nullptr; - if (!assets.empty()) { - mj_defaultVFS(&vfs); - vfs_ptr = &vfs; - for (const auto& asset : assets) { - std::string buffer_name = StripPath(asset.name); - const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( - vfs_ptr, buffer_name.c_str(), asset.content, asset.content_size); - if (vfs_error) { - mj_deleteVFS(vfs_ptr); - if (vfs_error == 2) { - throw py::value_error("Repeated file name in assets dict: " + - buffer_name); - } else { - throw py::value_error("Asset failed to load: " + buffer_name); - } - } - } - } - - raw::MjModel* model = loadfunc(filename.c_str(), vfs_ptr); - mj_deleteVFS(vfs_ptr); - if (model && !model->buffer) { - mj_deleteModel(model); - model = nullptr; - } - return model; -} - -MjModelWrapper MjModelWrapper::LoadXMLFile( - const std::string& filename, - const std::optional>& assets) { - const auto converted_assets = ConvertAssetsDict(assets); - raw::MjModel* model; - { - py::gil_scoped_release no_gil; - char error[1024]; - model = LoadModelFileImpl( - filename, converted_assets, - [&error](const char* filename, const mjVFS* vfs) { - return InterceptMjErrors(mj_loadXML)( - filename, vfs, error, sizeof(error)); - }); - if (!model) { - throw py::value_error(error); - } - } - return MjModelWrapper(model); -} - -MjModelWrapper MjModelWrapper::LoadBinaryFile( - const std::string& filename, - const std::optional>& assets) { - const auto converted_assets = ConvertAssetsDict(assets); - raw::MjModel* model; - { - py::gil_scoped_release no_gil; - model = LoadModelFileImpl( - filename, converted_assets, InterceptMjErrors(mj_loadModel)); - if (!model) { - throw py::value_error("mj_loadModel: failed to load from mjb"); - } - } - return MjModelWrapper(model); -} - -MjModelWrapper MjModelWrapper::LoadXML( - const std::string& xml, - const std::optional>& assets) { - auto converted_assets = ConvertAssetsDict(assets); - raw::MjModel* model; - { - py::gil_scoped_release no_gil; - std::string model_filename = "model_.xml"; - if (assets.has_value()) { - while (assets->find(model_filename) != assets->end()) { - model_filename = - model_filename.substr(0, model_filename.size() - 4) + "_.xml"; - } - } - converted_assets.emplace_back( - model_filename.c_str(), xml.c_str(), xml.length()); - char error[1024]; - model = LoadModelFileImpl( - model_filename, converted_assets, - [&error](const char* filename, const mjVFS* vfs) { - return InterceptMjErrors(mj_loadXML)( - filename, vfs, error, sizeof(error)); - }); - if (!model) { - throw py::value_error(error); - } - } - return MjModelWrapper(model); -} - -MjModelWrapper MjModelWrapper::WrapRawModel(raw::MjModel* m) { - return MjModelWrapper(m); -} - py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m, const MjDataWrapper& old_d) { raw::MjModel* m = static_cast(mju_malloc(sizeof(mjModel))); @@ -427,932 +82,8 @@ py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m, return py::make_tuple(m_pyobj, d_pyobj); } -namespace { -// A byte at the start of serialized mjModel structs, which can be incremented -// when we change the serialization logic to reject pickles from an unsupported -// future version. -constexpr static char kSerializationVersion = 1; - -void CheckInput(const std::istream& input, std::string class_name) { - if (input.fail()) { - throw py::value_error("Invalid serialized " + class_name + "."); - } -} - } // namespace -void MjModelWrapper::Serialize(std::ostream& output) const { - WriteChar(output, kSerializationVersion); - - int model_size = mj_sizeModel(get()); - WriteInt(output, model_size); - std::string buffer(model_size, 0); - mj_saveModel(get(), nullptr, buffer.data(), model_size); - WriteBytes(output, buffer.data(), model_size); -} - -std::unique_ptr MjModelWrapper::Deserialize( - std::istream& input) { - CheckInput(input, "mjModel"); - - char serializationVersion = ReadChar(input); - CheckInput(input, "mjModel"); - - if (serializationVersion != kSerializationVersion) { - throw py::value_error("Incompatible serialization version."); - } - - std::size_t model_size = ReadInt(input); - CheckInput(input, "mjModel"); - if (model_size < 0) { - throw py::value_error("Invalid serialized mjModel."); - } - std::string model_bytes(model_size, 0); - ReadBytes(input, model_bytes.data(), model_size); - CheckInput(input, "mjModel"); - - raw::MjModel* model = LoadModelFileImpl( - "model.mjb", - {{"model.mjb", model_bytes.data(), static_cast(model_size)}}, - InterceptMjErrors(mj_loadModel)); - if (!model) { - throw py::value_error("Invalid serialized mjModel."); - } - return std::unique_ptr(new MjModelWrapper(model)); -} - -// ==================== MJCONTACT ============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjContactWrapper::MjWrapper() - : WrapperBase(new raw::MjContact{}), - X(pos), - X(frame), - X(friction), - X(solref), - X(solreffriction), - X(solimp), - X(H), - X(geom), - X(flex), - X(elem), - X(vert) {} - -MjContactWrapper::MjWrapper(raw::MjContact* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(pos), - X(frame), - X(friction), - X(solref), - X(solreffriction), - X(solimp), - X(H), - X(geom), - X(flex), - X(elem), - X(vert) {} -#undef X - -MjContactWrapper::MjWrapper(const MjContactWrapper& other) - : MjContactWrapper() { - *this->ptr_ = *other.ptr_; -} - -MjContactList::MjStructList(raw::MjContact* ptr, int nconmax, - int* ncon, py::handle owner) - : StructListBase(ptr, nconmax, owner, /* lazy = */ true), - ncon_(ncon) {} - -// Slicing -MjContactList::MjStructList(MjContactList& other, py::slice slice) - : StructListBase(other, slice), - ncon_(other.ncon_) {} - -// ==================== MJDATA ================================================= -static void MjDataCapsuleDestructor(PyObject* pyobj) { - mj_deleteData( - static_cast(PyCapsule_GetPointer(pyobj, nullptr))); -} - -absl::flat_hash_map& -MjDataRawPointerMap() { - static auto* hash_map = - new absl::flat_hash_map(); - return *hash_map; -} - -MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept { - try { - auto& map = MjDataRawPointerMap(); - { - py::gil_scoped_acquire gil; - auto found = map.find(m); - return found != map.end() ? found->second : nullptr; - } - } catch (...) { - return nullptr; - } -} - -MjDataWrapper::MjWrapper(MjModelWrapper* model) - : WrapperBase(InterceptMjErrors(mj_makeData)(model->get()), - &MjDataCapsuleDestructor), -#undef MJ_M -#define MJ_M(x) model->get()->x -#define X(dtype, var, dim0, dim1) \ - var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)), - MJDATA_POINTERS -#undef MJ_M -#define MJ_M(x) (x) -#undef X - - contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), - -#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), - MJDATA_VECTOR -#undef X - model_(model), - model_ref_(py::cast(model_)), - indexer_(ptr_, model_->get(), owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; - } - if (!is_newly_inserted) { - throw UnexpectedError( - "MjDataRawPointerMap already contains this raw mjData*"); - } -} - -MjDataWrapper::MjWrapper(const MjDataWrapper& other) - : WrapperBase(other.Copy(), &MjDataCapsuleDestructor), -#undef MJ_M -#define MJ_M(x) other.model_->get()->x -#define X(dtype, var, dim0, dim1) \ - var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ - owner_)), - MJDATA_POINTERS -#undef MJ_M -#define MJ_M(x) (x) -#undef X - - contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), - -#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), - MJDATA_VECTOR -#undef X - model_(other.model_), - model_ref_(other.model_ref_), - indexer_(ptr_, model_->get(), owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; - } - if (!is_newly_inserted) { - throw UnexpectedError( - "MjDataRawPointerMap already contains this raw mjData*"); - } -} - -MjDataWrapper::MjWrapper(MjDataWrapper&& other) - : WrapperBase(other.ptr_, other.owner_), -#undef MJ_M -#define MJ_M(x) other.model_->get()->x -#define X(dtype, var, dim0, dim1) \ - var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ - owner_)), - MJDATA_POINTERS -#undef MJ_M -#define MJ_M(x) (x) -#undef X - - contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), - -#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), - MJDATA_VECTOR -#undef X - model_(other.model_), - model_ref_(std::move(other.model_ref_)), - indexer_(ptr_, model_->get(), owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = - MjDataRawPointerMap().insert_or_assign(ptr_, this).second; - } - if (is_newly_inserted) { - throw UnexpectedError( - "MjDataRawPointerMap does not contains the moved-from mjData*"); - } - other.ptr_ = nullptr; -} - -MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model) - : WrapperBase(other.Copy(), &MjDataCapsuleDestructor), -#undef MJ_M -#define MJ_M(x) other.model_->get()->x -#define X(dtype, var, dim0, dim1) \ - var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ - owner_)), - MJDATA_POINTERS -#undef MJ_M -#define MJ_M(x) (x) -#undef X - - contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), - -#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), - MJDATA_VECTOR -#undef X - model_(model), - model_ref_(py::cast(model_)), - indexer_(ptr_, model_->get(), owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; - } - if (!is_newly_inserted) { - throw UnexpectedError( - "MjDataRawPointerMap already contains this raw mjData*"); - } -} - -MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d) - : WrapperBase(d, &MjDataCapsuleDestructor), -#undef MJ_M -#define MJ_M(x) model->get()->x -#define X(dtype, var, dim0, dim1) \ - var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)), - MJDATA_POINTERS -#undef MJ_M -#define MJ_M(x) (x) -#undef X - - contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), - -#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), - MJDATA_VECTOR -#undef X - model_(model), - model_ref_(py::cast(model_)), - indexer_(ptr_, model_->get(), owner_) { - bool is_newly_inserted = false; - { - py::gil_scoped_acquire gil; - is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; - } - if (!is_newly_inserted) { - throw UnexpectedError( - "MjDataRawPointerMap already contains this raw mjData*"); - } -} - -MjDataWrapper::~MjWrapper() { - if (ptr_) { - bool erased = false; - { - py::gil_scoped_acquire gil; - erased = MjDataRawPointerMap().erase(ptr_); - } - if (!erased) { - std::cerr << "MjDataRawPointerMap does not contain this raw mjData*" - << std::endl; - std::terminate(); - } - } -} - -void MjDataWrapper::Serialize(std::ostream& output) const { - // TODO: Replace this custom serialization with a protobuf - WriteChar(output, kSerializationVersion); - - model_->Serialize(output); - - // Write struct and scalar fields -#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var)) - X(maxuse_stack); - X(maxuse_arena); - X(maxuse_con); - X(maxuse_efc); - X(solver); - X(timer); - X(warning); - X(ncon); - X(ne); - X(nf); - X(nJ); - X(nA); - X(nefc); - X(nisland); - X(time); - X(energy); -#undef X - - // Write buffer and arena contents - { - MJDATA_POINTERS_PREAMBLE((this->model_->get())) - -#define X(type, name, nr, nc) \ - WriteBytes(output, ptr_->name, sizeof(type)*(this->model_->get()->nr)*(nc)); - MJDATA_POINTERS -#undef X - -#undef MJ_M -#define MJ_M(x) this->model_->get()->x -#undef MJ_D -#define MJ_D(x) this->ptr_->x -#define X(type, name, nr, nc) \ - if ((nr) * (nc)) { \ - WriteBytes(output, ptr_->name, \ - ptr_->name ? sizeof(type) * (nr) * (nc) : 0); \ - } - - MJDATA_ARENA_POINTERS_CONTACT - MJDATA_ARENA_POINTERS_SOLVER - if (mj_isDual(this->model_->get())) { - MJDATA_ARENA_POINTERS_DUAL - } - if (this->ptr_->nisland) { - MJDATA_ARENA_POINTERS_ISLAND - } -#undef MJ_M -#define MJ_M(x) x -#undef MJ_D -#define MJ_D(x) x -#undef X - } -} - -MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { - char serializationVersion = ReadChar(input); - CheckInput(input, "mjData"); - - if (serializationVersion != kSerializationVersion) { - throw py::value_error("Incompatible serialization version."); - } - - // Read the model that was used to create the mjData. - std::unique_ptr m_wrapper = - MjModelWrapper::Deserialize(input); - raw::MjModel& m = *m_wrapper->get(); - - bool is_dual = mj_isDual(&m); - - raw::MjData* d = mj_makeData(&m); - if (!d) { - throw py::value_error("Failed to create mjData."); - } - - // Read structs and scalar fields -#define X(var) \ - ReadBytes(input, (void*) &d->var, sizeof(d->var)); \ - CheckInput(input, "mjData"); - - X(maxuse_stack); - X(maxuse_arena); - X(maxuse_con); - X(maxuse_efc); - X(solver); - X(timer); - X(warning); - X(ncon); - X(ne); - X(nf); - X(nJ); - X(nA); - X(nefc); - X(nisland); - X(time); - X(energy); -#undef X - - // Read buffer and arena contents - { - MJDATA_POINTERS_PREAMBLE((&m)) - -#define X(type, name, nr, nc) \ - ReadBytes(input, d->name, sizeof(type)*(m.nr)*(nc)); - MJDATA_POINTERS -#undef X - -#undef MJ_M -#define MJ_M(x) m.x -#undef MJ_D -#define MJ_D(x) d->x -// arena pointers might be null, so we need to check the size before allocating. -#define X(type, name, nr, nc) \ - if ((nr) * (nc)) { \ - std::size_t actual_nbytes = ReadInt(input); \ - if (actual_nbytes) { \ - if (actual_nbytes != sizeof(type) * (nr) * (nc)) { \ - input.setstate(input.rdstate() | std::ios_base::failbit); \ - } else { \ - d->name = static_castname)>( \ - mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type))); \ - input.read(reinterpret_cast(d->name), actual_nbytes); \ - } \ - } else { \ - d->name = nullptr; \ - } \ - } - - MJDATA_ARENA_POINTERS_CONTACT - MJDATA_ARENA_POINTERS_SOLVER - if (is_dual) { - MJDATA_ARENA_POINTERS_DUAL - } - if (d->nisland) { - MJDATA_ARENA_POINTERS_ISLAND - } -#undef MJ_M -#define MJ_M(x) x -#undef MJ_D -#define MJ_D(x) x -#undef X - } - CheckInput(input, "mjData"); - - // All bytes should have been used. - input.ignore(1); - if (!input.eof()) { - throw py::value_error("Invalid serialized mjData."); - } - - return MjDataWrapper(m_wrapper.release(), d); -} - -raw::MjData* MjDataWrapper::Copy() const { - const raw::MjModel* m = model_->get(); - return InterceptMjErrors(mj_copyData)(NULL, m, this->ptr_); -} - -// ==================== MJSTATISTIC ============================================ -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjStatisticWrapper::MjWrapper() - : WrapperBase(new raw::MjStatistic{}), - X(center) {} - -MjStatisticWrapper::MjWrapper(raw::MjStatistic* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(center) {} -#undef X - -MjStatisticWrapper::MjWrapper(const MjStatisticWrapper& other) - : MjStatisticWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJWARNINGSTAT ========================================== -MjWarningStatWrapper::MjWrapper() - : WrapperBase(new raw::MjWarningStat{}) {} - -MjWarningStatWrapper::MjWrapper(raw::MjWarningStat* ptr, py::handle owner) - : WrapperBase(ptr, owner) {} - -MjWarningStatWrapper::MjWrapper(const MjWarningStatWrapper& other) - : MjWarningStatWrapper() { - *this->ptr_ = *other.ptr_; -} - -#define X(type, var) \ - var(std::vector{num}, std::vector{sizeof(raw::MjWarningStat)}, \ - &ptr->var, owner) -MjWarningStatList::MjStructList(raw::MjWarningStat* ptr, int num, - py::handle owner) - : StructListBase(ptr, num, owner), X(int, lastinfo), X(int, number) {} -#undef X - -// Slicing -#define X(type, var) var(other.var[slice]) -MjWarningStatList::MjStructList(MjWarningStatList& other, py::slice slice) - : StructListBase(other, slice), - X(int, lastinfo), - X(int, number) {} -#undef X - -// ==================== MJTIMERSTAT ============================================ -MjTimerStatWrapper::MjWrapper() - : WrapperBase(new raw::MjTimerStat{}) {} - -MjTimerStatWrapper::MjWrapper(raw::MjTimerStat* ptr, py::handle owner) - : WrapperBase(ptr, owner) {} - -MjTimerStatWrapper::MjWrapper(const MjTimerStatWrapper& other) - : MjTimerStatWrapper() { - *this->ptr_ = *other.ptr_; -} - -#define X(type, var) \ - var(std::vector{num}, std::vector{sizeof(raw::MjTimerStat)}, \ - &ptr->var, owner) -MjTimerStatList::MjStructList(raw::MjTimerStat* ptr, int num, py::handle owner) - : StructListBase(ptr, num, owner), - X(mjtNum, duration), - X(int, number) {} -#undef X - -// Slicing -#define X(type, var) var(other.var[slice]) -MjTimerStatList::MjStructList(MjTimerStatList& other, py::slice slice) - : StructListBase(other, slice), - X(mjtNum, duration), - X(int, number) {} -#undef X - -// ==================== MJSOLVERSTAT =========================================== -MjSolverStatWrapper::MjWrapper() - : WrapperBase(new raw::MjSolverStat{}) {} - -MjSolverStatWrapper::MjWrapper(raw::MjSolverStat* ptr, py::handle owner) - : WrapperBase(ptr, owner) {} - -MjSolverStatWrapper::MjWrapper(const MjSolverStatWrapper& other) - : MjSolverStatWrapper() { - *this->ptr_ = *other.ptr_; -} - -#define X(type, var) \ - var(std::vector{num}, std::vector{sizeof(raw::MjSolverStat)}, \ - &ptr->var, owner) -MjSolverStatList::MjStructList(raw::MjSolverStat* ptr, int num, - py::handle owner) - : StructListBase(ptr, num, owner), - X(mjtNum, improvement), - X(mjtNum, gradient), - X(mjtNum, lineslope), - X(int, nactive), - X(int, nchange), - X(int, neval), - X(int, nupdate) {} -#undef X -#undef XN - -// Slicing -#define X(type, var) var(other.var[slice]) -MjSolverStatList::MjStructList(MjSolverStatList& other, py::slice slice) - : StructListBase(other, slice), - X(mjtNum, improvement), - X(mjtNum, gradient), - X(mjtNum, lineslope), - X(int, nactive), - X(int, nchange), - X(int, neval), - X(int, nupdate) {} -#undef X - -// ==================== MJVPERTURB ============================================= -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvPerturbWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjvPerturb* const pert = new raw::MjvPerturb; - mjv_defaultPerturb(pert); - return pert; - }()), - X(refpos), - X(refquat), - X(refselpos), - X(localpos) {} -#undef X - -MjvPerturbWrapper::MjWrapper(const MjvPerturbWrapper& other) - : MjvPerturbWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVCAMERA ============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvCameraWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjvCamera* const cam = new raw::MjvCamera; - mjv_defaultCamera(cam); - return cam; - }()), - X(lookat) {} -#undef X - -MjvCameraWrapper::MjWrapper(const MjvCameraWrapper& other) - : MjvCameraWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVGLCAMERA ============================================ -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvGLCameraWrapper::MjWrapper() - : WrapperBase(new raw::MjvGLCamera{}), - X(pos), - X(forward), - X(up) {} - -MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(pos), - X(forward), - X(up) {} -#undef X - -MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera&& other) - : MjvGLCameraWrapper() { - *this->ptr_ = other; -} - -MjvGLCameraWrapper::MjWrapper(const MjvGLCameraWrapper& other) - : MjvGLCameraWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVGEOM ================================================ -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvGeomWrapper::MjWrapper() - : WrapperBase(new raw::MjvGeom{}), - X(size), - X(pos), - mat([this]() { - static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9); - return InitPyArray(std::array{3, 3}, ptr_->mat, owner_); - }()), - X(rgba) { - mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr); -} - -MjvGeomWrapper::MjWrapper(raw::MjvGeom* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(size), - X(pos), - mat([this]() { - static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9); - return InitPyArray(std::array{3, 3}, ptr_->mat, owner_); - }()), - X(rgba) {} -#undef X - -MjvGeomWrapper::MjWrapper(const MjvGeomWrapper& other) - : MjvGeomWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVLIGHT =============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvLightWrapper::MjWrapper() - : WrapperBase(new raw::MjvLight{}), - X(pos), - X(dir), - X(attenuation), - X(ambient), - X(diffuse), - X(specular) {} - -MjvLightWrapper::MjWrapper(raw::MjvLight* ptr, py::handle owner) - : WrapperBase(ptr, owner), - X(pos), - X(dir), - X(attenuation), - X(ambient), - X(diffuse), - X(specular) {} -#undef X - -MjvLightWrapper::MjWrapper(const MjvLightWrapper& other) - : MjvLightWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVOPTION ============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvOptionWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjvOption* const opt = new raw::MjvOption; - mjv_defaultOption(opt); - return opt; - }()), - X(geomgroup), - X(sitegroup), - X(jointgroup), - X(tendongroup), - X(actuatorgroup), - X(flexgroup), - X(skingroup), - X(flags) {} -#undef X - -MjvOptionWrapper::MjWrapper(const MjvOptionWrapper& other) - : MjvOptionWrapper() { - *this->ptr_ = *other.ptr_; -} - -// ==================== MJVSCENE =============================================== -static void MjvSceneCapsuleDestructor(PyObject* pyobj) { - py::gil_scoped_acquire gil; - auto* scn = static_cast(PyCapsule_GetPointer(pyobj, nullptr)); - if (scn) { - mjv_freeScene(scn); - delete scn; - } -} - -#define X(var) var(InitPyArray(ptr_->var, owner_)) -#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_)) -MjvSceneWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjvScene *const scn = new raw::MjvScene; - mjv_defaultScene(scn); - InterceptMjErrors(mjv_makeScene)(nullptr, scn, 0); - return scn; - }(), &MjvSceneCapsuleDestructor), - nskinvert(0), - XN(geoms, 0), - XN(geomorder, 0), - XN(flexedgeadr, 0), - XN(flexedgenum, 0), - XN(flexvertadr, 0), - XN(flexvertnum, 0), - XN(flexfaceadr, 0), - XN(flexfacenum, 0), - XN(flexfaceused, 0), - XN(flexedge, 0), - XN(flexvert, 0), - XN(flexface, 0), - XN(flexnormal, 0), - XN(flextexcoord, 0), - XN(skinfacenum, 0), - XN(skinvertadr, 0), - XN(skinvertnum, 0), - XN(skinvert, 0), - XN(skinnormal, 0), - X(lights), - X(camera), - X(translate), - X(rotate), - X(flags), - X(framergb) {} - -#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_)) -MjvSceneWrapper::MjWrapper(const MjModelWrapper& model, int maxgeom) - : WrapperBase( - [maxgeom](const raw::MjModel* m) { - raw::MjvScene *const scn = new raw::MjvScene; - mjv_defaultScene(scn); - InterceptMjErrors(mjv_makeScene)(m, scn, maxgeom); - return scn; - }(model.get()), - &MjvSceneCapsuleDestructor), - nskinvert([](const raw::MjModel* m) { - int nskinvert = 0; - for (int i = 0; i < m->nskin; ++i) { - nskinvert += m->skin_vertnum[i]; - } - return nskinvert; - }(model.get())), - nflexface([](const raw::MjModel* m) { - int nflexface = 0; - int flexfacenum = 0; - for (int f=0; f < m->nflex; f++) { - if (m->flex_dim[f] == 0) { - // 1D : 0 - flexfacenum = 0; - } else if (m->flex_dim[f] == 2) { - // 2D: 2*fragments + 2*elements - flexfacenum = 2*m->flex_shellnum[f] + 2*m->flex_elemnum[f]; - } else { - // 3D: max(fragments, 4*maxlayer) - // find number of elements in biggest layer - int maxlayer = 0, layer = 0, nlayer = 1; - while (nlayer) { - nlayer = 0; - for (int e=0; e < m->flex_elemnum[f]; e++) { - if (m->flex_elemlayer[m->flex_elemadr[f]+e] == layer) { - nlayer++; - } - } - maxlayer = mjMAX(maxlayer, nlayer); - layer++; - } - flexfacenum = mjMAX(m->flex_shellnum[f], 4*maxlayer); - } - - // accumulate over flexes - nflexface += flexfacenum; - } - return nflexface; - }(model.get())), - nflexedge(model.get()->nflexedge), - nflexvert(model.get()->nflexvert), - XN(geoms, ptr_->maxgeom), - XN(geomorder, ptr_->maxgeom), - XN(flexedgeadr, ptr_->nflex), - XN(flexedgenum, ptr_->nflex), - XN(flexvertadr, ptr_->nflex), - XN(flexvertnum, ptr_->nflex), - XN(flexfaceadr, ptr_->nflex), - XN(flexfacenum, ptr_->nflex), - XN(flexfaceused, ptr_->nflex), - XN(flexedge, 2*nflexedge), - XN(flexvert, 3*nflexvert), - XN(flexface, 9*nflexface), - XN(flexnormal, 9*nflexface), - XN(flextexcoord, 6*nflexface), - XN(skinfacenum, ptr_->nskin), - XN(skinvertadr, ptr_->nskin), - XN(skinvertnum, ptr_->nskin), - XN(skinvert, 3*nskinvert), - XN(skinnormal, 3*nskinvert), - X(lights), - X(camera), - X(translate), - X(rotate), - X(flags), - X(framergb) {} -#undef X -#undef XN - -template -static T* MallocAndCopy(const T* src, int count) { - if (src) { - T* out = static_cast(mju_malloc(count * sizeof(T))); - std::memcpy(out, src, count * sizeof(T)); - return out; - } else { - return nullptr; - } -} - -MjvSceneWrapper::MjWrapper(const MjvSceneWrapper& other) - : MjvSceneWrapper() { - mjv_freeScene(ptr_); - *ptr_ = *other.ptr_; - -#define XN(var, n) \ - ptr_->var = MallocAndCopy(other.ptr_->var, n); \ - var = InitPyArray(std::array{n}, ptr_->var, owner_); - - XN(geoms, ptr_->ngeom); - XN(geomorder, ptr_->ngeom); - XN(flexedgeadr, ptr_->nflex); - XN(flexedgenum, ptr_->nflex); - XN(flexvertadr, ptr_->nflex); - XN(flexvertnum, ptr_->nflex); - XN(flexfaceadr, ptr_->nflex); - XN(flexfacenum, ptr_->nflex); - XN(flexfaceused, ptr_->nflex); - XN(flexedge, 2*nflexedge); - XN(flexvert, 3*nflexvert); - XN(flexface, 9*nflexface); - XN(flexnormal, 9*nflexface); - XN(flextexcoord, 6*nflexface); - XN(skinfacenum, ptr_->nskin); - XN(skinvertadr, ptr_->nskin); - XN(skinvertnum, ptr_->nskin); - XN(skinvert, 3*nskinvert); - XN(skinnormal, 3*nskinvert); - -#undef XN -} - -// ==================== MJVFIGURE ============================================== -#define X(var) var(InitPyArray(ptr_->var, owner_)) -MjvFigureWrapper::MjWrapper() - : WrapperBase([]() { - raw::MjvFigure* const fig = new raw::MjvFigure; - mjv_defaultFigure(fig); - return fig; - }()), - X(flg_ticklabel), - X(gridsize), - X(gridrgb), - X(figurergba), - X(panergba), - X(legendrgba), - X(textrgb), - X(linergb), - X(range), - X(highlight), - X(linepnt), - X(linedata), - X(xaxispixel), - X(yaxispixel), - X(xaxisdata), - X(yaxisdata), -#undef X - - linename([](raw::MjvFigure* ptr, py::handle owner) { -// Use a macro to help us static_assert that the array extents here are kept -// in sync with mjVisualize.h. -#define MAKE_STR_ARRAY(N1, N2) \ - static_assert( \ - std::is_same_v); \ - return py::array(py::dtype("|S" #N2), N1, ptr->linename, owner); - - MAKE_STR_ARRAY(mjMAXLINE, 100); - -#undef MAKE_STR_ARRAY - }(ptr_, owner_)) {} - -MjvFigureWrapper::MjWrapper(const MjvFigureWrapper& other) - : MjvFigureWrapper() { - *this->ptr_ = *other.ptr_; -} - PYBIND11_MODULE(_structs, m) { py::module_::import("mujoco._enums"); @@ -1420,8 +151,8 @@ PYBIND11_MODULE(_structs, m) { << self.attr("__class__").attr("__name__").cast(); #define X(type, var) \ - result << "\n " #var ": "; \ - StructReprImpl(self.attr(#var), result, 2); + result << "\n " #var ": "; \ + StructReprImpl(self.attr(#var), result, 2); X(raw::MjVisualGlobal, global_) X(raw::MjVisualQuality, quality) @@ -1438,10 +169,10 @@ PYBIND11_MODULE(_structs, m) { mjVisualGlobal.def("__copy__", [](const raw::MjVisualGlobal& other) { return raw::MjVisualGlobal(other); }); - mjVisualGlobal.def( - "__deepcopy__", [](const raw::MjVisualGlobal& other, py::dict) { - return raw::MjVisualGlobal(other); - }); + mjVisualGlobal.def("__deepcopy__", + [](const raw::MjVisualGlobal& other, py::dict) { + return raw::MjVisualGlobal(other); + }); DefineStructFunctions(mjVisualGlobal); #define X(var) mjVisualGlobal.def_readwrite(#var, &raw::MjVisualGlobal::var) X(orthographic); @@ -1462,10 +193,10 @@ PYBIND11_MODULE(_structs, m) { mjVisualQuality.def("__copy__", [](const raw::MjVisualQuality& other) { return raw::MjVisualQuality(other); }); - mjVisualQuality.def( - "__deepcopy__", [](const raw::MjVisualQuality& other, py::dict) { - return raw::MjVisualQuality(other); - }); + mjVisualQuality.def("__deepcopy__", + [](const raw::MjVisualQuality& other, py::dict) { + return raw::MjVisualQuality(other); + }); DefineStructFunctions(mjVisualQuality); #define X(var) mjVisualQuality.def_readwrite(#var, &raw::MjVisualQuality::var) X(shadowsize); @@ -1476,21 +207,20 @@ PYBIND11_MODULE(_structs, m) { #undef X py::class_ mjVisualHeadlight(mjVisual, "Headlight"); - mjVisualHeadlight.def( - "__copy__", [](const MjVisualHeadlightWrapper& other) { - return MjVisualHeadlightWrapper(other); - }); - mjVisualHeadlight.def( - "__deepcopy__", [](const MjVisualHeadlightWrapper& other, py::dict) { - return MjVisualHeadlightWrapper(other); - }); + mjVisualHeadlight.def("__copy__", [](const MjVisualHeadlightWrapper& other) { + return MjVisualHeadlightWrapper(other); + }); + mjVisualHeadlight.def("__deepcopy__", + [](const MjVisualHeadlightWrapper& other, py::dict) { + return MjVisualHeadlightWrapper(other); + }); DefineStructFunctions(mjVisualHeadlight); - #define X(var) \ - DefinePyArray(mjVisualHeadlight, #var, &MjVisualHeadlightWrapper::var) +#define X(var) \ + DefinePyArray(mjVisualHeadlight, #var, &MjVisualHeadlightWrapper::var) X(ambient); X(diffuse); X(specular); - #undef X +#undef X mjVisualHeadlight.def_property( "active", [](const MjVisualHeadlightWrapper& c) { return c.get()->active; }, @@ -1502,10 +232,9 @@ PYBIND11_MODULE(_structs, m) { mjVisualMap.def("__copy__", [](const raw::MjVisualMap& other) { return raw::MjVisualMap(other); }); - mjVisualMap.def( - "__deepcopy__", [](const raw::MjVisualMap& other, py::dict) { - return raw::MjVisualMap(other); - }); + mjVisualMap.def("__deepcopy__", [](const raw::MjVisualMap& other, py::dict) { + return raw::MjVisualMap(other); + }); DefineStructFunctions(mjVisualMap); #define X(var) mjVisualMap.def_readwrite(#var, &raw::MjVisualMap::var) X(stiffness); @@ -1527,10 +256,10 @@ PYBIND11_MODULE(_structs, m) { mjVisualScale.def("__copy__", [](const raw::MjVisualScale& other) { return raw::MjVisualScale(other); }); - mjVisualScale.def( - "__deepcopy__", [](const raw::MjVisualScale& other, py::dict) { - return raw::MjVisualScale(other); - }); + mjVisualScale.def("__deepcopy__", + [](const raw::MjVisualScale& other, py::dict) { + return raw::MjVisualScale(other); + }); DefineStructFunctions(mjVisualScale); #define X(var) mjVisualScale.def_readwrite(#var, &raw::MjVisualScale::var) X(forcewidth); @@ -1556,10 +285,10 @@ PYBIND11_MODULE(_structs, m) { mjVisualRgba.def("__copy__", [](const MjVisualRgbaWrapper& other) { return MjVisualRgbaWrapper(other); }); - mjVisualRgba.def( - "__deepcopy__", [](const MjVisualRgbaWrapper& other, py::dict) { - return MjVisualRgbaWrapper(other); - }); + mjVisualRgba.def("__deepcopy__", + [](const MjVisualRgbaWrapper& other, py::dict) { + return MjVisualRgbaWrapper(other); + }); DefineStructFunctions(mjVisualRgba); #define X(var) DefinePyArray(mjVisualRgba, #var, &MjVisualRgbaWrapper::var) X(fog); @@ -1607,26 +336,26 @@ PYBIND11_MODULE(_structs, m) { // ==================== MJMODEL ============================================== py::class_ mjModel(m, "MjModel"); mjModel.def_static( - "from_xml_string", &MjModelWrapper::LoadXML, - py::arg("xml"), py::arg_v("assets", py::none()), + "from_xml_string", &MjModelWrapper::LoadXML, py::arg("xml"), + py::arg_v("assets", py::none()), py::doc( -R"(Loads an MjModel from an XML string and an optional assets dictionary.)")); + R"(Loads an MjModel from an XML string and an optional assets dictionary.)")); mjModel.def_static("_from_model_ptr", [](uintptr_t addr) { return MjModelWrapper::WrapRawModel(reinterpret_cast(addr)); }); mjModel.def_static( - "from_xml_path", &MjModelWrapper::LoadXMLFile, - py::arg("filename"), py::arg_v("assets", py::none()), + "from_xml_path", &MjModelWrapper::LoadXMLFile, py::arg("filename"), + py::arg_v("assets", py::none()), py::doc( -R"(Loads an MjModel from an XML file and an optional assets dictionary. + R"(Loads an MjModel from an XML file and an optional assets dictionary. The filename for the XML can also refer to a key in the assets dictionary. This is useful for example when the XML is not available as a file on disk.)")); mjModel.def_static( - "from_binary_path", &MjModelWrapper::LoadBinaryFile, - py::arg("filename"), py::arg_v("assets", py::none()), + "from_binary_path", &MjModelWrapper::LoadBinaryFile, py::arg("filename"), + py::arg_v("assets", py::none()), py::doc( -R"(Loads an MjModel from an MJB file and an optional assets dictionary. + R"(Loads an MjModel from an MJB file and an optional assets dictionary. The filename for the MJB can also refer to a key in the assets dictionary. This is useful for example when the MJB is not available as a file on disk.)")); @@ -1686,30 +415,37 @@ This is useful for example when the MJB is not available as a file on disk.)")); return py::tuple(py::cast(fields)); }); -#define X(dtype, var, dim0, dim1) \ - if constexpr (std::string_view(#var) != "text_data" && \ - std::string_view(#var) != "names" && \ - std::string_view(#var) != "paths") { \ - DefinePyArray(mjModel, #var, &MjModelWrapper::var); \ +#define X(dtype, var, dim0, dim1) \ + if constexpr (std::string_view(#var) != "text_data" && \ + std::string_view(#var) != "names" && \ + std::string_view(#var) != "paths") { \ + DefinePyArray(mjModel, #var, &MjModelWrapper::var); \ } MJMODEL_POINTERS #undef X - mjModel.def_property_readonly( - "text_data", [](const MjModelWrapper& m) -> const auto& { - // Return the full bytes array of concatenated text data - return m.text_data_bytes; - }); - mjModel.def_property_readonly( - "names", [](const MjModelWrapper& m) -> const auto& { - // Return the full bytes array of concatenated names - return m.names_bytes; - }); - mjModel.def_property_readonly( - "paths", [](const MjModelWrapper& m) -> const auto& { - // Return the full bytes array of concatenated paths - return m.paths_bytes; - }); + mjModel.def_property_readonly("text_data", + [](const MjModelWrapper& m) -> const auto& { + // Return the full bytes array of concatenated + // text data + return m.text_data_bytes; + }); + mjModel.def_property_readonly("names", + [](const MjModelWrapper& m) -> const auto& { + // Return the full bytes array of concatenated + // names + return m.names_bytes; + }); + mjModel.def_property_readonly("paths", + [](const MjModelWrapper& m) -> const auto& { + // Return the full bytes array of concatenated + // paths + return m.paths_bytes; + }); + mjModel.def_property_readonly("signature", + [](const MjModelWrapper& m) -> const uint64_t& { + return m.get()->signature; + }); #define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \ mjModel.def( \ @@ -1717,22 +453,28 @@ This is useful for example when the MJB is not available as a file on disk.)")); [](MjModelWrapper& m, int i) -> auto& { return m.indexer().field(i); }, \ py::return_value_policy::reference_internal); \ mjModel.def( \ - #field, [](MjModelWrapper& m, std::string_view name) -> auto& { \ + #field, \ + [](MjModelWrapper& m, std::string_view name) -> auto& { \ return m.indexer().field##_by_name(name); \ }, \ py::return_value_policy::reference_internal, py::arg_v("name", "")); - MJMODEL_VIEW_GROUPS #undef XGROUP -#define XGROUP(spectype, field) \ - mjModel.def( \ - "bind_scalar", \ - [](MjModelWrapper& m, spectype& spec) -> auto& { \ - return m.indexer().field##_by_name(mjs_getString(spec.name)); \ - }, \ - py::return_value_policy::reference_internal, \ +#define XGROUP(spectype, field) \ + mjModel.def( \ + "bind_scalar", \ + [](MjModelWrapper& m, spectype& spec) -> auto& { \ + if (mjs_getSpec(spec.element)->element->signature != \ + m.get()->signature) { \ + throw py::value_error( \ + "The mjSpec does not match mjModel. Please recompile " \ + "the mjSpec."); \ + } \ + return m.indexer().field(mjs_getId(spec.element)); \ + }, \ + py::return_value_policy::reference_internal, \ py::arg_v("spec", py::none())); MJMODEL_BIND_GROUPS @@ -1744,7 +486,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); [](MjModelWrapper& m, int i) -> auto& { return m.indexer().field(i); }, \ py::return_value_policy::reference_internal); \ mjModel.def( \ - #altname, [](MjModelWrapper& m, std::string_view name) -> auto& { \ + #altname, \ + [](MjModelWrapper& m, std::string_view name) -> auto& { \ return m.indexer().field##_by_name(name); \ }, \ py::return_value_policy::reference_internal, py::arg_v("name", "")); @@ -1758,12 +501,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); py::class_ groupedViews(m, "_" #MjModelGroupedViews); \ FIELD_XMACROS \ groupedViews.def("__repr__", MjModelStructRepr); \ - groupedViews.def_property_readonly("id", [](GroupedViews& views) { \ - return views.index(); \ - }); \ - groupedViews.def_property_readonly("name", [](GroupedViews& views) { \ - return views.name(); \ - }); \ + groupedViews.def_property_readonly( \ + "id", [](GroupedViews& views) { return views.index(); }); \ + groupedViews.def_property_readonly( \ + "name", [](GroupedViews& views) { return views.name(); }); \ } #define X(type, prefix, var, dim0, dim1) \ groupedViews.def_property( \ @@ -1778,8 +519,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); { py::handle builtins(PyEval_GetBuiltins()); builtins[MjModelWrapper::kFromRawPointer] = - reinterpret_cast(reinterpret_cast( - &MjModelWrapper::FromRawPointer)); + reinterpret_cast( + reinterpret_cast(&MjModelWrapper::FromRawPointer)); } // ==================== MJWARNINGSTAT ======================================== @@ -1788,10 +529,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjWarningStat.def("__copy__", [](const MjWarningStatWrapper& other) { return MjWarningStatWrapper(other); }); - mjWarningStat.def( - "__deepcopy__", [](const MjWarningStatWrapper& other, py::dict) { - return MjWarningStatWrapper(other); - }); + mjWarningStat.def("__deepcopy__", + [](const MjWarningStatWrapper& other, py::dict) { + return MjWarningStatWrapper(other); + }); DefineStructFunctions(mjWarningStat); #define X(var) \ mjWarningStat.def_property( \ @@ -1804,10 +545,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); #undef X py::class_ mjWarningStatList(m, "_MjWarningStatList"); - mjWarningStatList.def( - "__getitem__", - &MjWarningStatList::operator[], - py::return_value_policy::reference); + mjWarningStatList.def("__getitem__", &MjWarningStatList::operator[], + py::return_value_policy::reference); mjWarningStatList.def( "__getitem__", [](MjWarningStatList& list, ::mjtWarning idx) { return list[idx]; }, @@ -1828,10 +567,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjTimerStat.def("__copy__", [](const MjTimerStatWrapper& other) { return MjTimerStatWrapper(other); }); - mjTimerStat.def( - "__deepcopy__", [](const MjTimerStatWrapper& other, py::dict) { - return MjTimerStatWrapper(other); - }); + mjTimerStat.def("__deepcopy__", + [](const MjTimerStatWrapper& other, py::dict) { + return MjTimerStatWrapper(other); + }); DefineStructFunctions(mjTimerStat); #define X(var) \ mjTimerStat.def_property( \ @@ -1844,10 +583,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); #undef X py::class_ mjTimerStatList(m, "_MjTimerStatList"); - mjTimerStatList.def( - "__getitem__", - &MjTimerStatList::operator[], - py::return_value_policy::reference); + mjTimerStatList.def("__getitem__", &MjTimerStatList::operator[], + py::return_value_policy::reference); mjTimerStatList.def( "__getitem__", [](MjTimerStatList& list, ::mjtTimer idx) { return list[idx]; }, @@ -1856,8 +593,7 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjTimerStatList.def("__len__", &MjTimerStatList::size); DefineStructFunctions(mjTimerStatList); -#define X(type, var) \ - mjTimerStatList.def_readonly(#var, &MjTimerStatList::var) +#define X(type, var) mjTimerStatList.def_readonly(#var, &MjTimerStatList::var) X(mjtNum, duration); X(int, number); #undef X @@ -1868,10 +604,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjSolverStat.def("__copy__", [](const MjSolverStatWrapper& other) { return MjSolverStatWrapper(other); }); - mjSolverStat.def( - "__deepcopy__", [](const MjSolverStatWrapper& other, py::dict) { - return MjSolverStatWrapper(other); - }); + mjSolverStat.def("__deepcopy__", + [](const MjSolverStatWrapper& other, py::dict) { + return MjSolverStatWrapper(other); + }); DefineStructFunctions(mjSolverStat); #define X(var) \ mjSolverStat.def_property( \ @@ -1890,13 +626,12 @@ This is useful for example when the MJB is not available as a file on disk.)")); py::class_ mjSolverStatList(m, "_MjSolverStatList"); mjSolverStatList.def("__getitem__", &MjSolverStatList::operator[], - py::return_value_policy::reference); + py::return_value_policy::reference); mjSolverStatList.def("__getitem__", &MjSolverStatList::Slice); mjSolverStatList.def("__len__", &MjSolverStatList::size); DefineStructFunctions(mjSolverStatList); -#define X(type, var) \ - mjSolverStatList.def_readonly(#var, &MjSolverStatList::var) +#define X(type, var) mjSolverStatList.def_readonly(#var, &MjSolverStatList::var) X(mjtNum, improvement); X(mjtNum, gradient); X(mjtNum, lineslope); @@ -1996,12 +731,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjData.def_property_readonly("_address", [](const MjDataWrapper& d) { return reinterpret_cast(d.get()); }); - mjData.def_property_readonly("model", [](const MjDataWrapper& d) { - return &d.model(); - }); - mjData.def("__copy__", [](const MjDataWrapper& other) { - return MjDataWrapper(other); - }); + mjData.def_property_readonly( + "model", [](const MjDataWrapper& d) { return &d.model(); }); + mjData.def("__copy__", + [](const MjDataWrapper& other) { return MjDataWrapper(other); }); mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict memo) { // Use copy.deepcopy(model) to make a model that Python is aware of. py::object new_model_py = @@ -2018,6 +751,9 @@ This is useful for example when the MJB is not available as a file on disk.)")); std::istringstream input(b, std::ios::in | std::ios::binary); return MjDataWrapper::Deserialize(input); })); + mjData.def_property_readonly( + "signature", + [](const MjDataWrapper& d) -> uint64_t { return d.get()->signature; }); #define X(type, var) \ mjData.def_property( \ @@ -2064,7 +800,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); [](MjDataWrapper& d, int i) -> auto& { return d.indexer().field(i); }, \ py::return_value_policy::reference_internal); \ mjData.def( \ - #field, [](MjDataWrapper& d, std::string_view name) -> auto& { \ + #field, \ + [](MjDataWrapper& d, std::string_view name) -> auto& { \ return d.indexer().field##_by_name(name); \ }, \ py::return_value_policy::reference_internal, py::arg_v("name", "")); @@ -2072,13 +809,19 @@ This is useful for example when the MJB is not available as a file on disk.)")); MJDATA_VIEW_GROUPS #undef XGROUP -#define XGROUP(spectype, field) \ - mjData.def( \ - "bind_scalar", \ - [](MjDataWrapper& d, spectype& spec) -> auto& { \ - return d.indexer().field##_by_name(mjs_getString(spec.name)); \ - }, \ - py::return_value_policy::reference_internal, \ +#define XGROUP(spectype, field) \ + mjData.def( \ + "bind_scalar", \ + [](MjDataWrapper& d, spectype& spec) -> auto& { \ + if (mjs_getSpec(spec.element)->element->signature != \ + d.get()->signature) { \ + throw py::value_error( \ + "The mjSpec does not match mjData. Please recompile " \ + "the mjSpec."); \ + } \ + return d.indexer().field(mjs_getId(spec.element)); \ + }, \ + py::return_value_policy::reference_internal, \ py::arg_v("spec", py::none())); MJDATA_BIND_GROUPS @@ -2090,7 +833,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); [](MjDataWrapper& d, int i) -> auto& { return d.indexer().field(i); }, \ py::return_value_policy::reference_internal); \ mjData.def( \ - #altname, [](MjDataWrapper& d, std::string_view name) -> auto& { \ + #altname, \ + [](MjDataWrapper& d, std::string_view name) -> auto& { \ return d.indexer().field##_by_name(name); \ }, \ py::return_value_policy::reference_internal, py::arg_v("name", "")); @@ -2104,12 +848,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); py::class_ groupedViews(m, "_" #MjDataGroupedViews); \ FIELD_XMACROS \ groupedViews.def("__repr__", MjDataStructRepr); \ - groupedViews.def_property_readonly("id", [](GroupedViews& views) { \ - return views.index(); \ - }); \ - groupedViews.def_property_readonly("name", [](GroupedViews& views) { \ - return views.name(); \ - }); \ + groupedViews.def_property_readonly( \ + "id", [](GroupedViews& views) { return views.index(); }); \ + groupedViews.def_property_readonly( \ + "name", [](GroupedViews& views) { return views.name(); }); \ } #define X(type, prefix, var, dim0, dim1) \ groupedViews.def_property( \ @@ -2123,9 +865,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); { py::handle builtins(PyEval_GetBuiltins()); - builtins[MjDataWrapper::kFromRawPointer] = - reinterpret_cast(reinterpret_cast( - &MjDataWrapper::FromRawPointer)); + builtins[MjDataWrapper::kFromRawPointer] = reinterpret_cast( + reinterpret_cast(&MjDataWrapper::FromRawPointer)); } // ==================== MJSTATISTIC ========================================== @@ -2134,10 +875,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjStatistic.def("__copy__", [](const MjStatisticWrapper& other) { return MjStatisticWrapper(other); }); - mjStatistic.def( - "__deepcopy__", [](const MjStatisticWrapper& other, py::dict) { - return MjStatisticWrapper(other); - }); + mjStatistic.def("__deepcopy__", + [](const MjStatisticWrapper& other, py::dict) { + return MjStatisticWrapper(other); + }); DefineStructFunctions(mjStatistic); #define X(var) \ @@ -2159,9 +900,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); // ==================== MJLROPT ============================================== py::class_ mjLROpt(m, "MjLROpt"); mjLROpt.def(py::init<>()); - mjLROpt.def("__copy__", [](const raw::MjLROpt& other) { - return raw::MjLROpt(other); - }); + mjLROpt.def("__copy__", + [](const raw::MjLROpt& other) { return raw::MjLROpt(other); }); mjLROpt.def("__deepcopy__", [](const raw::MjLROpt& other, py::dict) { return raw::MjLROpt(other); }); @@ -2246,11 +986,10 @@ This is useful for example when the MJB is not available as a file on disk.)")); mjvGLCamera.def("__copy__", [](const MjvGLCameraWrapper& other) { return MjvGLCameraWrapper(other); }); - mjvGLCamera.def( - "__deepcopy__", - [](const MjvGLCameraWrapper& other, py::dict) { - return MjvGLCameraWrapper(other); - }); + mjvGLCamera.def("__deepcopy__", + [](const MjvGLCameraWrapper& other, py::dict) { + return MjvGLCameraWrapper(other); + }); DefineStructFunctions(mjvGLCamera); #define X(var) \ mjvGLCamera.def_property( \ @@ -2334,9 +1073,12 @@ This is useful for example when the MJB is not available as a file on disk.)")); X(cutoff); X(exponent); X(headlight); - X(directional); + X(type); + X(texid); X(castshadow); X(bulbradius); + X(intensity); + X(range); #undef X #define X(var) DefinePyArray(mjvLight, #var, &MjvLightWrapper::var) @@ -2361,7 +1103,7 @@ This is useful for example when the MJB is not available as a file on disk.)")); #define X(var) \ mjvOption.def_property( \ #var, [](const MjvOptionWrapper& c) { return c.get()->var; }, \ - [](MjvOptionWrapper& c, decltype(raw::MjvOption::var) rhs) { \ + [](MjvOptionWrapper& c, decltype(raw::MjvOption::var) rhs) { \ c.get()->var = rhs; \ }) X(label); @@ -2384,8 +1126,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); // ==================== MJVSCENE ============================================= py::class_ mjvScene(m, "MjvScene"); mjvScene.def(py::init<>()); - mjvScene.def(py::init(), - py::arg("model"), py::arg("maxgeom")); + mjvScene.def(py::init(), py::arg("model"), + py::arg("maxgeom")); mjvScene.def("__copy__", [](const MjvSceneWrapper& other) { return MjvSceneWrapper(other); }); @@ -2512,11 +1254,9 @@ This is useful for example when the MJB is not available as a file on disk.)")); py::arg("cam1"), py::arg("cam2"), py::doc(python_traits::mjv_averageCamera::doc)); - m.def( - "_recompile_spec_addr", - [](uintptr_t spec_addr, const MjModelWrapper& m, const MjDataWrapper& d) { - return RecompileSpec(reinterpret_cast(spec_addr), m, d); - } - ); + m.def("_recompile_spec_addr", [](uintptr_t spec_addr, const MjModelWrapper& m, + const MjDataWrapper& d) { + return RecompileSpec(reinterpret_cast(spec_addr), m, d); + }); } // PYBIND11_MODULE NOLINT(readability/fn_size) } // namespace mujoco::python::_impl diff --git a/python/mujoco/structs_wrappers.cc b/python/mujoco/structs_wrappers.cc new file mode 100644 index 00000000..7c1ec9b0 --- /dev/null +++ b/python/mujoco/structs_wrappers.cc @@ -0,0 +1,1328 @@ +// Copyright 2022 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. + +#include + +#include +#include +#include +#include // NOLINT(build/c++11) +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include "errors.h" +#include "private.h" +#include "raw.h" +#include "serialization.h" +#include "structs.h" +#include +#include +#include +#include +#include +#include +#include + +namespace mujoco::python::_impl { + +namespace py = ::pybind11; + +namespace { +#define PTRDIFF(x, y) \ + reinterpret_cast(x) - reinterpret_cast(y) + +// Returns the shape of a NumPy array given the dimensions from an X Macro. +// If dim1 is a _literal_ constant 1, the resulting array is 1-dimensional of +// length dim0, otherwise the resulting array is 2-dimensional of shape +// (dim0, dim1). +#define X_ARRAY_SHAPE(dim0, dim1) XArrayShapeImpl(#dim1)((dim0), (dim1)) + +std::vector XArrayShapeImpl1D(int dim0, int dim1) { return {dim0}; } + +std::vector XArrayShapeImpl2D(int dim0, int dim1) { return {dim0, dim1}; } + +constexpr auto XArrayShapeImpl(const std::string_view dim1_str) { + if (dim1_str == "1") { + return XArrayShapeImpl1D; + } else { + return XArrayShapeImpl2D; + } +} + +inline std::size_t NConMax(const mjData* d) { + return d->narena / sizeof(mjContact); +} + +} // namespace + +// ==================== MJOPTION =============================================== +#define X(var, dim) , var(InitPyArray(std::array{dim}, ptr_->var, owner_)) +MjOptionWrapper::MjWrapper() + : WrapperBase([]() { + raw::MjOption* const opt = new raw::MjOption; + mj_defaultOption(opt); + return opt; + }()) MJOPTION_VECTORS {} + +MjOptionWrapper::MjWrapper(raw::MjOption* ptr, py::handle owner) + : WrapperBase(ptr, owner) MJOPTION_VECTORS {} +#undef X + +MjOptionWrapper::MjWrapper(const MjOptionWrapper& other) : MjOptionWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVISUAL =============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjVisualHeadlightWrapper::MjWrapper() + : WrapperBase(new raw::MjVisualHeadlight{}), + X(ambient), + X(diffuse), + X(specular) {} + +MjVisualHeadlightWrapper::MjWrapper(raw::MjVisualHeadlight* ptr, + py::handle owner) + : WrapperBase(ptr, owner), X(ambient), X(diffuse), X(specular) {} +#undef X + +MjVisualHeadlightWrapper::MjWrapper(const MjVisualHeadlightWrapper& other) + : MjVisualHeadlightWrapper() { + *this->ptr_ = *other.ptr_; +} + +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjVisualRgbaWrapper::MjWrapper() + : WrapperBase(new raw::MjVisualRgba{}), + X(fog), + X(haze), + X(force), + X(inertia), + X(joint), + X(actuator), + X(actuatornegative), + X(actuatorpositive), + X(com), + X(camera), + X(light), + X(selectpoint), + X(connect), + X(contactpoint), + X(contactforce), + X(contactfriction), + X(contacttorque), + X(contactgap), + X(rangefinder), + X(constraint), + X(slidercrank), + X(crankbroken), + X(frustum) {} + +MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner) + : WrapperBase(ptr, owner), + X(fog), + X(haze), + X(force), + X(inertia), + X(joint), + X(actuator), + X(actuatornegative), + X(actuatorpositive), + X(com), + X(camera), + X(light), + X(selectpoint), + X(connect), + X(contactpoint), + X(contactforce), + X(contactfriction), + X(contacttorque), + X(contactgap), + X(rangefinder), + X(constraint), + X(slidercrank), + X(crankbroken), + X(frustum) {} +#undef X + +MjVisualRgbaWrapper::MjWrapper(const MjVisualRgbaWrapper& other) + : MjVisualRgbaWrapper() { + *this->ptr_ = *other.ptr_; +} + +MjVisualWrapper::MjWrapper() + : WrapperBase(new raw::MjVisual{}), + headlight(&ptr_->headlight, owner_), + rgba(&ptr_->rgba, owner_) {} + +MjVisualWrapper::MjWrapper(raw::MjVisual* ptr, py::handle owner) + : WrapperBase(ptr, owner), + headlight(&ptr_->headlight, owner_), + rgba(&ptr_->rgba, owner_) {} + +MjVisualWrapper::MjWrapper(const MjVisualWrapper& other) : MjVisualWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJMODEL ================================================ +static void MjModelCapsuleDestructor(PyObject* pyobj) { + mj_deleteModel( + static_cast(PyCapsule_GetPointer(pyobj, nullptr))); +} + +static absl::flat_hash_map& +MjModelRawPointerMap() { + static auto* hash_map = + new absl::flat_hash_map(); + return *hash_map; +} + +MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept { + try { + auto& map = MjModelRawPointerMap(); + { + py::gil_scoped_acquire gil; + auto found = map.find(m); + return found != map.end() ? found->second : nullptr; + } + } catch (...) { + return nullptr; + } +} + +#undef MJ_M +#define MJ_M(x) ptr_->x +#define X(dtype, var, dim0, dim1) \ + , var(InitPyArray(X_ARRAY_SHAPE(ptr_->dim0, dim1), ptr_->var, owner_)) +MjModelWrapper::MjWrapper(raw::MjModel* ptr) + : WrapperBase(ptr, &MjModelCapsuleDestructor), + opt(&ptr->opt, owner_), + vis(&ptr->vis, owner_), + stat(&ptr->stat, owner_) MJMODEL_POINTERS, + text_data_bytes(ptr->text_data, ptr->ntextdata), + names_bytes(ptr->names, ptr->nnames), + paths_bytes(ptr->paths, ptr->npaths), + indexer_(ptr, owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second; + } + if (!is_newly_inserted) { + throw UnexpectedError( + "MjModelWrapper(mjModel*): MjModelRawPointerMap already contains this " + "raw mjModel*"); + } +} + +MjModelWrapper::MjWrapper(MjModelWrapper&& other) + : WrapperBase(other.ptr_, other.owner_), + opt(&ptr_->opt, owner_), + vis(&ptr_->vis, owner_), + stat(&ptr_->stat, owner_) MJMODEL_POINTERS, + text_data_bytes(ptr_->text_data, ptr_->ntextdata), + names_bytes(ptr_->names, ptr_->nnames), + paths_bytes(ptr_->paths, ptr_->npaths), + indexer_(ptr_, owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = + MjModelRawPointerMap().insert_or_assign(ptr_, this).second; + } + if (is_newly_inserted) { + throw UnexpectedError( + "MjModelRawPointerMap does not contains the moved-from mjModel*"); + } + other.ptr_ = nullptr; +} +#undef X +#undef MJ_M +#define MJ_M(x) x + +// Delegating to the MjModelWrapper::MjWrapper(raw::MjModel*) constructor, +// no need to modify MjModelRawPointerMap here. +MjModelWrapper::MjWrapper(const MjModelWrapper& other) + : MjModelWrapper(InterceptMjErrors(mj_copyModel)(NULL, other.get())) {} + +MjModelWrapper::~MjWrapper() { + if (ptr_) { + bool erased = false; + { + py::gil_scoped_acquire gil; + erased = MjModelRawPointerMap().erase(ptr_); + } + if (!erased) { + std::cerr << "MjModelRawPointerMap does not contain this raw mjModel*" + << std::endl; + std::terminate(); + } + } +} + +// Helper function for both LoadXMLFile and LoadBinaryFile. +// Creates a temporary MJB from the assets dictionary if one is supplied. +template +static raw::MjModel* LoadModelFileImpl(const std::string& filename, + const std::vector& assets, + LoadFunc&& loadfunc) { + mjVFS vfs; + mjVFS* vfs_ptr = nullptr; + if (!assets.empty()) { + mj_defaultVFS(&vfs); + vfs_ptr = &vfs; + for (const auto& asset : assets) { + std::string buffer_name = StripPath(asset.name); + const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( + vfs_ptr, buffer_name.c_str(), asset.content, asset.content_size); + if (vfs_error) { + mj_deleteVFS(vfs_ptr); + if (vfs_error == 2) { + throw py::value_error("Repeated file name in assets dict: " + + buffer_name); + } else { + throw py::value_error("Asset failed to load: " + buffer_name); + } + } + } + } + + raw::MjModel* model = loadfunc(filename.c_str(), vfs_ptr); + mj_deleteVFS(vfs_ptr); + if (model && !model->buffer) { + mj_deleteModel(model); + model = nullptr; + } + return model; +} + +MjModelWrapper MjModelWrapper::LoadXMLFile( + const std::string& filename, + const std::optional>& assets) { + const auto converted_assets = ConvertAssetsDict(assets); + raw::MjModel* model; + { + py::gil_scoped_release no_gil; + char error[1024]; + model = LoadModelFileImpl(filename, converted_assets, + [&error](const char* filename, const mjVFS* vfs) { + return InterceptMjErrors(mj_loadXML)( + filename, vfs, error, sizeof(error)); + }); + if (!model) { + throw py::value_error(error); + } + } + return MjModelWrapper(model); +} + +MjModelWrapper MjModelWrapper::LoadBinaryFile( + const std::string& filename, + const std::optional>& assets) { + const auto converted_assets = ConvertAssetsDict(assets); + raw::MjModel* model; + { + py::gil_scoped_release no_gil; + model = LoadModelFileImpl(filename, converted_assets, + InterceptMjErrors(mj_loadModel)); + if (!model) { + throw py::value_error("mj_loadModel: failed to load from mjb"); + } + } + return MjModelWrapper(model); +} + +MjModelWrapper MjModelWrapper::LoadXML( + const std::string& xml, + const std::optional>& assets) { + auto converted_assets = ConvertAssetsDict(assets); + raw::MjModel* model; + { + py::gil_scoped_release no_gil; + std::string model_filename = "model_.xml"; + if (assets.has_value()) { + while (assets->find(model_filename) != assets->end()) { + model_filename = + model_filename.substr(0, model_filename.size() - 4) + "_.xml"; + } + } + converted_assets.emplace_back(model_filename.c_str(), xml.c_str(), + xml.length()); + char error[1024]; + model = LoadModelFileImpl(model_filename, converted_assets, + [&error](const char* filename, const mjVFS* vfs) { + return InterceptMjErrors(mj_loadXML)( + filename, vfs, error, sizeof(error)); + }); + if (!model) { + throw py::value_error(error); + } + } + return MjModelWrapper(model); +} + +MjModelWrapper MjModelWrapper::WrapRawModel(raw::MjModel* m) { + return MjModelWrapper(m); +} + +py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m, + const MjDataWrapper& old_d) { + raw::MjModel* m = static_cast(mju_malloc(sizeof(mjModel))); + m->buffer = nullptr; + raw::MjData* d = mj_copyData(nullptr, old_m.get(), old_d.get()); + if (mj_recompile(spec, nullptr, m, d)) { + throw py::value_error(mjs_getError(spec)); + } + + py::object m_pyobj = py::cast((MjModelWrapper(m))); + py::object d_pyobj = + py::cast((MjDataWrapper(py::cast(m_pyobj), d))); + return py::make_tuple(m_pyobj, d_pyobj); +} + +namespace { +// A byte at the start of serialized mjModel structs, which can be incremented +// when we change the serialization logic to reject pickles from an unsupported +// future version. +constexpr static char kSerializationVersion = 1; + +void CheckInput(const std::istream& input, std::string class_name) { + if (input.fail()) { + throw py::value_error("Invalid serialized " + class_name + "."); + } +} + +} // namespace + +void MjModelWrapper::Serialize(std::ostream& output) const { + WriteChar(output, kSerializationVersion); + + int model_size = mj_sizeModel(get()); + WriteInt(output, model_size); + std::string buffer(model_size, 0); + mj_saveModel(get(), nullptr, buffer.data(), model_size); + WriteBytes(output, buffer.data(), model_size); +} + +std::unique_ptr MjModelWrapper::Deserialize( + std::istream& input) { + CheckInput(input, "mjModel"); + + char serializationVersion = ReadChar(input); + CheckInput(input, "mjModel"); + + if (serializationVersion != kSerializationVersion) { + throw py::value_error("Incompatible serialization version."); + } + + std::size_t model_size = ReadInt(input); + CheckInput(input, "mjModel"); + if (model_size < 0) { + throw py::value_error("Invalid serialized mjModel."); + } + std::string model_bytes(model_size, 0); + ReadBytes(input, model_bytes.data(), model_size); + CheckInput(input, "mjModel"); + + raw::MjModel* model = LoadModelFileImpl( + "model.mjb", + {{"model.mjb", model_bytes.data(), static_cast(model_size)}}, + InterceptMjErrors(mj_loadModel)); + if (!model) { + throw py::value_error("Invalid serialized mjModel."); + } + return std::unique_ptr(new MjModelWrapper(model)); +} + +// ==================== MJCONTACT ============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjContactWrapper::MjWrapper() + : WrapperBase(new raw::MjContact{}), + X(pos), + X(frame), + X(friction), + X(solref), + X(solreffriction), + X(solimp), + X(H), + X(geom), + X(flex), + X(elem), + X(vert) {} + +MjContactWrapper::MjWrapper(raw::MjContact* ptr, py::handle owner) + : WrapperBase(ptr, owner), + X(pos), + X(frame), + X(friction), + X(solref), + X(solreffriction), + X(solimp), + X(H), + X(geom), + X(flex), + X(elem), + X(vert) {} +#undef X + +MjContactWrapper::MjWrapper(const MjContactWrapper& other) + : MjContactWrapper() { + *this->ptr_ = *other.ptr_; +} + +MjContactList::MjStructList(raw::MjContact* ptr, int nconmax, int* ncon, + py::handle owner) + : StructListBase(ptr, nconmax, owner, /* lazy = */ true), ncon_(ncon) {} + +// Slicing +MjContactList::MjStructList(MjContactList& other, py::slice slice) + : StructListBase(other, slice), ncon_(other.ncon_) {} + +// ==================== MJDATA ================================================= +static void MjDataCapsuleDestructor(PyObject* pyobj) { + mj_deleteData( + static_cast(PyCapsule_GetPointer(pyobj, nullptr))); +} + +absl::flat_hash_map& MjDataRawPointerMap() { + static auto* hash_map = + new absl::flat_hash_map(); + return *hash_map; +} + +MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept { + try { + auto& map = MjDataRawPointerMap(); + { + py::gil_scoped_acquire gil; + auto found = map.find(m); + return found != map.end() ? found->second : nullptr; + } + } catch (...) { + return nullptr; + } +} + +namespace { +// default timer callback (seconds) +mjtNum GetTime() { + using Clock = std::chrono::steady_clock; + using Seconds = std::chrono::duration; + static const Clock::time_point tm_start = Clock::now(); + return Seconds(Clock::now() - tm_start).count(); +} +} // namespace + +MjDataWrapper::MjWrapper(MjModelWrapper* model) + : WrapperBase(InterceptMjErrors(mj_makeData)(model->get()), + &MjDataCapsuleDestructor), +#undef MJ_M +#define MJ_M(x) model->get()->x +#define X(dtype, var, dim0, dim1) \ + var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)), + MJDATA_POINTERS +#undef MJ_M +#define MJ_M(x) (x) +#undef X + + contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), + +#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), + MJDATA_VECTOR +#undef X + model_(model), + model_ref_(py::cast(model_)), + indexer_(ptr_, model_->get(), owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; + } + if (!is_newly_inserted) { + throw UnexpectedError( + "MjDataRawPointerMap already contains this raw mjData*"); + } + + // install default timer if not already installed + { + py::gil_scoped_acquire gil; + if (!mjcb_time) { + mjcb_time = GetTime; + } + } +} + +MjDataWrapper::MjWrapper(const MjDataWrapper& other) + : WrapperBase(other.Copy(), &MjDataCapsuleDestructor), +#undef MJ_M +#define MJ_M(x) other.model_->get()->x +#define X(dtype, var, dim0, dim1) \ + var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ + owner_)), + MJDATA_POINTERS +#undef MJ_M +#define MJ_M(x) (x) +#undef X + + contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), + +#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), + MJDATA_VECTOR +#undef X + model_(other.model_), + model_ref_(other.model_ref_), + indexer_(ptr_, model_->get(), owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; + } + if (!is_newly_inserted) { + throw UnexpectedError( + "MjDataRawPointerMap already contains this raw mjData*"); + } +} + +MjDataWrapper::MjWrapper(MjDataWrapper&& other) + : WrapperBase(other.ptr_, other.owner_), +#undef MJ_M +#define MJ_M(x) other.model_->get()->x +#define X(dtype, var, dim0, dim1) \ + var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ + owner_)), + MJDATA_POINTERS +#undef MJ_M +#define MJ_M(x) (x) +#undef X + + contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), + +#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), + MJDATA_VECTOR +#undef X + model_(other.model_), + model_ref_(std::move(other.model_ref_)), + indexer_(ptr_, model_->get(), owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = + MjDataRawPointerMap().insert_or_assign(ptr_, this).second; + } + if (is_newly_inserted) { + throw UnexpectedError( + "MjDataRawPointerMap does not contains the moved-from mjData*"); + } + other.ptr_ = nullptr; +} + +MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model) + : WrapperBase(other.Copy(), &MjDataCapsuleDestructor), +#undef MJ_M +#define MJ_M(x) other.model_->get()->x +#define X(dtype, var, dim0, dim1) \ + var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \ + owner_)), + MJDATA_POINTERS +#undef MJ_M +#define MJ_M(x) (x) +#undef X + + contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), + +#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), + MJDATA_VECTOR +#undef X + model_(model), + model_ref_(py::cast(model_)), + indexer_(ptr_, model_->get(), owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; + } + if (!is_newly_inserted) { + throw UnexpectedError( + "MjDataRawPointerMap already contains this raw mjData*"); + } +} + +MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d) + : WrapperBase(d, &MjDataCapsuleDestructor), +#undef MJ_M +#define MJ_M(x) model->get()->x +#define X(dtype, var, dim0, dim1) \ + var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)), + MJDATA_POINTERS +#undef MJ_M +#define MJ_M(x) (x) +#undef X + + contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)), + +#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)), + MJDATA_VECTOR +#undef X + model_(model), + model_ref_(py::cast(model_)), + indexer_(ptr_, model_->get(), owner_) { + bool is_newly_inserted = false; + { + py::gil_scoped_acquire gil; + is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; + } + if (!is_newly_inserted) { + throw UnexpectedError( + "MjDataRawPointerMap already contains this raw mjData*"); + } +} + +MjDataWrapper::~MjWrapper() { + if (ptr_) { + bool erased = false; + { + py::gil_scoped_acquire gil; + erased = MjDataRawPointerMap().erase(ptr_); + } + if (!erased) { + std::cerr << "MjDataRawPointerMap does not contain this raw mjData*" + << std::endl; + std::terminate(); + } + } +} + +void MjDataWrapper::Serialize(std::ostream& output) const { + // TODO: Replace this custom serialization with a protobuf + WriteChar(output, kSerializationVersion); + + model_->Serialize(output); + + // Write struct and scalar fields +#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var)) + X(maxuse_stack); + X(maxuse_arena); + X(maxuse_con); + X(maxuse_efc); + X(solver); + X(timer); + X(warning); + X(ncon); + X(ne); + X(nf); + X(nJ); + X(nA); + X(nefc); + X(nisland); + X(time); + X(energy); +#undef X + + // Write buffer and arena contents + { + MJDATA_POINTERS_PREAMBLE((this->model_->get())) + +#define X(type, name, nr, nc) \ + WriteBytes(output, ptr_->name, \ + sizeof(type) * (this->model_->get()->nr) * (nc)); + MJDATA_POINTERS +#undef X + +#undef MJ_M +#define MJ_M(x) this->model_->get()->x +#undef MJ_D +#define MJ_D(x) this->ptr_->x +#define X(type, name, nr, nc) \ + if ((nr) * (nc)) { \ + WriteBytes(output, ptr_->name, \ + ptr_->name ? sizeof(type) * (nr) * (nc) : 0); \ + } + + MJDATA_ARENA_POINTERS_CONTACT + MJDATA_ARENA_POINTERS_SOLVER + if (mj_isDual(this->model_->get())) { + MJDATA_ARENA_POINTERS_DUAL + } + if (this->ptr_->nisland) { + MJDATA_ARENA_POINTERS_ISLAND + } +#undef MJ_M +#define MJ_M(x) x +#undef MJ_D +#define MJ_D(x) x +#undef X + } +} + +MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { + char serializationVersion = ReadChar(input); + CheckInput(input, "mjData"); + + if (serializationVersion != kSerializationVersion) { + throw py::value_error("Incompatible serialization version."); + } + + // Read the model that was used to create the mjData. + std::unique_ptr m_wrapper = + MjModelWrapper::Deserialize(input); + raw::MjModel& m = *m_wrapper->get(); + + bool is_dual = mj_isDual(&m); + + raw::MjData* d = mj_makeData(&m); + if (!d) { + throw py::value_error("Failed to create mjData."); + } + + // Read structs and scalar fields +#define X(var) \ + ReadBytes(input, (void*)&d->var, sizeof(d->var)); \ + CheckInput(input, "mjData"); + + X(maxuse_stack); + X(maxuse_arena); + X(maxuse_con); + X(maxuse_efc); + X(solver); + X(timer); + X(warning); + X(ncon); + X(ne); + X(nf); + X(nJ); + X(nA); + X(nefc); + X(nisland); + X(time); + X(energy); +#undef X + + // Read buffer and arena contents + { + MJDATA_POINTERS_PREAMBLE((&m)) + +#define X(type, name, nr, nc) \ + ReadBytes(input, d->name, sizeof(type) * (m.nr) * (nc)); + MJDATA_POINTERS +#undef X + +#undef MJ_M +#define MJ_M(x) m.x +#undef MJ_D +#define MJ_D(x) d->x +// arena pointers might be null, so we need to check the size before allocating. +#define X(type, name, nr, nc) \ + if ((nr) * (nc)) { \ + std::size_t actual_nbytes = ReadInt(input); \ + if (actual_nbytes) { \ + if (actual_nbytes != sizeof(type) * (nr) * (nc)) { \ + input.setstate(input.rdstate() | std::ios_base::failbit); \ + } else { \ + d->name = static_castname)>( \ + mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type))); \ + input.read(reinterpret_cast(d->name), actual_nbytes); \ + } \ + } else { \ + d->name = nullptr; \ + } \ + } + + MJDATA_ARENA_POINTERS_CONTACT + MJDATA_ARENA_POINTERS_SOLVER + if (is_dual) { + MJDATA_ARENA_POINTERS_DUAL + } + if (d->nisland) { + MJDATA_ARENA_POINTERS_ISLAND + } +#undef MJ_M +#define MJ_M(x) x +#undef MJ_D +#define MJ_D(x) x +#undef X + } + CheckInput(input, "mjData"); + + // All bytes should have been used. + input.ignore(1); + if (!input.eof()) { + throw py::value_error("Invalid serialized mjData."); + } + + return MjDataWrapper(m_wrapper.release(), d); +} + +raw::MjData* MjDataWrapper::Copy() const { + const raw::MjModel* m = model_->get(); + return InterceptMjErrors(mj_copyData)(NULL, m, this->ptr_); +} + +// ==================== MJSTATISTIC ============================================ +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjStatisticWrapper::MjWrapper() + : WrapperBase(new raw::MjStatistic{}), X(center) {} + +MjStatisticWrapper::MjWrapper(raw::MjStatistic* ptr, py::handle owner) + : WrapperBase(ptr, owner), X(center) {} +#undef X + +MjStatisticWrapper::MjWrapper(const MjStatisticWrapper& other) + : MjStatisticWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJWARNINGSTAT ========================================== +MjWarningStatWrapper::MjWrapper() : WrapperBase(new raw::MjWarningStat{}) {} + +MjWarningStatWrapper::MjWrapper(raw::MjWarningStat* ptr, py::handle owner) + : WrapperBase(ptr, owner) {} + +MjWarningStatWrapper::MjWrapper(const MjWarningStatWrapper& other) + : MjWarningStatWrapper() { + *this->ptr_ = *other.ptr_; +} + +#define X(type, var) \ + var(std::vector{num}, std::vector{sizeof(raw::MjWarningStat)}, \ + &ptr->var, owner) +MjWarningStatList::MjStructList(raw::MjWarningStat* ptr, int num, + py::handle owner) + : StructListBase(ptr, num, owner), X(int, lastinfo), X(int, number) {} +#undef X + +// Slicing +#define X(type, var) var(other.var[slice]) +MjWarningStatList::MjStructList(MjWarningStatList& other, py::slice slice) + : StructListBase(other, slice), X(int, lastinfo), X(int, number) {} +#undef X + +// ==================== MJTIMERSTAT ============================================ +MjTimerStatWrapper::MjWrapper() : WrapperBase(new raw::MjTimerStat{}) {} + +MjTimerStatWrapper::MjWrapper(raw::MjTimerStat* ptr, py::handle owner) + : WrapperBase(ptr, owner) {} + +MjTimerStatWrapper::MjWrapper(const MjTimerStatWrapper& other) + : MjTimerStatWrapper() { + *this->ptr_ = *other.ptr_; +} + +#define X(type, var) \ + var(std::vector{num}, std::vector{sizeof(raw::MjTimerStat)}, \ + &ptr->var, owner) +MjTimerStatList::MjStructList(raw::MjTimerStat* ptr, int num, py::handle owner) + : StructListBase(ptr, num, owner), X(mjtNum, duration), X(int, number) {} +#undef X + +// Slicing +#define X(type, var) var(other.var[slice]) +MjTimerStatList::MjStructList(MjTimerStatList& other, py::slice slice) + : StructListBase(other, slice), X(mjtNum, duration), X(int, number) {} +#undef X + +// ==================== MJSOLVERSTAT =========================================== +MjSolverStatWrapper::MjWrapper() : WrapperBase(new raw::MjSolverStat{}) {} + +MjSolverStatWrapper::MjWrapper(raw::MjSolverStat* ptr, py::handle owner) + : WrapperBase(ptr, owner) {} + +MjSolverStatWrapper::MjWrapper(const MjSolverStatWrapper& other) + : MjSolverStatWrapper() { + *this->ptr_ = *other.ptr_; +} + +#define X(type, var) \ + var(std::vector{num}, std::vector{sizeof(raw::MjSolverStat)}, \ + &ptr->var, owner) +MjSolverStatList::MjStructList(raw::MjSolverStat* ptr, int num, + py::handle owner) + : StructListBase(ptr, num, owner), + X(mjtNum, improvement), + X(mjtNum, gradient), + X(mjtNum, lineslope), + X(int, nactive), + X(int, nchange), + X(int, neval), + X(int, nupdate) {} +#undef X +#undef XN + +// Slicing +#define X(type, var) var(other.var[slice]) +MjSolverStatList::MjStructList(MjSolverStatList& other, py::slice slice) + : StructListBase(other, slice), + X(mjtNum, improvement), + X(mjtNum, gradient), + X(mjtNum, lineslope), + X(int, nactive), + X(int, nchange), + X(int, neval), + X(int, nupdate) {} +#undef X + +// ==================== MJVPERTURB ============================================= +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvPerturbWrapper::MjWrapper() + : WrapperBase([]() { + raw::MjvPerturb* const pert = new raw::MjvPerturb; + mjv_defaultPerturb(pert); + return pert; + }()), + X(refpos), + X(refquat), + X(refselpos), + X(localpos) {} +#undef X + +MjvPerturbWrapper::MjWrapper(const MjvPerturbWrapper& other) + : MjvPerturbWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVCAMERA ============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvCameraWrapper::MjWrapper() + : WrapperBase([]() { + raw::MjvCamera* const cam = new raw::MjvCamera; + mjv_defaultCamera(cam); + return cam; + }()), + X(lookat) {} +#undef X + +MjvCameraWrapper::MjWrapper(const MjvCameraWrapper& other) + : MjvCameraWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVGLCAMERA ============================================ +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvGLCameraWrapper::MjWrapper() + : WrapperBase(new raw::MjvGLCamera{}), X(pos), X(forward), X(up) {} + +MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera* ptr, py::handle owner) + : WrapperBase(ptr, owner), X(pos), X(forward), X(up) {} +#undef X + +MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera&& other) : MjvGLCameraWrapper() { + *this->ptr_ = other; +} + +MjvGLCameraWrapper::MjWrapper(const MjvGLCameraWrapper& other) + : MjvGLCameraWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVGEOM ================================================ +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvGeomWrapper::MjWrapper() + : WrapperBase(new raw::MjvGeom{}), + X(size), + X(pos), + mat([this]() { + static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0]) * 9); + return InitPyArray(std::array{3, 3}, ptr_->mat, owner_); + }()), + X(rgba) { + mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr); +} + +MjvGeomWrapper::MjWrapper(raw::MjvGeom* ptr, py::handle owner) + : WrapperBase(ptr, owner), + X(size), + X(pos), + mat([this]() { + static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0]) * 9); + return InitPyArray(std::array{3, 3}, ptr_->mat, owner_); + }()), + X(rgba) {} +#undef X + +MjvGeomWrapper::MjWrapper(const MjvGeomWrapper& other) : MjvGeomWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVLIGHT =============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvLightWrapper::MjWrapper() + : WrapperBase(new raw::MjvLight{}), + X(pos), + X(dir), + X(attenuation), + X(ambient), + X(diffuse), + X(specular) {} + +MjvLightWrapper::MjWrapper(raw::MjvLight* ptr, py::handle owner) + : WrapperBase(ptr, owner), + X(pos), + X(dir), + X(attenuation), + X(ambient), + X(diffuse), + X(specular) {} +#undef X + +MjvLightWrapper::MjWrapper(const MjvLightWrapper& other) : MjvLightWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVOPTION ============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvOptionWrapper::MjWrapper() + : WrapperBase([]() { + raw::MjvOption* const opt = new raw::MjvOption; + mjv_defaultOption(opt); + return opt; + }()), + X(geomgroup), + X(sitegroup), + X(jointgroup), + X(tendongroup), + X(actuatorgroup), + X(flexgroup), + X(skingroup), + X(flags) {} +#undef X + +MjvOptionWrapper::MjWrapper(const MjvOptionWrapper& other) + : MjvOptionWrapper() { + *this->ptr_ = *other.ptr_; +} + +// ==================== MJVSCENE =============================================== +static void MjvSceneCapsuleDestructor(PyObject* pyobj) { + py::gil_scoped_acquire gil; + auto* scn = static_cast(PyCapsule_GetPointer(pyobj, nullptr)); + if (scn) { + mjv_freeScene(scn); + delete scn; + } +} + +#define X(var) var(InitPyArray(ptr_->var, owner_)) +#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_)) +MjvSceneWrapper::MjWrapper() + : WrapperBase( + []() { + raw::MjvScene* const scn = new raw::MjvScene; + mjv_defaultScene(scn); + InterceptMjErrors(mjv_makeScene)(nullptr, scn, 0); + return scn; + }(), + &MjvSceneCapsuleDestructor), + nskinvert(0), + XN(geoms, 0), + XN(geomorder, 0), + XN(flexedgeadr, 0), + XN(flexedgenum, 0), + XN(flexvertadr, 0), + XN(flexvertnum, 0), + XN(flexfaceadr, 0), + XN(flexfacenum, 0), + XN(flexfaceused, 0), + XN(flexedge, 0), + XN(flexvert, 0), + XN(flexface, 0), + XN(flexnormal, 0), + XN(flextexcoord, 0), + XN(skinfacenum, 0), + XN(skinvertadr, 0), + XN(skinvertnum, 0), + XN(skinvert, 0), + XN(skinnormal, 0), + X(lights), + X(camera), + X(translate), + X(rotate), + X(flags), + X(framergb) {} + +#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_)) +MjvSceneWrapper::MjWrapper(const MjModelWrapper& model, int maxgeom) + : WrapperBase( + [maxgeom](const raw::MjModel* m) { + raw::MjvScene* const scn = new raw::MjvScene; + mjv_defaultScene(scn); + InterceptMjErrors(mjv_makeScene)(m, scn, maxgeom); + return scn; + }(model.get()), + &MjvSceneCapsuleDestructor), + nskinvert([](const raw::MjModel* m) { + int nskinvert = 0; + for (int i = 0; i < m->nskin; ++i) { + nskinvert += m->skin_vertnum[i]; + } + return nskinvert; + }(model.get())), + nflexface([](const raw::MjModel* m) { + int nflexface = 0; + int flexfacenum = 0; + for (int f = 0; f < m->nflex; f++) { + if (m->flex_dim[f] == 0) { + // 1D : 0 + flexfacenum = 0; + } else if (m->flex_dim[f] == 2) { + // 2D: 2*fragments + 2*elements + flexfacenum = 2 * m->flex_shellnum[f] + 2 * m->flex_elemnum[f]; + } else { + // 3D: max(fragments, 4*maxlayer) + // find number of elements in biggest layer + int maxlayer = 0, layer = 0, nlayer = 1; + while (nlayer) { + nlayer = 0; + for (int e = 0; e < m->flex_elemnum[f]; e++) { + if (m->flex_elemlayer[m->flex_elemadr[f] + e] == layer) { + nlayer++; + } + } + maxlayer = mjMAX(maxlayer, nlayer); + layer++; + } + flexfacenum = mjMAX(m->flex_shellnum[f], 4 * maxlayer); + } + + // accumulate over flexes + nflexface += flexfacenum; + } + return nflexface; + }(model.get())), + nflexedge(model.get()->nflexedge), + nflexvert(model.get()->nflexvert), + XN(geoms, ptr_->maxgeom), + XN(geomorder, ptr_->maxgeom), + XN(flexedgeadr, ptr_->nflex), + XN(flexedgenum, ptr_->nflex), + XN(flexvertadr, ptr_->nflex), + XN(flexvertnum, ptr_->nflex), + XN(flexfaceadr, ptr_->nflex), + XN(flexfacenum, ptr_->nflex), + XN(flexfaceused, ptr_->nflex), + XN(flexedge, 2 * nflexedge), + XN(flexvert, 3 * nflexvert), + XN(flexface, 9 * nflexface), + XN(flexnormal, 9 * nflexface), + XN(flextexcoord, 6 * nflexface), + XN(skinfacenum, ptr_->nskin), + XN(skinvertadr, ptr_->nskin), + XN(skinvertnum, ptr_->nskin), + XN(skinvert, 3 * nskinvert), + XN(skinnormal, 3 * nskinvert), + X(lights), + X(camera), + X(translate), + X(rotate), + X(flags), + X(framergb) {} +#undef X +#undef XN + +template +static T* MallocAndCopy(const T* src, int count) { + if (src) { + T* out = static_cast(mju_malloc(count * sizeof(T))); + std::memcpy(out, src, count * sizeof(T)); + return out; + } else { + return nullptr; + } +} + +MjvSceneWrapper::MjWrapper(const MjvSceneWrapper& other) : MjvSceneWrapper() { + mjv_freeScene(ptr_); + *ptr_ = *other.ptr_; + +#define XN(var, n) \ + ptr_->var = MallocAndCopy(other.ptr_->var, n); \ + var = InitPyArray(std::array{n}, ptr_->var, owner_); + + XN(geoms, ptr_->ngeom); + XN(geomorder, ptr_->ngeom); + XN(flexedgeadr, ptr_->nflex); + XN(flexedgenum, ptr_->nflex); + XN(flexvertadr, ptr_->nflex); + XN(flexvertnum, ptr_->nflex); + XN(flexfaceadr, ptr_->nflex); + XN(flexfacenum, ptr_->nflex); + XN(flexfaceused, ptr_->nflex); + XN(flexedge, 2 * nflexedge); + XN(flexvert, 3 * nflexvert); + XN(flexface, 9 * nflexface); + XN(flexnormal, 9 * nflexface); + XN(flextexcoord, 6 * nflexface); + XN(skinfacenum, ptr_->nskin); + XN(skinvertadr, ptr_->nskin); + XN(skinvertnum, ptr_->nskin); + XN(skinvert, 3 * nskinvert); + XN(skinnormal, 3 * nskinvert); + +#undef XN +} + +// ==================== MJVFIGURE ============================================== +#define X(var) var(InitPyArray(ptr_->var, owner_)) +MjvFigureWrapper::MjWrapper() + : WrapperBase([]() { + raw::MjvFigure* const fig = new raw::MjvFigure; + mjv_defaultFigure(fig); + return fig; + }()), + X(flg_ticklabel), + X(gridsize), + X(gridrgb), + X(figurergba), + X(panergba), + X(legendrgba), + X(textrgb), + X(linergb), + X(range), + X(highlight), + X(linepnt), + X(linedata), + X(xaxispixel), + X(yaxispixel), + X(xaxisdata), + X(yaxisdata), +#undef X + + linename([](raw::MjvFigure* ptr, py::handle owner) { +// Use a macro to help us static_assert that the array extents here are kept +// in sync with mjVisualize.h. +#define MAKE_STR_ARRAY(N1, N2) \ + static_assert( \ + std::is_same_v); \ + return py::array(py::dtype("|S" #N2), N1, ptr->linename, owner); + MAKE_STR_ARRAY(mjMAXLINE, 100); + +#undef MAKE_STR_ARRAY + }(ptr_, owner_)) { +} + +MjvFigureWrapper::MjWrapper(const MjvFigureWrapper& other) + : MjvFigureWrapper() { + *this->ptr_ = *other.ptr_; +} + +} // namespace mujoco::python::_impl diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 65852c87..c1de94ec 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -23,7 +23,7 @@ import queue import sys import threading import time -from typing import Callable, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union import weakref import glfw @@ -115,9 +115,22 @@ class Handle: return sim.viewport return None - def set_figures(self, viewports_figures): + def set_figures( + self, viewports_figures: Union[Tuple[mujoco.MjrRect, mujoco.MjvFigure], + List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]] + ): + """Overlay figures on the viewer. + + Args: + viewports_figures: Single tuple or list of tuples of (viewport, figure) + viewport: Rectangle defining position and size of the figure + figure: MjvFigure object containing the figure data to display + """ sim = self._sim() if sim is not None: + # Convert single tuple to list if needed + if isinstance(viewports_figures, tuple): + viewports_figures = [viewports_figures] sim.set_figures(viewports_figures) def clear_figures(self): @@ -125,6 +138,73 @@ class Handle: if sim is not None: sim.clear_figures() + def set_texts(self, texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], + List[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]]]): + """Overlay text on the viewer. + + Args: + texts: Single tuple or list of tuples of (font, gridpos, text1, text2) + font: Font style from mujoco.mjtFontScale + gridpos: Position of text box from mujoco.mjtGridPos + text1: Left text column, defaults to empty string if None + text2: Right text column, defaults to empty string if None + """ + sim = self._sim() + if sim is not None: + # Convert single tuple to list if needed + if isinstance(texts, tuple): + texts = [texts] + + # Convert None values to empty strings + default_font = mujoco.mjtFontScale.mjFONTSCALE_150 + default_gridpos = mujoco.mjtGridPos.mjGRID_TOPLEFT + processed_texts = [( + default_font if font is None else font, + default_gridpos if gridpos is None else gridpos, + "" if text1 is None else text1, + "" if text2 is None else text2) + for font, gridpos, text1, text2 in texts] + + sim.set_texts(processed_texts) + + def clear_texts(self): + sim = self._sim() + if sim is not None: + sim.clear_texts() + + def set_images( + self, viewports_images: Union[Tuple[mujoco.MjrRect, np.ndarray], + List[Tuple[mujoco.MjrRect, np.ndarray]]] + ): + """Overlay images on the viewer. + + Args: + viewports_images: Single tuple or list of tuples of (viewport, image) + viewport: Rectangle defining position and size of the image + image: RGB image with shape (height, width, 3) + """ + sim = self._sim() + if sim is not None: + # Convert single tuple to list if needed + if isinstance(viewports_images, tuple): + viewports_images = [viewports_images] + + processed_images = [] + for viewport, image in viewports_images: + targ_shape = (viewport.height, viewport.width) + # Check if image is already the correct shape + if image.shape[:2] != targ_shape: + raise ValueError(f"Image shape {image.shape[:2]} does not match target shape {targ_shape}") + flipped = np.flip(image, axis=0) + contiguous = np.ascontiguousarray(flipped) + processed_images.append((viewport, contiguous)) + sim.set_images(processed_images) + + def clear_images(self): + sim = self._sim() + if sim is not None: + sim.clear_images() + def close(self): sim = self._sim() if sim is not None: diff --git a/python/pyproject.toml b/python/pyproject.toml index 98fc970e..c178ad79 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "mujoco" -version = "3.3.1" +version = "3.3.3" authors = [ {name = "Google DeepMind", email = "mujoco@deepmind.com"}, ] @@ -35,9 +35,9 @@ dynamic = ["readme", "scripts"] [project.urls] Homepage = "https://github.com/google-deepmind/mujoco" -Documentation = "https://mujoco.readthedocs.io/en/3.3.1" +Documentation = "https://mujoco.readthedocs.io/en/3.3.3" Repository = "https://github.com/google-deepmind/mujoco" -Changelog = "https://mujoco.readthedocs.io/en/3.3.1/changelog.html" +Changelog = "https://mujoco.readthedocs.io/en/3.3.3/changelog.html" [tool.setuptools] include-package-data = false diff --git a/python/rollout.ipynb b/python/rollout.ipynb index 61b21ffc..5df4ba50 100644 --- a/python/rollout.ipynb +++ b/python/rollout.ipynb @@ -5,7 +5,6 @@ "id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34", "metadata": { "editable": true, - "id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34", "tags": [] }, "source": [ @@ -42,7 +41,6 @@ "id": "5d8a6604-0948-4a42-a48d-249c7f0c462b", "metadata": { "editable": true, - "id": "5d8a6604-0948-4a42-a48d-249c7f0c462b", "tags": [] }, "source": [ @@ -55,7 +53,6 @@ "id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f", "metadata": { "editable": true, - "id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f", "tags": [ "hide-input" ] @@ -157,9 +154,7 @@ { "cell_type": "markdown", "id": "fc69d0f4", - "metadata": { - "id": "fc69d0f4" - }, + "metadata": {}, "source": [ "# Helper Functions" ] @@ -170,7 +165,6 @@ "id": "082482c7", "metadata": { "editable": true, - "id": "082482c7", "tags": [ "hide-input" ] @@ -273,9 +267,7 @@ { "cell_type": "markdown", "id": "c0570c2c", - "metadata": { - "id": "c0570c2c" - }, + "metadata": {}, "source": [ "# Using `rollout`\n", "\n", @@ -288,11 +280,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "849b93e5", - "metadata": { - "id": "849b93e5" - }, + "metadata": {}, "outputs": [], "source": [ "#@title Benchmarked models\n", @@ -378,9 +368,7 @@ { "cell_type": "markdown", "id": "55d171f7-541b-4441-aa18-da86d6716410", - "metadata": { - "id": "55d171f7-541b-4441-aa18-da86d6716410" - }, + "metadata": {}, "source": [ "## Usage\n", "\n", @@ -392,11 +380,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5", - "metadata": { - "id": "9cd2f94a-11df-4247-986c-5a56af69a1f5" - }, + "metadata": {}, "outputs": [], "source": [ "print(rollout.rollout.__doc__)" @@ -406,7 +392,6 @@ "cell_type": "markdown", "id": "b6f7a094-8352-4b07-99ee-5278e3036cd5", "metadata": { - "id": "b6f7a094-8352-4b07-99ee-5278e3036cd5", "tags": [] }, "source": [ @@ -418,11 +403,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe", - "metadata": { - "id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe" - }, + "metadata": {}, "outputs": [], "source": [ "nbatch = 100 # Simulate this many tops\n", @@ -454,20 +437,16 @@ { "cell_type": "markdown", "id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f", - "metadata": { - "id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f" - }, + "metadata": {}, "source": [ "Our model has an angular velocity sensor the middle of the top. Let's plot the response using the `sensordata` array that rollout returns." ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "957b8566-da31-410b-b385-e78241c5247a", - "metadata": { - "id": "957b8566-da31-410b-b385-e78241c5247a" - }, + "metadata": {}, "outputs": [], "source": [ "plt.figure(figsize=(12, 8))\n", @@ -483,9 +462,7 @@ { "cell_type": "markdown", "id": "58044bc1-f98c-4bbf-a703-40ba075552a0", - "metadata": { - "id": "58044bc1-f98c-4bbf-a703-40ba075552a0" - }, + "metadata": {}, "source": [ "### Example: different models\n", "100 gray tops is kind of boring. It would be better if they were colorful and different sizes!\n", @@ -497,11 +474,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2", - "metadata": { - "id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2" - }, + "metadata": {}, "outputs": [], "source": [ "# Make 100 tippe tops with different colors and sizes\n", @@ -569,20 +544,16 @@ { "cell_type": "markdown", "id": "cf485c08-72be-4169-89b6-9d93df8ebbe3", - "metadata": { - "id": "cf485c08-72be-4169-89b6-9d93df8ebbe3" - }, + "metadata": {}, "source": [ "Because the models are now different, the measurements of the gyro sensor are not consistent even though the initial state for each rollout was the same." ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03", - "metadata": { - "id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03" - }, + "metadata": {}, "outputs": [], "source": [ "plt.figure(figsize=(12, 8))\n", @@ -598,9 +569,7 @@ { "cell_type": "markdown", "id": "3841a669-6cd1-427e-a629-20a10a6e3a34", - "metadata": { - "id": "3841a669-6cd1-427e-a629-20a10a6e3a34" - }, + "metadata": {}, "source": [ "### Example: control inputs\n", "Open loop controls can be passed to `rollout` via the `control` argument. If passed, `nstep` no longer needs to be specified as it can be inferred from the size of `control`.\n", @@ -610,11 +579,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "2a184873-8d24-45da-b444-8d21f5dcd733", - "metadata": { - "id": "2a184873-8d24-45da-b444-8d21f5dcd733" - }, + "metadata": {}, "outputs": [], "source": [ "# Episode parameters.\n", @@ -665,9 +632,7 @@ { "cell_type": "markdown", "id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7", - "metadata": { - "id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7" - }, + "metadata": {}, "source": [ "`rollout`'s `control_spec` argument can be used to indicate `control` contains values for actuators, generalized forces, cartesian forces, mocap poses, and/or the activation/deactivation of equality constraints. Internally, this is managed through [mj_setState](https://mujoco.readthedocs.io/en/stable/APIreference/APIfunctions.html#mj-setstate) and `control_spec` corresponds to `mj_setState`'s `spec` argument.\n", "\n", @@ -676,11 +641,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5", - "metadata": { - "id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5" - }, + "metadata": {}, "outputs": [], "source": [ "xfrc_size = mujoco.mj_stateSize(humanoid_model, mujoco.mjtState.mjSTATE_XFRC_APPLIED)\n", @@ -713,9 +676,7 @@ { "cell_type": "markdown", "id": "0961c3ec-a691-4875-9a55-227a3d29c472", - "metadata": { - "id": "0961c3ec-a691-4875-9a55-227a3d29c472" - }, + "metadata": {}, "source": [ "# Advanced usage" ] @@ -723,9 +684,7 @@ { "cell_type": "markdown", "id": "VfYIyXWcLKfg", - "metadata": { - "id": "VfYIyXWcLKfg" - }, + "metadata": {}, "source": [ "## skip_checks\n", "\n", @@ -747,11 +706,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95", - "metadata": { - "id": "d02cc8e8-63cd-4852-ab3c-364a18025a95" - }, + "metadata": {}, "outputs": [], "source": [ "nbatch = 1000\n", @@ -795,9 +752,7 @@ { "cell_type": "markdown", "id": "92627030-4726-4689-be8b-f1ba75905104", - "metadata": { - "id": "92627030-4726-4689-be8b-f1ba75905104" - }, + "metadata": {}, "source": [ "As expected, as `nstep` increases, the benefits of using skip checks fades quickly. However, at low nstep and high batch sizes, it can make a significant difference.\n", "\n", @@ -807,9 +762,7 @@ { "cell_type": "markdown", "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c", - "metadata": { - "id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c" - }, + "metadata": {}, "source": [ "## Reusing threadpools (`Rollout` class)\n", "\n", @@ -820,11 +773,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49", - "metadata": { - "id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49" - }, + "metadata": {}, "outputs": [], "source": [ "nbatch = 100\n", @@ -862,9 +813,7 @@ { "cell_type": "markdown", "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671", - "metadata": { - "id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671" - }, + "metadata": {}, "source": [ "## Reusing threadpools (`rollout` method)\n", "\n", @@ -875,11 +824,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "b6aa6801", - "metadata": { - "id": "b6aa6801" - }, + "metadata": {}, "outputs": [], "source": [ "nbatch = 1000\n", @@ -899,20 +846,16 @@ { "cell_type": "markdown", "id": "144378d3", - "metadata": { - "id": "144378d3" - }, + "metadata": {}, "source": [ "Second, if `rollout` reuses the same threadpool between calls, it is no longer safe to call `rollout` from multiple threads. For example the following is not allowed (the offending lines are commented out to avoid crashing the interpreter):" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "7f46a6d8", - "metadata": { - "id": "7f46a6d8" - }, + "metadata": {}, "outputs": [], "source": [ "thread1 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n", @@ -928,9 +871,7 @@ { "cell_type": "markdown", "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a", - "metadata": { - "id": "78c1f864-5238-4e27-a7ec-d03c45484d9a" - }, + "metadata": {}, "source": [ "## chunk_size\n", "\n", @@ -941,11 +882,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "a1be8f93", - "metadata": { - "id": "a1be8f93" - }, + "metadata": {}, "outputs": [], "source": [ "nbatch = 100\n", @@ -989,9 +928,7 @@ { "cell_type": "markdown", "id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7", - "metadata": { - "id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7" - }, + "metadata": {}, "source": [ "## Warmstarting\n", "\n", @@ -1004,11 +941,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb", - "metadata": { - "id": "d4d9f660-f83c-432e-a579-124a7ecab4fb" - }, + "metadata": {}, "outputs": [], "source": [ "top_model_cg = copy.copy(top_model)\n", @@ -1069,9 +1004,7 @@ { "cell_type": "markdown", "id": "7c2cf4fa", - "metadata": { - "id": "7c2cf4fa" - }, + "metadata": {}, "source": [ "As expected, the middle animation (with warmstarting) matches the continuous rollout on the left. However, the model that did not use warmstarting diverged." ] @@ -1079,9 +1012,7 @@ { "cell_type": "markdown", "id": "7944637f", - "metadata": { - "id": "7944637f" - }, + "metadata": {}, "source": [ "# Benchmarks\n", "\n", @@ -1098,11 +1029,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "cb6355dd", - "metadata": { - "id": "cb6355dd" - }, + "metadata": {}, "outputs": [], "source": [ "def python_rollout(model, data, nbatch, nstep):\n", @@ -1114,9 +1043,7 @@ { "cell_type": "markdown", "id": "6fe4a78b", - "metadata": { - "id": "6fe4a78b" - }, + "metadata": {}, "source": [ "To run nbatch rollouts with `rollout`, we need to make an array of nbatch initial states to start the rollouts from.\n", "\n", @@ -1128,11 +1055,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "74f143e2", - "metadata": { - "id": "74f143e2" - }, + "metadata": {}, "outputs": [], "source": [ "def nthread_rollout(model, data, nbatch, nstep, nthread, rollout_):\n", @@ -1146,9 +1071,7 @@ { "cell_type": "markdown", "id": "b75dc44c", - "metadata": { - "id": "b75dc44c" - }, + "metadata": {}, "source": [ "Next, we benchmark the Python loop and `rollout` in both single threaded and multithreaded modes. The three benchmarks take about 2.5 minutes in total to run in total on an AMD 5800X3D." ] @@ -1158,8 +1081,7 @@ "execution_count": 0, "id": "0301e3ee", "metadata": { - "cellView": "form", - "id": "0301e3ee" + "cellView": "form" }, "outputs": [], "source": [ @@ -1282,20 +1204,16 @@ { "cell_type": "markdown", "id": "08fb0c12", - "metadata": { - "id": "08fb0c12" - }, + "metadata": {}, "source": [ "### Tippe Top Benchmark" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "f7e54830", - "metadata": { - "id": "f7e54830" - }, + "metadata": {}, "outputs": [], "source": [ "nominal_nbatch = 256 # Batch size to use when testing different nstep\n", @@ -1314,20 +1232,16 @@ { "cell_type": "markdown", "id": "edefb26e", - "metadata": { - "id": "edefb26e" - }, + "metadata": {}, "source": [ "### Humanoid Benchmark" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "c9e58c6c", - "metadata": { - "id": "c9e58c6c" - }, + "metadata": {}, "outputs": [], "source": [ "nominal_nbatch = 256 # Batch size to use when testing different nstep\n", @@ -1346,20 +1260,16 @@ { "cell_type": "markdown", "id": "468903bb", - "metadata": { - "id": "468903bb" - }, + "metadata": {}, "source": [ "### Humanoid100 Benchmark" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "83d775d4", - "metadata": { - "id": "83d775d4" - }, + "metadata": {}, "outputs": [], "source": [ "nominal_nbatch = 128 # Batch size to use when testing different nstep\n", @@ -1383,9 +1293,7 @@ { "cell_type": "markdown", "id": "d1133084", - "metadata": { - "id": "d1133084" - }, + "metadata": {}, "source": [ "# MJX versus `rollout`" ] @@ -1393,9 +1301,7 @@ { "cell_type": "markdown", "id": "c1638f2d", - "metadata": { - "id": "c1638f2d" - }, + "metadata": {}, "source": [ "Next we will benchmark `rollout` and MJX using the tippe top and humanoid models (humanoid100 is not supported by MJX).\n", "\n", @@ -1409,8 +1315,7 @@ "execution_count": 0, "id": "7c86d157", "metadata": { - "cellView": "form", - "id": "7c86d157" + "cellView": "form" }, "outputs": [], "source": [ @@ -1531,20 +1436,16 @@ { "cell_type": "markdown", "id": "a2dafd2e", - "metadata": { - "id": "a2dafd2e" - }, + "metadata": {}, "source": [ "### MJX Tippe Top Benchmark" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "98c580b0", - "metadata": { - "id": "98c580b0" - }, + "metadata": {}, "outputs": [], "source": [ "nominal_nbatch = 16384 # Batch size to use when testing different nstep\n", @@ -1560,20 +1461,16 @@ { "cell_type": "markdown", "id": "205da5da", - "metadata": { - "id": "205da5da" - }, + "metadata": {}, "source": [ "### MJX Humanoid Benchmark" ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "53166ae1", - "metadata": { - "id": "53166ae1" - }, + "metadata": {}, "outputs": [], "source": [ "nominal_nbatch = 4096 # Batch size to use when testing different nstep\n", @@ -1589,9 +1486,7 @@ { "cell_type": "markdown", "id": "fb2caa72", - "metadata": { - "id": "fb2caa72" - }, + "metadata": {}, "source": [ "### MJX Multiple Humanoids in one model\n", "\n", @@ -1599,16 +1494,14 @@ "\n", "Here we will produce a similar plot to compare MJX and with `rollout`. On a 5800X3D and 4090 the benchmark takes about 16.5 minutes to run.\n", "\n", - "**Note:** These results are not directly comparable since with the plot in the documentation because, in particular, the batch size was redued from 8192 to 4096 in order to fit the batch on a 4090." + "**Note:** These results are not directly comparable since with the plot in the documentation because, in particular, the batch size was reduced from 8192 to 4096 in order to fit the batch on a 4090." ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "3d6be608", - "metadata": { - "id": "3d6be608" - }, + "metadata": {}, "outputs": [], "source": [ "max_humanoids = 10\n", @@ -1656,11 +1549,9 @@ }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, "id": "b6c5fc2e", - "metadata": { - "id": "b6c5fc2e" - }, + "metadata": {}, "outputs": [], "source": [ "#@title Plot MJX nhumanoid benchmark\n", diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb index 0c7be7be..c0933859 100644 --- a/python/tutorial.ipynb +++ b/python/tutorial.ipynb @@ -1360,16 +1360,18 @@ "lines = ax[0,0].plot(sim_time, force)\n", "ax[0,0].set_title('contact force')\n", "ax[0,0].set_ylabel('Newton')\n", - "ax[0,0].legend(iter(lines), ('normal z', 'friction x', 'friction y'));\n", + "ax[0,0].legend(lines, ('normal z', 'friction x', 'friction y'));\n", "\n", "ax[1,0].plot(sim_time, acceleration)\n", "ax[1,0].set_title('acceleration')\n", "ax[1,0].set_ylabel('(meter,radian)/s/s')\n", + "ax[1,0].legend(['ax', 'ay', 'az', 'αx', 'αy', 'αz'])\n", "\n", "ax[2,0].plot(sim_time, velocity)\n", "ax[2,0].set_title('velocity')\n", "ax[2,0].set_ylabel('(meter,radian)/s')\n", "ax[2,0].set_xlabel('second')\n", + "ax[2,0].legend(['vx', 'vy', 'vz', 'ωx', 'ωy', 'ωz'])\n", "\n", "ax[0,1].plot(sim_time, ncon)\n", "ax[0,1].set_title('number of contacts')\n", diff --git a/sample/CMakeLists.txt b/sample/CMakeLists.txt index 8273ed11..f9457f2c 100644 --- a/sample/CMakeLists.txt +++ b/sample/CMakeLists.txt @@ -24,7 +24,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON) project( mujoco_samples - VERSION 3.3.1 + VERSION 3.3.3 DESCRIPTION "MuJoCo samples binaries" HOMEPAGE_URL "https://mujoco.org" ) diff --git a/sample/compile.cc b/sample/compile.cc index a5a572d8..14382393 100644 --- a/sample/compile.cc +++ b/sample/compile.cc @@ -18,7 +18,6 @@ #include #include #include -#include #include @@ -30,13 +29,12 @@ static constexpr char helpstring[] = " if infile is mjcf, compilation will be timed twice to measure the impact of caching\n\n" " Example: compile model.xml [model.mjb]\n"; -// timer +// timer (seconds) mjtNum gettm(void) { - using std::chrono::steady_clock; - using Microseconds = std::chrono::duration; - static steady_clock::time_point tm_start = steady_clock::now(); - auto elapsed = Microseconds(steady_clock::now() - tm_start); - return elapsed.count(); + using Clock = std::chrono::steady_clock; + using Seconds = std::chrono::duration; + static const Clock::time_point tm_start = Clock::now(); + return Seconds(Clock::now() - tm_start).count(); } // deallocate and print message @@ -140,12 +138,12 @@ int main(int argc, char** argv) { if (type1==typeXML) { double starttime = gettm(); m = mj_loadXML(argv[1], 0, error, 1000); - first = 1e-6 * (gettm() - starttime); + first = gettm() - starttime; if (m) { mj_deleteModel(m); starttime = gettm(); m = mj_loadXML(argv[1], 0, error, 1000); - second = 1e-6 * (gettm() - starttime); + second = gettm() - starttime; } } else { m = mj_loadModel(argv[1], 0); diff --git a/sample/testspeed.cc b/sample/testspeed.cc index 3791bf5b..8074f68e 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -36,13 +36,12 @@ int constraints[maxthread]; mjtNum iterations[maxthread]; mjtNum simtime[maxthread]; -// timer +// timer (microseconds) mjtNum gettm(void) { - using std::chrono::steady_clock; - using Microseconds = std::chrono::duration; - static steady_clock::time_point tm_start = steady_clock::now(); - auto elapsed = Microseconds(steady_clock::now() - tm_start); - return elapsed.count(); + using Clock = std::chrono::steady_clock; + using Microseconds = std::chrono::duration; + static const Clock::time_point tm_start = Clock::now(); + return Microseconds(Clock::now() - tm_start).count(); } // deallocate and print message @@ -98,8 +97,8 @@ void simulate(int id, int nstep, mjtNum* ctrl) { // accumulate statistics contacts[id] += d[id]->ncon; constraints[id] += d[id]->nefc; - int nisland = d[id]->solver_nisland; - if (nisland == 1) { + int nisland = mjMAX(1, mjMIN(d[id]->nisland, mjNISLAND)); + if (nisland == 1 || nisland == 0) { iterations[id] += d[id]->solver_niter[0]; } else { mjtNum niter = 0; diff --git a/simulate/CMakeLists.txt b/simulate/CMakeLists.txt index 89316dfe..96330194 100644 --- a/simulate/CMakeLists.txt +++ b/simulate/CMakeLists.txt @@ -29,7 +29,7 @@ set(MUJOCO_DEP_VERSION_lodepng project( mujoco_simulate - VERSION 3.3.1 + VERSION 3.3.3 DESCRIPTION "MuJoCo simulate binaries" HOMEPAGE_URL "https://mujoco.org" ) @@ -114,6 +114,7 @@ target_sources( target_compile_options(platform_ui_adapter PRIVATE ${MUJOCO_SIMULATE_COMPILE_OPTIONS}) if(APPLE) target_sources(platform_ui_adapter PUBLIC glfw_corevideo.h PRIVATE glfw_corevideo.mm) + set_source_files_properties(glfw_corevideo.mm PROPERTIES COMPILE_FLAGS -Wno-deprecated-declarations) target_link_libraries(platform_ui_adapter PUBLIC "-framework CoreVideo") endif() target_include_directories( diff --git a/simulate/main.cc b/simulate/main.cc index 3ebe873c..8d70e6e8 100644 --- a/simulate/main.cc +++ b/simulate/main.cc @@ -39,7 +39,7 @@ extern "C" { #if defined(__APPLE__) #include #endif - #include + #include #include #endif } diff --git a/simulate/simulate.cc b/simulate/simulate.cc index bf4b4a13..b6a1b9e4 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -298,7 +298,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) { memset(sim->figcost.linepnt, 0, mjMAXLINE*sizeof(int)); // number of islands that have diagnostics - int nisland = mjMIN(d->solver_nisland, mjNISLAND); + int nisland = mjMAX(1, mjMIN(d->nisland, mjNISLAND)); // iterate over islands for (int k=0; k < nisland; k++) { @@ -328,7 +328,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) { sim->figconstraint.linedata[start + 4][2*i] = i; // y - int nefc = nisland == 1 ? d->nefc : d->island_efcnum[k]; + int nefc = nisland == 1 ? d->nefc : d->island_nefc[k]; sim->figconstraint.linedata[start + 0][2*i+1] = nefc; const mjSolverStat* stat = d->solver + k*mjNSOLVER + i; sim->figconstraint.linedata[start + 1][2*i+1] = stat->nactive; @@ -413,7 +413,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) { static_cast(d->nefc), static_cast(sqrt_nnz), static_cast(d->ncon), - static_cast(solver_niter) + static_cast(solver_niter) / nisland }; // update figsize @@ -464,7 +464,7 @@ void InitializeSensor(mj::Simulate* sim) { // title mju::strcpy_arr(figsensor.title, "Sensor data"); - // y-tick nubmer format + // y-tick number format mju::strcpy_arr(figsensor.yformat, "%.1f"); // grid size @@ -546,6 +546,16 @@ void ShowFigure(mj::Simulate* sim, mjrRect viewport, mjvFigure* fig){ mjr_figure(viewport, fig, &sim->platform_ui->mjr_context()); } +void ShowOverlayText(mj::Simulate* sim, mjrRect viewport, int font, int gridpos, + std::string text1, std::string text2) { + mjr_overlay(font, gridpos, viewport, text1.c_str(), text2.c_str(), + &sim->platform_ui->mjr_context()); +} + +void ShowImage(mj::Simulate* sim, mjrRect viewport, const unsigned char* image) { + mjr_drawPixels(image, nullptr, viewport, &sim->platform_ui->mjr_context()); +} + // load state from history buffer static void LoadScrubState(mj::Simulate* sim) { // get index into circular buffer @@ -572,7 +582,7 @@ void UpdateInfoText(mj::Simulate* sim, const mjModel* m, const mjData* d, char tmp[20]; // number of islands with statistics - int nisland = mjMIN(d->solver_nisland, mjNISLAND); + int nisland = mjMAX(1, mjMIN(d->nisland, mjNISLAND)); // compute solver error (maximum over islands) mjtNum solerr = 0; @@ -679,7 +689,7 @@ void UpdateWatch(mj::Simulate* sim, const mjModel* m, const mjData* d) { // make physics section of UI void MakePhysicsSection(mj::Simulate* sim) { - mjOption* opt = sim->is_passive_ ? &sim->scnstate_.model.opt : &sim->m_->opt; + mjOption* opt = sim->is_passive_ ? &sim->m_passive_->opt : &sim->m_->opt; mjuiDef defPhysics[] = { {mjITEM_SECTION, "Physics", mjPRESERVE, nullptr, "AP"}, {mjITEM_SELECT, "Integrator", 2, &(opt->integrator), "Euler\nRK4\nimplicit\nimplicitfast"}, @@ -863,13 +873,13 @@ void MakeRenderingSection(mj::Simulate* sim, const mjModel* m) { // make visualization section of UI void MakeVisualizationSection(mj::Simulate* sim, const mjModel* m) { - mjStatistic* stat = sim->is_passive_ ? &sim->scnstate_.model.stat : &sim->m_->stat; - mjVisual* vis = sim->is_passive_ ? &sim->scnstate_.model.vis : &sim->m_->vis; + mjStatistic* stat = sim->is_passive_ ? &sim->m_passive_->stat : &sim->m_->stat; + mjVisual* vis = sim->is_passive_ ? &sim->m_passive_->vis : &sim->m_->vis; mjuiDef defVisualization[] = { {mjITEM_SECTION, "Visualization", mjPRESERVE, nullptr, "AV"}, {mjITEM_SEPARATOR, "Headlight", 1}, - {mjITEM_RADIO, "Active", 5, &(vis->headlight.active), "Off\nOn"}, + {mjITEM_RADIO, "Active", 2, &(vis->headlight.active), "Off\nOn"}, {mjITEM_EDITFLOAT, "Ambient", 2, &(vis->headlight.ambient), "3"}, {mjITEM_EDITFLOAT, "Diffuse", 2, &(vis->headlight.diffuse), "3"}, {mjITEM_EDITFLOAT, "Specular", 2, &(vis->headlight.specular), "3"}, @@ -882,7 +892,7 @@ void MakeVisualizationSection(mj::Simulate* sim, const mjModel* m) { {mjITEM_BUTTON, "Align", 2, nullptr, "CA"}, {mjITEM_SEPARATOR, "Global", 1}, {mjITEM_EDITNUM, "Extent", 2, &(stat->extent), "1"}, - {mjITEM_RADIO, "Inertia", 5, &(vis->global.ellipsoidinertia), "Box\nEllipsoid"}, + {mjITEM_RADIO, "Inertia", 2, &(vis->global.ellipsoidinertia), "Box\nEllipsoid"}, {mjITEM_RADIO, "BVH active", 5, &(vis->global.bvactive), "False\nTrue"}, {mjITEM_SEPARATOR, "Map", 1}, {mjITEM_EDITFLOAT, "Stiffness", 2, &(vis->map.stiffness), "1"}, @@ -1141,17 +1151,74 @@ void AlignAndScaleView(mj::Simulate* sim, const mjModel* m) { } -// copy qpos to clipboard as key -void CopyPose(mj::Simulate* sim, const mjModel* m, const mjData* d) { - char clipboard[5000] = ""); + + // qvel + mju::strcat_arr(clipboard, "\"\n qvel=\""); + for (int i = 0; i < m->nv; i++) { + mju::sprintf_arr(buf, format, d->qvel[i]); + if (i < m->nv-1) mju::strcat_arr(buf, " "); + mju::strcat_arr(clipboard, buf); + } + + // act + if (m->na > 0) { + mju::strcat_arr(clipboard, "\"\n act=\""); + for (int i = 0; i < m->na; i++) { + mju::sprintf_arr(buf, format, d->act[i]); + if (i < m->na-1) mju::strcat_arr(buf, " "); + mju::strcat_arr(clipboard, buf); + } + } + + // ctrl + if (m->nu > 0) { + mju::strcat_arr(clipboard, "\"\n ctrl=\""); + for (int i = 0; i < m->nu; i++) { + mju::sprintf_arr(buf, format, d->ctrl[i]); + if (i < m->nu-1) mju::strcat_arr(buf, " "); + mju::strcat_arr(clipboard, buf); + } + } + + if (m->nmocap > 0) { + // mocap_pos + mju::strcat_arr(clipboard, "\"\n mpos=\""); + for (int i = 0; i < 3*m->nmocap; i++) { + mju::sprintf_arr(buf, format, d->mocap_pos[i]); + if (i < 3*m->nmocap-1) mju::strcat_arr(buf, " "); + mju::strcat_arr(clipboard, buf); + } + + // mocap_quat + mju::strcat_arr(clipboard, "\"\n mquat=\""); + for (int i = 0; i < 4*m->nmocap; i++) { + mju::sprintf_arr(buf, format, d->mocap_quat[i]); + if (i < 4*m->nmocap-1) mju::strcat_arr(buf, " "); + mju::strcat_arr(clipboard, buf); + } + } + + mju::strcat_arr(clipboard, "\"\n/>"); // copy to clipboard sim->platform_ui->SetClipboardString(clipboard); @@ -1412,8 +1479,9 @@ void UiEvent(mjuiState* state) { sim->pending_.align = true; break; - case 4: // Copy pose - sim->pending_.copy_pose = true; + case 4: // Copy key + sim->pending_.copy_key = true; + sim->pending_.copy_key_full_precision = sim->platform_ui->IsShiftKeyPressed(); break; case 5: // Adjust key @@ -1435,7 +1503,7 @@ void UiEvent(mjuiState* state) { // physics section else if (it && it->sectionid==SECT_PHYSICS && sim->m_) { - mjOption* opt = sim->is_passive_ ? &sim->scnstate_.model.opt : &sim->m_->opt; + mjOption* opt = sim->is_passive_ ? &sim->m_passive_->opt : &sim->m_->opt; // update disable flags in mjOption opt->disableflags = 0; @@ -1697,15 +1765,14 @@ void UiEvent(mjuiState* state) { return; } + // local pointers used below + mjModel* model = sim->is_passive_ ? sim->m_passive_ : sim->m_; + mjData* data = sim->is_passive_ ? sim->d_passive_ : sim->d_; + // 3D scroll - if (state->type==mjEVENT_SCROLL && state->mouserect==3) { + if (state->type==mjEVENT_SCROLL && state->mouserect==3 && model) { // emulate vertical mouse motion = 2% of window height - if (sim->m_ && !sim->is_passive_) { - mjv_moveCamera(sim->m_, mjMOUSE_ZOOM, 0, -zoom_increment*state->sy, &sim->scn, &sim->cam); - } else { - mjv_moveCameraFromState( - &sim->scnstate_, mjMOUSE_ZOOM, 0, -zoom_increment*state->sy, &sim->scn, &sim->cam); - } + mjv_moveCamera(model, mjMOUSE_ZOOM, 0, -zoom_increment*state->sy, &sim->scn, &sim->cam); return; } @@ -1761,25 +1828,11 @@ void UiEvent(mjuiState* state) { // move perturb or camera mjrRect r = state->rect[3]; if (sim->pert.active) { - if (!sim->is_passive_) { - mjv_movePerturb( - sim->m_, sim->d_, action, state->dx / r.height, -state->dy / r.height, - &sim->scn, &sim->pert); - } else { - mjv_movePerturbFromState( - &sim->scnstate_, action, state->dx / r.height, -state->dy / r.height, - &sim->scn, &sim->pert); - } + mjv_movePerturb(model, data, action, state->dx / r.height, -state->dy / r.height, + &sim->scn, &sim->pert); } else { - if (!sim->is_passive_) { - mjv_moveCamera( - sim->m_, action, state->dx / r.height, -state->dy / r.height, - &sim->scn, &sim->cam); - } else { - mjv_moveCameraFromState( - &sim->scnstate_, action, state->dx / r.height, -state->dy / r.height, - &sim->scn, &sim->cam); - } + mjv_moveCamera(model, action, state->dx / r.height, -state->dy / r.height, + &sim->scn, &sim->cam); } return; } @@ -1813,10 +1866,11 @@ Simulate::Simulate(std::unique_ptr platform_ui, platform_ui(std::move(platform_ui)), uistate(this->platform_ui->state()) { mjv_defaultScene(&scn); - mjv_defaultSceneState(&scnstate_); } -// synchronize model and data + +//------------------------- Synchronize render and physics threads --------------------------------- + // operations which require holding the mutex, prevents racing with physics thread void Simulate::Sync() { MutexLock lock(this->mtx); @@ -1877,48 +1931,30 @@ void Simulate::Sync() { } } + // in passive mode, synchronize user's mjModel with changes made via the UI if (is_passive_) { - // synchronize m_->opt with changes made via the UI -#define X(name) \ - if (IsDifferent(scnstate_.model.opt.name, mjopt_prev_.name)) { \ - pending_.ui_update_physics = true; \ - Copy(m_->opt.name, scnstate_.model.opt.name); \ - } + // synchronize mjModel.opt + if (std::memcmp(&m_passive_->opt, &mjopt_prev_, sizeof(mjOption))) { + pending_.ui_update_physics = true; + m_->opt = m_passive_->opt; + } - X(timestep); - X(apirate); - X(impratio); - X(tolerance); - X(noslip_tolerance); - X(ccd_tolerance); - X(gravity); - X(wind); - X(magnetic); - X(density); - X(viscosity); - X(o_margin); - X(o_solref); - X(o_solimp); - X(o_friction); - X(integrator); - X(cone); - X(jacobian); - X(solver); - X(iterations); - X(noslip_iterations); - X(ccd_iterations); - X(disableflags); - X(enableflags); - X(disableactuator); - X(sdf_initpoints); - X(sdf_iterations); + // synchronize mjModel.vis + if (std::memcmp(&m_passive_->vis, &mjvis_prev_, sizeof(mjVisual))) { + pending_.ui_update_visualization = true; + m_->vis = m_passive_->vis; + } - #undef X + // synchronize mjModel.stat + if (std::memcmp(&m_passive_->stat, &mjstat_prev_, sizeof(mjStatistic))) { + pending_.ui_update_visualization = true; + m_->stat = m_passive_->stat; + } // synchronize number of mjWARN_VGEOMFULL warnings - if (scnstate_.data.warning[mjWARN_VGEOMFULL].number > warn_vgeomfull_prev_) { + if (d_passive_->warning[mjWARN_VGEOMFULL].number > warn_vgeomfull_prev_) { d_->warning[mjWARN_VGEOMFULL].number += - scnstate_.data.warning[mjWARN_VGEOMFULL].number - warn_vgeomfull_prev_; + d_passive_->warning[mjWARN_VGEOMFULL].number - warn_vgeomfull_prev_; } } @@ -1967,9 +2003,10 @@ void Simulate::Sync() { pending_.align = false; } - if (pending_.copy_pose) { - CopyPose(this, m_, d_); - pending_.copy_pose = false; + if (pending_.copy_key) { + CopyKey(this, m_, d_, pending_.copy_key_full_precision); + pending_.copy_key = false; + pending_.copy_key_full_precision = false; } if (pending_.load_from_history) { @@ -2066,25 +2103,19 @@ void Simulate::Sync() { pending_.select = false; } - // update scene + // update scene or sync data from user in passive mode if (!is_passive_) { mjv_updateScene(m_, d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn); } else { - mjv_updateSceneState(m_, d_, &this->opt, &scnstate_); + mjv_copyModel(m_passive_, m_); + mjv_copyData(d_passive_, m_passive_, d_); - // append geoms from user_scn to scnstate_ scratch space + // append geoms from user_scn to scratch space if (user_scn) { - int ngeom = user_scn->ngeom; - int maxgeom = scnstate_.scratch.maxgeom - scnstate_.scratch.ngeom; - if (ngeom > maxgeom) { - mj_warning(d_, mjWARN_VGEOMFULL, scnstate_.scratch.maxgeom); - ngeom = maxgeom; - } - if (ngeom > 0) { - std::memcpy(scnstate_.scratch.geoms + scnstate_.scratch.ngeom, - user_scn->geoms, - sizeof(mjvGeom) * ngeom); - scnstate_.scratch.ngeom += ngeom; + user_scn_geoms_.clear(); + user_scn_geoms_.reserve(user_scn->ngeom); + for (int i = 0; i < user_scn->ngeom; ++i) { + user_scn_geoms_.push_back(user_scn->geoms[i]); } } @@ -2100,8 +2131,10 @@ void Simulate::Sync() { Copy(user_scn_flags_prev_, user_scn->flags); } - mjopt_prev_ = scnstate_.model.opt; - warn_vgeomfull_prev_ = scnstate_.data.warning[mjWARN_VGEOMFULL].number; + mjopt_prev_ = m_passive_->opt; + mjvis_prev_ = m_passive_->vis; + mjstat_prev_ = m_passive_->stat; + warn_vgeomfull_prev_ = d_passive_->warning[mjWARN_VGEOMFULL].number; } // update settings @@ -2260,15 +2293,8 @@ void Simulate::LoadOnRenderThread() { } } - // re-create scene and context + // re-create scene mjv_makeScene(this->m_, &this->scn, kMaxGeom); - if (this->is_passive_) { - mjopt_prev_ = m_->opt; - opt_prev_ = opt; - cam_prev_ = cam; - warn_vgeomfull_prev_ = d_->warning[mjWARN_VGEOMFULL].number; - mjv_makeSceneState(this->m_, this->d_, &this->scnstate_, kMaxGeom); - } this->platform_ui->RefreshMjrContext(this->m_, 50*(this->font+1)); UiModify(&this->ui0, &this->uistate, &this->platform_ui->mjr_context()); @@ -2297,12 +2323,18 @@ void Simulate::LoadOnRenderThread() { mju::strcpy_arr(this->previous_filename, this->filename); } - // update scene + // update scene in managed mode, in passive mode copy data from user (update in RenderLoop) if (!is_passive_) { - mjv_updateScene(this->m_, this->d_, - &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn); + mjv_updateScene(this->m_, this->d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn); } else { - mjv_updateSceneState(this->m_, this->d_, &this->opt, &this->scnstate_); + mjopt_prev_ = m_->opt; + opt_prev_ = opt; + cam_prev_ = cam; + warn_vgeomfull_prev_ = d_->warning[mjWARN_VGEOMFULL].number; + + // full copy on init + m_passive_ = mj_copyModel(nullptr, m_); + d_passive_ = mj_copyData(nullptr, m_passive_, d_); } // set window title to model name @@ -2423,6 +2455,13 @@ void Simulate::Render() { pending_.ui_update_physics = false; } + if (pending_.ui_update_visualization) { + if (this->ui0_enable && this->ui0.sect[SECT_VISUALIZATION].state) { + mjui0_update_section(this, SECT_VISUALIZATION); + } + pending_.ui_update_visualization = false; + } + if (is_passive_) { if (this->ui0_enable && this->ui0.sect[SECT_RENDERING].state && (cam_prev_.type != cam.type || @@ -2593,10 +2632,38 @@ void Simulate::Render() { } // user figures + if (this->newfigurerequest.load() == 1) { + this->user_figures_.clear(); + std::swap(this->user_figures_, this->user_figures_new_); + int value = 1; + this->newfigurerequest.compare_exchange_strong(value, 0); + } for (auto& [viewport, figure] : this->user_figures_) { ShowFigure(this, viewport, &figure); } + // overlay text + if (this->newtextrequest.load() == 1) { + this->user_texts_.clear(); + std::swap(this->user_texts_, this->user_texts_new_); + int value = 1; + this->newtextrequest.compare_exchange_strong(value, 0); + } + for (auto& [font, gridpos, text1, text2] : this->user_texts_) { + ShowOverlayText(this, rect, font, gridpos, text1, text2); + } + + // user images + if (this->newimagerequest.load() == 1) { + this->user_images_.clear(); + std::swap(this->user_images_, this->user_images_new_); + int value = 1; + this->newimagerequest.compare_exchange_strong(value, 0); + } + for (auto& [viewport, image] : this->user_images_) { + ShowImage(this, viewport, image.get()); + } + // finalize this->platform_ui->SwapBuffers(); } @@ -2714,11 +2781,22 @@ void Simulate::RenderLoop() { } // update scene, doing a full sync if in fully managed mode - if (!this->is_passive_) { + if (!is_passive_) { Sync(); - } else { - scnstate_.data.warning[mjWARN_VGEOMFULL].number += mjv_updateSceneFromState( - &scnstate_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn); + } else if (m_passive_ && d_passive_) { + // the user has called Sync() in their code + mjv_updateScene(m_passive_, d_passive_, + &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn); + + // add user geoms to scene + int nusergeom = user_scn_geoms_.size(); + int ngeom = std::min(nusergeom, this->scn.maxgeom - this->scn.ngeom); + if (ngeom < nusergeom) { + mj_warning(d_passive_, mjWARN_VGEOMFULL, this->scn.maxgeom); + } + std::memcpy(this->scn.geoms + this->scn.ngeom, user_scn_geoms_.data(), + ngeom * sizeof(mjvGeom)); + this->scn.ngeom += ngeom; } } // MutexLock (unblocks simulation thread) @@ -2739,7 +2817,8 @@ void Simulate::RenderLoop() { const MutexLock lock(this->mtx); mjv_freeScene(&this->scn); if (is_passive_) { - mjv_freeSceneState(&scnstate_); + mj_deleteData(d_passive_); + mj_deleteModel(m_passive_); } this->exitrequest.store(2); diff --git a/simulate/simulate.h b/simulate/simulate.h index cd654192..38234e4c 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -129,8 +130,15 @@ class Simulate { std::vector ctrl_; std::vector ctrl_prev_; - mjvSceneState scnstate_; + // in passive mode the user owns m_ and d_, these "passive" instances are + // owned by Simulate, updated from the user by the Sync() method + mjModel* m_passive_ = nullptr; + mjData* d_passive_ = nullptr; + std::vector user_scn_geoms_; + mjOption mjopt_prev_; + mjVisual mjvis_prev_; + mjStatistic mjstat_prev_; mjvOption opt_prev_; mjvCamera cam_prev_; @@ -144,7 +152,8 @@ class Simulate { std::optional print_data; bool reset; bool align; - bool copy_pose; + bool copy_key; + bool copy_key_full_precision; bool load_from_history; bool load_key; bool save_key; @@ -155,6 +164,7 @@ class Simulate { bool ui_update_simulation; bool ui_update_physics; bool ui_update_rendering; + bool ui_update_visualization; bool ui_update_joint; bool ui_update_ctrl; bool ui_remake_ctrl; @@ -196,6 +206,9 @@ class Simulate { std::atomic_int droploadrequest = 0; std::atomic_int screenshotrequest = 0; std::atomic_int uiloadrequest = 0; + std::atomic_int newfigurerequest = 0; + std::atomic_int newtextrequest = 0; + std::atomic_int newimagerequest = 0; // loadrequest // 3: display a loading message @@ -249,10 +262,15 @@ class Simulate { mjvFigure figsize = {}; mjvFigure figsensor = {}; - // additional user-defined visualization geoms (used in passive mode) + // additional user-defined visualization mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; std::vector> user_figures_; + std::vector> user_figures_new_; + std::vector> user_texts_; + std::vector> user_texts_new_; + std::vector>> user_images_; + std::vector>> user_images_new_; // OpenGL rendering and UI int refresh_rate = 60; @@ -293,7 +311,7 @@ class Simulate { {mjITEM_BUTTON, "Reset", 2, nullptr, " #259"}, {mjITEM_BUTTON, "Reload", 5, nullptr, "CL"}, {mjITEM_BUTTON, "Align", 2, nullptr, "CA"}, - {mjITEM_BUTTON, "Copy pose", 2, nullptr, "CC"}, + {mjITEM_BUTTON, "Copy state", 2, nullptr, "CC"}, {mjITEM_SLIDERINT, "Key", 3, &this->key, "0 0"}, {mjITEM_BUTTON, "Load key", 3}, {mjITEM_BUTTON, "Save key", 3}, diff --git a/src/engine/CMakeLists.txt b/src/engine/CMakeLists.txt index 8957c4c6..ea4b2dd6 100644 --- a/src/engine/CMakeLists.txt +++ b/src/engine/CMakeLists.txt @@ -83,8 +83,6 @@ set(MUJOCO_ENGINE_SRCS engine_vis_init.h engine_vis_interact.c engine_vis_interact.h - engine_vis_state.c - engine_vis_state.h engine_vis_visualize.c engine_vis_visualize.h ) diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c index b724d280..16b3b05c 100644 --- a/src/engine/engine_collision_convex.c +++ b/src/engine/engine_collision_convex.c @@ -908,24 +908,31 @@ static void mju_rotateFrame(const mjtNum origin[3], const mjtNum rot[9], -// return true if multiccd can run in a single pass -static int singlePass(const mjCCDObj* obj1, const mjCCDObj* obj2) { +// return number of contacts supported by a single pass of narrowphase +static int maxContacts(const mjCCDObj* obj1, const mjCCDObj* obj2) { const mjModel* m = obj1->model; // single pass not supported for margins if (obj1->margin > 0 || obj2->margin > 0) { - return 0; + return 1; } - // supported geoms for single pass + // can return 8 contacts for box-box collision in one pass int type1 = m->geom_type[obj1->geom]; int type2 = m->geom_type[obj2->geom]; + if (type1 == mjGEOM_BOX && type2 == mjGEOM_BOX) { + return 8; + } + + // reduce mesh collisions to 4 contacts max if (type1 == mjGEOM_BOX || type1 == mjGEOM_MESH) { if (type2 == mjGEOM_BOX || type2 == mjGEOM_MESH) { - return 1; + return mjENABLED(mjENBL_MULTICCD) ? 4 : 1; } } - return 0; + + // not supported for other geom types + return 1; } @@ -937,17 +944,14 @@ int mjc_Convex(const mjModel* m, const mjData* d, mjCCDObj obj1, obj2; mjc_initCCDObj(&obj1, m, d, g1, margin); mjc_initCCDObj(&obj2, m, d, g2, margin); - int max_contacts = 1; - - if (mjENABLED(mjENBL_MULTICCD) && singlePass(&obj1, &obj2)) { - max_contacts = 4; - } + int max_contacts = maxContacts(&obj1, &obj2); // find initial contact int ncon = mjc_CCDIteration(m, d, &obj1, &obj2, con, max_contacts, margin); - // nativeccd supports multi Box-Box collision directly - if (!mjDISABLED(mjDSBL_NATIVECCD) && singlePass(&obj1, &obj2)) { + + // no additional contacts needed + if (!mjDISABLED(mjDSBL_NATIVECCD) && max_contacts > 1) { return ncon; } diff --git a/src/engine/engine_collision_convex.h b/src/engine/engine_collision_convex.h index c3a4b220..76df2a28 100644 --- a/src/engine/engine_collision_convex.h +++ b/src/engine/engine_collision_convex.h @@ -81,12 +81,12 @@ void mjc_pointSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]); void mjc_lineSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]); // pairwise geom collision functions using ccd -int mjc_PlaneConvex (const mjModel* m, const mjData* d, - mjContact* con, int g1, int g2, mjtNum margin); -int mjc_ConvexHField (const mjModel* m, const mjData* d, - mjContact* con, int g1, int g2, mjtNum margin); -int mjc_Convex (const mjModel* m, const mjData* d, - mjContact* con, int g1, int g2, mjtNum margin); +int mjc_PlaneConvex(const mjModel* m, const mjData* d, + mjContact* con, int g1, int g2, mjtNum margin); +int mjc_ConvexHField(const mjModel* m, const mjData* d, + mjContact* con, int g1, int g2, mjtNum margin); +MJAPI int mjc_Convex(const mjModel* m, const mjData* d, + mjContact* con, int g1, int g2, mjtNum margin); // geom-elem or elem-elem or vert-elem collision function using ccd int mjc_ConvexElem (const mjModel* m, const mjData* d, mjContact* con, diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c index fb13c666..a0a76004 100644 --- a/src/engine/engine_collision_gjk.c +++ b/src/engine/engine_collision_gjk.c @@ -14,10 +14,10 @@ #include "engine/engine_collision_gjk.h" -#include #include #include #include +#include #include #include @@ -25,14 +25,16 @@ #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" +#define mjMINVAL2 (mjMINVAL * mjMINVAL) +#define mjMAXVAL2 (mjMAXVAL * mjMAXVAL) + // subdistance algorithm for GJK that computes the barycentric coordinates of the point in a // simplex closest to the origin // implementation adapted from Montanari et al, ToG 2017 -static void subdistance(mjtNum lambda[4], int n, const mjtNum s1[3], const mjtNum s2[3], - const mjtNum s3[3], const mjtNum s4[3]); +static void subdistance(mjtNum lambda[4], int n, const Vertex simplex[4]); // compute the barycentric coordinates of the closest point to the origin in the n-simplex, -// where n = 3, 2, 1 respectively +// for n = 3, 2, 1 respectively static void S3D(mjtNum lambda[4], const mjtNum s1[3], const mjtNum s2[3], const mjtNum s3[3], const mjtNum s4[3]); static void S2D(mjtNum lambda[3], const mjtNum s1[3], const mjtNum s2[3], const mjtNum s3[3]); @@ -40,7 +42,7 @@ static void S1D(mjtNum lambda[2], const mjtNum s1[3], const mjtNum s2[3]); // compute the support point for GJK static void gjkSupport(Vertex* v, mjCCDObj* obj1, mjCCDObj* obj2, - const mjtNum x_k[3]); + const mjtNum x_k[3], mjtNum x_norm); // compute the linear combination of 1 - 4 3D vectors static inline void lincomb(mjtNum res[3], const mjtNum* coef, int n, const mjtNum v1[3], @@ -51,24 +53,24 @@ typedef struct { int verts[3]; // indices of the three vertices of the face in the polytope int adj[3]; // adjacent faces, one for each edge: [v1,v2], [v2,v3], [v3,v1] mjtNum v[3]; // projection of the origin on face, can be used as face normal - mjtNum dist; // norm of v; negative if deleted + mjtNum dist2; // squared norm of v; negative if deleted int index; // index in map; -1: not in map, -2: deleted from polytope } Face; // polytope used in the Expanding Polytope Algorithm (EPA) typedef struct { - Vertex* verts; // list of vertices that make up the polytope - int nverts; // number of vertices - Face* faces; // list of faces that make up the polytope - int nfaces; // number of faces - int maxfaces; // max number of faces that can be stored in polytope - Face** map; // linear map storing faces - int nmap; // number of faces in map - struct Horizon { // polytope boundary edges that can be seen from w - int* indices; // indices of faces on horizon - int* edges; // corresponding edge of each face on the horizon - int nedges; // number of edges in horizon - mjtNum* w; // point where horizon is created + Vertex* verts; // list of vertices that make up the polytope + int nverts; // number of vertices + Face* faces; // list of faces that make up the polytope + int nfaces; // number of faces + int maxfaces; // max number of faces that can be stored in polytope + Face** map; // linear map storing faces + int nmap; // number of faces in map + struct Horizon { // polytope boundary edges that can be seen from w + int* indices; // indices of faces on horizon + int* edges; // corresponding edge of each face on the horizon + int nedges; // number of edges in horizon + const mjtNum* w; // point where horizon is created } horizon; } Polytope; @@ -79,7 +81,7 @@ static int epaSupport(Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, // make copy of vertex in polytope and return its index static int insertVertex(Polytope* pt, const Vertex* v); -// attach a face to the polytope with the given vertex indices; return distance to origin +// attach a face to the polytope with the given vertex indices; return squared distance to origin static mjtNum attachFace(Polytope* pt, int v1, int v2, int v3, int adj1, int adj2, int adj3); // return 1 if objects are in contact; 0 if not; -1 if inconclusive @@ -180,21 +182,27 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { mjtNum cutoff2 = status->dist_cutoff * status->dist_cutoff; // if both geoms are discrete, finite convergence is guaranteed; set tolerance to 0 - mjtNum epsilon = discreteGeoms(obj1, obj2) ? 0 : status->tolerance * status->tolerance; + mjtNum epsilon = discreteGeoms(obj1, obj2) ? 0 : 0.5 * status->tolerance * status->tolerance; + mjtNum x_norm; // set initial guess sub3(x_k, x1_k, x2_k); for (; k < kmax; k++) { // compute the kth support point - gjkSupport(simplex + n, obj1, obj2, x_k); + x_norm = dot3(x_k, x_k); + if (x_norm < mjMINVAL2) { + break; + } + x_norm = mju_sqrt(x_norm); + gjkSupport(simplex + n, obj1, obj2, x_k, x_norm); mjtNum *s_k = simplex[n].vert; // stopping criteria using the Frank-Wolfe duality gap given by // |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - s_k) > mjtNum diff[3]; sub3(diff, x_k, s_k); - if (2*dot3(x_k, diff) < epsilon) { + if (dot3(x_k, diff) < epsilon) { if (!k) n = 1; break; } @@ -206,16 +214,16 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { status->gjk_iterations = k; status->nsimplex = 0; status->nx = 0; - status->dist = mjMAXVAL; + status->dist = mjMAX_LIMIT; return; } - } else if (status->dist_cutoff < mjMAXVAL) { + } else if (status->dist_cutoff < mjMAX_LIMIT) { mjtNum vs = dot3(x_k, s_k), vv = dot3(x_k, x_k); if (dot3(x_k, s_k) > 0 && (vs*vs / vv) >= cutoff2) { status->gjk_iterations = k; status->nsimplex = 0; status->nx = 0; - status->dist = mjMAXVAL; + status->dist = mjMAX_LIMIT; return; } } @@ -227,7 +235,7 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { int ret = gjkIntersect(status, obj1, obj2); if (ret != -1) { status->nx = 0; - status->dist = ret > 0 ? 0 : mjMAXVAL; + status->dist = ret > 0 ? 0 : mjMAX_LIMIT; return; } k = status->gjk_iterations; @@ -236,12 +244,12 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { // run the distance subalgorithm to compute the barycentric coordinates // of the closest point to the origin in the simplex - subdistance(lambda, n + 1, simplex[0].vert, simplex[1].vert, simplex[2].vert, simplex[3].vert); + subdistance(lambda, n + 1, simplex); // remove vertices from the simplex no longer needed n = 0; for (int i = 0; i < 4; i++) { - if (lambda[i] == 0) continue; + if (!lambda[i]) continue; simplex[n] = simplex[i]; lambda[n++] = lambda[i]; } @@ -251,7 +259,7 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { status->gjk_iterations = k; status->nsimplex = 0; status->nx = 0; - status->dist = mjMAXVAL; + status->dist = mjMAX_LIMIT; return; } @@ -283,58 +291,51 @@ static void gjk(mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { status->nx = 1; status->gjk_iterations = k; status->nsimplex = n; - status->dist = norm3(x_k); + status->dist = x_norm; } // compute the support point in obj1 and obj2 for Minkowski difference -static inline void support(mjtNum s1[3], mjtNum s2[3], mjCCDObj* obj1, mjCCDObj* obj2, +static inline void support(Vertex* v, mjCCDObj* obj1, mjCCDObj* obj2, const mjtNum dir[3], const mjtNum dir_neg[3]) { // obj1 - obj1->support(s1, obj1, dir); + obj1->support(v->vert1, obj1, dir); if (obj1->margin > 0 && obj1->geom >= 0) { mjtNum margin = 0.5 * obj1->margin; - s1[0] += dir[0] * margin; - s1[1] += dir[1] * margin; - s1[2] += dir[2] * margin; + v->vert1[0] += dir[0] * margin; + v->vert1[1] += dir[1] * margin; + v->vert1[2] += dir[2] * margin; } // obj2 - obj2->support(s2, obj2, dir_neg); + obj2->support(v->vert2, obj2, dir_neg); if (obj2->margin > 0 && obj2->geom >= 0) { mjtNum margin = 0.5 * obj2->margin; - s2[0] += dir_neg[0] * margin; - s2[1] += dir_neg[1] * margin; - s2[2] += dir_neg[2] * margin; + v->vert2[0] += dir_neg[0] * margin; + v->vert2[1] += dir_neg[1] * margin; + v->vert2[2] += dir_neg[2] * margin; } + + // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir) + sub3(v->vert, v->vert1, v->vert2); + + // copy vertex indices of discrete geoms + v->index1 = obj1->vertindex; + v->index2 = obj2->vertindex; } // compute the support points in obj1 and obj2 for the kth approximation point -static void gjkSupport(Vertex* v, mjCCDObj* obj1, mjCCDObj* obj2, - const mjtNum x_k[3]) { - mjtNum dir[3] = {-1, 0, 0}, dir_neg[3] = {1, 0, 0}; +static inline void gjkSupport(Vertex* v, mjCCDObj* obj1, mjCCDObj* obj2, + const mjtNum x_k[3], mjtNum x_norm) { + mjtNum dir[3], dir_neg[3]; // mjc_support requires a normalized direction - mjtNum norm = dot3(x_k, x_k); - if (norm > mjMINVAL*mjMINVAL) { - norm = 1/mju_sqrt(norm); - scl3(dir_neg, x_k, norm); - scl3(dir, dir_neg, -1); - } - - // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir) - support(v->vert1, v->vert2, obj1, obj2, dir, dir_neg); - sub3(v->vert, v->vert1, v->vert2); - // copy mesh indices - if (obj1->vertindex >= 0) { - v->index1 = obj1->vertindex; - } - if (obj2->vertindex >= 0) { - v->index2 = obj2->vertindex; - } + scl3(dir_neg, x_k, 1 / x_norm); + scl3(dir, dir_neg, -1); + support(v, obj1, obj2, dir, dir_neg); } @@ -354,16 +355,7 @@ static int epaSupport(Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, int n = pt->nverts++; Vertex* v = pt->verts + n; - - // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir) - support(v->vert1, v->vert2, obj1, obj2, dir, dir_neg); - sub3(v->vert, v->vert1, v->vert2); - if (obj1->vertindex >= 0) { - v->index1 = obj1->vertindex; - } - if (obj2->vertindex >= 0) { - v->index2 = obj2->vertindex; - } + support(v, obj1, obj2, dir, dir_neg); return n; } @@ -373,15 +365,7 @@ static int epaSupport(Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, static void gjkIntersectSupport(Vertex* v, mjCCDObj* obj1, mjCCDObj* obj2, const mjtNum dir[3]) { mjtNum dir_neg[3] = {-dir[0], -dir[1], -dir[2]}; - // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir) - support(v->vert1, v->vert2, obj1, obj2, dir, dir_neg); - sub3(v->vert, v->vert1, v->vert2); - if (obj1->vertindex >= 0) { - v->index1 = obj1->vertindex; - } - if (obj2->vertindex >= 0) { - v->index2 = obj2->vertindex; - } + support(v, obj1, obj2, dir, dir_neg); } @@ -393,13 +377,12 @@ static inline mjtNum signedDistance(mjtNum normal[3], const Vertex* v1, const Ve sub3(diff1, v3->vert, v1->vert); sub3(diff2, v2->vert, v1->vert); cross3(normal, diff1, diff2); - mjtNum norm = dot3(normal, normal); - if (norm > mjMINVAL*mjMINVAL && norm < mjMAXVAL*mjMAXVAL) { - norm = 1/mju_sqrt(norm); - scl3(normal, normal, norm); + mjtNum norm2 = dot3(normal, normal); + if (norm2 > mjMINVAL2 && norm2 < mjMAXVAL2) { + scl3(normal, normal, 1 / mju_sqrt(norm2)); return dot3(normal, v1->vert); } - return mjMAXVAL; // cannot recover normal (ignore face) + return mjMAX_LIMIT; // cannot recover normal (ignore face) } @@ -556,24 +539,33 @@ static inline int sameSign2(mjtNum a, mjtNum b) { // subdistance algorithm for GJK that computes the barycentric coordinates of the point in a // simplex closest to the origin // implementation adapted from Montanari et al, ToG 2017 -static inline void subdistance(mjtNum lambda[4], int n, const mjtNum s1[3], - const mjtNum s2[3], const mjtNum s3[3], const mjtNum s4[3]) { - lambda[0] = lambda[1] = lambda[2] = lambda[3] = 0; - if (n == 4) { - S3D(lambda, s1, s2, s3, s4); - } else if (n == 3) { - S2D(lambda, s1, s2, s3); - } else if (n == 2) { - S1D(lambda, s1, s2); - } else { +static inline void subdistance(mjtNum lambda[4], int n, const Vertex simplex[4]) { + memset(lambda, 0, 4 * sizeof(mjtNum)); + const mjtNum* s1 = simplex[0].vert; + const mjtNum* s2 = simplex[1].vert; + const mjtNum* s3 = simplex[2].vert; + const mjtNum* s4 = simplex[3].vert; + + switch (n) { + case 4: + S3D(lambda, s1, s2, s3, s4); + break; + case 3: + S2D(lambda, s1, s2, s3); + break; + case 2: + S1D(lambda, s1, s2); + break; + default: lambda[0] = 1; + break; } } -static void S3D(mjtNum lambda[4], const mjtNum s1[3], const mjtNum s2[3], const mjtNum s3[3], - const mjtNum s4[3]) { +static void S3D(mjtNum lambda[4], const mjtNum s1[3], const mjtNum s2[3], + const mjtNum s3[3], const mjtNum s4[3]) { // the matrix M is given by // [[ s1_x, s2_x, s3_x, s4_x ], // [ s1_y, s2_y, s3_y, s4_y ], @@ -607,7 +599,7 @@ static void S3D(mjtNum lambda[4], const mjtNum s1[3], const mjtNum s2[3], const } // find the smallest distance, and use the corresponding barycentric coordinates - mjtNum dmin = mjMAXVAL; + mjtNum dmin = mjMAX_LIMIT; if (!comp1) { mjtNum lambda_2d[3], x[3]; @@ -658,7 +650,6 @@ static void S3D(mjtNum lambda[4], const mjtNum s1[3], const mjtNum s2[3], const lambda[0] = lambda_2d[0]; lambda[1] = lambda_2d[1]; lambda[2] = lambda_2d[2]; - lambda[3] = 0; } } } @@ -758,7 +749,7 @@ static void S2D(mjtNum lambda[3], const mjtNum s1[3], const mjtNum s2[3], const } // find the smallest distance, and use the corresponding barycentric coordinates - mjtNum dmin = mjMAXVAL; + mjtNum dmin = mjMAX_LIMIT; if (!comp1) { mjtNum lambda_1d[2], x[3]; @@ -805,27 +796,29 @@ static void S1D(mjtNum lambda[2], const mjtNum s1[3], const mjtNum s2[3]) { projectOriginLine(p_o, s1, s2); // find the axis with the largest projection "shadow" of the simplex - mjtNum mu_max = 0; - int index; - for (int i = 0; i < 3; i++) { - mjtNum mu = s1[i] - s2[i]; - if (mju_abs(mu) >= mju_abs(mu_max)) { - mu_max = mu; - index = i; - } + mjtNum mu = s1[0] - s2[0]; + mjtNum mu_max = mu; + int index = 0; + + mu = s1[1] - s2[1]; + if (mju_abs(mu) >= mju_abs(mu_max)) { + mu_max = mu; + index = 1; + } + + mu = s1[2] - s2[2]; + if (mju_abs(mu) >= mju_abs(mu_max)) { + mu_max = mu; + index = 2; } mjtNum C1 = p_o[index] - s2[index]; mjtNum C2 = s1[index] - p_o[index]; - // inside the simplex - if (sameSign2(mu_max, C1) && sameSign2(mu_max, C2)) { - lambda[0] = C1 / mu_max; - lambda[1] = C2 / mu_max; - } else { - lambda[0] = 0; - lambda[1] = 1; - } + // determine if projection of origin lies inside 1-simplex + int same = sameSign2(mu_max, C1) && sameSign2(mu_max, C2); + lambda[0] = same ? C1 / mu_max : 0; + lambda[1] = same ? C2 / mu_max : 1; } @@ -833,22 +826,16 @@ static void S1D(mjtNum lambda[2], const mjtNum s1[3], const mjtNum s2[3]) { // replace a 3-simplex with one of its faces static inline void replaceSimplex3(Polytope* pt, mjCCDStatus* status, int v1, int v2, int v3) { + // reset status simplex status->nsimplex = 3; - Vertex* v = pt->verts; - copy3(status->simplex[0].vert1, v[v1].vert1); - copy3(status->simplex[1].vert1, v[v2].vert1); - copy3(status->simplex[2].vert1, v[v3].vert1); - - copy3(status->simplex[0].vert2, v[v1].vert2); - copy3(status->simplex[1].vert2, v[v2].vert2); - copy3(status->simplex[2].vert2, v[v3].vert2); - - copy3(status->simplex[0].vert, v[v1].vert); - copy3(status->simplex[1].vert, v[v2].vert); - copy3(status->simplex[2].vert, v[v3].vert); + status->simplex[0] = pt->verts[v1]; + status->simplex[1] = pt->verts[v2]; + status->simplex[2] = pt->verts[v3]; + // reset polytope pt->nfaces = 0; pt->nverts = 0; + pt->nmap = 0; } @@ -903,6 +890,26 @@ static void rotmat(mjtNum R[9], const mjtNum axis[3]) { +// return nonzero if the ray v1v2 intersects the triangle v3v4v5 +static inline int rayTriangle(const mjtNum v1[3], const mjtNum v2[3], const mjtNum v3[3], + const mjtNum v4[3], const mjtNum v5[3]) { + mjtNum diff12[3], diff13[3], diff14[3], diff15[3]; + sub3(diff12, v2, v1); + sub3(diff13, v3, v1); + sub3(diff14, v4, v1); + sub3(diff15, v5, v1); + + mjtNum vol1 = det3(diff13, diff14, diff12); + mjtNum vol2 = det3(diff14, diff15, diff12); + mjtNum vol3 = det3(diff15, diff13, diff12); + + if (vol1 >= 0 && vol2 >= 0 && vol3 >= 0) return 1; + if (vol1 <= 0 && vol2 <= 0 && vol3 <= 0) return -1; + return 0; +} + + + // create a polytope from a 1-simplex (returns 0 on success) static int polytope2(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) { mjtNum *v1 = status->simplex[0].vert, *v2 = status->simplex[1].vert; @@ -911,7 +918,7 @@ static int polytope2(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj sub3(diff, v2, v1); // find component with smallest magnitude (so cross product is largest) - mjtNum value = mjMAXVAL; + mjtNum value = mjMAX_LIMIT; int index = 0; for (int i = 0; i < 3; i++) { if (mju_abs(diff[i]) < value) { @@ -945,34 +952,34 @@ static int polytope2(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj mjtNum* v5 = pt->verts[v5i].vert; // build hexahedron - if (attachFace(pt, v1i, v3i, v4i, 1, 3, 2) < mjMINVAL) { + if (attachFace(pt, v1i, v3i, v4i, 1, 3, 2) < mjMINVAL2) { replaceSimplex3(pt, status, v1i, v3i, v4i); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v1i, v5i, v3i, 2, 4, 0) < mjMINVAL) { + if (attachFace(pt, v1i, v5i, v3i, 2, 4, 0) < mjMINVAL2) { replaceSimplex3(pt, status, v1i, v5i, v3i); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v1i, v4i, v5i, 0, 5, 1) < mjMINVAL) { + if (attachFace(pt, v1i, v4i, v5i, 0, 5, 1) < mjMINVAL2) { replaceSimplex3(pt, status, v1i, v4i, v5i); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v2i, v4i, v3i, 5, 0, 4) < mjMINVAL) { + if (attachFace(pt, v2i, v4i, v3i, 5, 0, 4) < mjMINVAL2) { replaceSimplex3(pt, status, v2i, v4i, v3i); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v2i, v3i, v5i, 3, 1, 5) < mjMINVAL) { + if (attachFace(pt, v2i, v3i, v5i, 3, 1, 5) < mjMINVAL2) { replaceSimplex3(pt, status, v2i, v3i, v5i); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v2i, v5i, v4i, 4, 2, 3) < mjMINVAL) { + if (attachFace(pt, v2i, v5i, v4i, 4, 2, 3) < mjMINVAL2) { replaceSimplex3(pt, status, v2i, v5i, v4i); return polytope3(pt, status, obj1, obj2); } - // check that origin is in the hexahedron - if (status->dist > 10*mjMINVAL && !testTetra(v1, v3, v4, v5) && !testTetra(v2, v3, v4, v5)) { - return mjEPA_P2_MISSING_ORIGIN; + // check hexahedron is convex + if (!rayTriangle(v1, v2, v3, v4, v5)) { + return mjEPA_P2_NONCONVEX; } for (int i = 0; i < 6; i++) { @@ -1101,28 +1108,26 @@ static int polytope3(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj } // create hexahedron for EPA - if (attachFace(pt, v4i, v1i, v2i, 1, 3, 2) < mjMINVAL) { + if (attachFace(pt, v4i, v1i, v2i, 1, 3, 2) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - if (attachFace(pt, v4i, v3i, v1i, 2, 4, 0) < mjMINVAL) { + if (attachFace(pt, v4i, v3i, v1i, 2, 4, 0) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - if (attachFace(pt, v4i, v2i, v3i, 0, 5, 1) < mjMINVAL) { + if (attachFace(pt, v4i, v2i, v3i, 0, 5, 1) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - if (attachFace(pt, v5i, v2i, v1i, 5, 0, 4) < mjMINVAL) { + if (attachFace(pt, v5i, v2i, v1i, 5, 0, 4) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - if (attachFace(pt, v5i, v1i, v3i, 3, 1, 5) < mjMINVAL) { + if (attachFace(pt, v5i, v1i, v3i, 3, 1, 5) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - if (attachFace(pt, v5i, v3i, v2i, 4, 2, 3) < mjMINVAL) { + if (attachFace(pt, v5i, v3i, v2i, 4, 2, 3) < mjMINVAL2) { return mjEPA_P3_ORIGIN_ON_FACE; } - - // if the origin is on the affine hull of any of the faces then the origin is not in the - // hexahedron or the hexahedron is degenerate + // populate face map for (int i = 0; i < 6; i++) { pt->map[i] = pt->faces + i; pt->faces[i].index = i; @@ -1141,19 +1146,19 @@ static int polytope4(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj int v4 = insertVertex(pt, status->simplex + 3); // if the origin is on a face, replace the 3-simplex with a 2-simplex - if (attachFace(pt, v1, v2, v3, 1, 3, 2) < mjMINVAL) { + if (attachFace(pt, v1, v2, v3, 1, 3, 2) < mjMINVAL2) { replaceSimplex3(pt, status, v1, v2, v3); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v1, v4, v2, 2, 3, 0) < mjMINVAL) { + if (attachFace(pt, v1, v4, v2, 2, 3, 0) < mjMINVAL2) { replaceSimplex3(pt, status, v1, v4, v2); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v1, v3, v4, 0, 3, 1) < mjMINVAL) { + if (attachFace(pt, v1, v3, v4, 0, 3, 1) < mjMINVAL2) { replaceSimplex3(pt, status, v1, v3, v4); return polytope3(pt, status, obj1, obj2); } - if (attachFace(pt, v4, v3, v2, 2, 0, 1) < mjMINVAL) { + if (attachFace(pt, v4, v3, v2, 2, 0, 1) < mjMINVAL2) { replaceSimplex3(pt, status, v4, v3, v2); return polytope3(pt, status, obj1, obj2); } @@ -1162,6 +1167,7 @@ static int polytope4(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj return mjEPA_P4_MISSING_ORIGIN; } + // populate face map for (int i = 0; i < 4; i++) { pt->map[i] = pt->faces + i; pt->faces[i].index = i; @@ -1171,16 +1177,10 @@ static int polytope4(Polytope* pt, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj } - // make a copy of vertex in polytope and return its index static inline int insertVertex(Polytope* pt, const Vertex* v) { int n = pt->nverts++; - Vertex* new_v = pt->verts + n; - copy3(new_v->vert1, v->vert1); - copy3(new_v->vert2, v->vert2); - new_v->index1 = v->index1; - new_v->index2 = v->index2; - sub3(new_v->vert, v->vert1, v->vert2); + pt->verts[n] = *v; return n; } @@ -1203,7 +1203,7 @@ static inline int maxFaces(Polytope* pt) { -// attach a face to the polytope with the given vertex indices; return distance to origin +// attach a face to the polytope with the given vertex indices; return squared distance to origin static inline mjtNum attachFace(Polytope* pt, int v1, int v2, int v3, int adj1, int adj2, int adj3) { Face* face = &pt->faces[pt->nfaces++]; @@ -1221,10 +1221,10 @@ static inline mjtNum attachFace(Polytope* pt, int v1, int v2, int v3, if (ret) { return 0; } - face->dist = norm3(face->v); + face->dist2 = dot3(face->v, face->v); face->index = -1; - return face->dist; + return face->dist2; } @@ -1248,10 +1248,8 @@ static inline int getEdge(Face* face, int vertex) { // recursive call to build horizon; return 1 if face is visible from w otherwise 0 static int horizonRec(Polytope* pt, Face* face, int e) { - mjtNum dist2 = face->dist * face->dist; - // v is visible from w so it is deleted and adjacent faces are checked - if (dot3(face->v, pt->horizon.w) >= dist2) { + if (dot3(face->v, pt->horizon.w) - face->dist2 > mjMINVAL) { deleteFace(pt, face); // recursively search the adjacent faces on the next two edges @@ -1331,44 +1329,68 @@ static void epaWitness(const Polytope* pt, const Face* face, mjtNum x1[3], mjtNu // return a face of the expanded polytope that best approximates the pentration depth // witness points are in status->{x1, x2} static Face* epa(mjCCDStatus* status, Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2) { - mjtNum tolerance = status->tolerance, lower, upper = FLT_MAX; - int k, kmax = status->max_iterations; + mjtNum upper = mjMAX_LIMIT, upper2 = mjMAX_LIMIT, lower2; Face* face = NULL, *pface = NULL; // face closest to origin + mjtNum tolerance = status->tolerance; + int discrete = discreteGeoms(obj1, obj2); + // tolerance is not used for discrete geoms + if (discrete && sizeof(mjtNum) == sizeof(double)) { + tolerance = mjMINVAL; + } + + int k, kmax = status->max_iterations; for (k = 0; k < kmax; k++) { pface = face; // find the face closest to the origin (lower bound for penetration depth) - lower = FLT_MAX; + lower2 = mjMAX_LIMIT; for (int i = 0; i < pt->nmap; i++) { - if (pt->map[i]->dist < lower) { + if (pt->map[i]->dist2 < lower2) { face = pt->map[i]; - lower = face->dist; + lower2 = face->dist2; } } // face not valid, return previous face - if (lower > upper || !face) { + if (lower2 > upper2 || !face) { face = pface; break; } // check if lower bound is 0 - if (lower <= 0) { + if (lower2 <= 0) { mju_warning("EPA: origin lies on affine hull of face"); break; } // compute support point w from the closest face's normal + mjtNum lower = mju_sqrt(lower2); int wi = epaSupport(pt, obj1, obj2, face->v, lower); - mjtNum* w = pt->verts[wi].vert; - mjtNum upper_k = dot3(face->v, w) / lower; // upper bound for kth iteration - if (upper_k < upper) upper = upper_k; + const Vertex* w = pt->verts + wi; + mjtNum upper_k = dot3(face->v, w->vert) / lower; // upper bound for kth iteration + if (upper_k < upper) { + upper = upper_k; + upper2 = upper * upper; + } if (upper - lower < tolerance) { break; } - pt->horizon.w = w; + // check if vertex w is a repeated support point + if (discrete) { + int i = 0, nverts = pt->nverts - 1; + for (; i < nverts; i++) { + if (w->index1 == pt->verts[i].index1 && w->index2 == pt->verts[i].index2) { + break; + } + } + if (i != nverts) { + break; + } + } + + pt->horizon.w = w->vert; horizon(pt, face); // unrecoverable numerical issue; at least one face was deleted so nedges is 3 or more @@ -1392,16 +1414,16 @@ static Face* epa(mjCCDStatus* status, Polytope* pt, mjCCDObj* obj1, mjCCDObj* ob int v1 = horFace->verts[horEdge], v2 = horFace->verts[(horEdge + 1) % 3]; horFace->adj[horEdge] = nfaces; - mjtNum dist = attachFace(pt, wi, v2, v1, nfaces + nedges - 1, horIndex, nfaces + 1); + mjtNum dist2 = attachFace(pt, wi, v2, v1, nfaces + nedges - 1, horIndex, nfaces + 1); // unrecoverable numerical issue - if (dist == 0) { + if (dist2 == 0) { face = NULL; break; } // store face in map - if (dist >= lower && dist <= upper) { + if (dist2 >= lower2 && dist2 <= upper2) { int i = pt->nmap++; pt->map[i] = &pt->faces[pt->nfaces - 1]; pt->map[i]->index = i; @@ -1417,16 +1439,16 @@ static Face* epa(mjCCDStatus* status, Polytope* pt, mjCCDObj* obj1, mjCCDObj* ob v1 = horFace->verts[horEdge]; v2 = horFace->verts[(horEdge + 1) % 3]; horFace->adj[horEdge] = cur; - dist = attachFace(pt, wi, v2, v1, cur - 1, horIndex, next); + dist2 = attachFace(pt, wi, v2, v1, cur - 1, horIndex, next); // unrecoverable numerical issue - if (dist == 0) { + if (dist2 == 0) { face = NULL; break; } // store face in map - if (dist >= lower && dist <= upper) { + if (dist2 >= lower2 && dist2 <= upper2) { int idx = pt->nmap++; pt->map[idx] = &pt->faces[pt->nfaces - 1]; pt->map[idx]->index = idx; @@ -1444,7 +1466,7 @@ static Face* epa(mjCCDStatus* status, Polytope* pt, mjCCDObj* obj1, mjCCDObj* ob if (face) { epaWitness(pt, face, status->x1, status->x2); status->nx = 1; - status->dist = -face->dist; + status->dist = -mju_sqrt(face->dist2); } else { status->nx = 0; status->dist = 0; @@ -1546,7 +1568,7 @@ static mjtNum planeNormal(mjtNum res[3], const mjtNum v1[3], const mjtNum v2[3], // find what side of a plane a point p lies static int halfspace(const mjtNum a[3], const mjtNum n[3], const mjtNum p[3]) { mjtNum diff[3] = {p[0] - a[0], p[1] - a[1], p[2] - a[2]}; - return dot3(diff, n) > 0; + return dot3(diff, n) > -mjMINVAL; } @@ -1557,7 +1579,7 @@ static mjtNum planeIntersect(mjtNum res[3], const mjtNum pn[3], mjtNum pd, mjtNum ab[3]; sub3(ab, b, a); mjtNum temp = dot3(pn, ab); - if (temp == 0.0) return mjMAXVAL; // parallel; no intersection + if (temp == 0.0) return mjMAX_LIMIT; // parallel; no intersection mjtNum t = (pd - dot3(pn, a)) / temp; if (t >= 0.0 && t <= 1.0) { res[0] = a[0] + t*ab[0]; @@ -1836,44 +1858,69 @@ static int meshEdgeNormals(mjtNum* res, mjtNum* endverts, int dim, mjCCDObj* obj +// try recovering box normal from collision normal +static int boxNormals2(mjtNum res[9], int resind[3], const mjtNum mat[9], const mjtNum n[3]) { + // list of box face normals + mjtNum normals[18] = {1, 0, 0, -1, 0, 0, + 0, 1, 0, 0, -1, 0, + 0, 0, 1, 0, 0, -1}; + + // get local coordinates of the normal + mjtNum local_n[3]; + local_n[0] = mat[0]*n[0] + mat[3]*n[1] + mat[6]*n[2]; + local_n[1] = mat[1]*n[0] + mat[4]*n[1] + mat[7]*n[2]; + local_n[2] = mat[2]*n[0] + mat[5]*n[1] + mat[8]*n[2]; + scl3(local_n, local_n, 1/mju_sqrt(dot3(local_n, local_n))); + + // determine if there is a side close to the normal + for (int i = 0; i < 6; i++) { + if (dot3(local_n, normals + 3*i) > mjFACE_TOL) { + globalcoord(res, mat, NULL, normals[3*i], normals[3*i + 1], normals[3*i + 2]); + resind[0] = i; + return 1; + } + } + return 0; +} + + + // compute possible face normals of a box given up to 3 vertices static int boxNormals(mjtNum res[9], int resind[3], int dim, mjCCDObj* obj, - int v1, int v2, int v3) { - // box data - int g = 3*obj->geom; - const mjtNum* mat = obj->data->geom_xmat + 3*g; - + int v1, int v2, int v3, const mjtNum dir[3]) { + const mjtNum* mat = obj->data->geom_xmat + 9*obj->geom; if (dim == 3) { + int c = 0; int x = ((v1 & 1) && (v2 & 1) && (v3 & 1)) - (!(v1 & 1) && !(v2 & 1) && !(v3 & 1)); int y = ((v1 & 2) && (v2 & 2) && (v3 & 2)) - (!(v1 & 2) && !(v2 & 2) && !(v3 & 2)); int z = ((v1 & 4) && (v2 & 4) && (v3 & 4)) - (!(v1 & 4) && !(v2 & 4) && !(v3 & 4)); globalcoord(res, mat, NULL, x, y, z); int sgn = x + y + z; - if (x) resind[0] = 0; - if (y) resind[0] = 2; - if (z) resind[0] = 4; + if (x) resind[c++] = 0; + if (y) resind[c++] = 2; + if (z) resind[c++] = 4; if (sgn == -1) resind[0]++; - return 1; + return c == 1 ? 1 : boxNormals2(res, resind, mat, dir); } if (dim == 2) { + int c = 0; int x = ((v1 & 1) && (v2 & 1)) - (!(v1 & 1) && !(v2 & 1)); int y = ((v1 & 2) && (v2 & 2)) - (!(v1 & 2) && !(v2 & 2)); int z = ((v1 & 4) && (v2 & 4)) - (!(v1 & 4) && !(v2 & 4)); if (x) { globalcoord(res, mat, NULL, x, 0, 0); - resind[0] = (x > 0) ? 0 : 1; + resind[c++] = (x > 0) ? 0 : 1; } if (y) { - int i = (x ? 1 : 0); - globalcoord(res + 3*i, mat, NULL, 0, y, 0); - resind[i] = (y > 0) ? 2 : 3; + globalcoord(res + 3*c, mat, NULL, 0, y, 0); + resind[c++] = (y > 0) ? 2 : 3; } if (z) { globalcoord(res + 3, mat, NULL, 0, 0, z); - resind[1] = (z > 0) ? 4 : 5; + resind[c++] = (z > 0) ? 4 : 5; } - return 2; + return c == 2 ? 2 : boxNormals2(res, resind, mat, dir); } if (dim == 1) { @@ -2086,14 +2133,18 @@ static void multicontact(Polytope* pt, Face* face, mjCCDStatus* status, mjtNum n1[3 * mjMAX_POLYVERT], n2[3 * mjMAX_POLYVERT]; // normals of possible face collisions int idx1[mjMAX_POLYVERT], idx2[mjMAX_POLYVERT]; // indices of faces + mjtNum dir[3], dir_neg[3]; + sub3(dir, status->x2, status->x1); + sub3(dir_neg, status->x1, status->x2); + // get all possible face normals for each geom if (obj1->geom_type == mjGEOM_BOX) { - nnorms1 = boxNormals(n1, idx1, nface1, obj1, v11i, v12i, v13i); + nnorms1 = boxNormals(n1, idx1, nface1, obj1, v11i, v12i, v13i, dir_neg); } else if (obj1->geom_type == mjGEOM_MESH) { nnorms1 = meshNormals(n1, idx1, nface1, obj1, v11i, v12i, v13i); } if (obj2->geom_type == mjGEOM_BOX) { - nnorms2 = boxNormals(n2, idx2, nface2, obj2, v21i, v22i, v23i); + nnorms2 = boxNormals(n2, idx2, nface2, obj2, v21i, v22i, v23i, dir); } else if (obj2->geom_type == mjGEOM_MESH) { nnorms2 = meshNormals(n2, idx2, nface2, obj2, v21i, v22i, v23i); } @@ -2160,25 +2211,24 @@ static void multicontact(Polytope* pt, Face* face, mjCCDStatus* status, // TODO(kylebayes): this approximates the contact direction, by scaling the face normal by the // single contact direction's magnitude. This is effective, but polygonClip should compute // this for each contact point. - mjtNum diff[3], approx_dir[3]; - sub3(diff, status->x2, status->x1); + mjtNum approx_dir[3]; // face1 is an edge; clip face1 against face2 if (edgecon1) { - scl3(approx_dir, n2 + 3*j, norm3(diff)); + scl3(approx_dir, n2 + 3*j, norm3(dir)); polygonClip(status, face2, nface2, face1, nface1, n2 + 3*j, approx_dir); return; } // face2 is an edge; clip face2 against face1 if (edgecon2) { - scl3(approx_dir, n1 + 3*j, -norm3(diff)); + scl3(approx_dir, n1 + 3*j, -norm3(dir)); polygonClip(status, face1, nface1, face2, nface2, n1 + 3*j, approx_dir); return; } // face-face collision - scl3(approx_dir, n2 + 3*j, norm3(diff)); + scl3(approx_dir, n2 + 3*j, norm3(dir)); polygonClip(status, face1, nface1, face2, nface2, n1 + 3*i, approx_dir); } @@ -2272,7 +2322,7 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m if (status->dist > status->tolerance) { inflate(status, full_margin1, full_margin2); if (status->dist > status->dist_cutoff) { - status->dist = mjMAXVAL; + status->dist = mjMAX_LIMIT; } return status->dist; } diff --git a/src/engine/engine_collision_gjk.h b/src/engine/engine_collision_gjk.h index 027929e3..55ee5f8b 100644 --- a/src/engine/engine_collision_gjk.h +++ b/src/engine/engine_collision_gjk.h @@ -15,6 +15,7 @@ #ifndef MUJOCO_SRC_ENGINE_ENGINE_COLLISION_GJK_H_ #define MUJOCO_SRC_ENGINE_ENGINE_COLLISION_GJK_H_ +#include #include #include @@ -27,6 +28,13 @@ extern "C" { #endif +// numerical max limit +#ifndef mjUSESINGLE + #define mjMAX_LIMIT DBL_MAX +#else + #define mjMAX_LIMIT FLT_MAX +#endif + // max number of EPA iterations #define mjMAX_EPA_ITERATIONS 170 @@ -44,7 +52,7 @@ typedef enum { mjEPA_NOCONTACT = -1, mjEPA_SUCCESS = 0, mjEPA_P2_INVALID_FACES, - mjEPA_P2_MISSING_ORIGIN, + mjEPA_P2_NONCONVEX, mjEPA_P2_ORIGIN_ON_FACE, mjEPA_P3_BAD_NORMAL, mjEPA_P3_INVALID_V4, diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c index c6fe8eba..591c96b8 100644 --- a/src/engine/engine_collision_sdf.c +++ b/src/engine/engine_collision_sdf.c @@ -158,7 +158,9 @@ static void geomGradient(mjtNum gradient[3], const mjModel* m, const mjData* d, e = mju_abs(x[2]); a[0] = c - size[0]; a[1] = e - size[1]; - mjtNum grada[3] = {x[0] / c, x[1] / c, x[2] / e}; + mjtNum grada[3] = {x[0] / mju_max(c, 1. / mjMAXVAL), + x[1] / mju_max(c, 1. / mjMAXVAL), + x[2] / mju_max(e, 1. / mjMAXVAL)}; int j = a[0] > a[1] ? 0 : 1; if (a[j] < 0) { gradient[0] = j == 0 ? grada[0] : 0; @@ -167,7 +169,7 @@ static void geomGradient(mjtNum gradient[3], const mjModel* m, const mjData* d, } else { b[0] = mju_max(a[0], 0); b[1] = mju_max(a[1], 0); - mjtNum bnorm = mju_norm(b, 2); + mjtNum bnorm = mju_max(mju_norm(b, 2), 1./mjMAXVAL); gradient[0] = grada[0] * b[0] / bnorm; gradient[1] = grada[1] * b[0] / bnorm; gradient[2] = grada[2] * b[1] / bnorm; @@ -263,7 +265,7 @@ void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, } // get sdf from geom id -static const mjpPlugin* getSDF(const mjModel* m, int id) { +const mjpPlugin* mjc_getSDF(const mjModel* m, int id) { int instance = m->geom_plugin[id]; const int nslot = mjp_pluginCount(); const int slot = m->plugin[instance]; @@ -583,7 +585,7 @@ int mjc_MeshSDF(const mjModel* m, const mjData* d, mjContact* con, int g1, int g // get sdf plugin int instance = m->geom_plugin[g2]; - const mjpPlugin* sdf_ptr = getSDF(m, g2); + const mjpPlugin* sdf_ptr = mjc_getSDF(m, g2); mjtGeom geomtype = mjGEOM_SDF; // copy into data @@ -725,12 +727,12 @@ int mjc_SDF(const mjModel* m, const mjData* d, mjContact* con, int g1, int g2, m mjtGeom geomtypes[2] = {m->geom_type[g2], m->geom_type[g1]}; instance[0] = m->geom_plugin[g2]; - sdf_ptr[0] = getSDF(m, g2); + sdf_ptr[0] = mjc_getSDF(m, g2); // get sdf plugins if (m->geom_type[g1] == mjGEOM_SDF) { instance[1] = m->geom_plugin[g1]; - sdf_ptr[1] = getSDF(m, g1); + sdf_ptr[1] = mjc_getSDF(m, g1); } else { instance[1] = g1; sdf_ptr[1] = NULL; diff --git a/src/engine/engine_collision_sdf.h b/src/engine/engine_collision_sdf.h index c6acbe1c..c0ea7cdb 100644 --- a/src/engine/engine_collision_sdf.h +++ b/src/engine/engine_collision_sdf.h @@ -19,27 +19,14 @@ #include #include #include +#include #ifdef __cplusplus extern "C" { #endif -typedef enum mjtSDFType_ { // signed distance function (SDF) type - mjSDFTYPE_SINGLE = 0, // single SDF - mjSDFTYPE_INTERSECTION, // max(A, B) - mjSDFTYPE_MIDSURFACE, // A - B - mjSDFTYPE_COLLISION, // A + B + abs(max(A, B)) -} mjtSDFType; - -struct mjSDF_ { - const mjpPlugin** plugin; - int* id; - mjtSDFType type; - mjtNum* relpos; - mjtNum* relmat; - mjtGeom* geomtype; -}; -typedef struct mjSDF_ mjSDF; +// get sdf from geom id +MJAPI const mjpPlugin* mjc_getSDF(const mjModel* m, int id); // signed distance function MJAPI mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]); diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 47918d48..3e625f49 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -378,50 +378,6 @@ void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* -// multiply Jacobian by vector, for one island -// flg_resunc and flg_vecunc denote whether res/vec are uncompressed -void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc) { - // no island, call regular function - if (island < 0) { - mj_mulJacVec(m, d, res, vec); - return; - } - - // sizes - int vecnnz = d->island_dofnum[island]; - int resnnz = d->island_efcnum[island]; - - // indices - int* vecind = d->island_dofind + d->island_dofadr[island]; - int* resind = d->island_efcind + d->island_efcadr[island]; - - // sparse Jacobian - if (mj_isSparse(m)) { - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int Jnnz = d->efc_J_rownnz[row]; - int Jrowadr = d->efc_J_rowadr[row]; - int* Jind = d->efc_J_colind + Jrowadr; - mjtNum* J = d->efc_J + Jrowadr; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse2(J, vec, Jnnz, Jind, vecnnz, vecind, flg_vecunc); - } - } - - // dense Jacobian - else { - int nv = m->nv; - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse(vec, d->efc_J + nv*row, vecnnz, vecind, flg_vecunc); - } - } -} - - - // multiply JacobianT by vector void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { // exit if no constraints @@ -443,50 +399,6 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* -// multiply Jacobian transpose by vector, for one island -// flg_resunc and flg_vecunc denote whether res/vec are uncompressed -void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc) { - // no island, call regular function - if (island < 0) { - mj_mulJacTVec(m, d, res, vec); - return; - } - - // sizes - int vecnnz = d->island_efcnum[island]; - int resnnz = d->island_dofnum[island]; - - // indices - int* vecind = d->island_efcind + d->island_efcadr[island]; - int* resind = d->island_dofind + d->island_dofadr[island]; - - // sparse Jacobian - if (mj_isSparse(m)) { - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int JTnnz = d->efc_JT_rownnz[row]; - int JTrowadr = d->efc_JT_rowadr[row]; - int* JTind = d->efc_JT_colind + JTrowadr; - mjtNum* JT = d->efc_JT + JTrowadr; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse2(JT, vec, JTnnz, JTind, vecnnz, vecind, flg_vecunc); - } - } - - // dense Jacobian - else { - int nefc = d->nefc; - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse(vec, d->efc_JT + nefc*row, vecnnz, vecind, flg_vecunc); - } - } -} - - - //--------------------- instantiate constraints by type -------------------------------------------- // equality constraints @@ -2084,7 +1996,7 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { if (mj_isSparse(m)) { // transpose mju_transposeSparse(d->efc_JT, d->efc_J, d->nefc, m->nv, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, + d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); @@ -2098,14 +2010,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { __msan_allocated_memory(d->efc_J_rowsuper, d->nefc); #endif // MEMORY_SANITIZER #endif // mjUSEAVX - - // supernodes of JT - mju_superSparse(m->nv, d->efc_JT_rowsuper, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind); - } else { - if (mjENABLED(mjENBL_ISLAND)) { - mju_transpose(d->efc_JT, d->efc_J, d->nefc, m->nv); - } } // compute diagApprox @@ -2167,11 +2071,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { continue; } - // traverse row j of M, marking new unique nonzeros - int nnzM = d->M_rownnz[j]; - int adrM = d->M_rowadr[j]; - for (int k=0; k < nnzM; k++) { - int c = d->M_colind[adrM + k]; + // traverse row j of C, marking new unique nonzeros + int nnzC = d->M_rownnz[j]; + int adrC = d->M_rowadr[j]; + for (int k=0; k < nnzC; k++) { + int c = d->M_colind[adrC + k]; if (marker[c] != r) { marker[c] = r; nnz++; @@ -2251,10 +2155,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { continue; } int j = B_colind[i]; - int adrM = d->M_rowadr[j]; - mju_addToSclSparseInc(B + adrB, d->qLD + adrM, + int adrC = d->M_rowadr[j]; + mju_addToSclSparseInc(B + adrB, d->qLD + adrC, nnzB, B_colind + adrB, - d->M_rownnz[j]-1, d->M_colind + adrM, -b); + d->M_rownnz[j]-1, d->M_colind + adrC, -b); } // B(r,:) <- sqrt(inv(D)) * B(r,:) @@ -2274,7 +2178,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { int* BT_colind = mjSTACKALLOC(d, nB, int); mjtNum* BT = mjSTACKALLOC(d, nB, mjtNum); mju_transposeSparse(BT, B, nefc, nv, - BT_rownnz, BT_rowadr, BT_colind, + BT_rownnz, BT_rowadr, BT_colind, NULL, B_rownnz, B_rowadr, B_colind); // allocate AR row nonzeros and addresses on arena @@ -2377,25 +2281,17 @@ void mj_referenceConstraint(const mjModel* m, mjData* d) { //---------------------------- update constraint state --------------------------------------------- -// compute efc_state, efc_force, qfrc_constraint, optionally restricted to one island -// island < 0: update all d->nefc constraints -// island >= 0: update only d->island_efcnum[island] constraints -// jar = Jac*qacc-aref is restricted to the island, in the above sense -// optional: cost(qacc) = shat(jar); cone Hessians -void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, - mjtNum cost[1], int flg_coneHessian, int island) { - int ne = d->ne, nf = d->nf; - const mjtNum *D = d->efc_D, *R = d->efc_R, *floss = d->efc_frictionloss; - mjtNum* force = d->efc_force; +// compute efc_state, efc_force +// optional: cost(qacc) = s_hat(jar); cone Hessians +void mj_constraintUpdate_impl(int ne, int nf, int nefc, + const mjtNum* D, const mjtNum* R, const mjtNum* floss, + const mjtNum* jar, const int* type, const int* id, + mjContact* contact, int* state, mjtNum* force, mjtNum cost[1], + int flg_coneHessian) { mjtNum s = 0; - int nefc = island < 0 ? d->nefc : d->island_efcnum[island]; - int* efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island]; - - // no constraints: clear qfrc_constraint and cost, return + // no constraints: clear cost, return if (!nefc) { - // can only occur for island == -1 - mju_zero(d->qfrc_constraint, m->nv); if (cost) { *cost = 0; } @@ -2403,55 +2299,49 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, } // compute unconstrained efc_force - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - force[i] = -D[i]*jar[c]; + for (int i=0; i < nefc; i++) { + force[i] = -D[i]*jar[i]; } // update constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // ==== equality if (i < ne) { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; continue; } // ==== friction if (i < ne + nf) { // linear negative - if (jar[c] <= -R[i]*floss[i]) { + if (jar[i] <= -R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[c]; + s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[i]; } force[i] = floss[i]; - - d->efc_state[i] = mjCNSTRSTATE_LINEARNEG; + state[i] = mjCNSTRSTATE_LINEARNEG; } // linear positive - else if (jar[c] >= R[i]*floss[i]) { + else if (jar[i] >= R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[c]; + s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[i]; } force[i] = -floss[i]; - - d->efc_state[i] = mjCNSTRSTATE_LINEARPOS; + state[i] = mjCNSTRSTATE_LINEARPOS; } // quadratic else { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } continue; } @@ -2459,36 +2349,35 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // ==== contact // non-negative constraint - if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) { + if (type[i] != mjCNSTR_CONTACT_ELLIPTIC) { // constraint is satisfied: no cost - if (jar[c] >= 0) { + if (jar[i] >= 0) { force[i] = 0; - d->efc_state[i] = mjCNSTRSTATE_SATISFIED; + state[i] = mjCNSTRSTATE_SATISFIED; } // quadratic else { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } } // contact with elliptic cone else { // get contact - mjContact* con = d->contact + d->efc_id[i]; + mjContact* con = contact + id[i]; mjtNum mu = con->mu, *friction = con->friction; int dim = con->dim; // map to regular dual cone space mjtNum U[6]; - U[0] = jar[c]*mu; + U[0] = jar[i]*mu; for (int j=1; j < dim; j++) { - U[j] = jar[c+j]*friction[j-1]; + U[j] = jar[i+j]*friction[j-1]; } // decompose into normal and tangent @@ -2498,19 +2387,17 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // top zone if (N >= mu*T || (T <= 0 && N >= 0)) { mju_zero(force+i, dim); - - d->efc_state[i] = mjCNSTRSTATE_SATISFIED; + state[i] = mjCNSTRSTATE_SATISFIED; } // bottom zone else if (mu*N+T <= 0 || (T <= 0 && N < 0)) { if (cost) { for (int j=0; j < dim; j++) { - s += 0.5*D[i+j]*jar[c+j]*jar[c+j]; + s += 0.5*D[i+j]*jar[i+j]*jar[i+j]; } } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } // middle zone @@ -2530,12 +2417,12 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, } // set state - d->efc_state[i] = mjCNSTRSTATE_CONE; + state[i] = mjCNSTRSTATE_CONE; // cone Hessian if (flg_coneHessian) { // get Hessian pointer - mjtNum* H = d->contact[d->efc_id[i]].H; + mjtNum* H = contact[id[i]].H; // set first row: (1, -mu/T * U) mjtNum scl = -mu/T; @@ -2546,10 +2433,11 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // set upper block: mu*N/T^3 * U*U' scl = mu*N/(T*T*T); - for (int k=1; k < dim; k++) + for (int k=1; k < dim; k++) { for (int j=k; j < dim; j++) { H[k*dim+j] = scl*U[j]*U[k]; } + } // add to diagonal: (mu^2 - mu*N/T) * I scl = mu*mu - mu*N/T; @@ -2576,19 +2464,14 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // replicate state in all cone dimensions for (int j=1; j < dim; j++) { - d->efc_state[i+j] = d->efc_state[i]; + state[i+j] = state[i]; } // advance to end of contact - c += (dim-1); + i += (dim-1); } } - // compute qfrc_constraint - int flg_vecunc = 1; - int flg_resunc = 1; - mj_mulJacTVec_island(m, d, d->qfrc_constraint, d->efc_force, island, flg_vecunc, flg_resunc); - // assign cost if (cost) { *cost = s; @@ -2598,8 +2481,11 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // compute efc_state, efc_force, qfrc_constraint -// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians +// optional: cost(qacc) = s_hat(jar) where jar = Jac*qacc-aref; cone Hessians void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian) { - mj_constraintUpdate_island(m, d, jar, cost, flg_coneHessian, -1); + mj_constraintUpdate_impl(d->ne, d->nf, d->nefc, d->efc_D, d->efc_R, d->efc_frictionloss, + jar, d->efc_type, d->efc_id, d->contact, d->efc_state, d->efc_force, + cost, flg_coneHessian); + mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force); } diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index 05c5fd57..a0a7c6ca 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -24,6 +24,7 @@ extern "C" { #endif + //-------------------------- Jacobian-related ------------------------------------------------------ // determine type of friction cone @@ -38,16 +39,9 @@ MJAPI int mj_isDual(const mjModel* m); // multiply Jacobian by vector MJAPI void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply Jacobian by vector, for one island -MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc); - // multiply JacobianT by vector MJAPI void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply JacobianT by vector, for one island -MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc); //-------------------------- utility functions ----------------------------------------------------- @@ -90,6 +84,7 @@ void mj_diagApprox(const mjModel* m, mjData* d); // compute efc_R, efc_D, efc_KDIP, adjust diagApprox void mj_makeImpedance(const mjModel* m, mjData* d); + //---------------------------- top-level API for constraint construction --------------------------- // main driver: call all functions above @@ -101,14 +96,19 @@ MJAPI void mj_projectConstraint(const mjModel* m, mjData* d); // compute efc_vel, efc_aref MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d); +// compute efc_state, efc_force +// optional: cost(qacc) = s_hat(jar); cone Hessians +MJAPI void mj_constraintUpdate_impl(int ne, int nf, int nefc, + const mjtNum* D, const mjtNum* R, const mjtNum* floss, + const mjtNum* jar, const int* type, const int* id, + mjContact* contact, int* state, mjtNum* force, mjtNum cost[1], + int flg_coneHessian); + // compute efc_state, efc_force, qfrc_constraint -// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians +// optional: cost(qacc) = s_hat(jar) where jar = Jac*qacc-aref; cone Hessians MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian); -// compute efc_state, efc_force, qfrc_constraint for one island -MJAPI void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, - mjtNum cost[1], int flg_coneHessian, int island); #ifdef __cplusplus } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 9b453c3e..d4deb5b3 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -861,6 +861,126 @@ void mj_tendon(const mjModel* m, mjData* d) { +// compute time derivative of dense tendon Jacobian for one tendon +void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) { + int nv = m->nv; + + // tendon id is invalid: return + if (id < 0 || id >= m->ntendon) { + return; + } + + // clear output + mju_zero(Jdot, nv); + + // fixed tendon has zero Jdot: return + int adr = m->tendon_adr[id]; + if (m->wrap_type[adr] == mjWRAP_JOINT) { + return; + } + + // allocate stack arrays + mj_markStack(d); + mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); + mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); + mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); + mjtNum* tmp = mjSTACKALLOC(d, nv, mjtNum); + + // process spatial tendon + mjtNum divisor = 1; + int wraptype, j = 0; + int num = m->tendon_num[id]; + while (j < num-1) { + // get 1st and 2nd object + int type0 = m->wrap_type[adr+j+0]; + int type1 = m->wrap_type[adr+j+1]; + int id0 = m->wrap_objid[adr+j+0]; + int id1 = m->wrap_objid[adr+j+1]; + + // pulley + if (type0 == mjWRAP_PULLEY || type1 == mjWRAP_PULLEY) { + // get divisor, insert obj=-2 + if (type0 == mjWRAP_PULLEY) { + divisor = m->wrap_prm[adr+j]; + } + + // move to next + j++; + continue; + } + + // init sequence; assume it starts with site + mjtNum wpnt[6]; + mju_copy3(wpnt, d->site_xpos+3*id0); + mjtNum vel[6]; + mj_objectVelocity(m, d, mjOBJ_SITE, id0, vel, /*flg_local=*/0); + mjtNum wvel[6] = {vel[3], vel[4], vel[5], 0, 0, 0}; + int wbody[2]; + wbody[0] = m->site_bodyid[id0]; + + // second object is geom: process site-geom-site + if (type1 == mjWRAP_SPHERE || type1 == mjWRAP_CYLINDER) { + // TODO(tassa) support geom wrapping (requires derivatives of mju_wrap) + mjERROR("geom wrapping not supported"); + } else { + wraptype = mjWRAP_NONE; + } + + // complete sequence + wbody[1] = m->site_bodyid[id1]; + mju_copy3(wpnt+3, d->site_xpos+3*id1); + mj_objectVelocity(m, d, mjOBJ_SITE, id1, vel, /*flg_local=*/0); + mju_copy3(wvel+3, vel+3); + + // accumulate moments if consecutive points are in different bodies + if (wbody[0] != wbody[1]) { + // dpnt = 3D position difference, normalize + mjtNum dpnt[3]; + mju_sub3(dpnt, wpnt+3, wpnt); + mjtNum norm = mju_normalize3(dpnt); + + // dvel = d / dt (dpnt) + mjtNum dvel[3]; + mju_sub3(dvel, wvel+3, wvel); + mjtNum dot = mju_dot3(dpnt, dvel); + mju_addToScl3(dvel, dpnt, -dot); + mju_scl3(dvel, dvel, norm > mjMINVAL ? 1/norm : 0); + + // TODO(tassa ) write sparse branch, requires mj_jacDotSparse + // if (mj_isSparse(m)) { ... } + + // get endpoint JacobianDots, subtract + mj_jacDot(m, d, jac1, 0, wpnt, wbody[0]); + mj_jacDot(m, d, jac2, 0, wpnt+3, wbody[1]); + mju_sub(jacdif, jac2, jac1, 3*nv); + + // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt + mju_mulMatTVec(tmp, jacdif, dpnt, 3, nv); + + // add to existing + mju_addToScl(Jdot, tmp, 1/divisor, nv); + + // get endpoint Jacobians, subtract + mj_jac(m, d, jac1, 0, wpnt, wbody[0]); + mj_jac(m, d, jac2, 0, wpnt+3, wbody[1]); + mju_sub(jacdif, jac2, jac1, 3*nv); + + // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt) + mju_mulMatTVec(tmp, jacdif, dvel, 3, nv); + + // add to existing + mju_addToScl(Jdot, tmp, 1/divisor, nv); + } + + // advance + j += (wraptype != mjWRAP_NONE ? 2 : 1); + } + + mj_freeStack(d); +} + + + // compute actuator/transmission lengths and moments void mj_transmission(const mjModel* m, mjData* d) { int nv = m->nv, nu = m->nu; @@ -1349,9 +1469,63 @@ void mj_transmission(const mjModel* m, mjData* d) { //-------------------------- inertia --------------------------------------------------------------- +// add tendon armature to qM +void mj_tendonArmature(const mjModel* m, mjData* d) { + int nv = m->nv, ntendon = m->ntendon, issparse = mj_isSparse(m); + + for (int k=0; k < ntendon; k++) { + mjtNum armature = m->tendon_armature[k]; + + if (!armature) { + continue; + } + + // dense + if (!issparse) { + mjtNum* ten_J = d->ten_J + nv*k; + for (int i=0; i < m->nv; i++) { + int Madr = m->dof_Madr[i]; + for (int j = i; j >= 0; j = m->dof_parentid[j]) { + d->qM[Madr++] += armature * ten_J[j] * ten_J[i]; + } + } + } + + // sparse + else { + // get sparse info for tendon k + int rowadr = d->ten_J_rowadr[k]; + int rownnz = d->ten_J_rownnz[k]; + const int* colind = d->ten_J_colind + rowadr; + mjtNum* ten_J = d->ten_J + rowadr; + + // iterate forward on nonzero rows i + for (int adr_i=0; adr_i < rownnz; adr_i++) { + int i = colind[adr_i]; + int Madr = m->dof_Madr[i]; + int adr_j = rownnz - 1; + + // iterate backward on ancestors of i, find matching column j + for (int j = i; j >= 0; j = m->dof_parentid[j]) { + // reduce adr_j until column index is no bigger than j + while (colind[adr_j] > j && adr_j >= 0) { + adr_j--; + } + + // found match, update qM + if (colind[adr_j] == j) { + d->qM[Madr++] += armature * ten_J[adr_j] * ten_J[adr_i]; + } + } + } + } + } +} + + + // composite rigid body inertia algorithm void mj_crb(const mjModel* m, mjData* d) { - TM_START; mjtNum buf[6]; mjtNum* crb = d->crb; int last_body = m->nbody - 1, nv = m->nv; @@ -1397,6 +1571,15 @@ void mj_crb(const mjModel* m, mjData* d) { d->qM[Madr_ij++] += mju_dot(d->cdof+6*j, buf, 6); } } +} + + + +void mj_makeM(const mjModel* m, mjData* d) { + TM_START; + mj_crb(m, d); + mj_tendonArmature(m, d); + mju_gather(d->M, d->qM, d->mapM2M, m->nC); TM_END(mjTIMER_POS_INERTIA); } @@ -1469,11 +1652,8 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd void mj_factorM(const mjModel* m, mjData* d) { TM_START; - int nM = m->nM; - for (int i=0; i < nM; i++) { - d->qLD[i] = d->qM[d->mapM2M[i]]; - } - mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + mju_copy(d->qLD, d->M, m->nC); + mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind); TM_ADD(mjTIMER_POS_INERTIA); } @@ -1481,32 +1661,25 @@ void mj_factorM(const mjModel* m, mjData* d) { // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv, - const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) { + const int* rownnz, const int* rowadr, const int* colind) { // backward loop over rows for (int k=nv-1; k >= 0; k--) { // get row k's address, diagonal index, inverse diagonal value - int rowadr_k = rowadr[k]; - int diag_k = rowadr_k + rownnz[k] - 1; - mjtNum invD = 1 / mat[diag_k]; + int start = rowadr[k]; + int diag = rownnz[k] - 1; + int end = start + diag; + mjtNum invD = 1 / mat[end]; if (diaginv) diaginv[k] = invD; - // skip if simple - if (diagnum[k]) { - continue; - } - - // update triangle above row k, inclusive - for (int adr=diag_k - 1; adr >= rowadr_k; adr--) { - // tmp = L(k, i) / L(k, k) - mjtNum tmp = mat[adr] * invD; - + // update triangle above row k + for (int adr=end - 1; adr >= start; adr--) { // update row i < k: L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k) int i = colind[adr]; - mju_addToScl(mat + rowadr[i], mat + rowadr_k, -tmp, rownnz[i]); - - // update ith element of row k: L(k, i) /= L(k, k) - mat[adr] = tmp; + mju_addToScl(mat + rowadr[i], mat + start, -mat[adr] * invD, rownnz[i]); } + + // update row k: L(k, :) /= L(k, k) + mju_scl(mat + start, mat + start, invD, diag); } } @@ -1626,12 +1799,12 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n, // in-place sparse backsubstitution: x = inv(L'*D*L)*x -void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n, - const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) { +void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, + const int* rownnz, const int* rowadr, const int* colind) { // x <- L^-T x for (int i=nv-1; i > 0; i--) { // skip diagonal rows - if (diagnum[i]) { + if (rownnz[i] == 1) { continue; } @@ -1642,7 +1815,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=start; adr < end; adr++) { - x[colind[adr]] -= qLDs[adr] * x_i; + x[colind[adr]] -= qLD[adr] * x_i; } } } @@ -1655,7 +1828,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, mjtNum x_i; if ((x_i = x[i+offset])) { for (int adr=start; adr < end; adr++) { - x[offset + colind[adr]] -= qLDs[adr] * x_i; + x[offset + colind[adr]] -= qLD[adr] * x_i; } } } @@ -1682,8 +1855,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, // x <- L^-1 x for (int i=1; i < nv; i++) { // skip diagonal rows - if (diagnum[i]) { - i += diagnum[i] - 1; // iterating forward: skip ahead, adjust i + if (rownnz[i] == 1) { continue; } @@ -1693,13 +1865,13 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, // one vector if (n == 1) { - x[i] -= mju_dotSparse(qLDs+adr, x, d, colind+adr, /*flg_unc1=*/0); + x[i] -= mju_dotSparse(qLD+adr, x, d, colind+adr); } // multiple vectors else { for (int offset=0; offset < n*nv; offset+=nv) { - x[i+offset] -= mju_dotSparse(qLDs+adr, x+offset, d, colind+adr, /*flg_unc1=*/0); + x[i+offset] -= mju_dotSparse(qLD+adr, x+offset, d, colind+adr); } } } @@ -1715,66 +1887,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { mju_copy(x, y, n*m->nv); } mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); -} - - -// in-place sparse backsubstitution for one island: x = inv(L'*D*L)*x -// L is in lower triangle of qLD; D is on diagonal of qLD -void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) { - // if no islands, call mj_solveLD - const mjtNum* qLD = d->qLD; - const mjtNum* qLDiagInv = d->qLDiagInv; - if (island < 0) { - mj_solveLD(x, qLD, qLDiagInv, m->nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); - return; - } - - // local copies of key variables - const int* rownnz = d->M_rownnz; - const int* rowadr = d->M_rowadr; - const int* colind = d->M_colind; - const int* diagnum = m->dof_simplenum; - - // local constants: island specific - int ndof = d->island_dofnum[island]; - const int* dofind = d->island_dofind + d->island_dofadr[island]; - const int* islandind = d->dof_islandind; - - // x <- inv(L') * x; skip simple, exploit sparsity of input vector - for (int k=ndof-1; k >= 0; k--) { - int i = dofind[k]; - mjtNum x_k; - if (!diagnum[i] && (x_k = x[k])) { - int start = rowadr[i]; - int end = start + rownnz[i] - 1; - for (int adr=end-1; adr >= start; adr--) { - x[islandind[colind[adr]]] -= qLD[adr] * x_k; - } - } - } - - // x <- inv(D) * x - for (int k=ndof-1; k >= 0; k--) { - x[k] *= qLDiagInv[dofind[k]]; // x(i) /= L(i,i) - } - - // x <- inv(L) * x; skip simple - for (int k=0; k < ndof; k++) { - int i = dofind[k]; - - // skip diagonal rows - if (diagnum[i]) { - continue; - } - - int start = rowadr[i]; - int end = start + rownnz[i] - 1; - for (int adr=end-1; adr >= start; adr--) { - x[k] -= x[islandind[colind[adr]]] * qLD[adr]; - } - } + d->M_rownnz, d->M_rowadr, d->M_colind); } @@ -2200,3 +2313,53 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { mju_addTo(d->cfrc_int+6*m->body_parentid[j], d->cfrc_int+6*j, 6); } } + + + +// add bias force due to tendon armature +void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { + int ntendon = m->ntendon, nv = m->nv, issparse = mj_isSparse(m); + mjtNum* ten_Jdot = NULL; + mj_markStack(d); + + // add bias term due to tendon armature + for (int i=0; i < ntendon; i++) { + mjtNum armature = m->tendon_armature[i]; + + // no armature: skip + if (!armature) { + continue; + } + + // allocate if required + if (!ten_Jdot) { + ten_Jdot = mjSTACKALLOC(d, nv, mjtNum); + } + + // get dense d/dt(tendon Jacobian) for tendon i + mj_tendonDot(m, d, i, ten_Jdot); + + // add bias term: qfrc += ten_J * armature * dot(ten_Jdot, qvel) + mjtNum coef = armature * mju_dot(ten_Jdot, d->qvel, nv); + + if (coef) { + // dense + if (!issparse) { + mju_addToScl(qfrc, d->ten_J + nv*i, coef, nv); + } + + // sparse + else { + int nnz = d->ten_J_rownnz[i]; + int adr = d->ten_J_rowadr[i]; + const int* colind = d->ten_J_colind + adr; + const mjtNum* ten_J = d->ten_J + adr; + for (int j=0; j < nnz; j++) { + qfrc[colind[j]] += coef * ten_J[j]; + } + } + } + } + + mj_freeStack(d); +} diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 7f1b9fb3..6a672b3a 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -39,6 +39,9 @@ MJAPI void mj_flex(const mjModel* m, mjData* d); // compute tendon lengths, velocities and moment arms MJAPI void mj_tendon(const mjModel* m, mjData* d); +// compute time derivative of dense tendon Jacobian for one tendon +MJAPI void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot); + // compute actuator transmission lengths and moments MJAPI void mj_transmission(const mjModel* m, mjData* d); @@ -48,13 +51,19 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d); // composite rigid body inertia algorithm MJAPI void mj_crb(const mjModel* m, mjData* d); +// add tendon armature to qM +MJAPI void mj_tendonArmature(const mjModel* m, mjData* d); + +// make inertia matrix +MJAPI void mj_makeM(const mjModel* m, mjData* d); + // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd (legacy implementation) MJAPI void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv); // sparse L'*D*L factorizaton of inertia-like matrix MJAPI void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv, - const int* rownnz, const int* rowadr, const int* diagnum, const int* colind); + const int* rownnz, const int* rowadr, const int* colind); // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd MJAPI void mj_factorM(const mjModel* m, mjData* d); @@ -65,15 +74,12 @@ MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n, // in-place sparse backsubstitution: x = inv(L'*D*L)*x // handle n vectors at once -MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n, - const int* rownnz, const int* rowadr, const int* diagnum, const int* colind); +MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, + const int* rownnz, const int* rowadr, const int* colind); // sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); -// sparse backsubstitution for one island: x = inv(L'*D*L)*x, use factorization in d -MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island); - // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, const mjtNum* sqrtInvD, int n); @@ -96,6 +102,12 @@ MJAPI void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result); // RNE with complete data: compute cacc, cfrc_ext, cfrc_int MJAPI void mj_rnePostConstraint(const mjModel* m, mjData* d); + +//-------------------------- tendon bias ----------------------------------------------------------- + +// add bias force due to tendon armature +MJAPI void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc); + #ifdef __cplusplus } #endif diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index a61d48a0..16d33cae 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -114,15 +114,15 @@ typedef struct mjFwdPositionArgs_ mjFwdPositionArgs; // wrapper for mj_crb and mj_factorM void* mj_inertialThreaded(void* args) { mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args; - mj_crb(forward_args->m, forward_args->d); // timed internally (POS_INERTIA) - mj_factorM(forward_args->m, forward_args->d); // timed internally (POS_INERTIA) + mj_makeM(forward_args->m, forward_args->d); + mj_factorM(forward_args->m, forward_args->d); return NULL; } // wrapper for mj_collision void* mj_collisionThreaded(void* args) { mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args; - mj_collision(forward_args->m, forward_args->d); // timed internally (POS_COLLISION) + mj_collision(forward_args->m, forward_args->d); return NULL; } @@ -142,9 +142,12 @@ void mj_fwdPosition(const mjModel* m, mjData* d) { // no threadpool: inertia and collision on main thread if (!d->threadpool) { - mj_crb(m, d); // timed internally (POS_INERTIA) - mj_factorM(m, d); // timed internally (POS_INERTIA) - mj_collision(m, d); // timed internally (POS_COLLISION) + // inertia, timed internally (POS_INERTIA) + mj_makeM(m, d); + mj_factorM(m, d); + + // collision, timed internally (POS_COLLISION) + mj_collision(m, d); } // have threadpool: inertia and collision on separate threads @@ -222,6 +225,9 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) { // compute qfrc_bias with abbreviated RNE (without acceleration) mj_rne(m, d, 0, d->qfrc_bias); + // add bias force due to tendon armature + mj_tendonBias(m, d, d->qfrc_bias); + TM_END(mjTIMER_VELOCITY); } @@ -271,7 +277,7 @@ static void clampVec(mjtNum* vec, const mjtNum* range, const mjtByte* limited, i // (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot) void mj_fwdActuation(const mjModel* m, mjData* d) { TM_START; - int nv = m->nv, nu = m->nu; + int nv = m->nv, nu = m->nu, ntendon = m->ntendon; mjtNum gain, bias, tau; mjtNum *prm, *force = d->actuator_force; @@ -284,6 +290,9 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { return; } + // any tendon transmission targets with force limits + int tendon_frclimited = 0; + // local, clamped copy of ctrl mj_markStack(d); mjtNum *ctrl = mjSTACKALLOC(d, nu, mjtNum); @@ -379,6 +388,11 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { continue; } + // check for tendon transmission with force limits + if (ntendon && !tendon_frclimited && m->actuator_trntype[i] == mjTRN_TENDON) { + tendon_frclimited = m->tendon_actfrclimited[m->actuator_trnid[2*i]]; + } + // extract gain info prm = m->actuator_gainprm + mjNGAIN*i; @@ -474,6 +488,38 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { } } + // clamp tendon total actuator force + if (tendon_frclimited) { + // compute total force for each tendon + mjtNum* tendon_total_force = mjSTACKALLOC(d, ntendon, mjtNum); + mju_zero(tendon_total_force, ntendon); + for (int i=0; i < nu; i++) { + if (m->actuator_trntype[i] == mjTRN_TENDON) { + int tendon_id = m->actuator_trnid[2*i]; + if (m->tendon_actfrclimited[tendon_id]) { + tendon_total_force[tendon_id] += force[i]; + } + } + } + + // scale tendon actuator forces if limited and outside range + for (int i=0; i < nu; i++) { + if (m->actuator_trntype[i] != mjTRN_TENDON) { + continue; + } + int tendon_id = m->actuator_trnid[2*i]; + mjtNum tendon_force = tendon_total_force[tendon_id]; + if (m->tendon_actfrclimited[tendon_id] && tendon_force) { + const mjtNum* range = m->tendon_actfrcrange + 2 * tendon_id; + if (tendon_force < range[0]) { + force[i] *= range[0] / tendon_force; + } else if (tendon_force > range[1]) { + force[i] *= range[1] / tendon_force; + } + } + } + } + // clamp actuator_force clampVec(force, m->actuator_forcerange, m->actuator_forcelimited, nu, NULL); @@ -584,6 +630,15 @@ static void warmstart(const mjModel* m, mjData* d) { } } + // have island structure: unconstrained qacc = qacc_smooth + if (d->nisland > 0) { + // loop over unconstrained dofs in map_idof2dof[nidof, nv) + for (int i=d->nidof; i < nv; i++) { + int dof = d->map_idof2dof[i]; + d->qacc[dof] = d->qacc_smooth[dof]; + } + } + mj_freeStack(d); } @@ -604,32 +659,35 @@ struct mjSolIslandArgs_ { }; typedef struct mjSolIslandArgs_ mjSolIslandArgs; -// extract arguments, pass to solver -void* mj_solCG_island_wrapper(void* args) { +// extract arguments, pass to CG solver +static void* CG_wrapper(void* args) { mjSolIslandArgs* solargs = (mjSolIslandArgs*) args; mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations); return NULL; } - - +// extract arguments, pass to Newton solver +static void* Newton_wrapper(void* args) { + mjSolIslandArgs* solargs = (mjSolIslandArgs*) args; + mj_solNewton_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations); + return NULL; +} // CG solver, multi-threaded over islands -void mj_solCG_island_multithreaded(const mjModel* m, mjData* d) { +static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) { mj_markStack(d); // allocate array of arguments to be passed to threads - mjSolIslandArgs* sol_cg_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs); + mjSolIslandArgs* sol_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs); mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask); - for (int island = 0; island < d->nisland; ++island) { - sol_cg_island_args[island].m = m; - sol_cg_island_args[island].d = d; - sol_cg_island_args[island].island = island; + sol_island_args[island].m = m; + sol_island_args[island].d = d; + sol_island_args[island].island = island; mju_defaultTask(&tasks[island]); - tasks[island].func = mj_solCG_island_wrapper; - tasks[island].args = &sol_cg_island_args[island]; + tasks[island].func = flg_Newton ? Newton_wrapper : CG_wrapper; + tasks[island].args = &sol_island_args[island]; mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]); } @@ -668,24 +726,43 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) { mju_zeroInt(d->solver_niter, mjNISLAND); // check if islands are supported - int islands_supported = mjENABLED(mjENBL_ISLAND) && - d->nisland > 0 && - m->opt.solver == mjSOL_CG && - m->opt.noslip_iterations == 0; + int islands_supported = mjENABLED(mjENBL_ISLAND) && + nisland > 0 && + m->opt.noslip_iterations == 0 && + (m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON); // run solver over constraint islands if (islands_supported) { - // no threadpool, loop over islands + int nidof = d->nidof; + + // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island) + mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof); + mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof); + mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof); + mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof); + mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc); + mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc); + + // solve per island, with or without threads if (!d->threadpool) { + // no threadpool, loop over islands for (int island=0; island < nisland; island++) { - mj_solCG_island(m, d, island, m->opt.iterations); + if (m->opt.solver == mjSOL_NEWTON) { + mj_solNewton_island(m, d, island, m->opt.iterations); + } else { + mj_solCG_island(m, d, island, m->opt.iterations); + } } + } else { + // have threadpool, solve using threads + solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON); } - else { - // solve using threads - mj_solCG_island_multithreaded(m, d); - } - d->solver_nisland = nisland; + + + // copy back solver outputs (scatter dofs since ni <= nv) + mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof); + mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof); + mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc); } // run solver over all constraints @@ -706,9 +783,6 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) { default: mjERROR("unknown solver type %d", m->opt.solver); } - - // one (monolithic) island - d->solver_nisland = 1; } // save result for next step warmstart @@ -770,7 +844,7 @@ static void mj_advance(const mjModel* m, mjData* d, // Euler integrator, semi-implicit in velocity, possibly skipping factorisation void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nM = m->nM; + int nv = m->nv, nC = m->nC; mj_markStack(d); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum); @@ -795,22 +869,20 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { else { if (!skipfactor) { // qH = M + h*diag(B) - for (int i=0; i < nM; i++) { - d->qH[i] = d->qM[d->mapM2M[i]]; - } + mju_copy(d->qH, d->M, nC); for (int i=0; i < nv; i++) { d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i]; } // factorize in-place - mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind); } // solve mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); mju_copy(qacc, qfrc, m->nv); mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); } // advance state and time @@ -939,7 +1011,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { // fully implicit in velocity, possibly skipping factorization void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nM = m->nM, nD = m->nD; + int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC; mj_markStack(d); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); @@ -954,10 +1026,8 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // compute analytical derivative qDeriv mjd_smooth_vel(m, d, /* flg_bias = */ 1); - // set qLU = qM - for (int i=0; i < nD; i++) { - d->qLU[i] = d->qM[d->mapM2D[i]]; - } + // gather qLU <- qM (lower to full) + mju_gather(d->qLU, d->qM, d->mapM2D, nD); // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); @@ -977,28 +1047,24 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // compute analytical derivative qDeriv; skip rne derivative mjd_smooth_vel(m, d, /* flg_bias = */ 0); - // modified mass matrix MhB = qDeriv[Lower] + // modified mass matrix: gather MhB <- qDeriv (full to lower) mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum); - for (int i=0; i < nM; i++) { - MhB[i] = d->qDeriv[d->mapD2M[i]]; - } + mju_gather(MhB, d->qDeriv, d->mapD2M, nM); // set MhB = M - dt*qDeriv mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); - // copy into qH - for (int i=0; i < nM; i++) { - d->qH[i] = MhB[d->mapM2M[i]]; - } + // gather qH <- MhB (legacy to CSR) + mju_gather(d->qH, MhB, d->mapM2M, nC); // factorize in-place - mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind); } // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc mju_copy(qacc, qfrc, nv); mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); } else { mjERROR("integrator must be implicit or implicitfast"); diff --git a/src/engine/engine_inverse.c b/src/engine/engine_inverse.c index f372947a..de8b9a52 100644 --- a/src/engine/engine_inverse.c +++ b/src/engine/engine_inverse.c @@ -31,6 +31,7 @@ #include "engine/engine_support.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" +#include "engine/engine_util_misc.h" #include "engine/engine_util_sparse.h" // position-dependent computations @@ -45,8 +46,8 @@ void mj_invPosition(const mjModel* m, mjData* d) { mj_tendon(m, d); TM_END(mjTIMER_POS_KINEMATICS); - mj_crb(m, d); // timed internally (POS_INERTIA) - mj_factorM(m, d); // timed internally (POS_INERTIA) + mj_makeM(m, d); // timed internally (POS_INERTIA) + mj_factorM(m, d); // timed internally (POS_INERTIA) mj_collision(m, d); // timed internally (POS_COLLISION) @@ -114,10 +115,8 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) { // compute qDeriv mjd_smooth_vel(m, d, /* flg_bias = */ 1); - // set qLU = qM - for (int i=0; i < nD; i++) { - d->qLU[i] = d->qM[d->mapM2D[i]]; - } + // gather qLU <- qM (lower to full) + mju_gather(d->qLU, d->qM, d->mapM2D, nD); // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index f925410c..9a9ad366 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1139,9 +1139,6 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src, const int* rownnz; const int* rowadr; if (reduced && !upper) { - rownnz = d->C_rownnz; - rowadr = d->C_rowadr; - } else if (!reduced && !upper) { rownnz = d->M_rownnz; rowadr = d->M_rowadr; } else if (!reduced && upper) { @@ -1251,24 +1248,13 @@ static void makeDofDofmaps(const mjModel* m, mjData* d) { } // make mapM2C - for (int i=0; i < nC; i++) d->mapM2C[i] = -1; - copyM2Sparse(m, d, d->mapM2C, M, /*reduced=*/1, /*upper=*/0); + for (int i=0; i < nC; i++) d->mapM2M[i] = -1; + copyM2Sparse(m, d, d->mapM2M, M, /*reduced=*/1, /*upper=*/0); // check that all indices are filled in for (int i=0; i < nC; i++) { - if (d->mapM2C[i] < 0) { - mjERROR("unassigned index in mapM2C"); - } - } - - // make mapM2M - for (int i=0; i < nM; i++) d->mapM2M[i] = -1; - copyM2Sparse(m, d, d->mapM2M, M, /*reduced=*/0, /*upper=*/0); - - // check that all indices are filled in - for (int i=0; i < nM; i++) { if (d->mapM2M[i] < 0) { - mjERROR("unassigned index in mapM2M"); + mjERROR("unassigned index in mapM2C"); } } @@ -1429,8 +1415,9 @@ mjData* mj_makeData(const mjModel* m) { -// copy mjData, if dest==NULL create new data -mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { +// copy mjData, if dest==NULL create new data; +// flg_all 1: copy all fields, 0: skip fields not required for visualization +mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int flg_all) { void* save_buffer; void* save_arena; @@ -1475,10 +1462,25 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { // copy buffer { MJDATA_POINTERS_PREAMBLE(m) - #define X(type, name, nr, nc) \ - memcpy((char*)dest->name, (const char*)src->name, sizeof(type)*(m->nr)*nc); - MJDATA_POINTERS - #undef X + if (flg_all) { + #define X(type, name, nr, nc) \ + memcpy((char*)dest->name, (const char*)src->name, sizeof(type)*(m->nr)*nc); + MJDATA_POINTERS + #undef X + } else { + // redefine XNV to nothing + #undef XNV + #define XNV(type, name, nr, nc) + + #define X(type, name, nr, nc) \ + memcpy((char*)dest->name, (const char*)src->name, sizeof(type)*(m->nr)*nc); + MJDATA_POINTERS + #undef X + + // redefine XNV to be the same as X + #undef XNV + #define XNV X + } } @@ -1488,7 +1490,8 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { #undef MJ_M #define MJ_M(n) (m->n) - #define X(type, name, nr, nc) \ + if (flg_all) { + #define X(type, name, nr, nc) \ if (src->name) { \ dest->name = (type*)((char*)dest->arena + PTRDIFF(src->name, src->arena)); \ ASAN_UNPOISON_MEMORY_REGION(dest->name, sizeof(type) * nr * nc); \ @@ -1496,8 +1499,28 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { } else { \ dest->name = NULL; \ } - MJDATA_ARENA_POINTERS - #undef X + MJDATA_ARENA_POINTERS + #undef X + } else { + // redefine XNV to nothing + #undef XNV + #define XNV(type, name, nr, nc) + + #define X(type, name, nr, nc) \ + if (src->name) { \ + dest->name = (type*)((char*)dest->arena + PTRDIFF(src->name, src->arena)); \ + ASAN_UNPOISON_MEMORY_REGION(dest->name, sizeof(type) * nr * nc); \ + memcpy((char*)dest->name, (const char*)src->name, sizeof(type) * nr * nc); \ + } else { \ + dest->name = NULL; \ + } + MJDATA_ARENA_POINTERS + #undef X + + // redefine XNV to be the same as X + #undef XNV + #define XNV X + } #undef MJ_M #define MJ_M(n) n @@ -1529,6 +1552,14 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { } +mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { + return mj_copyDataVisual(dest, m, src, /*flg_all=*/1); +} + + +mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src) { + return mj_copyDataVisual(dest, m, src, /*flg_all=*/0); +} static void maybe_lock_alloc_mutex(mjData* d) { if (d->threadpool != 0) { @@ -1904,7 +1935,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { memset(d->warning, 0, mjNWARNING*sizeof(mjWarningStat)); memset(d->timer, 0, mjNTIMER*sizeof(mjTimerStat)); memset(d->solver, 0, mjNSOLVER*mjNISLAND*sizeof(mjSolverStat)); - d->solver_nisland = 0; mju_zeroInt(d->solver_niter, mjNISLAND); mju_zeroInt(d->solver_nnz, mjNISLAND); mju_zero(d->solver_fwdinv, 2); @@ -1918,6 +1948,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->nJ = 0; d->nA = 0; d->nisland = 0; + d->nidof = 0; // clear global properties d->time = 0; @@ -1989,9 +2020,8 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { makeBSparse(m, d); checkDBSparse(m, d); - // make M, C - makeDofDofSparse(m, d, d->M_rownnz, d->M_rowadr, NULL, d->M_colind, /*reduced=*/0, /*upper=*/0); - makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, NULL, d->C_colind, /*reduced=*/1, /*upper=*/0); + // make C + makeDofDofSparse(m, d, d->M_rownnz, d->M_rowadr, NULL, d->M_colind, /*reduced=*/1, /*upper=*/0); // make index mappings: mapM2D, mapD2M, mapM2C, mapM2M makeDofDofmaps(m, d); @@ -2014,6 +2044,9 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { } } } + + // copy signature from model + d->signature = m->signature; } @@ -2076,6 +2109,7 @@ static int sensorSize(mjtSensor sensor_type, int sensor_dim) { case mjSENS_ACTUATORVEL: case mjSENS_ACTUATORFRC: case mjSENS_JOINTACTFRC: + case mjSENS_TENDONACTFRC: case mjSENS_JOINTLIMITPOS: case mjSENS_JOINTLIMITVEL: case mjSENS_JOINTLIMITFRC: @@ -2140,6 +2174,7 @@ static int numObjects(const mjModel* m, mjtObj objtype) { case mjOBJ_DEFAULT: case mjOBJ_FRAME: case mjOBJ_UNKNOWN: + case mjOBJ_MODEL: return -1; case mjOBJ_BODY: case mjOBJ_XBODY: @@ -2371,7 +2406,7 @@ const char* mj_validateReferences(const mjModel* m) { } } for (int i=0; i < m->ntex; i++) { - int tex_adr = m->tex_adr[i] + 3*m->tex_height[i]*m->tex_width[i]; + int tex_adr = m->tex_adr[i] + m->tex_nchannel[i]*m->tex_height[i]*m->tex_width[i]; if (tex_adr > m->ntexdata || m->tex_adr[i] < 0) { return "Invalid model: tex_adr out of bounds."; } diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 809d9e83..bdf0e2ce 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -99,6 +99,9 @@ MJAPI void mj_makeRawData(mjData** dest, const mjModel* m); // m is only required to contain the size fields from MJMODEL_INTS. MJAPI mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src); +// copy mjData, skip large arrays not required for abstract visualization +MJAPI mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src); + // set data to defaults MJAPI void mj_resetData(const mjModel* m, mjData* d); diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 9a8be763..05364373 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -16,6 +16,7 @@ #include #include +#include #include #include @@ -26,12 +27,65 @@ #include "engine/engine_support.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" +#include "engine/engine_util_sparse.h" #ifdef MEMORY_SANITIZER #include #endif +//-------------------------- local utilities ------------------------------------------------------- + +// clear island-related arena pointers in mjData +static void clearIsland(mjData* d, size_t parena) { +#define X(type, name, nr, nc) d->name = NULL; + MJDATA_ARENA_POINTERS_ISLAND +#undef X + d->nefc = 0; + d->nisland = 0; + d->nidof = 0; + d->parena = parena; + + // poison remaining memory +#ifdef ADDRESS_SANITIZER + ASAN_POISON_MEMORY_REGION( + (char*)d->arena + d->parena, d->narena - d->pstack - d->parena); +#endif +} + + + +// allocate island arrays on arena, return 1 on success, 0 on failure +static int arenaAllocIsland(const mjModel* m, mjData* d) { +#undef MJ_M +#define MJ_M(n) m->n +#undef MJ_D +#define MJ_D(n) d->n + + size_t parena_old = d->parena; + +#define X(type, name, nr, nc) \ + d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \ + if (!d->name) { \ + mj_warning(d, mjWARN_CNSTRFULL, d->narena); \ + clearIsland(d, parena_old); \ + return 0; \ + } + + MJDATA_ARENA_POINTERS_ISLAND + +#undef X + +#undef MJ_M +#define MJ_M(n) n +#undef MJ_D +#define MJ_D(n) n + return 1; +} + + + +//-------------------------- flood-fill and graph construction ------------------------------------ // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: @@ -87,54 +141,6 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons -// clear island-related arena pointers in mjData -static void clearIsland(mjData* d, size_t parena) { -#define X(type, name, nr, nc) d->name = NULL; - MJDATA_ARENA_POINTERS_ISLAND -#undef X - d->nefc = 0; - d->nisland = 0; - d->parena = parena; - - // poison remaining memory -#ifdef ADDRESS_SANITIZER - ASAN_POISON_MEMORY_REGION( - (char*)d->arena + d->parena, d->narena - d->pstack - d->parena); -#endif -} - - - -// allocate island arrays on arena, return 1 on success, 0 on failure -static int arenaAllocIsland(const mjModel* m, mjData* d) { -#undef MJ_M -#define MJ_M(n) m->n -#undef MJ_D -#define MJ_D(n) d->n - - size_t parena_old = d->parena; - -#define X(type, name, nr, nc) \ - d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \ - if (!d->name) { \ - mj_warning(d, mjWARN_CNSTRFULL, d->narena); \ - clearIsland(d, parena_old); \ - return 0; \ - } - - MJDATA_ARENA_POINTERS_ISLAND - -#undef X - -#undef MJ_M -#define MJ_M(n) n -#undef MJ_D -#define MJ_D(n) n - return 1; -} - - - // return upper bound on number of tree-tree edges static int countMaxEdge(const mjModel* m, const mjData* d) { int nedge_max = 0; @@ -411,14 +417,17 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg +//-------------------------- main entry-point ----------------------------------------------------- + // discover islands: -// nisland, island_dofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext +// nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext void mj_island(const mjModel* m, mjData* d) { int nv = m->nv, nefc = d->nefc, ntree=m->ntree; // no constraints: quick return if (!nefc || m->nflex) { // TODO: add flex support to island discovery d->nisland = 0; + d->nidof = 0; return; } @@ -454,86 +463,178 @@ void mj_island(const mjModel* m, mjData* d) { int* stack = mjSTACKALLOC(d, nedge, int); d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack); + // no islands found: quick return + if (!d->nisland) { + d->nidof = 0; + mj_freeStack(d); + return; + } + + // count ni: total number of dofs in islands + int nidof = 0; + for (int i=0; i < nv; i++) { + nidof += (tree_island[m->dof_treeid[i]] >= 0); + } + d->nidof = nidof; + // allocate island arrays on arena if (!arenaAllocIsland(m, d)) { mj_freeStack(d); return; } - int nisland = d->nisland; // local copy + // local copy + int nisland = d->nisland; - // compute dof_island, island_dofnum - int num_dof_unc = 0; // number of unconstrained dofs - mju_zeroInt(d->island_dofnum, nisland); + + // ------------------------------------- degrees of freedom -------------------------------------- + + // compute dof_island, island_nv + mju_zeroInt(d->island_nv, nisland); for (int i=0; i < nv; i++) { - // dof_island - int island = tree_island[m->dof_treeid[i]]; + // assign dofs to islands + int island = tree_island[m->dof_treeid[i]]; // -1 if unconstrained d->dof_island[i] = island; - // island_dofnum + // increment island_nv if (island >= 0) { - d->island_dofnum[island]++; - } else { - num_dof_unc++; + d->island_nv[island]++; } } - // compute island_dofadr - if (nisland) d->island_dofadr[0] = 0; + // compute island_idofadr (cumsum of island_nv) + d->island_idofadr[0] = 0; for (int i=1; i < nisland; i++) { - d->island_dofadr[i] = d->island_dofadr[i-1] + d->island_dofnum[i-1]; + d->island_idofadr[i] = d->island_idofadr[i-1] + d->island_nv[i-1]; } - // reset island_dofnum - mju_zeroInt(d->island_dofnum, nisland); - - // compute dof_islandind, island_dofind - int num_dof_island = 0; - for (int i=0; i < nv; i++) { - int island = d->dof_island[i]; + // compute dof <-> idof maps + int* island_nv2 = mjSTACKALLOC(d, nisland + 1, int); // last element counts unconstrained dofs + mju_zeroInt(island_nv2, nisland + 1); + for (int dof=0; dof < nv; dof++) { + int island = d->dof_island[dof]; + int idof; if (island >= 0) { - d->island_dofind[d->island_dofadr[island] + d->island_dofnum[island]] = i; - d->dof_islandind[i] = d->island_dofnum[island]++; - num_dof_island++; + // constrained dof + idof = d->island_idofadr[island] + island_nv2[island]++; } else { - d->dof_islandind[i] = -1; + // unconstrained dof + idof = nidof + island_nv2[nisland]++; } + + d->map_dof2idof[dof] = idof; + d->map_idof2dof[idof] = dof; // only the first ni elements of map_idof2dof are in some island } - // sanity check, SHOULD NOT OCCUR - if (num_dof_island + num_dof_unc != nv) { - mjERROR("not all islands assigned to dofs"); + // SHOULD NOT OCCUR + if (!mju_compare(island_nv2, d->island_nv, nisland)) mjERROR("island_nv miscount"); + if (nidof + island_nv2[nisland] != nv) mjERROR("miscount of unconstrained dofs"); + + // compute island_dofadr (used for visualization) + for (int i=0; i < nisland; i++) { + d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]]; } - // finalize dof_islandind: set remaining indices to -1 - for (int i=num_dof_island; i < nv; i++) { - d->island_dofind[i] = -1; - } + // inertia: block-diagonalize both iLD <- qLD and iM <- qM + mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind, + d->qLD, d->M_rownnz, d->M_rowadr, d->M_colind, + nidof, nisland, + d->map_idof2dof, d->map_dof2idof, + d->island_idofadr, d->island_idofadr, + d->iM, d->M); + mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof); - // compute efc_island, island_efcnum - mju_zeroInt(d->island_efcnum, nisland); + + // ------------------------------------- constraints --------------------------------------------- + + // compute efc_island, island_{ne,nf,nefc} + mju_zeroInt(d->island_ne, nisland); + mju_zeroInt(d->island_nf, nisland); + mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { int tree[2]; treeFirst(m, d, tree, i); int island = tree_island[tree[0]]; d->efc_island[i] = island; - d->island_efcnum[island]++; + d->island_nefc[island]++; + switch (d->efc_type[i]) { + case mjCNSTR_EQUALITY: + d->island_ne[island]++; + break; + case mjCNSTR_FRICTION_DOF: + case mjCNSTR_FRICTION_TENDON: + d->island_nf[island]++; + break; + default: + break; + } } - // compute island_efcadr - if (nisland) d->island_efcadr[0] = 0; + // compute island_iefcadr (cumsum of island_nefc) + d->island_iefcadr[0] = 0; for (int i=1; i < nisland; i++) { - d->island_efcadr[i] = d->island_efcadr[i-1] + d->island_efcnum[i-1]; + d->island_iefcadr[i] = d->island_iefcadr[i-1] + d->island_nefc[i-1]; } - // reset island_efcnum - mju_zeroInt(d->island_efcnum, nisland); - - // compute efc_islandind - for (int i=0; i < nefc; i++) { - int island = d->efc_island[i]; - d->island_efcind[d->island_efcadr[island] + (d->island_efcnum[island]++)] = i; + // compute efc <-> iefc maps + int* island_nefc2 = island_nv2; // reuse island_nv2 + mju_zeroInt(island_nefc2, nisland); + for (int c=0; c < nefc; c++) { + int island = d->efc_island[c]; + int ic = d->island_iefcadr[island] + island_nefc2[island]++; + d->map_efc2iefc[c] = ic; + d->map_iefc2efc[ic] = c; } + // SHOULD NOT OCCUR + if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); + + // dense: block-diagonalize Jacobian + if (!mj_isSparse(m)) { + mju_blockDiag(d->iefc_J, d->efc_J, + nv, nidof, nisland, + d->map_iefc2efc, d->map_idof2dof, + d->island_nefc, d->island_nv, + d->island_iefcadr, d->island_idofadr); + } + + // sparse + else { + // block-diagonalize Jacobian + mju_blockDiagSparse(d->iefc_J, d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind, + d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + nefc, nisland, + d->map_iefc2efc, d->map_dof2idof, + d->island_iefcadr, d->island_idofadr, NULL, NULL); + + // recompute rowsuper per island + for (int island=0; island < nisland; island++) { + int adr = d->island_iefcadr[island]; + mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr, + d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind); + } + + // block-diagonalize Jacobian-transpose + mju_blockDiagSparse(d->iefc_JT, d->iefc_JT_rownnz, d->iefc_JT_rowadr, d->iefc_JT_colind, + d->efc_JT, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, + nidof, nisland, + d->map_idof2dof, d->map_efc2iefc, + d->island_idofadr, d->island_iefcadr, NULL, NULL); + + // recompute rowsuper per island + for (int island=0; island < nisland; island++) { + int adr = d->island_idofadr[island]; + mju_superSparse(d->island_nv[island], d->iefc_JT_rowsuper + adr, + d->iefc_JT_rownnz + adr, d->iefc_JT_rowadr + adr, d->iefc_JT_colind); + } + } + + // copy position-dependent efc vectors required by solver + mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc); + mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc); + mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc); + mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc); + mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc); + mj_freeStack(d); } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 40296906..5b0d1b5b 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -84,7 +84,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) { case mjJNT_BALL: { - // convert quatertion difference into angular "velocity" + // convert quaternion difference into angular "velocity" mjtNum dif[3], quat[4]; mju_copy4(quat, d->qpos+padr); mju_normalize4(quat); @@ -116,9 +116,50 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // flex elasticity for (int f=0; f < m->nflex; f++) { mjtNum* k = m->flex_stiffness + 21*m->flex_elemadr[f]; + mjtNum* b = m->flex_bending + 16*m->flex_edgeadr[f]; int dim = m->flex_dim[f]; - if (dim == 1 || m->flex_rigid[f] || k[0] == 0) { + if (dim == 1 || m->flex_rigid[f]) { + continue; + } + + // add bending forces to qfrc_spring + if (dim == 2) { + mjtNum* xpos = d->flexvert_xpos + 3*m->flex_vertadr[f]; + int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f]; + + for (int e = 0; e < m->flex_edgenum[f]; e++) { + const int* edge = m->flex_edge + 2*(e+m->flex_edgeadr[f]); + const int* flap = m->flex_edgeflap + 2*(e+m->flex_edgeadr[f]); + int v[4] = {edge[0], edge[1], flap[0], flap[1]}; + if (v[3] == -1) { + // skip boundary edges + continue; + } + mjtNum force[12] = {0}; + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 4; j++) { + for (int x = 0; x < 3; x++) { + force[3*i+x] += b[16*e+4*i+j] * xpos[3*v[j]+x]; + } + } + } + + // TODO: add damping + + // insert into global force + for (int i = 0; i < 4; i++) { + int bid = bodyid[v[i]]; + int body_dofnum = m->body_dofnum[bid]; + int body_dofadr = m->body_dofadr[bid]; + for (int x = 0; x < body_dofnum; x++) { + d->qfrc_spring[body_dofadr+x] -= force[3*i+x]; + } + } + } + } + + if (k[0] == 0) { continue; } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index fcea96b2..d6e5204e 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1008,9 +1008,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena // SOLVER STAT if (d->nefc) { fprintf(fp, "SOLVER STAT\n"); - fprintf(fp, " solver_nisland = %d\n", d->solver_nisland); printVector(" solver_fwdinv = ", d->solver_fwdinv, 2, fp, float_format); - int nisland_stat = mjMIN(d->solver_nisland, mjNISLAND); + int nisland_stat = mjMAX(1, mjMIN(d->nisland, mjNISLAND)); for (int island=0; island < nisland_stat; island++) { int niter_stat = mjMIN(mjNSOLVER, d->solver_niter[island]); if (niter_stat) { @@ -1124,9 +1123,9 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz, d->moment_rowadr, d->moment_colind, fp, float_format); printArray("CRB", m->nbody, 10, d->crb, fp, float_format); - printInertia("QM", d->qM, m, fp, float_format); - + printSparse("M", d->M, m->nv, d->M_rownnz, + d->M_rowadr, d->M_colind, fp, float_format); printSparse("QLD", d->qLD, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind, fp, float_format); printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); @@ -1162,7 +1161,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena fprintf(fp, "\n\n"); // M sparse structure - mj_printSparsity("M: inertia matrix", m->nv, m->nv, d->M_rowadr, NULL, d->M_rownnz, + mj_printSparsity("M: reduced inertia matrix", m->nv, m->nv, d->M_rowadr, NULL, d->M_rownnz, NULL, d->M_colind, fp); fprintf(fp, NAME_FORMAT, "M_rownnz"); @@ -1171,58 +1170,27 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena } fprintf(fp, "\n\n"); - // M_rowadr + // C_rowadr fprintf(fp, NAME_FORMAT, "M_rowadr"); for (int i = 0; i < m->nv; i++) { fprintf(fp, " %d", d->M_rowadr[i]); } fprintf(fp, "\n\n"); - // M_colind + // C_colind fprintf(fp, NAME_FORMAT, "M_colind"); - for (int i = 0; i < m->nM; i++) { + for (int i = 0; i < m->nC; i++) { fprintf(fp, " %d", d->M_colind[i]); } fprintf(fp, "\n\n"); // mapM2M fprintf(fp, NAME_FORMAT, "mapM2M"); - for (int i = 0; i < m->nM; i++) { + for (int i = 0; i < m->nC; i++) { fprintf(fp, " %d", d->mapM2M[i]); } fprintf(fp, "\n\n"); - // C sparse structure - mj_printSparsity("C: reduced dof-dof matrix", m->nv, m->nv, d->C_rowadr, NULL, d->C_rownnz, - NULL, d->C_colind, fp); - - fprintf(fp, NAME_FORMAT, "C_rownnz"); - for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->C_rownnz[i]); - } - fprintf(fp, "\n\n"); - - // C_rowadr - fprintf(fp, NAME_FORMAT, "C_rowadr"); - for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->C_rowadr[i]); - } - fprintf(fp, "\n\n"); - - // C_colind - fprintf(fp, NAME_FORMAT, "C_colind"); - for (int i = 0; i < m->nC; i++) { - fprintf(fp, " %d", d->C_colind[i]); - } - fprintf(fp, "\n\n"); - - // mapM2C - fprintf(fp, NAME_FORMAT, "mapM2C"); - for (int i = 0; i < m->nC; i++) { - fprintf(fp, " %d", d->mapM2C[i]); - } - fprintf(fp, "\n\n"); - // D sparse structure mj_printSparsity("D: dof-dof matrix", m->nv, m->nv, d->D_rowadr, d->D_diag, d->D_rownnz, NULL, d->D_colind, fp); @@ -1393,27 +1361,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFNUM"); + fprintf(fp, NAME_FORMAT, "ISLAND_NV"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_dofnum[i]); + fprintf(fp, " %d", d->island_nv[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFADR"); + fprintf(fp, NAME_FORMAT, "ISLAND_IDOFADR"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_dofadr[i]); + fprintf(fp, " %d", d->island_idofadr[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFIND"); + fprintf(fp, NAME_FORMAT, "MAP_IDOF2DOF"); for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->island_dofind[i]); - } - fprintf(fp, "\n\n"); + int dof = d->map_idof2dof[i]; + if (i > 0) { + int dofprev = d->map_idof2dof[i-1]; - fprintf(fp, NAME_FORMAT, "DOF_ISLANDIND"); - for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->dof_islandind[i]); + // print '|' at island boundaries + if (d->dof_island[dof] != d->dof_island[dofprev]) { + fprintf(fp, " |"); + } + } + fprintf(fp, " %d", dof); } fprintf(fp, "\n\n"); @@ -1423,21 +1394,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCNUM"); + fprintf(fp, NAME_FORMAT, "ISLAND_NEFC"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_efcnum[i]); + fprintf(fp, " %d", d->island_nefc[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCADR"); + fprintf(fp, NAME_FORMAT, "ISLAND_IEFCADR"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_efcadr[i]); + fprintf(fp, " %d", d->island_iefcadr[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCIND"); + fprintf(fp, NAME_FORMAT, "MAP_IEFC2EFC"); for (int i = 0; i < d->nefc; i++) { - fprintf(fp, " %d", d->island_efcind[i]); + int efc = d->map_iefc2efc[i]; + if (i > 0) { + int efcprev = d->map_iefc2efc[i-1]; + + // print '|' at island boundaries + if (d->efc_island[efc] != d->efc_island[efcprev]) { + fprintf(fp, " |"); + } + } + fprintf(fp, " %d", efc); } fprintf(fp, "\n\n"); } diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index baad81c4..485474c1 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -698,8 +698,8 @@ void mj_sensorVel(const mjModel* m, mjData* d) { // acceleration/force-dependent sensors void mj_sensorAcc(const mjModel* m, mjData* d) { int rootid, bodyid, objtype, objid, adr, nusersensor = 0; - int ne = d->ne, nf = d->nf, nefc = d->nefc; - mjtNum tmp[6], conforce[6], conray[3]; + int ne = d->ne, nf = d->nf, nefc = d->nefc, nu = m->nu; + mjtNum tmp[6], conforce[6], conray[3], frc; mjContact* con; // disabled sensors: return @@ -825,6 +825,16 @@ void mj_sensorAcc(const mjModel* m, mjData* d) { d->sensordata[adr] = d->qfrc_actuator[m->jnt_dofadr[objid]]; break; + case mjSENS_TENDONACTFRC: // tendonactfrc + frc = 0.0; + for (int j=0; j < nu; j++) { + if (m->actuator_trntype[j] == mjTRN_TENDON && m->actuator_trnid[2*j] == objid) { + frc += d->actuator_force[j]; + } + } + d->sensordata[adr] = frc; + break; + case mjSENS_JOINTLIMITFRC: // jointlimitfrc d->sensordata[adr] = 0; for (int j=ne+nf; j < nefc; j++) { @@ -928,9 +938,7 @@ void mj_energyPos(const mjModel* m, mjData* d) { switch ((mjtJoint) m->jnt_type[i]) { case mjJNT_FREE: - mju_copy4(quat, d->qpos+padr); - mju_normalize4(quat); - mju_sub3(dif, quat, m->qpos_spring+padr); + mju_sub3(dif, d->qpos+padr, m->qpos_spring+padr); d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif); // continue with rotations @@ -938,7 +946,7 @@ void mj_energyPos(const mjModel* m, mjData* d) { mjFALLTHROUGH; case mjJNT_BALL: - // covert quatertion difference into angular "velocity" + // convert quaternion difference into angular "velocity" mju_copy4(quat, d->qpos+padr); mju_normalize4(quat); mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr); diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 73f8919e..f737688c 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -102,10 +102,10 @@ static void set0(mjModel* m, mjData* d) { memset(m->flex_rigid, 0, m->nflex); // run remaining computations - mj_crb(m, d); + mj_tendon(m, d); + mj_makeM(m, d); mj_factorM(m, d); mj_flex(m, d); - mj_tendon(m, d); mj_transmission(m, d); // restore flex rigidity diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index f23d8bf5..635fa6af 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -187,8 +187,7 @@ static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int res[j] = d->efc_b[i+j] + mju_dotSparse(d->efc_AR + d->efc_AR_rowadr[i+j], d->efc_force, d->efc_AR_rownnz[i+j], - d->efc_AR_colind + d->efc_AR_rowadr[i+j], - /*flg_unc1=*/0); + d->efc_AR_colind + d->efc_AR_rowadr[i+j]); } } @@ -765,14 +764,54 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) { // CG context struct _mjCGContext { - int flg_Newton; // 1: Newton, 0: CG - - // island-related + int is_sparse; // 1: sparse, 0: dense + int is_elliptic; // 1: elliptic, 0: pyramidal int island; // current island index, -1 if monolithic + + // sizes int nv; // number of dofs - int nefc; // number of constraints - int* dofind; // dof indices of this island, NULL if monolithic - int* efcind; // constraint indices of this island, NULL if monolithic + int ne; // number of equalities + int nf; // number of friction constraints + int nefc; // number of all constraints + + // contact array + mjContact* contact; + + // dof arrays + const mjtNum* qfrc_smooth; + const mjtNum* qacc_smooth; + mjtNum* qfrc_constraint; + mjtNum* qacc; + + // inertia + const int* M_rownnz; + const int* M_rowadr; + const int* M_colind; + const mjtNum* M; + const mjtNum* qLD; + const mjtNum* qLDiagInv; + + // efc arrays + const mjtNum* efc_D; + const mjtNum* efc_R; + const mjtNum* efc_frictionloss; + const mjtNum* efc_aref; + const int* efc_id; + const int* efc_type; + mjtNum* efc_force; + int* efc_state; + + // Jacobians + const int* J_rownnz; + const int* J_rowadr; + const int* J_rowsuper; + const int* J_colind; + const int* JT_rownnz; + const int* JT_rowadr; + const int* JT_rowsuper; + const int* JT_colind; + const mjtNum* J; + const mjtNum* JT; // common arrays (CGallocate) mjtNum* Jaref; // Jac*qacc - aref (nefc x 1) @@ -786,14 +825,15 @@ struct _mjCGContext { // Newton arrays, known-size (CGallocate) mjtNum* D; // constraint inertia (nefc x 1) - mjtNum* C; // reduced sparse inertia matrix (nC x 1) int* H_rowadr; // Hessian row addresses (nv x 1) int* H_rownnz; // Hessian row nonzeros (nv x 1) int* H_lowernnz; // Hessian lower triangle row nonzeros (nv x 1) int* L_rownnz; // Hessian factor row nonzeros (nv x 1) int* L_rowadr; // Hessian factor row addresses (nv x 1) + int* buf_ind; // index buffer for sparse addition (nv x 1) + mjtNum* buf_val; // value buffer for sparse addition (nv x 1) - // Newton arrays, computed-size (HessianMake) + // Newton arrays, computed-size (MakeHessian) int nH; // number of nonzeros in Hessian H int* H_colind; // Hessian column indices (nH x 1) mjtNum* H; // Hessian (nH x 1) @@ -818,23 +858,125 @@ struct _mjCGContext { typedef struct _mjCGContext mjCGContext; -// allocate fixed-size arrays in mjCGContext -// mj_{mark/free}Stack in calling function! -static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, - int island, int flg_Newton) { + +// set sizes and pointers to mjData arrays in mjCGContext +static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int island) { // clear everything memset(ctx, 0, sizeof(mjCGContext)); - // get sizes - int nv = island < 0 ? m->nv : d->island_dofnum[island]; - int nefc = island < 0 ? d->nefc : d->island_efcnum[island]; - - // island-related + // globals + ctx->is_sparse = mj_isSparse(m); + ctx->is_elliptic = (m->opt.cone == mjCONE_ELLIPTIC); + ctx->contact = d->contact; ctx->island = island; - ctx->nv = nv; - ctx->nefc = nefc; - ctx->dofind = island < 0 ? NULL : d->island_dofind + d->island_dofadr[island]; - ctx->efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island]; + + // set sizes and pointers (monolithic) + if (island < 0) { + // sizes + ctx->nv = m->nv; + ctx->ne = d->ne; + ctx->nf = d->nf; + ctx->nefc = d->nefc; + + // dof arrays + ctx->qfrc_smooth = d->qfrc_smooth; + ctx->qfrc_constraint = d->qfrc_constraint; + ctx->qacc_smooth = d->qacc_smooth; + ctx->qacc = d->qacc; + + // inertia + ctx->M_rownnz = d->M_rownnz; + ctx->M_rowadr = d->M_rowadr; + ctx->M_colind = d->M_colind; + ctx->M = d->M; + ctx->qLD = d->qLD; + ctx->qLDiagInv = d->qLDiagInv; + + // efc arrays + ctx->efc_D = d->efc_D; + ctx->efc_R = d->efc_R; + ctx->efc_frictionloss = d->efc_frictionloss; + ctx->efc_aref = d->efc_aref; + ctx->efc_id = d->efc_id; + ctx->efc_type = d->efc_type; + ctx->efc_force = d->efc_force; + ctx->efc_state = d->efc_state; + + // Jacobians + ctx->J = d->efc_J; + if (ctx->is_sparse) { + ctx->J_rownnz = d->efc_J_rownnz; + ctx->J_rowadr = d->efc_J_rowadr; + ctx->J_rowsuper = d->efc_J_rowsuper; + ctx->J_colind = d->efc_J_colind; + ctx->JT_rownnz = d->efc_JT_rownnz; + ctx->JT_rowadr = d->efc_JT_rowadr; + ctx->JT_rowsuper = d->efc_JT_rowsuper; + ctx->JT_colind = d->efc_JT_colind; + ctx->JT = d->efc_JT; + } + } + + // set sizes and pointers (per-island) + else { + // sizes + ctx->nv = d->island_nv[island]; + ctx->ne = d->island_ne[island]; + ctx->nf = d->island_nf[island]; + ctx->nefc = d->island_nefc[island]; + + // dof arrays + int idofadr = d->island_idofadr[island]; + ctx->qfrc_smooth = d->ifrc_smooth + idofadr; + ctx->qfrc_constraint = d->ifrc_constraint + idofadr; + ctx->qacc_smooth = d->iacc_smooth + idofadr; + ctx->qacc = d->iacc + idofadr; + + // inertia + ctx->M_rownnz = d->iM_rownnz + idofadr; + ctx->M_rowadr = d->iM_rowadr + idofadr; + ctx->M_colind = d->iM_colind; + ctx->M = d->iM; + ctx->qLD = d->iLD; + ctx->qLDiagInv = d->iLDiagInv + idofadr; + + // efc arrays + int iefcadr = d->island_iefcadr[island]; + ctx->efc_D = d->iefc_D + iefcadr; + ctx->efc_R = d->iefc_R + iefcadr; + ctx->efc_frictionloss = d->iefc_frictionloss + iefcadr; + ctx->efc_aref = d->iefc_aref + iefcadr; + ctx->efc_id = d->iefc_id + iefcadr; + ctx->efc_type = d->iefc_type + iefcadr; + ctx->efc_force = d->iefc_force + iefcadr; + ctx->efc_state = d->iefc_state + iefcadr; + + // Jacobians + if (!ctx->is_sparse) { + ctx->J = d->iefc_J + d->nidof * iefcadr; + } else { + ctx->J_rownnz = d->iefc_J_rownnz + iefcadr; + ctx->J_rowadr = d->iefc_J_rowadr + iefcadr; + ctx->J_rowsuper = d->iefc_J_rowsuper + iefcadr; + ctx->J_colind = d->iefc_J_colind; + ctx->JT_rownnz = d->iefc_JT_rownnz + idofadr; + ctx->JT_rowadr = d->iefc_JT_rowadr + idofadr; + ctx->JT_rowsuper = d->iefc_JT_rowsuper + idofadr; + ctx->JT_colind = d->iefc_JT_colind; + ctx->J = d->iefc_J; + ctx->JT = d->iefc_JT; + } + } +} + + + +// allocate fixed-size arrays in mjCGContext +// mj_{mark/free}Stack in calling function! +static void CGallocate(mjData* d, mjCGContext* ctx, int flg_Newton) { + // local sizes + int nv = ctx->nv; + int nefc = ctx->nefc; // common arrays ctx->Jaref = mjSTACKALLOC(d, nefc, mjtNum); @@ -847,18 +989,18 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, ctx->quad = mjSTACKALLOC(d, nefc*3, mjtNum); // Newton only, known-size arrays - ctx->flg_Newton = flg_Newton; if (flg_Newton) { - ctx->D = mjSTACKALLOC(d, nefc, mjtNum); + ctx->D = mjSTACKALLOC(d, nefc, mjtNum); // sparse Newton only - if (mj_isSparse(m)) { - ctx->C = mjSTACKALLOC(d, m->nC, mjtNum); + if (ctx->is_sparse) { ctx->H_rowadr = mjSTACKALLOC(d, nv, int); ctx->H_rownnz = mjSTACKALLOC(d, nv, int); ctx->H_lowernnz = mjSTACKALLOC(d, nv, int); ctx->L_rownnz = mjSTACKALLOC(d, nv, int); ctx->L_rowadr = mjSTACKALLOC(d, nv, int); + ctx->buf_val = mjSTACKALLOC(d, nv, mjtNum); + ctx->buf_ind = mjSTACKALLOC(d, nv, int); } } } @@ -866,28 +1008,35 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, // update efc_force, qfrc_constraint, cost-related -static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) { +static void CGupdateConstraint(mjCGContext* ctx, int flg_HessianCone) { int nefc = ctx->nefc, nv = ctx->nv; - const int* dofind = ctx->dofind; - const int* efcind = ctx->efcind; // update constraints - mj_constraintUpdate_island(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton, ctx->island); + mj_constraintUpdate_impl(ctx->ne, ctx->nf, ctx->nefc, ctx->efc_D, ctx->efc_R, + ctx->efc_frictionloss, ctx->Jaref, ctx->efc_type, ctx->efc_id, + ctx->contact, ctx->efc_state, ctx->efc_force, + &(ctx->cost), flg_HessianCone); + + // compute qfrc_constraint (dense or sparse) + if (!ctx->is_sparse) { + mju_mulMatTVec(ctx->qfrc_constraint, ctx->J, ctx->efc_force, nefc, nv); + } else { + mju_mulMatVecSparse(ctx->qfrc_constraint, ctx->JT, ctx->efc_force, nv, + ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper); + } // count active and cone ctx->nactive = 0; ctx->ncone = 0; - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - ctx->nactive += (d->efc_state[i] != mjCNSTRSTATE_SATISFIED); - ctx->ncone += (d->efc_state[i] == mjCNSTRSTATE_CONE); + for (int i=0; i < nefc; i++) { + ctx->nactive += (ctx->efc_state[i] != mjCNSTRSTATE_SATISFIED); + ctx->ncone += (ctx->efc_state[i] == mjCNSTRSTATE_CONE); } // add Gauss cost, set in quadratic[0] mjtNum Gauss = 0; - for (int c=0; c < nv; c++) { - int i = dofind ? dofind[c] : c; - Gauss += 0.5 * (ctx->Ma[c] - d->qfrc_smooth[i]) * (d->qacc[i] - d->qacc_smooth[i]); + for (int i=0; i < nv; i++) { + Gauss += 0.5 * (ctx->Ma[i] - ctx->qfrc_smooth[i]) * (ctx->qacc[i] - ctx->qacc_smooth[i]); } ctx->quadGauss[0] = Gauss; @@ -895,22 +1044,19 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) { } -// TODO(tassa): Restore mjData const-ness. + // update grad, Mgrad -static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) { +static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) { int nv = ctx->nv; - const int* dofind = ctx->dofind; // grad = M*qacc - qfrc_smooth - qfrc_constraint - for (int c=0; c < nv; c++) { - int i = dofind ? dofind[c] : c; - ctx->grad[c] = ctx->Ma[c] - d->qfrc_smooth[i] - d->qfrc_constraint[i]; + for (int i=0; i < nv; i++) { + ctx->grad[i] = ctx->Ma[i] - ctx->qfrc_smooth[i] - ctx->qfrc_constraint[i]; } // Newton: Mgrad = H \ grad - // TODO: b/295296178 - add island support to Newton solver - if (ctx->flg_Newton) { - if (mj_isSparse(m)) { + if (flg_Newton) { + if (ctx->is_sparse) { mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L), ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind); } else { @@ -921,44 +1067,32 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) { // CG: Mgrad = M \ grad else { mju_copy(ctx->Mgrad, ctx->grad, nv); - mj_solveM_island(m, d, ctx->Mgrad, ctx->island); + mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); } } // prepare quadratic polynomials and contact cone quantities -static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { - int nv = ctx->nv, nefc = ctx->nefc, island = ctx->island; - const int* dofind = ctx->dofind; - const int* efcind = ctx->efcind; +static void CGprepare(mjCGContext* ctx) { + int nv = ctx->nv, nefc = ctx->nefc; const mjtNum* v = ctx->search; // Gauss: alpha^2*0.5*v'*M*v + alpha*v'*(Ma-qfrc_smooth) + 0.5*(a-qacc_smooth)'*(Ma-qfrc_smooth) // quadGauss[0] already computed in CGupdateConstraint - mjtNum v_dot_smooth; - if (island < 0) { - v_dot_smooth = mju_dot(d->qfrc_smooth, v, nv); - } else { - v_dot_smooth = 0; - for (int c=0; c < nv; c++) { - v_dot_smooth += d->qfrc_smooth[dofind[c]] * v[c]; - } - } - ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - v_dot_smooth; + ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(ctx->qfrc_smooth, v, nv); ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv); // process constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // pointers to numeric data - const mjtNum* Jv = ctx->Jv + c; - const mjtNum* Jaref = ctx->Jaref + c; - const mjtNum* D = d->efc_D + i; + const mjtNum* Jv = ctx->Jv + i; + const mjtNum* Jaref = ctx->Jaref + i; + const mjtNum* D = ctx->efc_D + i; // pointer to this quadratic - mjtNum* quad = ctx->quad + 3*c; + mjtNum* quad = ctx->quad + 3*i; // init with scalar quadratic mjtNum DJ0 = D[0]*Jaref[0]; @@ -967,12 +1101,12 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { quad[2] = Jv[0]*D[0]*Jv[0]; // elliptic cone: extra processing - if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { + if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // extract contact info - mjContact* con = d->contact + d->efc_id[i]; + const mjContact* con = ctx->contact + ctx->efc_id[i]; int dim = con->dim; mjtNum U[6], V[6], UU = 0, UV = 0, VV = 0, mu = con->mu; - mjtNum* friction = con->friction; + const mjtNum* friction = con->friction; // complete vector quadratic (for bottom zone) for (int j=1; j < dim; j++) { @@ -1006,7 +1140,7 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu))); // advance to next constraint - c += (dim-1); + i += (dim-1); } // apply scaling @@ -1028,9 +1162,8 @@ typedef struct _mjCGPnt mjCGPnt; // evaluate linesearch cost, return first and second derivatives -static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* p) { - int ne = d->ne, nf = d->nf, nefc = ctx->nefc; - const int* efcind = ctx->efcind; +static void CGeval(mjCGContext* ctx, mjCGPnt* p) { + int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc; // clear result mjtNum cost = 0, alpha = p->alpha; @@ -1041,26 +1174,24 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* mju_copy3(quadTotal, ctx->quadGauss); // process constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // equality if (i < ne) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); continue; } // friction if (i < ne + nf) { // search point, friction loss, bound (Rf) - mjtNum start = ctx->Jaref[c], dir = ctx->Jv[c]; + mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i]; mjtNum x = start + alpha*dir; - mjtNum f = d->efc_frictionloss[i]; - mjtNum Rf = d->efc_R[i]*f; + mjtNum f = ctx->efc_frictionloss[i]; + mjtNum Rf = ctx->efc_R[i]*f; // -bound < x < bound : quadratic if (-Rf < x && x < Rf) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); } // x < -bound : linear negative @@ -1078,10 +1209,10 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* } // limit and contact - if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone + if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone // extract contact info - mjContact* con = d->contact + d->efc_id[i]; - mjtNum* quad = ctx->quad + 3*c; + const mjContact* con = ctx->contact + ctx->efc_id[i]; + mjtNum* quad = ctx->quad + 3*i; int dim = con->dim; mjtNum mu = con->mu; @@ -1137,14 +1268,14 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* } // advance to next constraint - c += (dim-1); + i += (dim-1); } else { // inequality // search point - mjtNum x = ctx->Jaref[c] + alpha*ctx->Jv[c]; + mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i]; // active if (x < 0) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); } } } @@ -1170,7 +1301,7 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* // update bracket point given 3 candidate points -static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, +static int updateBracket(mjCGContext* ctx, mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) { int flag = 0; for (int i=0; i < 3; i++) { @@ -1192,7 +1323,7 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, // compute next point if updated if (flag) { pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1]; - CGeval(m, d, ctx, pnext); + CGeval(ctx, pnext); } return flag; @@ -1201,8 +1332,8 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, // line search -static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { - int nv = ctx->nv; +static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations) { + int nv = ctx->nv, nefc = ctx->nefc; mjCGPnt p0, p1, p2, pmid, p1next, p2next; // clear results @@ -1218,23 +1349,31 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // compute scaled gradtol and slope scaling - mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm / ctx->scale; + mjtNum gtol = tolerance * snorm / ctx->scale; mjtNum slopescl = ctx->scale / snorm; - // compute Mv, Jv - mj_mulM_island(m, d, ctx->Mv, ctx->search, ctx->island, /*flg_vecunc=*/0); - mj_mulJacVec_island(m, d, ctx->Jv, ctx->search, ctx->island, /*flg_resunc=*/0, /*flg_vecunc=*/0); + // compute Mv = M * v + mju_mulSymVecSparse(ctx->Mv, ctx->M, ctx->search, nv, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); + + // compute Jv = J * search (dense or sparse) + if (!ctx->is_sparse) { + mju_mulMatVec(ctx->Jv, ctx->J, ctx->search, nefc, nv); + } else { + mju_mulMatVecSparse(ctx->Jv, ctx->J, ctx->search, nefc, + ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->J_rowsuper); + } // prepare quadratics and cones - CGprepare(m, d, ctx); + CGprepare(ctx); // init at alpha = 0, save p0.alpha = 0; - CGeval(m, d, ctx, &p0); + CGeval(ctx, &p0); // always attempt one Newton step p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1]; - CGeval(m, d, ctx, &p1); + CGeval(ctx, &p1); if (p0.cost < p1.cost) { p1 = p0; } @@ -1289,14 +1428,14 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { // one-sided search int p2update = 0; - while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < m->opt.ls_iterations) { + while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < ls_iterations) { // save current p2 = p1; p2update = 1; // move to Newton point w.r.t current p1.alpha -= p1.deriv[0]/p1.deriv[1]; - CGeval(m, d, ctx, &p1); + CGeval(ctx, &p1); // check for convergence if (mju_abs(p1.deriv[0]) < gtol) { @@ -1306,7 +1445,7 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // check for failure to bracket - if (ctx->LSiter >= m->opt.ls_iterations) { + if (ctx->LSiter >= ls_iterations) { ctx->LSresult = 3; // could not bracket ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; return p1.alpha; @@ -1322,13 +1461,13 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { // compute next-points for bracket p2next = p1; p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1]; - CGeval(m, d, ctx, &p1next); + CGeval(ctx, &p1next); // bracketed search - while (ctx->LSiter < m->opt.ls_iterations) { + while (ctx->LSiter < ls_iterations) { // evaluate at midpoint pmid.alpha = 0.5*(p1.alpha + p2.alpha); - CGeval(m, d, ctx, &pmid); + CGeval(ctx, &pmid); // make list of candidates mjCGPnt candidates[3] = {p1next, p2next, pmid}; @@ -1349,8 +1488,8 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // update brackets - int b1 = updateBracket(m, d, ctx, &p1, candidates, &p1next); - int b2 = updateBracket(m, d, ctx, &p2, candidates, &p2next); + int b1 = updateBracket(ctx, &p1, candidates, &p1next); + int b2 = updateBracket(ctx, &p2, candidates, &p2next); // no update possible: numerical accuracy reached, use midpoint if (!b1 && !b2) { @@ -1384,34 +1523,30 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { // allocate and compute Hessian given efc_state // mj_{mark/free}Stack in caller function! -static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { - int nv = m->nv, nefc = d->nefc; +static void MakeHessian(mjData* d, mjCGContext* ctx) { + int nv = ctx->nv, nefc = ctx->nefc; // compute constraint inertia for (int i=0; i < nefc; i++) { - ctx->D[i] = d->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? d->efc_D[i] : 0; + ctx->D[i] = ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? ctx->efc_D[i] : 0; } // sparse - if (mj_isSparse(m)) { - // copy values of reduced sparse inertia matrix C - for (int i=0; i < m->nC; i++) { - ctx->C[i] = d->qM[d->mapM2C[i]]; - } - + if (ctx->is_sparse) { // initialize Hessian rowadr, rownnz mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv, - d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, - d->efc_JT_rowsuper, d, /*flg_upper=*/0); + ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, + ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, + ctx->JT_rowsuper, d, /*flg_upper=*/0); // add nC to Hessian total nonzeros (unavoidable overcounting since H_colind is still unknown) - ctx->nH = m->nC + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1]; + ctx->nH = ctx->M_rowadr[nv - 1] + ctx->M_rownnz[nv - 1] + + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1]; // shift H row addresses to make room for C int shift = 0; for (int r = 0; r < nv - 1; r++) { - shift += d->C_rownnz[r]; + shift += ctx->M_rownnz[r]; ctx->H_rowadr[r + 1] += shift; } @@ -1420,15 +1555,16 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { ctx->H = mjSTACKALLOC(d, ctx->nH, mjtNum); // compute H = J'*D*J - mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, ctx->D, nefc, nv, + mju_sqrMatTDSparse(ctx->H, ctx->J, ctx->JT, ctx->D, nefc, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, - d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, + ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, NULL, + ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d, /*diagind=*/NULL); // add mass matrix: H = J'*D*J + C - mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, - ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind); + mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, + ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, + ctx->buf_val, ctx->buf_ind); // transiently compute H'; mju_cholFactorNNZ is memory-contiguous in upper triangle layout mj_markStack(d); @@ -1436,7 +1572,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { int* HT_rowadr = mjSTACKALLOC(d, nv, int); int* HT_colind = mjSTACKALLOC(d, ctx->nH, int); mju_transposeSparse(NULL, NULL, nv, nv, - HT_rownnz, HT_rowadr, HT_colind, + HT_rownnz, HT_rowadr, HT_colind, NULL, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind); // count total and row non-zeros of reverse-Cholesky factor L @@ -1452,7 +1588,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { // allocate L_colind, L, Lcone ctx->L_colind = mjSTACKALLOC(d, ctx->nL, int); ctx->L = mjSTACKALLOC(d, ctx->nL, mjtNum); - if (m->opt.cone == mjCONE_ELLIPTIC) { + if (ctx->is_elliptic) { ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum); } @@ -1482,47 +1618,49 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { // allocate L, Lcone ctx->nL = nv*nv; ctx->L = mjSTACKALLOC(d, ctx->nL, mjtNum); - if (m->opt.cone == mjCONE_ELLIPTIC) { + if (ctx->is_elliptic) { ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum); } // compute H = M + J'*D*J - mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv); - mj_addMDense(m, d, ctx->L); + mju_sqrMatTD_impl(ctx->L, ctx->J, ctx->D, nefc, nv, /*flg_upper=*/ 0); + mju_addToSymSparse(ctx->L, ctx->M, ctx->nv, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, + /*flg_upper=*/ 0); } } // forward declaration of HessianCone (readability) -static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx); +static void HessianCone(mjData* d, mjCGContext* ctx); // factorize Hessian: L = chol(H), maybe (re)compute H given efc_state -static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx, - int flg_recompute) { - int nv = m->nv, nefc = d->nefc; +static void FactorizeHessian(mjData* d, mjCGContext* ctx, int flg_recompute) { + int nv = ctx->nv, nefc = ctx->nefc; // maybe compute constraint inertia if (flg_recompute) { for (int i=0; i < nefc; i++) { - ctx->D[i] = d->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? d->efc_D[i] : 0; + ctx->D[i] = ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? ctx->efc_D[i] : 0; } } // sparse - if (mj_isSparse(m)) { + if (ctx->is_sparse) { // maybe compute H = M + J'*D*J if (flg_recompute) { // compute H = J'*D*J - mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, ctx->D, nefc, nv, + mju_sqrMatTDSparse(ctx->H, ctx->J, ctx->JT, ctx->D, nefc, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, - d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, + ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, NULL, + ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d, /*diagind=*/NULL); // add mass matrix: H = J'*D*J + C - mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, - ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind); + mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, + ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, + ctx->buf_val, ctx->buf_ind); } // copy H lower-triangle into L, fill-in already accounted for @@ -1552,8 +1690,10 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx, else { // maybe compute H = M + J'*D*J if (flg_recompute) { - mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv); - mj_addMDense(m, d, ctx->L); + mju_sqrMatTD_impl(ctx->L, ctx->J, ctx->D, nefc, nv, /*flg_upper=*/ 0); + mju_addToSymSparse(ctx->L, ctx->M, ctx->nv, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, + /*flg_upper=*/ 0); } // factorize H @@ -1562,7 +1702,7 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx, // add cones to factor if present if (ctx->ncone) { - HessianCone(m, d, ctx); + HessianCone(d, ctx); } // mark full update @@ -1572,8 +1712,8 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx, // elliptic case: Hcone = H + cone_contributions -static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { - int nv = m->nv, nefc = d->nefc; +static void HessianCone(mjData* d, mjCGContext* ctx) { + int nv = ctx->nv, nefc = ctx->nefc; mjtNum local[36]; // start with Hcone = H @@ -1588,8 +1728,8 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { // add contributions for (int i=0; i < nefc; i++) { - if (d->efc_state[i] == mjCNSTRSTATE_CONE) { - mjContact* con = d->contact + d->efc_id[i]; + if (ctx->efc_state[i] == mjCNSTRSTATE_CONE) { + mjContact* con = ctx->contact + ctx->efc_id[i]; int dim = con->dim; // Cholesky of local Hessian @@ -1597,15 +1737,15 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { mju_cholFactor(local, dim, mjMINVAL); // sparse - if (mj_isSparse(m)) { + if (ctx->is_sparse) { // get nnz for row i (same for all rows in contact) - const int nnz = d->efc_J_rownnz[i]; + const int nnz = ctx->J_rownnz[i]; // compute LTJ = L'*J for this contact mju_zero(LTJ, dim*nnz); for (int r=0; r < dim; r++) { for (int c=0; c <= r; c++) { - mju_addToScl(LTJ+c*nnz, d->efc_J+d->efc_J_rowadr[i+r], local[r*dim+c], nnz); + mju_addToScl(LTJ+c*nnz, ctx->J+ctx->J_rowadr[i+r], local[r*dim+c], nnz); } } @@ -1613,7 +1753,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { for (int r=0; r < dim; r++) { // copy data for this row mju_copy(LTJ_row, LTJ+r*nnz, nnz); - mju_copyInt(LTJ_ind, d->efc_J_colind+d->efc_J_rowadr[i+r], nnz); + mju_copyInt(LTJ_ind, ctx->J_colind+ctx->J_rowadr[i+r], nnz); // update mju_cholUpdateSparse(ctx->Lcone, LTJ_row, nv, 1, @@ -1627,7 +1767,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { mju_zero(LTJ, dim*nv); for (int r=0; r < dim; r++) { for (int c=0; c <= r; c++) { - mju_addToScl(LTJ+c*nv, d->efc_J+(i+r)*nv, local[r*dim+c], nv); + mju_addToScl(LTJ+c*nv, ctx->J+(i+r)*nv, local[r*dim+c], nv); } } @@ -1651,8 +1791,8 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { // incremental update to Hessian factor due to changes in efc_state -static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, const int* oldstate) { - int rank, nv = m->nv, nefc = d->nefc; +static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate) { + int rank, nv = ctx->nv, nefc = ctx->nefc; mj_markStack(d); // local space @@ -1667,32 +1807,32 @@ static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, co int flag_update = -1; // add quad - if (oldstate[i] != mjCNSTRSTATE_QUADRATIC && d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) { + if (oldstate[i] != mjCNSTRSTATE_QUADRATIC && ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC) { flag_update = 1; } // subtract quad - else if (oldstate[i] == mjCNSTRSTATE_QUADRATIC && d->efc_state[i] != mjCNSTRSTATE_QUADRATIC) { + else if (oldstate[i] == mjCNSTRSTATE_QUADRATIC && ctx->efc_state[i] != mjCNSTRSTATE_QUADRATIC) { flag_update = 0; } // perform update if flagged if (flag_update != -1) { // update with vec = J(i,:)*sqrt(D[i])) - if (mj_isSparse(m)) { + if (ctx->is_sparse) { // get nnz and adr of row i - const int nnz = d->efc_J_rownnz[i], adr = d->efc_J_rowadr[i]; + const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i]; // scale vec, copy colind - mju_scl(vec, d->efc_J+adr, mju_sqrt(d->efc_D[i]), nnz); - mju_copyInt(vec_ind, d->efc_J_colind+adr, nnz); + mju_scl(vec, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz); + mju_copyInt(vec_ind, ctx->J_colind+adr, nnz); // sparse update or downdate rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, vec_ind, d); } else { - mju_scl(vec, d->efc_J+i*nv, mju_sqrt(d->efc_D[i]), nv); + mju_scl(vec, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv); rank = mju_cholUpdate(ctx->L, vec, nv, flag_update); } ctx->nupdate++; @@ -1700,7 +1840,7 @@ static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, co // recompute H directly if accuracy lost if (rank < nv) { mj_freeStack(d); - FactorizeHessian(m, d, ctx, /*flg_recompute=*/1); + FactorizeHessian(d, ctx, /*flg_recompute=*/1); // nothing else to do return; @@ -1710,7 +1850,7 @@ static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, co // add cones if present if (ctx->ncone) { - HessianCone(m, d, ctx); + HessianCone(d, ctx); } mj_freeStack(d); @@ -1726,14 +1866,13 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, mjCGContext ctx; mj_markStack(d); - // allocate context - CGallocate(m, d, &ctx, island, flg_Newton); + // make context + CGpointers(m, d, &ctx, island); + CGallocate(d, &ctx, flg_Newton); // local copies int nv = ctx.nv; int nefc = ctx.nefc; - const int* dofind = ctx.dofind; - const int* efcind = ctx.efcind; // allocate local storage if (!flg_Newton) { @@ -1743,27 +1882,28 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } int* oldstate = mjSTACKALLOC(d, nefc, int); - // initialize matrix-vector products - int flg_vecunc = 1; // d->qacc is uncompressed - mj_mulM_island(m, d, ctx.Ma, d->qacc, island, flg_vecunc); - int flg_resunc = 0; // ctx.Jaref is compressed - mj_mulJacVec_island(m, d, ctx.Jaref, d->qacc, island, flg_resunc, flg_vecunc); - if (island < 0) { - mju_subFrom(ctx.Jaref, d->efc_aref, nefc); + // compute Ma = M * qacc + mju_mulSymVecSparse(ctx.Ma, ctx.M, ctx.qacc, nv, + ctx.M_rownnz, ctx.M_rowadr, ctx.M_colind); + + + // compute Jaref = J * qacc - aref (dense or sparse) + if (!ctx.is_sparse) { + mju_mulMatVec(ctx.Jaref, ctx.J, ctx.qacc, nefc, nv); } else { - for (int c=0; c < nefc; c++) { - ctx.Jaref[c] -= d->efc_aref[efcind[c]]; - } + mju_mulMatVecSparse(ctx.Jaref, ctx.J, ctx.qacc, nefc, + ctx.J_rownnz, ctx.J_rowadr, ctx.J_colind, ctx.J_rowsuper); } + mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc); // first update - CGupdateConstraint(m, d, &ctx); + CGupdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); if (flg_Newton) { // compute and factorize Hessian - MakeHessian(m, d, &ctx); - FactorizeHessian(m, d, &ctx, /*flg_recompute=*/0); + MakeHessian(d, &ctx); + FactorizeHessian(d, &ctx, /*flg_recompute=*/0); } - CGupdateGradient(m, d, &ctx); + CGupdateGradient(&ctx, flg_Newton); // start both with preconditioned gradient mju_scl(ctx.search, ctx.Mgrad, -1, nv); @@ -1774,8 +1914,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv)); } else { mjtNum island_inertia = 0; - for (int c=0; c < nv; c++) { - island_inertia += d->qM[m->dof_Madr[dofind[c]]]; + for (int i=0; i < nv; i++) { + int diag_i = ctx.M_rowadr[i] + ctx.M_rownnz[i] - 1; + island_inertia += ctx.M[diag_i]; } scale = 1 / island_inertia; } @@ -1784,7 +1925,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, // main loop while (iter < maxiter) { // perform linesearch - alpha = CGsearch(m, d, &ctx); + alpha = CGsearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations); // no improvement: done if (alpha == 0) { @@ -1792,13 +1933,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } // move to new solution - if (island < 0) { - mju_addToScl(d->qacc, ctx.search, alpha, nv); - } else { - for (int c=0; c < nv; c++) { - d->qacc[dofind[c]] += alpha * ctx.search[c]; - } - } + mju_addToScl(ctx.qacc, ctx.search, alpha, nv); mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv); mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc); @@ -1807,27 +1942,20 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, mju_copy(gradold, ctx.grad, nv); mju_copy(Mgradold, ctx.Mgrad, nv); } - if (island < 0) { - mju_copyInt(oldstate, d->efc_state, nefc); - } else { - for (int c=0; c < nefc; c++) { - oldstate[c] = d->efc_state[efcind[c]]; - } - } + mju_copyInt(oldstate, ctx.efc_state, nefc); mjtNum oldcost = ctx.cost; // update - CGupdateConstraint(m, d, &ctx); + CGupdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); if (flg_Newton) { - HessianIncremental(m, d, &ctx, oldstate); + HessianIncremental(d, &ctx, oldstate); } - CGupdateGradient(m, d, &ctx); + CGupdateGradient(&ctx, flg_Newton); // count state changes int nchange = 0; - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - nchange += (d->efc_state[i] != oldstate[c]); + for (int i=0; i < nefc; i++) { + nchange += (ctx.efc_state[i] != oldstate[i]); } // scale improvement, gradient, save stats @@ -1859,8 +1987,8 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } // update - for (int c=0; c < nv; c++) { - ctx.search[c] = -ctx.Mgrad[c] + beta*ctx.search[c]; + for (int i=0; i < nv; i++) { + ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i]; } } } @@ -1910,3 +2038,10 @@ void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter) { void mj_solNewton(const mjModel* m, mjData* d, int maxiter) { mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/1); } + + + +// Newton entry point (one island) +void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter) { + mj_solCGNewton(m, d, island, maxiter, /*flg_Newton=*/1); +} diff --git a/src/engine/engine_solver.h b/src/engine/engine_solver.h index 6947489a..7ee007de 100644 --- a/src/engine/engine_solver.h +++ b/src/engine/engine_solver.h @@ -32,9 +32,13 @@ void mj_solCG(const mjModel* m, mjData* d, int maxiter); // Newton solver void mj_solNewton(const mjModel* m, mjData* d, int maxiter); + //------------------------------ per-island solvers ------------------------------------------------ // CG solver void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter); +// Newton entry point +void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter); + #endif // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_ diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 4f5d06af..2173879a 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -41,8 +41,8 @@ //-------------------------- Constants ------------------------------------------------------------- - #define mjVERSION 331 -#define mjVERSIONSTRING "3.3.1" + #define mjVERSION 333 +#define mjVERSIONSTRING "3.3.3" // names of disable flags const char* mjDISABLESTRING[mjNDISABLE] = { @@ -972,14 +972,9 @@ void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) { -// multiply vector by inertia matrix -void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { - int nv = m->nv; - const mjtNum* M = d->qM; - const int* Madr = m->dof_Madr; - const int* parentid = m->dof_parentid; - const int* simplenum = m->dof_simplenum; - +// multiply vector by inertia matrix (implementation) +void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M, + const int* Madr, const int* parentid, const int* simplenum) { mju_zero(res, nv); for (int i=0; i < nv; i++) { @@ -1031,64 +1026,9 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) -// multiply vector by inertia matrix for one dof island -void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_vecunc) { - // if no island, call regular function - if (island < 0) { - mj_mulM(m, d, res, vec); - return; - } - - // local constants: general - const mjtNum* M = d->qM; - const int* Madr = m->dof_Madr; - const int* parentid = m->dof_parentid; - const int* simplenum = m->dof_simplenum; - - // local constants: island specific - int ndof = d->island_dofnum[island]; - const int* dofind = d->island_dofind + d->island_dofadr[island]; - const int* islandind = d->dof_islandind; - - mju_zero(res, ndof); - - for (int k=0; k < ndof; k++) { - // address in full dof vector - int i = dofind[k]; - - // address in M - int adr = Madr[i]; - - // diagonal - if (flg_vecunc) { - res[k] = M[adr]*vec[i]; - } else { - res[k] = M[adr]*vec[k]; - } - - // simple dof: continue - if (simplenum[i]) { - continue; - } - - // off-diagonal - int j = parentid[i]; - while (j >= 0) { - adr++; - int l = islandind[j]; - if (flg_vecunc) { - res[k] += M[adr]*vec[j]; - res[l] += M[adr]*vec[i]; - } else { - res[k] += M[adr]*vec[l]; - res[l] += M[adr]*vec[k]; - } - - // advance to parent - j = parentid[j]; - } - } +// multiply vector by inertia matrix +void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { + mj_mulM_impl(res, vec, m->nv, d->qM, m->dof_Madr, m->dof_parentid, m->dof_simplenum); } @@ -1108,7 +1048,7 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) // non-simple: add off-diagonals if (!m->dof_simplenum[i]) { int adr = d->M_rowadr[i]; - res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr, /*flg_unc1=*/0); + res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr); } } @@ -1122,77 +1062,26 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) // add inertia matrix to destination matrix -// destination can be sparse uncompressed, or dense when all int* are NULL +// destination can be sparse or dense when all int* are NULL void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind) { + int nv = m->nv; // sparse if (rownnz && rowadr && colind) { - int nC = m->nC; mj_markStack(d); + mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum); + int* buf_ind = mjSTACKALLOC(d, nv, int); - // create reduced sparse inertia matrix C - mjtNum* C = mjSTACKALLOC(d, nC, mjtNum); - for (int i=0; i < nC; i++) { - C[i] = d->qM[d->mapM2C[i]]; - } + mju_addToMatSparse(dst, rownnz, rowadr, colind, nv, + d->M, d->M_rownnz, d->M_rowadr, d->M_colind, + buf_val, buf_ind); - mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, - d->C_rownnz, d->C_rowadr, d->C_colind); mj_freeStack(d); } // dense else { - mj_addMDense(m, d, dst); - } -} - - - -// add inertia matrix to sparse destination matrix -void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst, - int* rownnz, int* rowadr, int* colind, mjtNum* M, - int* M_rownnz, int* M_rowadr, int* M_colind) { - int nv = m->nv; - - mj_markStack(d); - int* buf_ind = mjSTACKALLOC(d, nv, int); - mjtNum* sparse_buf = mjSTACKALLOC(d, nv, mjtNum); - - // add to destination - for (int i=0; i < nv; i++) { - rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1, - rownnz[i], M_rownnz[i], colind + rowadr[i], - M_colind + M_rowadr[i], sparse_buf, buf_ind); - } - mj_freeStack(d); -} - - - -// add inertia matrix to dense destination matrix -void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) { - int nv = m->nv; - - for (int i = 0; i < nv; i++) { - int adr = m->dof_Madr[i]; - int j = i; - while (j >= 0) { - // add - dst[i*nv+j] += d->qM[adr]; - if (j < i) { - dst[j*nv+i] += d->qM[adr]; - } - - // only diagonal if simplenum - if (m->dof_simplenum[i]) { - break; - } - - // advance - j = m->dof_parentid[j]; - adr++; - } + mju_addToSymSparse(dst, d->M, nv, d->M_rownnz, d->M_rowadr, d->M_colind, /*flg_upper=*/ 1); } } diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index c188ea38..3d5f14c3 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -120,29 +120,21 @@ MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body); // convert sparse inertia matrix M into full matrix MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M); +// multiply vector by inertia matrix (implementation) +MJAPI void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M, + const int* Madr, const int* parentid, const int* simplenum); + // multiply vector by inertia matrix MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply vector by inertia matrix for one dof island -MJAPI void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_vecunc); - // multiply vector by (inertia matrix)^(1/2) MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); // add inertia matrix to destination matrix -// destination can be sparse uncompressed, or dense when all int* are NULL +// destination can be sparse or dense when all int* are NULL MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind); -// add inertia matrix to sparse destination matrix -MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst, - int* rownnz, int* rowadr, int* colind, mjtNum* M, - int* M_rownnz, int* M_rowadr, int* M_colind); - -// add inertia matrix to dense destination matrix -MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst); - //-------------------------- perturbations --------------------------------------------------------- diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c index c870c6fb..95055d72 100644 --- a/src/engine/engine_util_blas.c +++ b/src/engine/engine_util_blas.c @@ -843,9 +843,9 @@ void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, } - -// compute M'*diag*M (diag=NULL: compute M'*M) -void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) { +// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional +void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, + int nr, int nc, int flg_upper) { mjtNum tmp; // half of MatMat routine: only lower triangle @@ -870,15 +870,23 @@ void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, in } } - // make symmetric - for (int i=0; i < nc; i++) { - for (int j=i+1; j < nc; j++) { - res[i*nc+j] = res[j*nc+i]; + // flg_upper is set: make symmetric + if (flg_upper) { + for (int i=0; i < nc; i++) { + for (int j=i+1; j < nc; j++) { + res[i*nc+j] = res[j*nc+i]; + } } } } +// compute M'*diag*M (diag=NULL: compute M'*M) +void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) { + mju_sqrMatTD_impl(res, mat, diag, nr, nc, /*flg_upper=*/ 1); +} + + // multiply matrices, first argument transposed void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, diff --git a/src/engine/engine_util_blas.h b/src/engine/engine_util_blas.h index 619a3225..4b2519c9 100644 --- a/src/engine/engine_util_blas.h +++ b/src/engine/engine_util_blas.h @@ -223,6 +223,10 @@ MJAPI void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, MJAPI void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, int r1, int c1, int c2); +// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional +void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc, + int flg_upper); + // compute M'*diag*M (diag=NULL: compute M'*M) MJAPI void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc); diff --git a/src/engine/engine_util_errmem.c b/src/engine/engine_util_errmem.c index 7c63c1ce..04e98009 100644 --- a/src/engine/engine_util_errmem.c +++ b/src/engine/engine_util_errmem.c @@ -97,7 +97,7 @@ void mju_writeLog(const char* type, const char* msg) { // get time time(&rawtime); -#if defined(_POSIX_C_SOURCE) || defined(__APPLE__) || defined(__STDC_VERSION_TIME_H__) +#if defined(_POSIX_C_SOURCE) || defined(__APPLE__) || defined(__STDC_VERSION_TIME_H__) || defined(__EMSCRIPTEN__) localtime_r(&rawtime, &timeinfo); #elif _MSC_VER localtime_s(&timeinfo, &rawtime); @@ -123,10 +123,9 @@ void mju_error_raw(const char* msg) { } else { // write to log and console mju_writeLog("ERROR", msg); - printf("ERROR: %s\n\nPress Enter to exit ...", msg); + printf("ERROR: %s\n\n", msg); - // pause, exit - getchar(); + // exit exit(EXIT_FAILURE); } } diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index 528b8f21..03435f48 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -107,12 +107,12 @@ static mjtNum wrap_circle(mjtNum pnt[4], const mjtNum end[4], const mjtNum* side return -1; } + mjtNum sqrt0 = mju_sqrt(sqlen0 - sqrad); + mjtNum sqrt1 = mju_sqrt(sqlen1 - sqrad); + // construct the two solutions, compute goodness mjtNum sol[2][2][2], good[2]; for (int i=0; i < 2; i++) { - mjtNum sqrt0 = mju_sqrt(sqlen0 - sqrad); - mjtNum sqrt1 = mju_sqrt(sqlen1 - sqrad); - int sgn = (i == 0 ? 1 : -1); sol[i][0][0] = (end[0]*sqrad + sgn*radius*end[1]*sqrt0)/sqlen0; @@ -1390,6 +1390,42 @@ void mju_n2d(double* res, const mjtNum* vec, int n) { +// gather +void mju_gather(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[i] = vec[ind[i]]; + } +} + + + +// scatter +void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[ind[i]] = vec[i]; + } +} + + + +// gather integers +void mju_gatherInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[i] = vec[ind[i]]; + } +} + + + +// scatter integers +void mju_scatterInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[ind[i]] = vec[i]; + } +} + + + // insertion sort, increasing order void mju_insertionSort(mjtNum* list, int n) { for (int i=1; i < n; i++) { diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index e266f814..f7f23782 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -156,6 +156,18 @@ MJAPI void mju_d2n(mjtNum* res, const double* vec, int n); // convert from mjtNum to double MJAPI void mju_n2d(double* res, const mjtNum* vec, int n); +// gather mjtNums +MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n); + +// scatter mjtNums +MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n); + +// gather integers +MJAPI void mju_gatherInt(int* res, const int* vec, const int* ind, int n); + +// scatter integers +MJAPI void mju_scatterInt(int* res, const int* vec, const int* ind, int n); + // insertion sort, increasing order MJAPI void mju_insertionSort(mjtNum* list, int n); diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 83fdb634..4d9bde30 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -191,6 +191,60 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, +// precount row non-zeros of reverse-Cholesky factor L, return total non-zeros +// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package' +// reads pattern from upper triangle +int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind, + int n, mjData* d) { + mj_markStack(d); + int* parent = mjSTACKALLOC(d, n, int); + int* flag = mjSTACKALLOC(d, n, int); + + // loop over rows in reverse order + for (int r = n - 1; r >= 0; r--) { + parent[r] = -1; + flag[r] = r; + L_rownnz[r] = 1; // start with 1 for diagonal + + // loop over non-zero columns of upper triangle + int start = rowadr[r]; + int end = start + rownnz[r]; + for (int c = start; c < end; c++) { + int i = colind[c]; + + // skip lower triangle + if (i <= r) { + continue; + } + + // traverse from i to ancestor, stop when row is flagged + while (flag[i] != r) { + // if not yet set, set parent to current row + if (parent[i] == -1) { + parent[i] = r; + } + + // increment non-zeros, flag row i, advance to parent + L_rownnz[i]++; + flag[i] = r; + i = parent[i]; + } + } + } + + mj_freeStack(d); + + // sum up all row non-zeros + int nnz = 0; + for (int r = 0; r < n; r++) { + nnz += L_rownnz[r]; + } + + return nnz; +} + + + // sparse reverse-order Cholesky solve void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n, const int* rownnz, const int* rowadr, const int* colind) { @@ -221,7 +275,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int // x(i) -= sum_j L(i,j)*x(j), j=0:i-1 if (nnz > 1) { - res[i] -= mju_dotSparse(mat+adr, res, nnz-1, colind+adr, /*flg_unc1=*/0); + res[i] -= mju_dotSparse(mat+adr, res, nnz-1, colind+adr); // modulo AVX, the above line does // for (int j=0; j 0) { int adr = rowadr[i] + d1; - res[i] -= mju_dotSparse(LU+adr, res, nnz, colind+adr, /*flg_unc1=*/0); + res[i] -= mju_dotSparse(LU+adr, res, nnz, colind+adr); } } @@ -666,7 +720,7 @@ void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n, int d = diag[i]; int adr = rowadr[i]; if (d > 0) { - res[i] -= mju_dotSparse(LU+adr, res, d, colind+adr, /*flg_unc1=*/0); + res[i] -= mju_dotSparse(LU+adr, res, d, colind+adr); } // divide by diagonal element of L diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 6bbdc6bb..6a772e2c 100644 --- a/src/engine/engine_util_solve.h +++ b/src/engine/engine_util_solve.h @@ -38,6 +38,10 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, int* rownnz, const int* rowadr, int* colind, mjData* d); +// precount row non-zeros of reverse-Cholesky factor L, return total +MJAPI int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, + const int* colind, int n, mjData* d); + // sparse reverse-order Cholesky solve void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n, const int* rownnz, const int* rowadr, const int* colind); diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index abb95be0..c9cf3bab 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -18,6 +18,7 @@ #include #include +#include #include // IWYU pragma: keep #include #include "engine/engine_io.h" @@ -60,9 +61,8 @@ void mju_dotSparseX3(mjtNum* res0, mjtNum* res1, mjtNum* res2, // dot-product, both vectors are sparse -// flg_unc2: is vec2 memory layout uncompressed -mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1, int nnz2, - const int* ind2, int flg_unc2) { +mjtNum mju_dotSparse2(const mjtNum* vec1, const int* ind1, int nnz1, + const mjtNum* vec2, const int* ind2, int nnz2) { int i1 = 0, i2 = 0; mjtNum res = 0; @@ -77,12 +77,7 @@ mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const in // match: accumulate result, advance both if (adr1 == adr2) { - if (flg_unc2) { - res += vec1[i1++] * vec2[adr2]; - i2++; - } else { - res += vec1[i1++] * vec2[i2++]; - } + res += vec1[i1++] * vec2[i2++]; } // otherwise advance smaller @@ -161,7 +156,7 @@ void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, #else // regular sparse dot-product for (int r=0; r < nr; r++) { - res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc1=*/0); + res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]); } #endif // mjUSEAVX } @@ -192,6 +187,72 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int } +// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, + const mjtNum* M, const int* M_rownnz, const int* M_rowadr, + const int* M_colind, + mjtNum* buf_val, int* buf_ind) { + for (int i=0; i < nr; i++) { + rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1, + rownnz[i], M_rownnz[i], colind + rowadr[i], + M_colind + M_rowadr[i], buf_val, buf_ind); + } +} + + +// add symmetric matrix (lower triangle) to dense matrix, upper triangle optional +void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind, int flg_upper) { + for (int i=0; i < n; i++) { + int start = rowadr[i]; + int end = start + rownnz[i]; + for (int adr=start; adr < end; adr++) { + mjtNum val = mat[adr]; + int j = colind[adr]; + + // lower + diagonal + res[i*n + j] += val; + + // strict upper + if (flg_upper && j < i) { + res[j*n + i] += val; + } + } + } +} + + + +// multiply symmetric matrix (only lower triangle represented) by vector: +// res = (mat + strict_upper(mat')) * vec +void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat, + const mjtNum* restrict vec, int n, + const int* restrict rownnz, const int* restrict rowadr, + const int* restrict colind) { + // clear res + mju_zero(res, n); + + // multiply + for (int i=0; i < n; i++) { + int adr = rowadr[i]; + int diag = rownnz[i] - 1; + const mjtNum* row = mat + adr; + + // diagonal + res[i] = row[diag] * vec[i]; + + // off-diagonals + const int* ind = colind + adr; + for (int k=diag-1; k >= 0; k--) { + int j = ind[k]; + mjtNum val = row[k]; + res[i] += val * vec[j]; // strict lower + res[j] += val * vec[i]; // strict upper + } + } +} + + // count the number of non-zeros in the sum of two sparse vectors int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) { @@ -470,9 +531,9 @@ int mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, in -// transpose sparse matrix +// transpose sparse matrix, optionally compute row supernodes void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc, - int* res_rownnz, int* res_rowadr, int* res_colind, + int* res_rownnz, int* res_rowadr, int* res_colind, int* res_rowsuper, const int* rownnz, const int* rowadr, const int* colind) { // clear number of non-zeros for each row of transposed mju_zeroInt(res_rownnz, nc); @@ -486,22 +547,40 @@ void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc, } } + // init res_rowsuper + if (res_rowsuper) { + for (int i = 0; i < nc - 1; i++) { + res_rowsuper[i] = (res_rownnz[i] == res_rownnz[i + 1]); + } + res_rowsuper[nc - 1] = 0; + } + // compute the row addresses for the transposed matrix res_rowadr[0] = 0; for (int i = 1; i < nc; i++) { res_rowadr[i] = res_rowadr[i-1] + res_rownnz[i-1]; } - // iterate through each non-zero entry of mat + // iterate through each row (column) of mat (res) for (int r = 0; r < nr; r++) { + int c_prev = -1; int start = rowadr[r]; int end = start + rownnz[r]; for (int i = start; i < end; i++) { // swap rows with columns and increment res_rowadr - int c = res_rowadr[colind[i]]++; - res_colind[c] = r; + int c = colind[i]; + int adr = res_rowadr[c]++; + res_colind[adr] = r; if (res) { - res[c] = mat[i]; + res[adr] = mat[i]; + } + + // mark non-supernodes + if (res_rowsuper) { + if (c > 0 && c != c_prev + 1 && res_rowsuper[c - 1]) { + res_rowsuper[c - 1] = 0; + } + c_prev = c; } } } @@ -510,8 +589,16 @@ void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc, for (int i = nc-1; i > 0; i--) { res_rowadr[i] = res_rowadr[i-1]; } - res_rowadr[0] = 0; + + // accumulate supernodes + if (res_rowsuper) { + for (int i = nc - 2; i >= 0; i--) { + if (res_rowsuper[i]) { + res_rowsuper[i] += res_rowsuper[i + 1]; + } + } + } } @@ -651,8 +738,10 @@ void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) { -// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout -// res_rowadr is required to be precomputed +// max number of supernodes handled +#define mjMAXSUPER 8 + +// compute sparse M'*diag*M (diag=NULL: compute M'*M), res_rowadr must be precomputed void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag, int nr, int nc, int* res_rownnz, const int* res_rowadr, int* res_colind, @@ -661,6 +750,206 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, const int* rownnzT, const int* rowadrT, const int* colindT, const int* rowsuperT, mjData* d, int* diagind) { + mj_markStack(d); + + // reinterpret transposed matrices as compressed sparse column + const mjtNum* mat_csc = matT; + const int* colnnz = rownnzT; + const int* coladr = rowadrT; + const int* rowind = colindT; + const int* colsuper = rowsuperT; + const mjtNum* matT_csc = mat; + const int* colnnzT = rownnz; + const int* coladrT = rowadr; + const int* rowindT = colind; + // rowsuper is unused + + // marker[i] = 1 if row i is set in current column + int* marker = mjSTACKALLOC(d, nc, int); + mju_zeroInt(marker, nc); + + // dense buffer (considered column-major) containing up to mjMAXSUPER columns + mjtNum* buffer = mjSTACKALLOC(d, nc*mjMAXSUPER, mjtNum); + + // dense index vector of the current column (unsorted) + int* buffer_idx = mjSTACKALLOC(d, nc, int); + + // rowstart[i]: address of first row in column mat'[:, i] with index > current column + int* rowstart = mjSTACKALLOC(d, nr, int); + mju_zeroInt(rowstart, nr); + + // clear res_rownnz + mju_zeroInt(res_rownnz, nc); + + // construct res[lower+diagonal], by column + for (int c=0; c < nc; c++) { + int buffer_nnz = 0; + + // prepare column c of mat + int nnz = colnnz[c]; + int adr = coladr[c]; + const int* ind = rowind + adr; + + // val: array of ns > 0 column pointers with identical pattern to c + const mjtNum* val[mjMAXSUPER]; + + // first column is c + int ns = 1; + val[0] = mat_csc + adr; + + // add c's supernodes, if any + int cs; + if (colsuper && (cs = colsuper[c])) { + ns += mjMIN(cs, mjMAXSUPER - 1); + for (int s=1; s < ns; s++) { + val[s] = mat_csc + coladr[c + s]; + } + } + + // diagonal special-case: dense dot product of column c, with/out diag + mjtNum diag_c[mjMAXSUPER]; + if (diag) { + for (int s=0; s < ns; s++) { + mjtNum ds = 0; + for (int k=0; k < nnz; k++) { + ds += (val[s][k] * val[s][k]) * diag[ind[k]]; + } + diag_c[s] = ds; + } + } else { + for (int s=0; s < ns; s++) { + diag_c[s] = mju_dot(val[s], val[s], nnz); + } + } + + // in the strict lower triangle, compute + // res[:, c] = mat' * mat[:, c] = sum_r(diag[r] * mat'[:, r] * mat[:, c]) + for (int i=0; i < nnz; i++) { + // prepare column r of mat' + int r = ind[i]; + int adrT = coladrT[r]; + int nnzT = colnnzT[r]; + const int* indT = rowindT + adrT; + const mjtNum* valT = matT_csc + adrT; + + // get v[s] = diag[r] * mat[r, c + s] for s in [0, ns) + mjtNum v[mjMAXSUPER]; + if (diag) { + mjtNum diag_r = diag[r]; + for (int s=0; s < ns; s++) { + v[s] = diag_r * val[s][i]; + } + } else { + for (int s=0; s < ns; s++) { + v[s] = val[s][i]; + } + } + + // gather to dense buffer columns: buffer[:, s] += mat'[:, r] * v[s] + for (int k=rowstart[r]; k < nnzT; k++) { + int j = indT[k]; + + // if j is not in the strict lower triangle, increment rowstart and continue + if (j <= c) { + rowstart[r]++; + continue; + } + + // first nonzero in row j: mark and set value + if (!marker[j]) { + // mark j and save it + marker[j] = 1; + buffer_idx[buffer_nnz++] = j; + + // set value + mjtNum vk = valT[k]; + for (int s=0; s < ns; s++) { + buffer[s*nc + j] = vk * v[s]; + } + } + + // otherwise existing nonzero in row j: add to value + else { + mjtNum vk = valT[k]; + for (int s=0; s < ns; s++) { + buffer[s*nc + j] += vk * v[s]; + } + } + } + } + + // scatter to res from dense buffer: res[:, c + s] = buffer[:, s] for s in [0, ns) + + // write values under diagonal + for (int i=0; i < buffer_nnz; i++) { + int j = buffer_idx[i]; + marker[j] = 0; + int adr_j = res_rowadr[j] + res_rownnz[j]; + + // truncate row to strict lower triangle + int lower = j - c; + int nm = mjMIN(ns, lower); + + // increment nonzeros + res_rownnz[j] += nm; + + // write value + for (int s=0; s < nm; s++) { + res[adr_j + s] = buffer[s*nc + j]; + } + + // write index + for (int s=0; s < nm; s++) { + res_colind[adr_j + s] = c + s; + } + } + + // write diagonal value + for (int s=0; s < ns; s++) { + int adr_s = res_rowadr[c + s] + res_rownnz[c + s]++; + res_colind[adr_s] = c + s; + res[adr_s] = diag_c[s]; + } + + // supernode: skip ahead if ns > 1 + c += ns - 1; + } + + // upper triangle requested: save diagonal indices and fill + if (diagind) { + // save diagonal indices + for (int i=0; i < nc; i++) { + diagind[i] = res_rowadr[i] + res_rownnz[i] - 1; + } + + // fill upper triangle + for (int i=0; i < nc; i++) { + int start = res_rowadr[i]; + int end = start + res_rownnz[i] - 1; + for (int j=start; j < end; j++) { + int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++; + res[adr] = res[j]; + res_colind[adr] = i; + } + } + } + + mj_freeStack(d); +} + +#undef mjMAXSUPER + + + +// legacy row-based implementation (reference) +void mju_sqrMatTDSparse_row(mjtNum* res, const mjtNum* mat, const mjtNum* matT, + const mjtNum* diag, int nr, int nc, + int* res_rownnz, const int* res_rowadr, int* res_colind, + const int* rownnz, const int* rowadr, + const int* colind, const int* rowsuper, + const int* rownnzT, const int* rowadrT, + const int* colindT, const int* rowsuperT, + mjData* d, int* diagind) { // allocate space for accumulation buffer and matT mj_markStack(d); @@ -796,54 +1085,102 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, mj_freeStack(d); } -// precount row non-zeros of reverse-Cholesky factor L, return total non-zeros -// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package' -// reads pattern from upper triangle -int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind, - int n, mjData* d) { - mj_markStack(d); - int* parent = mjSTACKALLOC(d, n, int); - int* flag = mjSTACKALLOC(d, n, int); - // loop over rows in reverse order - for (int r = n - 1; r >= 0; r--) { - parent[r] = -1; - flag[r] = r; - L_rownnz[r] = 1; // start with 1 for diagonal - // loop over non-zero columns of upper triangle - int start = rowadr[r]; - int end = start + rownnz[r]; - for (int c = start; c < end; c++) { - int i = colind[c]; - - // skip lower triangle - if (i <= r) { - continue; - } - - // traverse from i to ancestor, stop when row is flagged - while (flag[i] != r) { - // if not yet set, set parent to current row - if (parent[i] == -1) { - parent[i] = r; - } - - // increment non-zeros, flag row i, advance to parent - L_rownnz[i]++; - flag[i] = r; - i = parent[i]; +// block-diagonalize a dense matrix +// res output matrix +// mat input matrix +// nc_mat number of columns in mat +// nc_res number of columns in res +// nb number of blocks +// perm_r reverse permutation of rows (res -> mat) +// perm_c reverse permutation of columns (res -> mat) +// block_nr number of rows in each block +// block_nc number of columns in each block +// block_r first row of each block +// block_c first column of each block +void mju_blockDiag(mjtNum* restrict res, const mjtNum* restrict mat, + int nc_mat, int nc_res, int nb, + const int* restrict perm_r, const int* restrict perm_c, + const int* restrict block_nr, const int* restrict block_nc, + const int* restrict block_r, const int* restrict block_c) { + for (int b=0; b < nb; b++) { + int bnr = block_nr[b]; + int bnc = block_nc[b]; + const int* adr_r = perm_r + block_r[b]; + const int* adr_c = perm_c + block_c[b]; + int adr = nc_res * block_r[b]; + for (int r = 0; r < bnr; r++) { + for (int c = 0; c < bnc; c++) { + res[adr++] = mat[nc_mat * adr_r[r] + adr_c[c]]; } } } - - mj_freeStack(d); - - // sum up all row non-zeros - int nnz = 0; - for (int r = 0; r < n; r++) { - nnz += L_rownnz[r]; - } - - return nnz; } + + +// block-diagonalize a sparse matrix +// res values of the target matrix res +// res_rownnz number of non-zeros in each row of res +// res_rowadr row address of each non-zero in res +// res_colind column index of each non-zero in res +// mat values of the source matrix mat +// mat_rownnz number of non-zeros in each row of mat +// mat_rowadr row address of each non-zero in mat +// mat_colind column index of each non-zero in mat +// nr number of rows in mat/res +// nb number of blocks +// perm_r reverse permutation of rows (res -> mat) +// perm_c forward permutation of columns (mat -> res) +// block_r first row of each block in res +// block_c first column of each block in res +// mat2 optional additional source matrix (same structure as mat) +// res2 optional additional target matrix (same structure as res) +void mju_blockDiagSparse(mjtNum* restrict res, int* restrict res_rownnz, + int* restrict res_rowadr, int* restrict res_colind, + const mjtNum* restrict mat, const int* restrict rownnz, + const int* restrict rowadr, const int* restrict colind, + int nr, int nb, + const int* restrict perm_r, const int* restrict perm_c, + const int* restrict block_r, const int* restrict block_c, + mjtNum* restrict res2, const mjtNum* restrict mat2) { + int block = 0; + int col_offset = block_c[block]; + int row_next = block + 1 < nb ? block_r[block + 1] : nr; + for (int r=0; r < nr; r++) { + // row k in mat goes to row r in res + int k = perm_r[r]; + + // rownnz + int nnz = rownnz[k]; + res_rownnz[r] = nnz; + + // rowadr + int res_adr = (r == 0) ? 0 : res_rowadr[r-1] + res_rownnz[r-1]; + res_rowadr[r] = res_adr; + + // colind + int* res_colind_r = res_colind + res_adr; + mjtNum* res_r = res + res_adr; + int mat_adr = rowadr[k]; + const int* colind_k = colind + mat_adr; + const mjtNum* mat_k = mat + mat_adr; + for (int j=0; j < nnz; j++) { + res_colind_r[j] = perm_c[colind_k[j]] - col_offset; + } + + // values (dense copy: partial order within block is guaranteed) + mju_copy(res_r, mat_k, nnz); + if (mat2 && res2) { + mju_copy(res2 + res_adr, mat2 + mat_adr, nnz); + } + + // end of block reached: update block counter, column offset, next row + if (r + 1 >= row_next && block + 1 < nb ) { + block++; + col_offset = block_c[block]; + row_next = block + 1 < nb ? block_r[block + 1] : nr; + } + } +} + diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 92c2dad0..f6c2f7c2 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -28,9 +28,9 @@ extern "C" { //------------------------------ sparse operations ------------------------------------------------- -// dot-product, both vectors are sparse, vec2 can be uncompressed -MJAPI mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1, - int nnz2, const int* ind2, int flg_unc2); +// dot-product, both vectors are sparse +MJAPI mjtNum mju_dotSparse2(const mjtNum* vec1, const int* ind1, int nnz1, + const mjtNum* vec2, const int* ind2, int nnz2); // convert matrix from dense to sparse // nnz is size of res and colind, return 1 if too small, 0 otherwise @@ -50,6 +50,22 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); +// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, + const mjtNum* M, const int* M_rownnz, const int* M_rowadr, + const int* M_colind, + mjtNum* buf_val, int* buf_ind); + +// add symmetric matrix (only lower triangle represented) to dense matrix +MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind, + int flg_upper); + +// multiply symmetric matrix (only lower triangle represented) by vector: +// res = (mat + strict_upper(mat')) * vec +MJAPI void mju_mulSymVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n, + const int* rownnz, const int* rowadr, const int* colind); + // compress sparse matrix, remove elements with abs(value) <= minval, return total non-zeros MJAPI int mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind, mjtNum minval); @@ -75,9 +91,9 @@ int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum s int mju_addChains(int* res, int n, int NV1, int NV2, const int* chain1, const int* chain2); -// transpose sparse matrix +// transpose sparse matrix, optionally compute row supernodes MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc, - int* res_rownnz, int* res_rowadr, int* res_colind, + int* res_rownnz, int* res_rowadr, int* res_colind, int* res_rowsuper, const int* rownnz, const int* rowadr, const int* colind); // construct row supernodes @@ -94,6 +110,16 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT const int* colindT, const int* rowsuperT, mjData* d, int* diagind); +// LEGACY: row-based implementation +MJAPI void mju_sqrMatTDSparse_row(mjtNum* res, const mjtNum* mat, const mjtNum* matT, + const mjtNum* diag, int nr, int nc, + int* res_rownnz, const int* res_rowadr, int* res_colind, + const int* rownnz, const int* rowadr, + const int* colind, const int* rowsuper, + const int* rownnzT, const int* rowadrT, + const int* colindT, const int* rowsuperT, + mjData* d, int* diagind); + // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr, const int* rownnz, const int* rowadr, const int* colind, @@ -103,19 +129,29 @@ MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr, // precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc); -// precount row non-zeros of reverse-Cholesky factor L, return total -MJAPI int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, - const int* colind, int n, mjData* d); +// block-diagonalize a dense matrix +MJAPI void mju_blockDiag(mjtNum* res, const mjtNum* mat, + int nc_mat, int nc_res, int nb, + const int* perm_r, const int* perm_c, + const int* block_nr, const int* block_nc, + const int* blockadr_r, const int* blockadr_c); + +// block-diagonalize a sparse matrix +MJAPI void mju_blockDiagSparse( + mjtNum* res, int* res_rownnz, int* res_rowadr, int* res_colind, + const mjtNum* mat, const int* rownnz, const int* rowadr, const int* colind, + int nr, int nb, + const int* perm_r, const int* perm_c, + const int* block_r, const int* block_c, + mjtNum* res2, const mjtNum* mat2); // ------------------------------ inlined functions ------------------------------------------------ // dot-product, first vector is sparse -// flg_unc1: is vec1 memory layout uncompressed static inline -mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1, - int flg_unc1) { +mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1) { #ifdef mjUSEAVX - return mju_dotSparse_avx(vec1, vec2, nnz1, ind1, flg_unc1); + return mju_dotSparse_avx(vec1, vec2, nnz1, ind1); #else int i = 0; mjtNum res = 0; @@ -125,33 +161,18 @@ mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int mjtNum res2 = 0; mjtNum res3 = 0; - - if (flg_unc1) { - for (; i <= n_4; i+=4) { - res0 += vec1[ind1[i+0]] * vec2[ind1[i+0]]; - res1 += vec1[ind1[i+1]] * vec2[ind1[i+1]]; - res2 += vec1[ind1[i+2]] * vec2[ind1[i+2]]; - res3 += vec1[ind1[i+3]] * vec2[ind1[i+3]]; - } - } else { - for (; i <= n_4; i+=4) { - res0 += vec1[i+0] * vec2[ind1[i+0]]; - res1 += vec1[i+1] * vec2[ind1[i+1]]; - res2 += vec1[i+2] * vec2[ind1[i+2]]; - res3 += vec1[i+3] * vec2[ind1[i+3]]; - } + for (; i <= n_4; i+=4) { + res0 += vec1[i+0] * vec2[ind1[i+0]]; + res1 += vec1[i+1] * vec2[ind1[i+1]]; + res2 += vec1[i+2] * vec2[ind1[i+2]]; + res3 += vec1[i+3] * vec2[ind1[i+3]]; } + res = (res0 + res2) + (res1 + res3); // scalar part - if (flg_unc1) { - for (; i < nnz1; i++) { - res += vec1[ind1[i]] * vec2[ind1[i]]; - } - } else { - for (; i < nnz1; i++) { - res += vec1[i] * vec2[ind1[i]]; - } + for (; i < nnz1; i++) { + res += vec1[i] * vec2[ind1[i]]; } return res; diff --git a/src/engine/engine_util_sparse_avx.h b/src/engine/engine_util_sparse_avx.h index 708db7c7..bdc9e931 100644 --- a/src/engine/engine_util_sparse_avx.h +++ b/src/engine/engine_util_sparse_avx.h @@ -30,10 +30,8 @@ //------------------------------ sparse operations using avx --------------------------------------- // dot-product, first vector is sparse -// flg_unc1: is vec1 memory layout uncompressed static inline -mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1, - int flg_unc1) { +mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1) { int i = 0; mjtNum res = 0; int nnz1_4 = nnz1 - 4; @@ -48,43 +46,22 @@ mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const vec2[ind1[2]], vec2[ind1[1]], vec2[ind1[0]]); - if (flg_unc1) { - val1 = _mm256_set_pd(vec1[ind1[3]], - vec1[ind1[2]], - vec1[ind1[1]], - vec1[ind1[0]]); - } else { - val1 = _mm256_loadu_pd(vec1); - } + + val1 = _mm256_loadu_pd(vec1); + sum = _mm256_mul_pd(val1, val2); i = 4; // parallel computation - if (flg_unc1) { - while (i<=nnz1_4) { - val1 = _mm256_set_pd(vec1[ind1[i+3]], - vec1[ind1[i+2]], - vec1[ind1[i+1]], - vec1[ind1[i+0]]); - val2 = _mm256_set_pd(vec2[ind1[i+3]], - vec2[ind1[i+2]], - vec2[ind1[i+1]], - vec2[ind1[i+0]]); - prod = _mm256_mul_pd(val1, val2); - sum = _mm256_add_pd(sum, prod); - i += 4; - } - } else { - while (i<=nnz1_4) { - val1 = _mm256_loadu_pd(vec1+i); - val2 = _mm256_set_pd(vec2[ind1[i+3]], - vec2[ind1[i+2]], - vec2[ind1[i+1]], - vec2[ind1[i+0]]); - prod = _mm256_mul_pd(val1, val2); - sum = _mm256_add_pd(sum, prod); - i += 4; - } + while (i<=nnz1_4) { + val1 = _mm256_loadu_pd(vec1+i); + val2 = _mm256_set_pd(vec2[ind1[i+3]], + vec2[ind1[i+2]], + vec2[ind1[i+1]], + vec2[ind1[i+0]]); + prod = _mm256_mul_pd(val1, val2); + sum = _mm256_add_pd(sum, prod); + i += 4; } // reduce @@ -96,14 +73,8 @@ mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const } // scalar part - if (flg_unc1) { - for (; i < nnz1; i++) { - res += vec1[ind1[i]] * vec2[ind1[i]]; - } - } else { - for (; i < nnz1; i++) { - res += vec1[i] * vec2[ind1[i]]; - } + for (; i < nnz1; i++) { + res += vec1[i] * vec2[ind1[i]]; } return res; @@ -209,7 +180,7 @@ void mju_mulMatVecSparse_avx(mjtNum* res, const mjtNum* mat, const mjtNum* vec, if (!rowsuper) { // regular sparse dot-product for (int r=0; r0) { - res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc1=*/0); + res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]); r++; rs--; @@ -243,7 +214,7 @@ void mju_mulMatVecSparse_avx(mjtNum* res, const mjtNum* mat, const mjtNum* vec, } else { - res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc1=*/0); + res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]); } } } diff --git a/src/engine/engine_util_spatial.c b/src/engine/engine_util_spatial.c index c170d923..dda47cf6 100644 --- a/src/engine/engine_util_spatial.c +++ b/src/engine/engine_util_spatial.c @@ -532,7 +532,7 @@ void mju_makeFrame(mjtNum frame[9]) { } // if yaxis undefined, set yaxis to (0,1,0) if possible, otherwise (0,0,1) - if (mju_norm3(frame+3) < 0.5) { + if (mju_dot3(frame+3, frame+3) < 0.25) { mju_zero3(frame+3); if (frame[1] < 0.5 && frame[1] > -0.5) { diff --git a/src/engine/engine_vis_state.c b/src/engine/engine_vis_state.c deleted file mode 100644 index 9a0ea210..00000000 --- a/src/engine/engine_vis_state.c +++ /dev/null @@ -1,415 +0,0 @@ -// 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. - -#include "engine/engine_vis_state.h" - -#include - -#include -#include -#include -#include -#include -#include "engine/engine_core_constraint.h" -#include "engine/engine_plugin.h" -#include "engine/engine_support.h" -#include "engine/engine_util_errmem.h" -#include "engine/engine_vis_init.h" -#include "engine/engine_vis_interact.h" -#include "engine/engine_vis_visualize.h" - -#ifdef MEMORY_SANITIZER - #include -#endif - -// this source file needs to treat XMJV differently from other X macros -#undef XMJV - - - -// round size up to multiples of 64-byte cache lines -static inline size_t roundUpToCacheLine(size_t n) { - return 64 * ((n / 64) + (n % 64 ? 1 : 0)); -} - - - -// set default scene -void mjv_defaultSceneState(mjvSceneState* scnstate) { - memset(scnstate, 0, sizeof(mjvSceneState)); - mjv_defaultScene(&scnstate->scratch); -} - - - -// allocate and init scene state -void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnstate, int maxgeom) { - mjv_freeScene(&scnstate->scratch); - mju_free(scnstate->buffer); - -#ifdef MEMORY_SANITIZER - __msan_allocated_memory(scnstate, sizeof(mjvSceneState)); - mjv_defaultScene(&scnstate->scratch); -#endif - - scnstate->nbuffer = 0; - scnstate->maxgeom = maxgeom; - -#define X(var) -#define XMJV(var) scnstate->model.var = m->var; - MJMODEL_INTS -#undef XMJV -#undef X - -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - scnstate->nbuffer += roundUpToCacheLine(sizeof(dtype) * m->dim0 * dim1); - MJMODEL_POINTERS -#undef XMJV -#undef X - -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - scnstate->nbuffer += roundUpToCacheLine(sizeof(dtype) * m->dim0 * dim1); - MJDATA_POINTERS -#undef XMJV -#undef X - - // create an arena in the scnstate, to allow visualization code to use the stack. - // TODO: Consider allocating way less than narena, since stack allocations in - // visualization code are much smaller than the arena space required by the model, - // typically. - scnstate->nbuffer += roundUpToCacheLine(m->narena); - // buffer space required for contacts - int condimmax = mj_isPyramidal(m) ? 10 : 6; - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->contact) * maxgeom); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->efc_force) * maxgeom * condimmax); - - // buffer space required for islands - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofadr) * m->ntree); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofind) * m->nv); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->dof_island) * m->nv); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->efc_island) * maxgeom * condimmax); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->tendon_efcadr) * m->ntendon); - - scnstate->buffer = mju_malloc(scnstate->nbuffer); - - char* ptr = scnstate->buffer; - -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - scnstate->model.var = (dtype*)ptr; \ - ptr += roundUpToCacheLine(sizeof(dtype) * m->dim0 * dim1); - MJMODEL_POINTERS -#undef XMJV -#undef X - -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - scnstate->data.var = (dtype*)ptr; \ - ptr += roundUpToCacheLine(sizeof(dtype) * m->dim0 * dim1); - MJDATA_POINTERS -#undef XMJV -#undef X - - scnstate->model.narena = m->narena; - scnstate->data.arena = (void*)ptr; - ptr += roundUpToCacheLine(m->narena); - - scnstate->data.contact = (mjContact*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.contact) * scnstate->maxgeom); - - scnstate->data.efc_force = (mjtNum*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.efc_force) * scnstate->maxgeom * condimmax); - - scnstate->data.island_dofadr = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofadr) * scnstate->model.ntree); - - scnstate->data.island_dofind = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofind) * scnstate->model.nv); - - scnstate->data.dof_island = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.dof_island) * scnstate->model.nv); - - scnstate->data.efc_island = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.efc_island) * scnstate->maxgeom * condimmax); - - scnstate->data.tendon_efcadr = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.tendon_efcadr) * m->ntendon); - - // should not occur - if (ptr - (char*)scnstate->buffer != scnstate->nbuffer) { - mjERROR("mjvSceneState buffer is not fully used"); - } - - mjv_makeScene(m, &scnstate->scratch, maxgeom); -} - - - -// free scene state -void mjv_freeSceneState(mjvSceneState* scnstate) { - mjv_freeScene(&scnstate->scratch); - mju_free(scnstate->buffer); - mjv_defaultSceneState(scnstate); -} - - - -// shallow copy scene state into model and data for use with mjv functions -void mjv_assignFromSceneState(const mjvSceneState* scnstate, mjModel* m, mjData* d) { - if (m) { - memset(m, 0, sizeof(mjModel)); - -#ifdef MEMORY_SANITIZER - // Tell msan to treat the entire buffer as uninitialized - __msan_allocated_memory(m, sizeof(mjModel)); -#endif - -#define X(var) -#define XMJV(var) m->var = scnstate->model.var; - MJMODEL_INTS -#undef XMJV -#undef X - - m->opt = scnstate->model.opt; - m->vis = scnstate->model.vis; - m->stat = scnstate->model.stat; - m->narena = scnstate->model.narena; - -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) m->var = scnstate->model.var; - MJMODEL_POINTERS -#undef XMJV -#undef X - } - - if (d) { - memset(d, 0, sizeof(mjData)); - -#ifdef MEMORY_SANITIZER - // Tell msan to treat the entire buffer as uninitialized - __msan_allocated_memory(d, sizeof(mjData)); -#endif - - memcpy(d->warning, scnstate->data.warning, sizeof(d->warning)); - d->threadpool = 0; - d->nefc = scnstate->data.nefc; - d->ncon = scnstate->data.ncon; - d->nisland = scnstate->data.nisland; - d->time = scnstate->data.time; - d->narena = scnstate->model.narena; - d->arena = scnstate->data.arena; - d->parena = 0; - d->pbase = 0; - d->pstack = 0; - - #define X(dtype, var, dim0, dim1) - #define XMJV(dtype, var, dim0, dim1) d->var = scnstate->data.var; - MJDATA_POINTERS - #undef XMJV - #undef X - - d->contact = scnstate->data.contact; - d->efc_force = scnstate->data.efc_force; - d->island_dofadr = scnstate->data.island_dofadr; - d->island_dofind = scnstate->data.island_dofind; - d->dof_island = scnstate->data.dof_island; - d->efc_island = scnstate->data.efc_island; - d->tendon_efcadr = scnstate->data.tendon_efcadr; - } -} - - - -// update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings -int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt, - const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn) { - // shallow-copy scnstate pointers into mjModel and mjData - mjModel m; - mjData d; - mjv_assignFromSceneState(scnstate, &m, &d); - - // save the number of mjWARN_VGEOMFULL warnings before the scene update - int warning_start = d.warning[mjWARN_VGEOMFULL].number; - - // copy mjvGeoms added by plugins - int nplugingeom = scnstate->scratch.ngeom; - if (nplugingeom > scn->maxgeom) { - mj_warning(&d, mjWARN_VGEOMFULL, scn->maxgeom); - scn->ngeom = scn->maxgeom; - } else { - scn->ngeom = nplugingeom; - } - memcpy(scn->geoms, scnstate->scratch.geoms, sizeof(mjvGeom) * scn->ngeom); - - // add all categories - mjv_addGeoms(&m, &d, opt, pert, catmask, scn); - - // update camera - mjv_updateCamera(&m, &d, cam, scn); - - // add lights - mjv_makeLights(&m, &d, scn); - - // update flexes - if (opt->flags[mjVIS_FLEXVERT] || opt->flags[mjVIS_FLEXEDGE] || - opt->flags[mjVIS_FLEXFACE] || opt->flags[mjVIS_FLEXSKIN]) { - mjv_updateActiveFlex(&m, &d, scn, opt); - } - - // update skins - if (opt->flags[mjVIS_SKIN]) { - mjv_updateActiveSkin(&m, &d, scn, opt); - } - - // return the number of new mjWARN_VGEOMFULL warnings generated - return d.warning[mjWARN_VGEOMFULL].number - warning_start; -} - - - -// update a scene state from model and data -void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt, - mjvSceneState* scnstate) { - // Check that mjModel sizes haven't changed. -#define X(var) -#define XMJV(var) \ - if (scnstate->model.var != m->var) { \ - mjERROR("m->%s changed: %d vs %d", #var, scnstate->model.var, m->var); \ - } - MJMODEL_INTS -#undef XMJV -#undef X - - // Update plugin visualization cache. - scnstate->scratch.ngeom = 0; - if (m->nplugin) { - const int nslot = mjp_pluginCount(); - // iterate over plugins, call visualize if defined - for (int i=0; i < m->nplugin; i++) { - const int slot = m->plugin[i]; - const mjpPlugin* plugin = mjp_getPluginAtSlotUnsafe(slot, nslot); - if (!plugin) { - mjERROR("invalid plugin slot: %d", slot); - } - if (plugin->visualize) { - plugin->visualize(m, d, opt, &scnstate->scratch, i); - } - } - } - - // Copy variable-sized arrays in mjModel. -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - memcpy(scnstate->model.var, m->var, sizeof(dtype) * m->dim0 * dim1); - MJMODEL_POINTERS -#undef XMJV -#undef X - - scnstate->model.opt = m->opt; - scnstate->model.vis = m->vis; - scnstate->model.stat = m->stat; - - // Copy mjData variables. - memcpy(scnstate->data.warning, d->warning, sizeof(d->warning)); - scnstate->data.time = d->time; - - // Copy variable-sized arrays in mjData. -#define X(dtype, var, dim0, dim1) -#define XMJV(dtype, var, dim0, dim1) \ - memcpy(scnstate->data.var, d->var, sizeof(dtype) * m->dim0 * dim1); - MJDATA_POINTERS -#undef XMJV -#undef X - - // Copy contacts. - { - if (d->ncon > scnstate->maxgeom) { - mj_warning(d, mjWARN_VGEOMFULL, scnstate->maxgeom); - scnstate->data.ncon = scnstate->maxgeom; - } else { - scnstate->data.ncon = d->ncon; - } - memcpy(scnstate->data.contact, d->contact, sizeof(*d->contact) * scnstate->data.ncon); - } - - // Copy only the entries in efc_force and efc_island that correspond to contacts. - { - scnstate->data.nefc = 0; - for (int i = 0; i < scnstate->data.ncon; ++i) { - const mjContact* con = &d->contact[i]; - scnstate->data.nefc += con->dim; - } - scnstate->data.nefc += scnstate->model.ntendon; - - int efc_address = 0; - int ispyramid = mj_isPyramidal(m); - for (int i = 0; i < scnstate->data.ncon; ++i) { - mjContact* con = &scnstate->data.contact[i]; - int dim = con->dim; - if (ispyramid && dim > 1){ - dim = 2*(dim - 1); - } - for (int j = 0; j < dim; ++j) { - scnstate->data.efc_force[efc_address + j] = d->efc_force[con->efc_address + j]; - if (d->nisland) { - scnstate->data.efc_island[efc_address + j] = d->efc_island[con->efc_address + j]; - } - } - con->efc_address = efc_address; - efc_address += dim; - } - if (d->nisland) { - for (int i = 0; i < scnstate->model.ntendon; ++i) { - int efcadr = d->tendon_efcadr[i]; - if (efcadr != -1) { - scnstate->data.efc_island[efcadr] = d->efc_island[efcadr]; - } - } - } - } - - // Copy island data. - scnstate->data.nisland = d->nisland; - if (d->nisland) { - memcpy(scnstate->data.island_dofadr, d->island_dofadr, sizeof(*d->island_dofadr) * d->nisland); - memcpy(scnstate->data.island_dofind, d->island_dofind, sizeof(*d->island_dofind) * m->nv); - memcpy(scnstate->data.dof_island, d->dof_island, sizeof(*d->dof_island) * m->nv); - memcpy(scnstate->data.tendon_efcadr, d->tendon_efcadr, sizeof(*d->tendon_efcadr) * m->ntendon); - } -} - - - -// move camera with mouse given a scene state; action is mjtMouse -void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvCamera* cam) { - mjModel m; - mjv_assignFromSceneState(scnstate, &m, NULL); - mjv_moveCamera(&m, action, reldx, reldy, scn, cam); -} - - - -// move perturb object with mouse given a scene state; action is mjtMouse -void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvPerturb* pert) { - mjModel m; - mjData d; - mjv_assignFromSceneState(scnstate, &m, &d); - mjv_movePerturb(&m, &d, action, reldx, reldy, scn, pert); -} diff --git a/src/engine/engine_vis_state.h b/src/engine/engine_vis_state.h deleted file mode 100644 index 7e1abfbe..00000000 --- a/src/engine/engine_vis_state.h +++ /dev/null @@ -1,63 +0,0 @@ -// 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. - -#ifndef MUJOCO_SRC_ENGINE_ENGINE_VIS_STATE_H_ -#define MUJOCO_SRC_ENGINE_ENGINE_VIS_STATE_H_ - -#include -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif -// set default scene state -MJAPI void mjv_defaultSceneState(mjvSceneState* scnstate); - -// allocate and init scene state -MJAPI void mjv_makeSceneState(const mjModel* m, const mjData* d, - mjvSceneState* scnstate, int maxgeom); - -// free scene state -MJAPI void mjv_freeSceneState(mjvSceneState* scnstate); - -// shallow copy scene state into model and data for use with mjv functions -void mjv_assignFromSceneState(const mjvSceneState* scnstate, mjModel* m, mjData* d); - -// update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings -MJAPI int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt, - const mjvPerturb* pert, mjvCamera* cam, int catmask, - mjvScene* scn); - -// update a scene state from model and data -MJAPI void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt, - mjvSceneState* scnstate); - -// move camera with mouse given a scene state; action is mjtMouse -MJAPI void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvCamera* cam); - -// move perturb object with mouse given a scene state; action is mjtMouse -MJAPI void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action, - mjtNum reldx, mjtNum reldy, - const mjvScene* scn, mjvPerturb* pert); - -#ifdef __cplusplus -} -#endif - -#endif // MUJOCO_SRC_ENGINE_ENGINE_VIS_STATE_H_ diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c index 1e665a2c..adcc410c 100644 --- a/src/engine/engine_vis_visualize.c +++ b/src/engine/engine_vis_visualize.c @@ -91,9 +91,9 @@ static void makeLabel(const mjModel* m, mjtObj type, int id, char* label) { // assign pseudo-random rgba to constraint island using Halton sequence static void islandColor(float rgba[4], int islanddofadr) { - rgba[0] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 2); - rgba[1] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 3); - rgba[2] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 5); + rgba[0] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 2); + rgba[1] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 3); + rgba[2] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 5); rgba[3] = 1; } @@ -152,7 +152,7 @@ static void addContactGeom(const mjModel* m, mjData* d, const mjtByte* flags, // override standard colors if visualizing islands if (vopt->flags[mjVIS_ISLAND] && d->nisland && efc_adr >= 0) { // set color using island's first dof - islandColor(thisgeom->rgba, d->island_dofind[d->island_dofadr[d->efc_island[efc_adr]]]); + islandColor(thisgeom->rgba, d->island_dofadr[d->efc_island[efc_adr]]); } // otherwise regular colors (different for included and excluded contacts) @@ -1344,7 +1344,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, int island = d->dof_island[m->body_dofadr[weld_id]]; if (island > -1) { // color using island's first dof - islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]); + islandColor(rgba_island, d->island_dofadr[island]); } } } @@ -1835,7 +1835,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, if (d->tendon_efcadr[i] != -1) { // set color using island's first dof int island = d->efc_island[d->tendon_efcadr[i]]; - islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]); + islandColor(rgba_island, d->island_dofadr[island]); } } setMaterial(m, thisgeom, tendon_matid, rgba, vopt->flags); @@ -2143,8 +2143,12 @@ void mjv_makeLights(const mjModel* m, const mjData* d, mjvScene* scn) { // set default properties memset(thislight, 0, sizeof(mjvLight)); thislight->headlight = 1; - thislight->directional = 1; + thislight->texid = -1; + thislight->type = mjLIGHT_DIRECTIONAL; thislight->castshadow = 0; + thislight->bulbradius = 0.02; + thislight->intensity = 0; + thislight->range = 10; // compute head position and gaze direction in model space mjtNum hpos[3], hfwd[3]; @@ -2169,10 +2173,13 @@ void mjv_makeLights(const mjModel* m, const mjData* d, mjvScene* scn) { // copy properties memset(thislight, 0, sizeof(mjvLight)); - thislight->directional = m->light_directional[i]; + thislight->type = m->light_type[i]; + thislight->texid = m->light_texid[i]; thislight->castshadow = m->light_castshadow[i]; thislight->bulbradius = m->light_bulbradius[i]; - if (!thislight->directional) { + thislight->intensity = m->light_intensity[i]; + thislight->range = m->light_range[i]; + if (thislight->type == mjLIGHT_SPOT) { f2f(thislight->attenuation, m->light_attenuation+3*i, 3); thislight->exponent = m->light_exponent[i]; thislight->cutoff = m->light_cutoff[i]; diff --git a/src/experimental/usd/mjcPhysics/actuatorAPI.cpp b/src/experimental/usd/mjcPhysics/actuatorAPI.cpp new file mode 100644 index 00000000..fb43eee7 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/actuatorAPI.cpp @@ -0,0 +1,391 @@ +// Copyright 2025 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. + +#include "./actuatorAPI.h" + +#include "pxr/usd/sdf/assetPath.h" +#include "pxr/usd/sdf/types.h" +#include "pxr/usd/usd/schemaRegistry.h" +#include "pxr/usd/usd/typed.h" + +PXR_NAMESPACE_OPEN_SCOPE + +// Register the schema with the TfType system. +TF_REGISTRY_FUNCTION(TfType) { + TfType::Define >(); +} + +/* virtual */ +MjcPhysicsActuatorAPI::~MjcPhysicsActuatorAPI() {} + +/* static */ +MjcPhysicsActuatorAPI MjcPhysicsActuatorAPI::Get(const UsdStagePtr &stage, + const SdfPath &path) { + if (!stage) { + TF_CODING_ERROR("Invalid stage"); + return MjcPhysicsActuatorAPI(); + } + return MjcPhysicsActuatorAPI(stage->GetPrimAtPath(path)); +} + +/* virtual */ +UsdSchemaKind MjcPhysicsActuatorAPI::_GetSchemaKind() const { + return MjcPhysicsActuatorAPI::schemaKind; +} + +/* static */ +bool MjcPhysicsActuatorAPI::CanApply(const UsdPrim &prim, std::string *whyNot) { + return prim.CanApplyAPI(whyNot); +} + +/* static */ +MjcPhysicsActuatorAPI MjcPhysicsActuatorAPI::Apply(const UsdPrim &prim) { + if (prim.ApplyAPI()) { + return MjcPhysicsActuatorAPI(prim); + } + return MjcPhysicsActuatorAPI(); +} + +/* static */ +const TfType &MjcPhysicsActuatorAPI::_GetStaticTfType() { + static TfType tfType = TfType::Find(); + return tfType; +} + +/* static */ +bool MjcPhysicsActuatorAPI::_IsTypedSchema() { + static bool isTyped = _GetStaticTfType().IsA(); + return isTyped; +} + +/* virtual */ +const TfType &MjcPhysicsActuatorAPI::_GetTfType() const { + return _GetStaticTfType(); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcCtrlLimitedAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcCtrlLimited); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcCtrlLimitedAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcCtrlLimited, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcForceLimitedAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcForceLimited); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcForceLimitedAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcForceLimited, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcActLimitedAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcActLimited); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcActLimitedAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcActLimited, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcCtrlRangeMinAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcCtrlRangeMin); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcCtrlRangeMinAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcCtrlRangeMin, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcCtrlRangeMaxAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcCtrlRangeMax); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcCtrlRangeMaxAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcCtrlRangeMax, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcForceRangeMinAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcForceRangeMin); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcForceRangeMinAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcForceRangeMin, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcForceRangeMaxAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcForceRangeMax); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcForceRangeMaxAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcForceRangeMax, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcActRangeMinAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcActRangeMin); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcActRangeMinAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcActRangeMin, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcActRangeMaxAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcActRangeMax); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcActRangeMaxAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcActRangeMax, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcLengthRangeMinAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcLengthRangeMin); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcLengthRangeMinAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcLengthRangeMin, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcLengthRangeMaxAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcLengthRangeMax); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcLengthRangeMaxAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcLengthRangeMax, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcGearAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcGear); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcGearAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcGear, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcCrankLengthAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcCrankLength); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcCrankLengthAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcCrankLength, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcJointInParentAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcJointInParent); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcJointInParentAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcJointInParent, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcActDimAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcActDim); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcActDimAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcActDim, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcDynTypeAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcDynType); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcDynTypeAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcDynType, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcGainTypeAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcGainType); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcGainTypeAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcGainType, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcBiasTypeAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcBiasType); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcBiasTypeAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcBiasType, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcDynPrmAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcDynPrm); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcDynPrmAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcDynPrm, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcGainPrmAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcGainPrm); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcGainPrmAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcGainPrm, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcBiasPrmAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcBiasPrm); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcBiasPrmAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcBiasPrm, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsActuatorAPI::GetMjcActEarlyAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcActEarly); +} + +UsdAttribute MjcPhysicsActuatorAPI::CreateMjcActEarlyAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcActEarly, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdRelationship MjcPhysicsActuatorAPI::GetMjcRefSiteRel() const { + return GetPrim().GetRelationship(MjcPhysicsTokens->mjcRefSite); +} + +UsdRelationship MjcPhysicsActuatorAPI::CreateMjcRefSiteRel() const { + return GetPrim().CreateRelationship(MjcPhysicsTokens->mjcRefSite, + /* custom = */ false); +} + +UsdRelationship MjcPhysicsActuatorAPI::GetMjcCrankSiteRel() const { + return GetPrim().GetRelationship(MjcPhysicsTokens->mjcCrankSite); +} + +UsdRelationship MjcPhysicsActuatorAPI::CreateMjcCrankSiteRel() const { + return GetPrim().CreateRelationship(MjcPhysicsTokens->mjcCrankSite, + /* custom = */ false); +} + +UsdRelationship MjcPhysicsActuatorAPI::GetMjcSliderSiteRel() const { + return GetPrim().GetRelationship(MjcPhysicsTokens->mjcSliderSite); +} + +UsdRelationship MjcPhysicsActuatorAPI::CreateMjcSliderSiteRel() const { + return GetPrim().CreateRelationship(MjcPhysicsTokens->mjcSliderSite, + /* custom = */ false); +} + +namespace { +static inline TfTokenVector _ConcatenateAttributeNames( + const TfTokenVector &left, const TfTokenVector &right) { + TfTokenVector result; + result.reserve(left.size() + right.size()); + result.insert(result.end(), left.begin(), left.end()); + result.insert(result.end(), right.begin(), right.end()); + return result; +} +} // namespace + +/*static*/ +const TfTokenVector &MjcPhysicsActuatorAPI::GetSchemaAttributeNames( + bool includeInherited) { + static TfTokenVector localNames = { + MjcPhysicsTokens->mjcCtrlLimited, MjcPhysicsTokens->mjcForceLimited, + MjcPhysicsTokens->mjcActLimited, MjcPhysicsTokens->mjcCtrlRangeMin, + MjcPhysicsTokens->mjcCtrlRangeMax, MjcPhysicsTokens->mjcForceRangeMin, + MjcPhysicsTokens->mjcForceRangeMax, MjcPhysicsTokens->mjcActRangeMin, + MjcPhysicsTokens->mjcActRangeMax, MjcPhysicsTokens->mjcLengthRangeMin, + MjcPhysicsTokens->mjcLengthRangeMax, MjcPhysicsTokens->mjcGear, + MjcPhysicsTokens->mjcCrankLength, MjcPhysicsTokens->mjcJointInParent, + MjcPhysicsTokens->mjcActDim, MjcPhysicsTokens->mjcDynType, + MjcPhysicsTokens->mjcGainType, MjcPhysicsTokens->mjcBiasType, + MjcPhysicsTokens->mjcDynPrm, MjcPhysicsTokens->mjcGainPrm, + MjcPhysicsTokens->mjcBiasPrm, MjcPhysicsTokens->mjcActEarly, + }; + static TfTokenVector allNames = _ConcatenateAttributeNames( + UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames); + + if (includeInherited) + return allNames; + else + return localNames; +} + +PXR_NAMESPACE_CLOSE_SCOPE + +// ===================================================================== // +// Feel free to add custom code below this line. It will be preserved by +// the code generator. +// +// Just remember to wrap code in the appropriate delimiters: +// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'. +// ===================================================================== // +// --(BEGIN CUSTOM CODE)-- diff --git a/src/experimental/usd/mjcPhysics/actuatorAPI.h b/src/experimental/usd/mjcPhysics/actuatorAPI.h new file mode 100644 index 00000000..49e6721a --- /dev/null +++ b/src/experimental/usd/mjcPhysics/actuatorAPI.h @@ -0,0 +1,828 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_GENERATED_ACTUATORAPI_H +#define MJCPHYSICS_GENERATED_ACTUATORAPI_H + +/// \file mjcPhysics/actuatorAPI.h + +#include "./api.h" +#include "./tokens.h" +#include "pxr/base/gf/matrix4d.h" +#include "pxr/base/gf/vec3d.h" +#include "pxr/base/gf/vec3f.h" +#include "pxr/base/tf/token.h" +#include "pxr/base/tf/type.h" +#include "pxr/base/vt/value.h" +#include "pxr/pxr.h" +#include "pxr/usd/usd/apiSchemaBase.h" +#include "pxr/usd/usd/prim.h" +#include "pxr/usd/usd/stage.h" + +PXR_NAMESPACE_OPEN_SCOPE + +class SdfAssetPath; + +// -------------------------------------------------------------------------- // +// PHYSICSACTUATORAPI // +// -------------------------------------------------------------------------- // + +/// \class MjcPhysicsActuatorAPI +/// +/// API describing a Mujoco actuator. +/// +/// For any described attribute \em Fallback \em Value or \em Allowed \em Values +/// below that are text/tokens, the actual token is published and defined in +/// \ref MjcPhysicsTokens. So to set an attribute to the value "rightHanded", +/// use MjcPhysicsTokens->rightHanded as the value. +/// +class MjcPhysicsActuatorAPI : public UsdAPISchemaBase { + public: + /// Compile time constant representing what kind of schema this class is. + /// + /// \sa UsdSchemaKind + static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI; + + /// Construct a MjcPhysicsActuatorAPI on UsdPrim \p prim . + /// Equivalent to MjcPhysicsActuatorAPI::Get(prim.GetStage(), prim.GetPath()) + /// for a \em valid \p prim, but will not immediately throw an error for + /// an invalid \p prim + explicit MjcPhysicsActuatorAPI(const UsdPrim &prim = UsdPrim()) + : UsdAPISchemaBase(prim) {} + + /// Construct a MjcPhysicsActuatorAPI on the prim held by \p schemaObj . + /// Should be preferred over MjcPhysicsActuatorAPI(schemaObj.GetPrim()), + /// as it preserves SchemaBase state. + explicit MjcPhysicsActuatorAPI(const UsdSchemaBase &schemaObj) + : UsdAPISchemaBase(schemaObj) {} + + /// Destructor. + MJCPHYSICS_API + virtual ~MjcPhysicsActuatorAPI(); + + /// Return a vector of names of all pre-declared attributes for this schema + /// class and all its ancestor classes. Does not include attributes that + /// may be authored by custom/extended methods of the schemas involved. + MJCPHYSICS_API + static const TfTokenVector &GetSchemaAttributeNames( + bool includeInherited = true); + + /// Return a MjcPhysicsActuatorAPI holding the prim adhering to this + /// schema at \p path on \p stage. If no prim exists at \p path on + /// \p stage, or if the prim at that path does not adhere to this schema, + /// return an invalid schema object. This is shorthand for the following: + /// + /// \code + /// MjcPhysicsActuatorAPI(stage->GetPrimAtPath(path)); + /// \endcode + /// + MJCPHYSICS_API + static MjcPhysicsActuatorAPI Get(const UsdStagePtr &stage, + const SdfPath &path); + + /// Returns true if this single-apply API schema can be applied to + /// the given \p prim. If this schema can not be a applied to the prim, + /// this returns false and, if provided, populates \p whyNot with the + /// reason it can not be applied. + /// + /// Note that if CanApply returns false, that does not necessarily imply + /// that calling Apply will fail. Callers are expected to call CanApply + /// before calling Apply if they want to ensure that it is valid to + /// apply a schema. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static bool CanApply(const UsdPrim &prim, std::string *whyNot = nullptr); + + /// Applies this single-apply API schema to the given \p prim. + /// This information is stored by adding "PhysicsActuatorAPI" to the + /// token-valued, listOp metadata \em apiSchemas on the prim. + /// + /// \return A valid MjcPhysicsActuatorAPI object is returned upon success. + /// An invalid (or empty) MjcPhysicsActuatorAPI object is returned upon + /// failure. See \ref UsdPrim::ApplyAPI() for conditions + /// resulting in failure. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static MjcPhysicsActuatorAPI Apply(const UsdPrim &prim); + + protected: + /// Returns the kind of schema this class belongs to. + /// + /// \sa UsdSchemaKind + MJCPHYSICS_API + UsdSchemaKind _GetSchemaKind() const override; + + private: + // needs to invoke _GetStaticTfType. + friend class UsdSchemaRegistry; + MJCPHYSICS_API + static const TfType &_GetStaticTfType(); + + static bool _IsTypedSchema(); + + // override SchemaBase virtuals. + MJCPHYSICS_API + const TfType &_GetTfType() const override; + + public: + // --------------------------------------------------------------------- // + // MJCCTRLLIMITED + // --------------------------------------------------------------------- // + /// If true, the control input to this actuator is automatically clamped to + /// ctrlrange at runtime. If false, control input clamping is disabled. If + /// 'auto' and autolimits is set in compiler, control clamping will + /// automatically be set to true if ctrlrange is defined without explicitly + /// setting this attribute to 'true'. Note that control input clamping can + /// also be globally disabled with the clampctrl attribute of option/flag. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:ctrlLimited = "auto"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | false, true, auto | + MJCPHYSICS_API + UsdAttribute GetMjcCtrlLimitedAttr() const; + + /// See GetMjcCtrlLimitedAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcCtrlLimitedAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCFORCELIMITED + // --------------------------------------------------------------------- // + /// If true, the force output of this actuator is automatically clamped to + /// forcerange at runtime. If false, force clamping is disabled. If 'auto' and + /// autolimits is set in compiler, force clamping will automatically be set to + /// true if forcerange is defined without explicitly setting this attribute to + /// 'true'. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:forceLimited = "auto"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | false, true, auto | + MJCPHYSICS_API + UsdAttribute GetMjcForceLimitedAttr() const; + + /// See GetMjcForceLimitedAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcForceLimitedAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCACTLIMITED + // --------------------------------------------------------------------- // + /// If true, the internal state (activation) associated with this actuator is + /// automatically clamped to actrange at runtime. If false, activation + /// clamping is disabled. If 'auto' and autolimits is set in compiler, + /// activation clamping will automatically be set to true if actrange is + /// defined without explicitly setting this attribute to 'true'. See the + /// Activation clamping section for more details. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:actLimited = "auto"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | false, true, auto | + MJCPHYSICS_API + UsdAttribute GetMjcActLimitedAttr() const; + + /// See GetMjcActLimitedAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcActLimitedAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCCTRLRANGEMIN + // --------------------------------------------------------------------- // + /// Minimum range for clamping the control input. The first value must be + /// smaller than the second value. Setting this attribute without specifying + /// ctrllimited is an error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:ctrlRange:min = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcCtrlRangeMinAttr() const; + + /// See GetMjcCtrlRangeMinAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcCtrlRangeMinAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCCTRLRANGEMAX + // --------------------------------------------------------------------- // + /// Maximum range for clamping the control input. The first value must be + /// smaller than the second value. Setting this attribute without specifying + /// ctrllimited is an error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:ctrlRange:max = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcCtrlRangeMaxAttr() const; + + /// See GetMjcCtrlRangeMaxAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcCtrlRangeMaxAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCFORCERANGEMIN + // --------------------------------------------------------------------- // + /// Minimum range for clamping the force output. The first value must be no + /// greater than the second value. Setting this attribute without specifying + /// forcelimited is an error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:forceRange:min = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcForceRangeMinAttr() const; + + /// See GetMjcForceRangeMinAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcForceRangeMinAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCFORCERANGEMAX + // --------------------------------------------------------------------- // + /// Maximum range for clamping the force output. The first value must be no + /// greater than the second value. Setting this attribute without specifying + /// forcelimited is an error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:forceRange:max = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcForceRangeMaxAttr() const; + + /// See GetMjcForceRangeMaxAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcForceRangeMaxAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCACTRANGEMIN + // --------------------------------------------------------------------- // + /// Minimum range for clamping the activation state. The first value must be + /// no greater than the second value. See the Activation clamping section for + /// more details. Setting this attribute without specifying actlimited is an + /// error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:actRange:min = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcActRangeMinAttr() const; + + /// See GetMjcActRangeMinAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcActRangeMinAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCACTRANGEMAX + // --------------------------------------------------------------------- // + /// Maximum range for clamping the activation state. The first value must be + /// no greater than the second value. See the Activation clamping section for + /// more details. Setting this attribute without specifying actlimited is an + /// error if autolimits is 'false' in compiler. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:actRange:max = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcActRangeMaxAttr() const; + + /// See GetMjcActRangeMaxAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcActRangeMaxAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCLENGTHRANGEMIN + // --------------------------------------------------------------------- // + /// Minimum range of feasible lengths of the actuator’s transmission. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:lengthRange:min = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcLengthRangeMinAttr() const; + + /// See GetMjcLengthRangeMinAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcLengthRangeMinAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCLENGTHRANGEMAX + // --------------------------------------------------------------------- // + /// Maximum range of feasible lengths of the actuator’s transmission. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:lengthRange:max = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcLengthRangeMaxAttr() const; + + /// See GetMjcLengthRangeMaxAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcLengthRangeMaxAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCGEAR + // --------------------------------------------------------------------- // + /// This attribute scales the length (and consequently moment arms, velocity + /// and force) of the actuator, for all transmission types. It is different + /// from the gain in the force generation mechanism, because the gain only + /// scales the force output and does not affect the length, moment arms and + /// velocity. For actuators with scalar transmission, only the first element + /// of this vector is used. The remaining elements are needed for joint, + /// jointinparent and site transmissions where this attribute is used to + /// specify 3D force and torque axes. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:gear = [1, 0, 0, 0, 0, 0]` | + /// | C++ Type | VtArray | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->DoubleArray | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcGearAttr() const; + + /// See GetMjcGearAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcGearAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCCRANKLENGTH + // --------------------------------------------------------------------- // + /// Used only for the slider-crank transmission type. Specifies the length of + /// the connecting rod. The compiler expects this value to be positive when a + /// slider-crank transmission is present. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:crankLength = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcCrankLengthAttr() const; + + /// See GetMjcCrankLengthAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcCrankLengthAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCJOINTINPARENT + // --------------------------------------------------------------------- // + /// If true and applied to ball and free joints, the 3d rotation axis given by + /// gear is defined in the parent frame (which is the world frame for free + /// joints) rather than the child frame. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:jointInParent = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcJointInParentAttr() const; + + /// See GetMjcJointInParentAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcJointInParentAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCACTDIM + // --------------------------------------------------------------------- // + /// Dimension of the activation state. The default value of -1 instructs the + /// compiler to set the dimension according to the dyntype. Values larger than + /// 1 are only allowed for user-defined activation dynamics, as native types + /// require dimensions of only 0 or 1. For activation dimensions bigger than + /// 1, the last element is used to generate force. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:actDim = -1` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcActDimAttr() const; + + /// See GetMjcActDimAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcActDimAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCDYNTYPE + // --------------------------------------------------------------------- // + /// Activation dynamics type for the actuator. The available dynamics types + /// were already described in the Actuation model section. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:dynType = "none"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | none, integrator, filter, + /// filterexact, muscle, user | + MJCPHYSICS_API + UsdAttribute GetMjcDynTypeAttr() const; + + /// See GetMjcDynTypeAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcDynTypeAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCGAINTYPE + // --------------------------------------------------------------------- // + /// The gain and bias together determine the output of the force generation + /// mechanism, which is currently assumed to be affine. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:gainType = "fixed"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | fixed, affine, muscle, user | + MJCPHYSICS_API + UsdAttribute GetMjcGainTypeAttr() const; + + /// See GetMjcGainTypeAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcGainTypeAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCBIASTYPE + // --------------------------------------------------------------------- // + /// The gain and bias together determine the output of the force generation + /// mechanism, which is currently assumed to be affine. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:biasType = "none"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | none, affine, muscle, user | + MJCPHYSICS_API + UsdAttribute GetMjcBiasTypeAttr() const; + + /// See GetMjcBiasTypeAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcBiasTypeAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCDYNPRM + // --------------------------------------------------------------------- // + /// Activation dynamics parameters. The built-in activation types (except for + /// muscle) use only the first parameter, but we provide additional parameters + /// in case user callbacks implement a more elaborate model. The length of + /// this array is not enforced by the parser, so the user can enter as many + /// parameters as needed. These defaults are not compatible with muscle + /// actuators. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:dynPrm = [1, 0, 0, 0, 0, 0, 0, 0, 0, + /// 0]` | | C++ Type | VtArray | | \ref Usd_Datatypes "Usd Type" | + /// SdfValueTypeNames->DoubleArray | | \ref SdfVariability "Variability" | + /// SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcDynPrmAttr() const; + + /// See GetMjcDynPrmAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcDynPrmAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCGAINPRM + // --------------------------------------------------------------------- // + /// Gain parameters. The built-in gain types (except for muscle) use only the + /// first parameter, but we provide additional parameters in case user + /// callbacks implement a more elaborate model. The length of this array is + /// not enforced by the parser, so the user can enter as many parameters as + /// needed. These defaults are not compatible with muscle actuators. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:gainPrm = [1, 0, 0, 0, 0, 0, 0, 0, + /// 0, 0]` | | C++ Type | VtArray | | \ref Usd_Datatypes "Usd Type" | + /// SdfValueTypeNames->DoubleArray | | \ref SdfVariability "Variability" | + /// SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcGainPrmAttr() const; + + /// See GetMjcGainPrmAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcGainPrmAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCBIASPRM + // --------------------------------------------------------------------- // + /// Bias parameters. The affine bias type uses three parameters. The length of + /// this array is not enforced by the parser, so the user can enter as many + /// parameters as needed. These defaults are not compatible with muscle + /// actuators. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:biasPrm = [0, 0, 0, 0, 0, 0, 0, 0, + /// 0, 0]` | | C++ Type | VtArray | | \ref Usd_Datatypes "Usd Type" | + /// SdfValueTypeNames->DoubleArray | | \ref SdfVariability "Variability" | + /// SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcBiasPrmAttr() const; + + /// See GetMjcBiasPrmAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcBiasPrmAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCACTEARLY + // --------------------------------------------------------------------- // + /// If true, force computation will use the next value of the activation + /// variable rather than the current one. Setting this flag reduces the delay + /// between the control and accelerations by one time-step. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:actEarly = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMjcActEarlyAttr() const; + + /// See GetMjcActEarlyAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMjcActEarlyAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MJCREFSITE + // --------------------------------------------------------------------- // + /// When applied to a site, measure the translation and rotation w.r.t the + /// frame of the refsite. In this case the actuator does have length and + /// position actuators can be used to directly control an end effector, see + /// refsite.xml example model. As above, the length is the dot product of the + /// gear vector and the frame difference. So gear='0 1 0 0 0 0' means + /// 'Y-offset of site in the refsite frame', while gear='0 0 0 0 0 1' means + /// rotation 'Z- rotation of site in the refsite frame'. It is recommended to + /// use a normalized gear vector with nonzeros in only the first 3 or the last + /// 3 elements of gear, so the actuator length will be in either length units + /// or radians, respectively. As with ball joints (see joint above), for + /// rotations which exceed a total angle of pi will wrap around, so tighter + /// limits are recommended. + /// + MJCPHYSICS_API + UsdRelationship GetMjcRefSiteRel() const; + + /// See GetMjcRefSiteRel(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create + MJCPHYSICS_API + UsdRelationship CreateMjcRefSiteRel() const; + + public: + // --------------------------------------------------------------------- // + // MJCCRANKSITE + // --------------------------------------------------------------------- // + /// If specified, the actuator acts on a slider-crank mechanism which is + /// implicitly determined by the actuator (i.e., it is not a separate model + /// element). The target site corresponds to the pin joining the crank and the + /// connecting rod. + /// + MJCPHYSICS_API + UsdRelationship GetMjcCrankSiteRel() const; + + /// See GetMjcCrankSiteRel(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create + MJCPHYSICS_API + UsdRelationship CreateMjcCrankSiteRel() const; + + public: + // --------------------------------------------------------------------- // + // MJCSLIDERSITE + // --------------------------------------------------------------------- // + /// Used only for the slider-crank transmission type. The target site is the + /// pin joining the slider and the connecting rod. The slider moves along the + /// z-axis of the slidersite frame. Therefore the site should be oriented as + /// needed when it is defined in the kinematic tree; its orientation cannot be + /// changed in the actuator definition. + /// + MJCPHYSICS_API + UsdRelationship GetMjcSliderSiteRel() const; + + /// See GetMjcSliderSiteRel(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create + MJCPHYSICS_API + UsdRelationship CreateMjcSliderSiteRel() const; + + public: + // ===================================================================== // + // Feel free to add custom code below this line, it will be preserved by + // the code generator. + // + // Just remember to: + // - Close the class declaration with }; + // - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE + // - Close the include guard with #endif + // ===================================================================== // + // --(BEGIN CUSTOM CODE)-- +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/mjcPhysics/api.h b/src/experimental/usd/mjcPhysics/api.h new file mode 100644 index 00000000..dc5476c9 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/api.h @@ -0,0 +1,42 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_API_H +#define MJCPHYSICS_API_H + +#include "pxr/base/arch/export.h" + +#if defined(PXR_STATIC) +#define MJCPHYSICS_API +#define MJCPHYSICS_API_TEMPLATE_CLASS(...) +#define MJCPHYSICS_API_TEMPLATE_STRUCT(...) +#define MJCPHYSICS_LOCAL +#else +#if defined(MJCPHYSICS_EXPORTS) +#define MJCPHYSICS_API ARCH_EXPORT +#define MJCPHYSICS_API_TEMPLATE_CLASS(...) \ + ARCH_EXPORT_TEMPLATE(class, __VA_ARGS__) +#define MJCPHYSICS_API_TEMPLATE_STRUCT(...) \ + ARCH_EXPORT_TEMPLATE(struct, __VA_ARGS__) +#else +#define MJCPHYSICS_API ARCH_IMPORT +#define MJCPHYSICS_API_TEMPLATE_CLASS(...) \ + ARCH_IMPORT_TEMPLATE(class, __VA_ARGS__) +#define MJCPHYSICS_API_TEMPLATE_STRUCT(...) \ + ARCH_IMPORT_TEMPLATE(struct, __VA_ARGS__) +#endif +#define MJCPHYSICS_LOCAL ARCH_HIDDEN +#endif + +#endif diff --git a/src/experimental/usd/mjcPhysics/collisionAPI.cpp b/src/experimental/usd/mjcPhysics/collisionAPI.cpp new file mode 100644 index 00000000..ff113849 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/collisionAPI.cpp @@ -0,0 +1,124 @@ +// Copyright 2025 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. + +#include "./collisionAPI.h" + +#include "pxr/usd/sdf/assetPath.h" +#include "pxr/usd/sdf/types.h" +#include "pxr/usd/usd/schemaRegistry.h" +#include "pxr/usd/usd/typed.h" + +PXR_NAMESPACE_OPEN_SCOPE + +// Register the schema with the TfType system. +TF_REGISTRY_FUNCTION(TfType) { + TfType::Define >(); +} + +/* virtual */ +MjcPhysicsCollisionAPI::~MjcPhysicsCollisionAPI() {} + +/* static */ +MjcPhysicsCollisionAPI MjcPhysicsCollisionAPI::Get(const UsdStagePtr &stage, + const SdfPath &path) { + if (!stage) { + TF_CODING_ERROR("Invalid stage"); + return MjcPhysicsCollisionAPI(); + } + return MjcPhysicsCollisionAPI(stage->GetPrimAtPath(path)); +} + +/* virtual */ +UsdSchemaKind MjcPhysicsCollisionAPI::_GetSchemaKind() const { + return MjcPhysicsCollisionAPI::schemaKind; +} + +/* static */ +bool MjcPhysicsCollisionAPI::CanApply(const UsdPrim &prim, + std::string *whyNot) { + return prim.CanApplyAPI(whyNot); +} + +/* static */ +MjcPhysicsCollisionAPI MjcPhysicsCollisionAPI::Apply(const UsdPrim &prim) { + if (prim.ApplyAPI()) { + return MjcPhysicsCollisionAPI(prim); + } + return MjcPhysicsCollisionAPI(); +} + +/* static */ +const TfType &MjcPhysicsCollisionAPI::_GetStaticTfType() { + static TfType tfType = TfType::Find(); + return tfType; +} + +/* static */ +bool MjcPhysicsCollisionAPI::_IsTypedSchema() { + static bool isTyped = _GetStaticTfType().IsA(); + return isTyped; +} + +/* virtual */ +const TfType &MjcPhysicsCollisionAPI::_GetTfType() const { + return _GetStaticTfType(); +} + +UsdAttribute MjcPhysicsCollisionAPI::GetShellInertiaAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcShellinertia); +} + +UsdAttribute MjcPhysicsCollisionAPI::CreateShellInertiaAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcShellinertia, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +namespace { +static inline TfTokenVector _ConcatenateAttributeNames( + const TfTokenVector &left, const TfTokenVector &right) { + TfTokenVector result; + result.reserve(left.size() + right.size()); + result.insert(result.end(), left.begin(), left.end()); + result.insert(result.end(), right.begin(), right.end()); + return result; +} +} // namespace + +/*static*/ +const TfTokenVector &MjcPhysicsCollisionAPI::GetSchemaAttributeNames( + bool includeInherited) { + static TfTokenVector localNames = { + MjcPhysicsTokens->mjcShellinertia, + }; + static TfTokenVector allNames = _ConcatenateAttributeNames( + UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames); + + if (includeInherited) + return allNames; + else + return localNames; +} + +PXR_NAMESPACE_CLOSE_SCOPE + +// ===================================================================== // +// Feel free to add custom code below this line. It will be preserved by +// the code generator. +// +// Just remember to wrap code in the appropriate delimiters: +// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'. +// ===================================================================== // +// --(BEGIN CUSTOM CODE)-- diff --git a/src/experimental/usd/mjcPhysics/collisionAPI.h b/src/experimental/usd/mjcPhysics/collisionAPI.h new file mode 100644 index 00000000..e3681526 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/collisionAPI.h @@ -0,0 +1,185 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_GENERATED_COLLISIONAPI_H +#define MJCPHYSICS_GENERATED_COLLISIONAPI_H + +/// \file mjcPhysics/collisionAPI.h + +#include "./api.h" +#include "./tokens.h" +#include "pxr/base/gf/matrix4d.h" +#include "pxr/base/gf/vec3d.h" +#include "pxr/base/gf/vec3f.h" +#include "pxr/base/tf/token.h" +#include "pxr/base/tf/type.h" +#include "pxr/base/vt/value.h" +#include "pxr/pxr.h" +#include "pxr/usd/usd/apiSchemaBase.h" +#include "pxr/usd/usd/prim.h" +#include "pxr/usd/usd/stage.h" + +PXR_NAMESPACE_OPEN_SCOPE + +class SdfAssetPath; + +// -------------------------------------------------------------------------- // +// COLLISIONAPI // +// -------------------------------------------------------------------------- // + +/// \class MjcPhysicsCollisionAPI +/// +/// API describing a Mujoco collider. +/// +class MjcPhysicsCollisionAPI : public UsdAPISchemaBase { + public: + /// Compile time constant representing what kind of schema this class is. + /// + /// \sa UsdSchemaKind + static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI; + + /// Construct a MjcPhysicsCollisionAPI on UsdPrim \p prim . + /// Equivalent to MjcPhysicsCollisionAPI::Get(prim.GetStage(), prim.GetPath()) + /// for a \em valid \p prim, but will not immediately throw an error for + /// an invalid \p prim + explicit MjcPhysicsCollisionAPI(const UsdPrim &prim = UsdPrim()) + : UsdAPISchemaBase(prim) {} + + /// Construct a MjcPhysicsCollisionAPI on the prim held by \p schemaObj . + /// Should be preferred over MjcPhysicsCollisionAPI(schemaObj.GetPrim()), + /// as it preserves SchemaBase state. + explicit MjcPhysicsCollisionAPI(const UsdSchemaBase &schemaObj) + : UsdAPISchemaBase(schemaObj) {} + + /// Destructor. + MJCPHYSICS_API + virtual ~MjcPhysicsCollisionAPI(); + + /// Return a vector of names of all pre-declared attributes for this schema + /// class and all its ancestor classes. Does not include attributes that + /// may be authored by custom/extended methods of the schemas involved. + MJCPHYSICS_API + static const TfTokenVector &GetSchemaAttributeNames( + bool includeInherited = true); + + /// Return a MjcPhysicsCollisionAPI holding the prim adhering to this + /// schema at \p path on \p stage. If no prim exists at \p path on + /// \p stage, or if the prim at that path does not adhere to this schema, + /// return an invalid schema object. This is shorthand for the following: + /// + /// \code + /// MjcPhysicsCollisionAPI(stage->GetPrimAtPath(path)); + /// \endcode + /// + MJCPHYSICS_API + static MjcPhysicsCollisionAPI Get(const UsdStagePtr &stage, + const SdfPath &path); + + /// Returns true if this single-apply API schema can be applied to + /// the given \p prim. If this schema can not be a applied to the prim, + /// this returns false and, if provided, populates \p whyNot with the + /// reason it can not be applied. + /// + /// Note that if CanApply returns false, that does not necessarily imply + /// that calling Apply will fail. Callers are expected to call CanApply + /// before calling Apply if they want to ensure that it is valid to + /// apply a schema. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static bool CanApply(const UsdPrim &prim, std::string *whyNot = nullptr); + + /// Applies this single-apply API schema to the given \p prim. + /// This information is stored by adding "CollisionAPI" to the + /// token-valued, listOp metadata \em apiSchemas on the prim. + /// + /// \return A valid MjcPhysicsCollisionAPI object is returned upon success. + /// An invalid (or empty) MjcPhysicsCollisionAPI object is returned upon + /// failure. See \ref UsdPrim::ApplyAPI() for conditions + /// resulting in failure. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static MjcPhysicsCollisionAPI Apply(const UsdPrim &prim); + + protected: + /// Returns the kind of schema this class belongs to. + /// + /// \sa UsdSchemaKind + MJCPHYSICS_API + UsdSchemaKind _GetSchemaKind() const override; + + private: + // needs to invoke _GetStaticTfType. + friend class UsdSchemaRegistry; + MJCPHYSICS_API + static const TfType &_GetStaticTfType(); + + static bool _IsTypedSchema(); + + // override SchemaBase virtuals. + MJCPHYSICS_API + const TfType &_GetTfType() const override; + + public: + // --------------------------------------------------------------------- // + // SHELLINERTIA + // --------------------------------------------------------------------- // + /// Enables handling of the inertia assuming mass is concentrated on the + /// surface. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:shellinertia = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetShellInertiaAttr() const; + + /// See GetShellInertiaAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateShellInertiaAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // ===================================================================== // + // Feel free to add custom code below this line, it will be preserved by + // the code generator. + // + // Just remember to: + // - Close the class declaration with }; + // - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE + // - Close the include guard with #endif + // ===================================================================== // + // --(BEGIN CUSTOM CODE)-- +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/mjcPhysics/generatedSchema.usda b/src/experimental/usd/mjcPhysics/generatedSchema.usda new file mode 100644 index 00000000..b911cf59 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/generatedSchema.usda @@ -0,0 +1,336 @@ +#usda 1.0 +( + "WARNING: THIS FILE IS GENERATED BY usdGenSchema. DO NOT EDIT." +) + +class "SceneAPI" ( + doc = "API providing global simulation options for Mujoco." +) +{ + uniform bool mjc:flag:actuation = 1 ( + displayName = "Actuation Forces Toggle" + doc = "Enables all standard computations related to actuator forces, including actuator dynamics." + ) + uniform bool mjc:flag:autoreset = 1 ( + displayName = "Automatic Simulation Reset Toggle" + doc = "Enables the automatic resetting of the simulation state when numerical issues are detected." + ) + uniform bool mjc:flag:clampctrl = 1 ( + displayName = "Control Input Clamping Toggle" + doc = "Enables the clamping of control inputs to all actuators, according to actuator-specific attributes." + ) + uniform bool mjc:flag:constraint = 1 ( + displayName = "Constraint Solver Toggle" + doc = "Enables constraint solver." + ) + uniform bool mjc:flag:contact = 1 ( + displayName = "Contact Constraints and Collision Detection Toggle" + doc = "Enables collision detection and all standard computations related to contact constraints." + ) + uniform bool mjc:flag:energy = 0 ( + displayName = "Energy Computation Toggle" + doc = "Enables the computation of potential and kinetic energy (mjData.energy[0,1])." + ) + uniform bool mjc:flag:equality = 1 ( + displayName = "Equality Constraints Toggle" + doc = "Enables all standard computations related to equality constraints." + ) + uniform bool mjc:flag:eulerdamp = 1 ( + displayName = "Euler Integrator Damping Toggle" + doc = "Enables implicit integration with respect to joint damping in the Euler integrator." + ) + uniform bool mjc:flag:filterparent = 1 ( + displayName = "Parent-Child Contact Filtering Toggle" + doc = "Enables the filtering of contact pairs where the two geoms belong to a parent and child body." + ) + uniform bool mjc:flag:frictionloss = 1 ( + displayName = "Friction Loss Constraints Toggle" + doc = "Enables all standard computations related to friction loss constraints." + ) + uniform bool mjc:flag:fwdinv = 0 ( + displayName = "Forward/Inverse Dynamics Comparison Toggle" + doc = "Enables the automatic comparison of forward and inverse dynamics." + ) + uniform bool mjc:flag:gravity = 1 ( + displayName = "Gravity Toggle" + doc = "Enables the application of gravitational acceleration as defined in mjOption." + ) + uniform bool mjc:flag:invdiscrete = 0 ( + displayName = "Discrete-Time Inverse Dynamics Toggle" + doc = "Enables discrete-time inverse dynamics with mj_inverse for integrators other than RK4." + ) + uniform bool mjc:flag:island = 0 ( + displayName = "Constraint Island Discovery Toggle" + doc = "Enables the discovery of constraint islands." + ) + uniform bool mjc:flag:limit = 1 ( + displayName = "Joint and Tendon Limit Constraints Toggle" + doc = "Enables all standard computations related to joint and tendon limit constraints." + ) + uniform bool mjc:flag:midphase = 1 ( + displayName = "Mid-Phase Collision Filtering Toggle" + doc = "Enables mid-phase collision filtering using a static AABB bounding volume hierarchy (BVH)." + ) + uniform bool mjc:flag:multiccd = 0 ( + displayName = "Multiple Contact Collision Detection (CCD) Toggle" + doc = "Enables multiple-contact collision detection for geom pairs using a general-purpose convex-convex collider." + ) + uniform bool mjc:flag:nativeccd = 1 ( + displayName = "Native Convex Collision Detection Toggle" + doc = "Enables the native convex collision detection pipeline instead of using the libccd library." + ) + uniform bool mjc:flag:override = 0 ( + displayName = "Contact Override Mechanism Toggle" + doc = "Enables the contact override mechanism." + ) + uniform bool mjc:flag:passive = 1 ( + displayName = "Passive Forces Toggle" + doc = "Enables the simulation of joint and tendon spring-dampers, fluid dynamics forces, and custom passive forces." + ) + uniform bool mjc:flag:refsafe = 1 ( + displayName = "Solver Reference Safety Mechanism Toggle" + doc = "Enables a safety mechanism that prevents instabilities due to solref[0] being too small compared to the simulation timestep." + ) + uniform bool mjc:flag:sensor = 1 ( + displayName = "Sensor Computations Toggle" + doc = "Enables all computations related to sensors." + ) + uniform bool mjc:flag:warmstart = 1 ( + displayName = "Solver Warm-Starting Toggle" + doc = "Enables warm-starting of the constraint solver, using the solution from the previous time step to initialize the iterative optimization." + ) + uniform int[] mjc:option:actuatorgroupdisable ( + displayName = "Actuator Group Disable" + doc = "List of actuator groups to disable." + ) + uniform double mjc:option:apirate = 100 ( + displayName = "ApiRate" + doc = """Determines the rate (in Hz) at which an external API allows + the update function to be executed.""" + ) + uniform int mjc:option:ccd_iterations = 50 ( + displayName = "CCD Iterations" + doc = "Maximum number of iterations of the algorithm used for convex collisions." + ) + uniform double mjc:option:ccd_tolerance = 0.000001 ( + displayName = "CCD Tolerance" + doc = """Tolerance threshold used for early termination of the convex + collision algorithm.""" + ) + uniform token mjc:option:cone = "pyramidal" ( + allowedTokens = ["pyramidal", "elliptic"] + displayName = "Friction Cone Type" + doc = "The type of contact friction cone." + ) + uniform double mjc:option:density = 0 ( + displayName = "Density" + doc = "Density of medium." + ) + uniform double mjc:option:impratio = 1 ( + displayName = "Impedance Ratio" + doc = """Ratio of frictional-to-normal constraint impedance for elliptic + friction cones.""" + ) + uniform token mjc:option:integrator = "euler" ( + allowedTokens = ["euler", "rk4", "implicit", "implicitfast"] + displayName = "Integrator" + doc = "Numerical integrator to be used." + ) + uniform int mjc:option:iterations = 100 ( + displayName = "Solver Iterations" + doc = "Maximum number of iterations of the constraint solver." + ) + uniform token mjc:option:jacobian = "auto" ( + allowedTokens = ["auto", "dense", "sparse"] + displayName = "Jacobian Type" + doc = "The type of constraint Jacobian and matrices computed from it." + ) + uniform int mjc:option:ls_iterations = 50 ( + displayName = "Linesearch Iterations" + doc = """Maximum number of linesearch iterations performed by CG/Newton + constraint solvers.""" + ) + uniform double mjc:option:ls_tolerance = 0.01 ( + displayName = "Linesearch Tolerance" + doc = "Tolerance threshold used for early termination of the linesearch algorithm." + ) + uniform double3 mjc:option:magnetic = (0, -0.5, 0) ( + displayName = "Magnetic Flux" + doc = "Global magnetic flux." + ) + uniform int mjc:option:noslip_iterations = 0 ( + displayName = "Noslip Iterations" + doc = "Maximum number of iterations of the Noslip solver." + ) + uniform double mjc:option:noslip_tolerance = 0.000001 ( + displayName = "Noslip Tolerance" + doc = "Tolerance threshold used for early termination of the Noslip solver." + ) + uniform double[] mjc:option:o_friction = [1, 1, 0.005, 0.0001, 0.0001] ( + displayName = "Contact Override Friction" + doc = """Replaces the friction parameter of all active contact pairs when + Contact override is enabled.""" + ) + uniform double mjc:option:o_margin = 0 ( + displayName = "Contact Override Margin" + doc = """Replaces the margin parameter of all active contact pairs when + Contact override is enabled.""" + ) + uniform double[] mjc:option:o_solimp = [0.9, 0.95, 0.001, 0.5, 2] ( + displayName = "Contact Override SolImp" + doc = """Replaces the solimp parameter of all active contact pairs when + Contact override is enabled.""" + ) + uniform double[] mjc:option:o_solref = [0.02, 1] ( + displayName = "Contact Override SolRef" + doc = """Replaces the solref parameter of all active contact pairs when + Contact override is enabled.""" + ) + uniform int mjc:option:sdf_initpoints = 40 ( + displayName = "SDF Initial Points" + doc = """Number of starting points used for finding contacts with Signed + Distance Field collisions.""" + ) + uniform int mjc:option:sdf_iterations = 10 ( + displayName = "SDF Iterations" + doc = """Number of iterations used for Signed Distance Field collisions + (per initial point).""" + ) + uniform token mjc:option:solver = "newton" ( + allowedTokens = ["pgs", "cg", "newton"] + displayName = "Solver" + doc = "Constraint solver algorithm to be used." + ) + uniform double mjc:option:timestep = 0.002 ( + displayName = "Timestep" + doc = "Controls the timestep in seconds used by MuJoCo." + ) + uniform double mjc:option:tolerance = 1e-8 ( + displayName = "Solver Tolerance" + doc = """Tolerance threshold used for early termination of the iterative + solver.""" + ) + uniform double mjc:option:viscosity = 0 ( + displayName = "Viscosity" + doc = "Viscosity of medium." + ) + uniform double3 mjc:option:wind = (0, 0, 0) ( + displayName = "Wind Velocity" + doc = "Velocity vector of medium (i.e. wind)." + ) +} + +class "SiteAPI" ( + doc = "API describing a Mujoco site." +) +{ +} + +class "CollisionAPI" ( + doc = "API describing a Mujoco collider." +) +{ + uniform bool mjc:shellinertia = 0 ( + displayName = "Shell Inertia" + doc = "Enables handling of the inertia assuming mass is concentrated on the surface." + ) +} + +class "MeshCollisionAPI" ( + doc = "API describing a Mujoco collider." +) +{ + uniform token mjc:inertia = "legacy" ( + allowedTokens = ["legacy", "convex", "exact", "shell"] + displayName = "Inertia" + doc = "Controls how a mesh is used when mass and inertia are inferred from geometry." + ) +} + +class "PhysicsActuatorAPI" ( + doc = "API describing a Mujoco actuator." +) +{ + uniform int mjc:actDim = -1 ( + doc = "Dimension of the activation state. The default value of -1 instructs the compiler to set the dimension according to the dyntype. Values larger than 1 are only allowed for user-defined activation dynamics, as native types require dimensions of only 0 or 1. For activation dimensions bigger than 1, the last element is used to generate force." + ) + uniform bool mjc:actEarly = 0 ( + doc = "If true, force computation will use the next value of the activation variable rather than the current one. Setting this flag reduces the delay between the control and accelerations by one time-step." + ) + uniform token mjc:actLimited = "auto" ( + allowedTokens = ["false", "true", "auto"] + doc = "If true, the internal state (activation) associated with this actuator is automatically clamped to actrange at runtime. If false, activation clamping is disabled. If 'auto' and autolimits is set in compiler, activation clamping will automatically be set to true if actrange is defined without explicitly setting this attribute to 'true'. See the Activation clamping section for more details." + ) + uniform double mjc:actRange:max = 0 ( + doc = "Maximum range for clamping the activation state. The first value must be no greater than the second value. See the Activation clamping section for more details. Setting this attribute without specifying actlimited is an error if autolimits is 'false' in compiler." + ) + uniform double mjc:actRange:min = 0 ( + doc = "Minimum range for clamping the activation state. The first value must be no greater than the second value. See the Activation clamping section for more details. Setting this attribute without specifying actlimited is an error if autolimits is 'false' in compiler." + ) + uniform double[] mjc:biasPrm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Bias parameters. The affine bias type uses three parameters. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + uniform token mjc:biasType = "none" ( + allowedTokens = ["none", "affine", "muscle", "user"] + doc = "The gain and bias together determine the output of the force generation mechanism, which is currently assumed to be affine." + ) + uniform double mjc:crankLength = 0 ( + doc = "Used only for the slider-crank transmission type. Specifies the length of the connecting rod. The compiler expects this value to be positive when a slider-crank transmission is present." + ) + rel mjc:crankSite ( + doc = "If specified, the actuator acts on a slider-crank mechanism which is implicitly determined by the actuator (i.e., it is not a separate model element). The target site corresponds to the pin joining the crank and the connecting rod." + ) + uniform token mjc:ctrlLimited = "auto" ( + allowedTokens = ["false", "true", "auto"] + doc = "If true, the control input to this actuator is automatically clamped to ctrlrange at runtime. If false, control input clamping is disabled. If 'auto' and autolimits is set in compiler, control clamping will automatically be set to true if ctrlrange is defined without explicitly setting this attribute to 'true'. Note that control input clamping can also be globally disabled with the clampctrl attribute of option/flag." + ) + uniform double mjc:ctrlRange:max = 0 ( + doc = "Maximum range for clamping the control input. The first value must be smaller than the second value. Setting this attribute without specifying ctrllimited is an error if autolimits is 'false' in compiler." + ) + uniform double mjc:ctrlRange:min = 0 ( + doc = "Minimum range for clamping the control input. The first value must be smaller than the second value. Setting this attribute without specifying ctrllimited is an error if autolimits is 'false' in compiler." + ) + uniform double[] mjc:dynPrm = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Activation dynamics parameters. The built-in activation types (except for muscle) use only the first parameter, but we provide additional parameters in case user callbacks implement a more elaborate model. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + uniform token mjc:dynType = "none" ( + allowedTokens = ["none", "integrator", "filter", "filterexact", "muscle", "user"] + doc = "Activation dynamics type for the actuator. The available dynamics types were already described in the Actuation model section." + ) + uniform token mjc:forceLimited = "auto" ( + allowedTokens = ["false", "true", "auto"] + doc = "If true, the force output of this actuator is automatically clamped to forcerange at runtime. If false, force clamping is disabled. If 'auto' and autolimits is set in compiler, force clamping will automatically be set to true if forcerange is defined without explicitly setting this attribute to 'true'." + ) + uniform double mjc:forceRange:max = 0 ( + doc = "Maximum range for clamping the force output. The first value must be no greater than the second value. Setting this attribute without specifying forcelimited is an error if autolimits is 'false' in compiler." + ) + uniform double mjc:forceRange:min = 0 ( + doc = "Minimum range for clamping the force output. The first value must be no greater than the second value. Setting this attribute without specifying forcelimited is an error if autolimits is 'false' in compiler." + ) + uniform double[] mjc:gainPrm = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Gain parameters. The built-in gain types (except for muscle) use only the first parameter, but we provide additional parameters in case user callbacks implement a more elaborate model. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + uniform token mjc:gainType = "fixed" ( + allowedTokens = ["fixed", "affine", "muscle", "user"] + doc = "The gain and bias together determine the output of the force generation mechanism, which is currently assumed to be affine." + ) + uniform double[] mjc:gear = [1, 0, 0, 0, 0, 0] ( + doc = "This attribute scales the length (and consequently moment arms, velocity and force) of the actuator, for all transmission types. It is different from the gain in the force generation mechanism, because the gain only scales the force output and does not affect the length, moment arms and velocity. For actuators with scalar transmission, only the first element of this vector is used. The remaining elements are needed for joint, jointinparent and site transmissions where this attribute is used to specify 3D force and torque axes." + ) + uniform bool mjc:jointInParent = 0 ( + doc = "If true and applied to ball and free joints, the 3d rotation axis given by gear is defined in the parent frame (which is the world frame for free joints) rather than the child frame." + ) + uniform double mjc:lengthRange:max = 0 ( + doc = "Maximum range of feasible lengths of the actuator’s transmission." + ) + uniform double mjc:lengthRange:min = 0 ( + doc = "Minimum range of feasible lengths of the actuator’s transmission." + ) + rel mjc:refSite ( + doc = "When applied to a site, measure the translation and rotation w.r.t the frame of the refsite. In this case the actuator does have length and position actuators can be used to directly control an end effector, see refsite.xml example model. As above, the length is the dot product of the gear vector and the frame difference. So gear='0 1 0 0 0 0' means 'Y-offset of site in the refsite frame', while gear='0 0 0 0 0 1' means rotation 'Z- rotation of site in the refsite frame'. It is recommended to use a normalized gear vector with nonzeros in only the first 3 or the last 3 elements of gear, so the actuator length will be in either length units or radians, respectively. As with ball joints (see joint above), for rotations which exceed a total angle of pi will wrap around, so tighter limits are recommended." + ) + rel mjc:sliderSite ( + doc = "Used only for the slider-crank transmission type. The target site is the pin joining the slider and the connecting rod. The slider moves along the z-axis of the slidersite frame. Therefore the site should be oriented as needed when it is defined in the kinematic tree; its orientation cannot be changed in the actuator definition." + ) +} + diff --git a/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp b/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp new file mode 100644 index 00000000..50c579dd --- /dev/null +++ b/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp @@ -0,0 +1,126 @@ +// Copyright 2025 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. + +#include "./meshCollisionAPI.h" + +#include "pxr/usd/sdf/assetPath.h" +#include "pxr/usd/sdf/types.h" +#include "pxr/usd/usd/schemaRegistry.h" +#include "pxr/usd/usd/typed.h" + +PXR_NAMESPACE_OPEN_SCOPE + +// Register the schema with the TfType system. +TF_REGISTRY_FUNCTION(TfType) { + TfType::Define >(); +} + +/* virtual */ +MjcPhysicsMeshCollisionAPI::~MjcPhysicsMeshCollisionAPI() {} + +/* static */ +MjcPhysicsMeshCollisionAPI MjcPhysicsMeshCollisionAPI::Get( + const UsdStagePtr &stage, const SdfPath &path) { + if (!stage) { + TF_CODING_ERROR("Invalid stage"); + return MjcPhysicsMeshCollisionAPI(); + } + return MjcPhysicsMeshCollisionAPI(stage->GetPrimAtPath(path)); +} + +/* virtual */ +UsdSchemaKind MjcPhysicsMeshCollisionAPI::_GetSchemaKind() const { + return MjcPhysicsMeshCollisionAPI::schemaKind; +} + +/* static */ +bool MjcPhysicsMeshCollisionAPI::CanApply(const UsdPrim &prim, + std::string *whyNot) { + return prim.CanApplyAPI(whyNot); +} + +/* static */ +MjcPhysicsMeshCollisionAPI MjcPhysicsMeshCollisionAPI::Apply( + const UsdPrim &prim) { + if (prim.ApplyAPI()) { + return MjcPhysicsMeshCollisionAPI(prim); + } + return MjcPhysicsMeshCollisionAPI(); +} + +/* static */ +const TfType &MjcPhysicsMeshCollisionAPI::_GetStaticTfType() { + static TfType tfType = TfType::Find(); + return tfType; +} + +/* static */ +bool MjcPhysicsMeshCollisionAPI::_IsTypedSchema() { + static bool isTyped = _GetStaticTfType().IsA(); + return isTyped; +} + +/* virtual */ +const TfType &MjcPhysicsMeshCollisionAPI::_GetTfType() const { + return _GetStaticTfType(); +} + +UsdAttribute MjcPhysicsMeshCollisionAPI::GetInertiaAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcInertia); +} + +UsdAttribute MjcPhysicsMeshCollisionAPI::CreateInertiaAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcInertia, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +namespace { +static inline TfTokenVector _ConcatenateAttributeNames( + const TfTokenVector &left, const TfTokenVector &right) { + TfTokenVector result; + result.reserve(left.size() + right.size()); + result.insert(result.end(), left.begin(), left.end()); + result.insert(result.end(), right.begin(), right.end()); + return result; +} +} // namespace + +/*static*/ +const TfTokenVector &MjcPhysicsMeshCollisionAPI::GetSchemaAttributeNames( + bool includeInherited) { + static TfTokenVector localNames = { + MjcPhysicsTokens->mjcInertia, + }; + static TfTokenVector allNames = _ConcatenateAttributeNames( + UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames); + + if (includeInherited) + return allNames; + else + return localNames; +} + +PXR_NAMESPACE_CLOSE_SCOPE + +// ===================================================================== // +// Feel free to add custom code below this line. It will be preserved by +// the code generator. +// +// Just remember to wrap code in the appropriate delimiters: +// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'. +// ===================================================================== // +// --(BEGIN CUSTOM CODE)-- diff --git a/src/experimental/usd/mjcPhysics/meshCollisionAPI.h b/src/experimental/usd/mjcPhysics/meshCollisionAPI.h new file mode 100644 index 00000000..55c8c27e --- /dev/null +++ b/src/experimental/usd/mjcPhysics/meshCollisionAPI.h @@ -0,0 +1,191 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H +#define MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H + +/// \file mjcPhysics/meshCollisionAPI.h + +#include "./api.h" +#include "./tokens.h" +#include "pxr/base/gf/matrix4d.h" +#include "pxr/base/gf/vec3d.h" +#include "pxr/base/gf/vec3f.h" +#include "pxr/base/tf/token.h" +#include "pxr/base/tf/type.h" +#include "pxr/base/vt/value.h" +#include "pxr/pxr.h" +#include "pxr/usd/usd/apiSchemaBase.h" +#include "pxr/usd/usd/prim.h" +#include "pxr/usd/usd/stage.h" + +PXR_NAMESPACE_OPEN_SCOPE + +class SdfAssetPath; + +// -------------------------------------------------------------------------- // +// MESHCOLLISIONAPI // +// -------------------------------------------------------------------------- // + +/// \class MjcPhysicsMeshCollisionAPI +/// +/// API describing a Mujoco collider. +/// +/// For any described attribute \em Fallback \em Value or \em Allowed \em Values +/// below that are text/tokens, the actual token is published and defined in +/// \ref MjcPhysicsTokens. So to set an attribute to the value "rightHanded", +/// use MjcPhysicsTokens->rightHanded as the value. +/// +class MjcPhysicsMeshCollisionAPI : public UsdAPISchemaBase { + public: + /// Compile time constant representing what kind of schema this class is. + /// + /// \sa UsdSchemaKind + static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI; + + /// Construct a MjcPhysicsMeshCollisionAPI on UsdPrim \p prim . + /// Equivalent to MjcPhysicsMeshCollisionAPI::Get(prim.GetStage(), + /// prim.GetPath()) for a \em valid \p prim, but will not immediately throw an + /// error for an invalid \p prim + explicit MjcPhysicsMeshCollisionAPI(const UsdPrim &prim = UsdPrim()) + : UsdAPISchemaBase(prim) {} + + /// Construct a MjcPhysicsMeshCollisionAPI on the prim held by \p schemaObj . + /// Should be preferred over MjcPhysicsMeshCollisionAPI(schemaObj.GetPrim()), + /// as it preserves SchemaBase state. + explicit MjcPhysicsMeshCollisionAPI(const UsdSchemaBase &schemaObj) + : UsdAPISchemaBase(schemaObj) {} + + /// Destructor. + MJCPHYSICS_API + virtual ~MjcPhysicsMeshCollisionAPI(); + + /// Return a vector of names of all pre-declared attributes for this schema + /// class and all its ancestor classes. Does not include attributes that + /// may be authored by custom/extended methods of the schemas involved. + MJCPHYSICS_API + static const TfTokenVector &GetSchemaAttributeNames( + bool includeInherited = true); + + /// Return a MjcPhysicsMeshCollisionAPI holding the prim adhering to this + /// schema at \p path on \p stage. If no prim exists at \p path on + /// \p stage, or if the prim at that path does not adhere to this schema, + /// return an invalid schema object. This is shorthand for the following: + /// + /// \code + /// MjcPhysicsMeshCollisionAPI(stage->GetPrimAtPath(path)); + /// \endcode + /// + MJCPHYSICS_API + static MjcPhysicsMeshCollisionAPI Get(const UsdStagePtr &stage, + const SdfPath &path); + + /// Returns true if this single-apply API schema can be applied to + /// the given \p prim. If this schema can not be a applied to the prim, + /// this returns false and, if provided, populates \p whyNot with the + /// reason it can not be applied. + /// + /// Note that if CanApply returns false, that does not necessarily imply + /// that calling Apply will fail. Callers are expected to call CanApply + /// before calling Apply if they want to ensure that it is valid to + /// apply a schema. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static bool CanApply(const UsdPrim &prim, std::string *whyNot = nullptr); + + /// Applies this single-apply API schema to the given \p prim. + /// This information is stored by adding "MeshCollisionAPI" to the + /// token-valued, listOp metadata \em apiSchemas on the prim. + /// + /// \return A valid MjcPhysicsMeshCollisionAPI object is returned upon + /// success. An invalid (or empty) MjcPhysicsMeshCollisionAPI object is + /// returned upon failure. See \ref UsdPrim::ApplyAPI() for conditions + /// resulting in failure. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static MjcPhysicsMeshCollisionAPI Apply(const UsdPrim &prim); + + protected: + /// Returns the kind of schema this class belongs to. + /// + /// \sa UsdSchemaKind + MJCPHYSICS_API + UsdSchemaKind _GetSchemaKind() const override; + + private: + // needs to invoke _GetStaticTfType. + friend class UsdSchemaRegistry; + MJCPHYSICS_API + static const TfType &_GetStaticTfType(); + + static bool _IsTypedSchema(); + + // override SchemaBase virtuals. + MJCPHYSICS_API + const TfType &_GetTfType() const override; + + public: + // --------------------------------------------------------------------- // + // INERTIA + // --------------------------------------------------------------------- // + /// Controls how a mesh is used when mass and inertia are inferred from + /// geometry. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:inertia = "legacy"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | legacy, convex, exact, shell | + MJCPHYSICS_API + UsdAttribute GetInertiaAttr() const; + + /// See GetInertiaAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateInertiaAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // ===================================================================== // + // Feel free to add custom code below this line, it will be preserved by + // the code generator. + // + // Just remember to: + // - Close the class declaration with }; + // - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE + // - Close the include guard with #endif + // ===================================================================== // + // --(BEGIN CUSTOM CODE)-- +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/mjcPhysics/plugInfo.json b/src/experimental/usd/mjcPhysics/plugInfo.json new file mode 100644 index 00000000..908651ba --- /dev/null +++ b/src/experimental/usd/mjcPhysics/plugInfo.json @@ -0,0 +1,68 @@ +#Portions of this file auto - generated by usdGenSchema. +#Edits will survive regeneration except for comments and +#changes to types with autoGenerated = true. +{ + "Plugins": [ + { + "Info": { + "Types": { + "MjcPhysicsActuatorAPI": { + "alias": { + "UsdSchemaBase": "PhysicsActuatorAPI" + }, + "autoGenerated": true, + "bases": [ + "UsdAPISchemaBase" + ], + "schemaKind": "singleApplyAPI" + }, + "MjcPhysicsCollisionAPI": { + "alias": { + "UsdSchemaBase": "CollisionAPI" + }, + "autoGenerated": true, + "bases": [ + "UsdAPISchemaBase" + ], + "schemaKind": "singleApplyAPI" + }, + "MjcPhysicsMeshCollisionAPI": { + "alias": { + "UsdSchemaBase": "MeshCollisionAPI" + }, + "autoGenerated": true, + "bases": [ + "UsdAPISchemaBase" + ], + "schemaKind": "singleApplyAPI" + }, + "MjcPhysicsSceneAPI": { + "alias": { + "UsdSchemaBase": "SceneAPI" + }, + "autoGenerated": true, + "bases": [ + "UsdAPISchemaBase" + ], + "schemaKind": "singleApplyAPI" + }, + "MjcPhysicsSiteAPI": { + "alias": { + "UsdSchemaBase": "SiteAPI" + }, + "autoGenerated": true, + "bases": [ + "UsdAPISchemaBase" + ], + "schemaKind": "singleApplyAPI" + } + } + }, + "LibraryPath": "", + "Name": "mjcPhysics", + "ResourcePath": "", + "Root": ".", + "Type": "library" + } + ] +} diff --git a/src/experimental/usd/mjcPhysics/sceneAPI.cpp b/src/experimental/usd/mjcPhysics/sceneAPI.cpp new file mode 100644 index 00000000..7bce5c2d --- /dev/null +++ b/src/experimental/usd/mjcPhysics/sceneAPI.cpp @@ -0,0 +1,701 @@ +// Copyright 2025 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. + +#include "./sceneAPI.h" + +#include "pxr/usd/sdf/assetPath.h" +#include "pxr/usd/sdf/types.h" +#include "pxr/usd/usd/schemaRegistry.h" +#include "pxr/usd/usd/typed.h" + +PXR_NAMESPACE_OPEN_SCOPE + +// Register the schema with the TfType system. +TF_REGISTRY_FUNCTION(TfType) { + TfType::Define >(); +} + +/* virtual */ +MjcPhysicsSceneAPI::~MjcPhysicsSceneAPI() {} + +/* static */ +MjcPhysicsSceneAPI MjcPhysicsSceneAPI::Get(const UsdStagePtr &stage, + const SdfPath &path) { + if (!stage) { + TF_CODING_ERROR("Invalid stage"); + return MjcPhysicsSceneAPI(); + } + return MjcPhysicsSceneAPI(stage->GetPrimAtPath(path)); +} + +/* virtual */ +UsdSchemaKind MjcPhysicsSceneAPI::_GetSchemaKind() const { + return MjcPhysicsSceneAPI::schemaKind; +} + +/* static */ +bool MjcPhysicsSceneAPI::CanApply(const UsdPrim &prim, std::string *whyNot) { + return prim.CanApplyAPI(whyNot); +} + +/* static */ +MjcPhysicsSceneAPI MjcPhysicsSceneAPI::Apply(const UsdPrim &prim) { + if (prim.ApplyAPI()) { + return MjcPhysicsSceneAPI(prim); + } + return MjcPhysicsSceneAPI(); +} + +/* static */ +const TfType &MjcPhysicsSceneAPI::_GetStaticTfType() { + static TfType tfType = TfType::Find(); + return tfType; +} + +/* static */ +bool MjcPhysicsSceneAPI::_IsTypedSchema() { + static bool isTyped = _GetStaticTfType().IsA(); + return isTyped; +} + +/* virtual */ +const TfType &MjcPhysicsSceneAPI::_GetTfType() const { + return _GetStaticTfType(); +} + +UsdAttribute MjcPhysicsSceneAPI::GetTimestepAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionTimestep); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateTimestepAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionTimestep, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetApiRateAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionApirate); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateApiRateAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionApirate, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetImpRatioAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionImpratio); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateImpRatioAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionImpratio, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetWindAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionWind); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateWindAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionWind, SdfValueTypeNames->Double3, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetMagneticAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionMagnetic); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateMagneticAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionMagnetic, SdfValueTypeNames->Double3, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetDensityAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionDensity); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateDensityAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionDensity, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetViscosityAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionViscosity); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateViscosityAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionViscosity, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetOMarginAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionO_margin); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateOMarginAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionO_margin, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetOSolRefAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionO_solref); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateOSolRefAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionO_solref, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetOSolImpAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionO_solimp); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateOSolImpAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionO_solimp, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetOFrictionAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionO_friction); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateOFrictionAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionO_friction, SdfValueTypeNames->DoubleArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetIntegratorAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionIntegrator); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateIntegratorAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionIntegrator, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetConeAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionCone); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateConeAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionCone, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetJacobianAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionJacobian); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateJacobianAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionJacobian, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetSolverAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionSolver); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateSolverAttr(VtValue const &defaultValue, + bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionSolver, SdfValueTypeNames->Token, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetIterationsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionIterations); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateIterationsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionIterations, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetToleranceAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionTolerance); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateToleranceAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionTolerance, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetLSIterationsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionLs_iterations); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateLSIterationsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionLs_iterations, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetLSToleranceAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionLs_tolerance); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateLSToleranceAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionLs_tolerance, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetNoslipIterationsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionNoslip_iterations); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateNoslipIterationsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionNoslip_iterations, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetNoslipToleranceAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionNoslip_tolerance); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateNoslipToleranceAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionNoslip_tolerance, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetCCDIterationsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionCcd_iterations); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateCCDIterationsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionCcd_iterations, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetCCDToleranceAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionCcd_tolerance); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateCCDToleranceAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionCcd_tolerance, SdfValueTypeNames->Double, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetSDFIterationsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionSdf_iterations); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateSDFIterationsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionSdf_iterations, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetSDFInitPointsAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcOptionSdf_initpoints); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateSDFInitPointsAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionSdf_initpoints, SdfValueTypeNames->Int, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetActuatorGroupDisableAttr() const { + return GetPrim().GetAttribute( + MjcPhysicsTokens->mjcOptionActuatorgroupdisable); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateActuatorGroupDisableAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcOptionActuatorgroupdisable, + SdfValueTypeNames->IntArray, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetConstraintFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagConstraint); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateConstraintFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagConstraint, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetEqualityFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagEquality); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateEqualityFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagEquality, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetFrictionLossFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagFrictionloss); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateFrictionLossFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagFrictionloss, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetLimitFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagLimit); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateLimitFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagLimit, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetContactFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagContact); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateContactFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagContact, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetPassiveFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagPassive); +} + +UsdAttribute MjcPhysicsSceneAPI::CreatePassiveFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagPassive, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetGravityFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagGravity); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateGravityFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagGravity, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetClampCtrlFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagClampctrl); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateClampCtrlFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagClampctrl, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetWarmStartFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagWarmstart); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateWarmStartFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagWarmstart, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetFilterParentFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagFilterparent); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateFilterParentFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagFilterparent, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetActuationFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagActuation); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateActuationFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagActuation, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetRefSafeFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagRefsafe); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateRefSafeFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagRefsafe, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetSensorFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagSensor); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateSensorFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagSensor, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetMidPhaseFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagMidphase); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateMidPhaseFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagMidphase, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetNativeCCDFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagNativeccd); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateNativeCCDFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagNativeccd, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetEulerDampFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagEulerdamp); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateEulerDampFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagEulerdamp, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetAutoResetFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagAutoreset); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateAutoResetFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagAutoreset, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetOverrideFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagOverride); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateOverrideFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagOverride, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetEnergyFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagEnergy); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateEnergyFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagEnergy, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetFwdinvFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagFwdinv); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateFwdinvFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagFwdinv, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetInvDiscreteFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagInvdiscrete); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateInvDiscreteFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagInvdiscrete, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetMultiCCDFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagMulticcd); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateMultiCCDFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagMulticcd, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +UsdAttribute MjcPhysicsSceneAPI::GetIslandFlagAttr() const { + return GetPrim().GetAttribute(MjcPhysicsTokens->mjcFlagIsland); +} + +UsdAttribute MjcPhysicsSceneAPI::CreateIslandFlagAttr( + VtValue const &defaultValue, bool writeSparsely) const { + return UsdSchemaBase::_CreateAttr( + MjcPhysicsTokens->mjcFlagIsland, SdfValueTypeNames->Bool, + /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely); +} + +namespace { +static inline TfTokenVector _ConcatenateAttributeNames( + const TfTokenVector &left, const TfTokenVector &right) { + TfTokenVector result; + result.reserve(left.size() + right.size()); + result.insert(result.end(), left.begin(), left.end()); + result.insert(result.end(), right.begin(), right.end()); + return result; +} +} // namespace + +/*static*/ +const TfTokenVector &MjcPhysicsSceneAPI::GetSchemaAttributeNames( + bool includeInherited) { + static TfTokenVector localNames = { + MjcPhysicsTokens->mjcOptionTimestep, + MjcPhysicsTokens->mjcOptionApirate, + MjcPhysicsTokens->mjcOptionImpratio, + MjcPhysicsTokens->mjcOptionWind, + MjcPhysicsTokens->mjcOptionMagnetic, + MjcPhysicsTokens->mjcOptionDensity, + MjcPhysicsTokens->mjcOptionViscosity, + MjcPhysicsTokens->mjcOptionO_margin, + MjcPhysicsTokens->mjcOptionO_solref, + MjcPhysicsTokens->mjcOptionO_solimp, + MjcPhysicsTokens->mjcOptionO_friction, + MjcPhysicsTokens->mjcOptionIntegrator, + MjcPhysicsTokens->mjcOptionCone, + MjcPhysicsTokens->mjcOptionJacobian, + MjcPhysicsTokens->mjcOptionSolver, + MjcPhysicsTokens->mjcOptionIterations, + MjcPhysicsTokens->mjcOptionTolerance, + MjcPhysicsTokens->mjcOptionLs_iterations, + MjcPhysicsTokens->mjcOptionLs_tolerance, + MjcPhysicsTokens->mjcOptionNoslip_iterations, + MjcPhysicsTokens->mjcOptionNoslip_tolerance, + MjcPhysicsTokens->mjcOptionCcd_iterations, + MjcPhysicsTokens->mjcOptionCcd_tolerance, + MjcPhysicsTokens->mjcOptionSdf_iterations, + MjcPhysicsTokens->mjcOptionSdf_initpoints, + MjcPhysicsTokens->mjcOptionActuatorgroupdisable, + MjcPhysicsTokens->mjcFlagConstraint, + MjcPhysicsTokens->mjcFlagEquality, + MjcPhysicsTokens->mjcFlagFrictionloss, + MjcPhysicsTokens->mjcFlagLimit, + MjcPhysicsTokens->mjcFlagContact, + MjcPhysicsTokens->mjcFlagPassive, + MjcPhysicsTokens->mjcFlagGravity, + MjcPhysicsTokens->mjcFlagClampctrl, + MjcPhysicsTokens->mjcFlagWarmstart, + MjcPhysicsTokens->mjcFlagFilterparent, + MjcPhysicsTokens->mjcFlagActuation, + MjcPhysicsTokens->mjcFlagRefsafe, + MjcPhysicsTokens->mjcFlagSensor, + MjcPhysicsTokens->mjcFlagMidphase, + MjcPhysicsTokens->mjcFlagNativeccd, + MjcPhysicsTokens->mjcFlagEulerdamp, + MjcPhysicsTokens->mjcFlagAutoreset, + MjcPhysicsTokens->mjcFlagOverride, + MjcPhysicsTokens->mjcFlagEnergy, + MjcPhysicsTokens->mjcFlagFwdinv, + MjcPhysicsTokens->mjcFlagInvdiscrete, + MjcPhysicsTokens->mjcFlagMulticcd, + MjcPhysicsTokens->mjcFlagIsland, + }; + static TfTokenVector allNames = _ConcatenateAttributeNames( + UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames); + + if (includeInherited) + return allNames; + else + return localNames; +} + +PXR_NAMESPACE_CLOSE_SCOPE + +// ===================================================================== // +// Feel free to add custom code below this line. It will be preserved by +// the code generator. +// +// Just remember to wrap code in the appropriate delimiters: +// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'. +// ===================================================================== // +// --(BEGIN CUSTOM CODE)-- diff --git a/src/experimental/usd/mjcPhysics/sceneAPI.h b/src/experimental/usd/mjcPhysics/sceneAPI.h new file mode 100644 index 00000000..2ac745c7 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/sceneAPI.h @@ -0,0 +1,1379 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_GENERATED_SCENEAPI_H +#define MJCPHYSICS_GENERATED_SCENEAPI_H + +/// \file mjcPhysics/sceneAPI.h + +#include "./api.h" +#include "./tokens.h" +#include "pxr/base/gf/matrix4d.h" +#include "pxr/base/gf/vec3d.h" +#include "pxr/base/gf/vec3f.h" +#include "pxr/base/tf/token.h" +#include "pxr/base/tf/type.h" +#include "pxr/base/vt/value.h" +#include "pxr/pxr.h" +#include "pxr/usd/usd/apiSchemaBase.h" +#include "pxr/usd/usd/prim.h" +#include "pxr/usd/usd/stage.h" + +PXR_NAMESPACE_OPEN_SCOPE + +class SdfAssetPath; + +// -------------------------------------------------------------------------- // +// SCENEAPI // +// -------------------------------------------------------------------------- // + +/// \class MjcPhysicsSceneAPI +/// +/// API providing global simulation options for Mujoco. +/// +/// For any described attribute \em Fallback \em Value or \em Allowed \em Values +/// below that are text/tokens, the actual token is published and defined in +/// \ref MjcPhysicsTokens. So to set an attribute to the value "rightHanded", +/// use MjcPhysicsTokens->rightHanded as the value. +/// +class MjcPhysicsSceneAPI : public UsdAPISchemaBase { + public: + /// Compile time constant representing what kind of schema this class is. + /// + /// \sa UsdSchemaKind + static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI; + + /// Construct a MjcPhysicsSceneAPI on UsdPrim \p prim . + /// Equivalent to MjcPhysicsSceneAPI::Get(prim.GetStage(), prim.GetPath()) + /// for a \em valid \p prim, but will not immediately throw an error for + /// an invalid \p prim + explicit MjcPhysicsSceneAPI(const UsdPrim &prim = UsdPrim()) + : UsdAPISchemaBase(prim) {} + + /// Construct a MjcPhysicsSceneAPI on the prim held by \p schemaObj . + /// Should be preferred over MjcPhysicsSceneAPI(schemaObj.GetPrim()), + /// as it preserves SchemaBase state. + explicit MjcPhysicsSceneAPI(const UsdSchemaBase &schemaObj) + : UsdAPISchemaBase(schemaObj) {} + + /// Destructor. + MJCPHYSICS_API + virtual ~MjcPhysicsSceneAPI(); + + /// Return a vector of names of all pre-declared attributes for this schema + /// class and all its ancestor classes. Does not include attributes that + /// may be authored by custom/extended methods of the schemas involved. + MJCPHYSICS_API + static const TfTokenVector &GetSchemaAttributeNames( + bool includeInherited = true); + + /// Return a MjcPhysicsSceneAPI holding the prim adhering to this + /// schema at \p path on \p stage. If no prim exists at \p path on + /// \p stage, or if the prim at that path does not adhere to this schema, + /// return an invalid schema object. This is shorthand for the following: + /// + /// \code + /// MjcPhysicsSceneAPI(stage->GetPrimAtPath(path)); + /// \endcode + /// + MJCPHYSICS_API + static MjcPhysicsSceneAPI Get(const UsdStagePtr &stage, const SdfPath &path); + + /// Returns true if this single-apply API schema can be applied to + /// the given \p prim. If this schema can not be a applied to the prim, + /// this returns false and, if provided, populates \p whyNot with the + /// reason it can not be applied. + /// + /// Note that if CanApply returns false, that does not necessarily imply + /// that calling Apply will fail. Callers are expected to call CanApply + /// before calling Apply if they want to ensure that it is valid to + /// apply a schema. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static bool CanApply(const UsdPrim &prim, std::string *whyNot = nullptr); + + /// Applies this single-apply API schema to the given \p prim. + /// This information is stored by adding "SceneAPI" to the + /// token-valued, listOp metadata \em apiSchemas on the prim. + /// + /// \return A valid MjcPhysicsSceneAPI object is returned upon success. + /// An invalid (or empty) MjcPhysicsSceneAPI object is returned upon + /// failure. See \ref UsdPrim::ApplyAPI() for conditions + /// resulting in failure. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static MjcPhysicsSceneAPI Apply(const UsdPrim &prim); + + protected: + /// Returns the kind of schema this class belongs to. + /// + /// \sa UsdSchemaKind + MJCPHYSICS_API + UsdSchemaKind _GetSchemaKind() const override; + + private: + // needs to invoke _GetStaticTfType. + friend class UsdSchemaRegistry; + MJCPHYSICS_API + static const TfType &_GetStaticTfType(); + + static bool _IsTypedSchema(); + + // override SchemaBase virtuals. + MJCPHYSICS_API + const TfType &_GetTfType() const override; + + public: + // --------------------------------------------------------------------- // + // TIMESTEP + // --------------------------------------------------------------------- // + /// Controls the timestep in seconds used by MuJoCo. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:timestep = 0.002` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetTimestepAttr() const; + + /// See GetTimestepAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateTimestepAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // APIRATE + // --------------------------------------------------------------------- // + /// Determines the rate (in Hz) at which an external API allows + /// the update function to be executed. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:apirate = 100` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetApiRateAttr() const; + + /// See GetApiRateAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateApiRateAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // IMPRATIO + // --------------------------------------------------------------------- // + /// Ratio of frictional-to-normal constraint impedance for elliptic + /// friction cones. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:impratio = 1` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetImpRatioAttr() const; + + /// See GetImpRatioAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateImpRatioAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // WIND + // --------------------------------------------------------------------- // + /// Velocity vector of medium (i.e. wind). + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double3 mjc:option:wind = (0, 0, 0)` | + /// | C++ Type | GfVec3d | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double3 | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetWindAttr() const; + + /// See GetWindAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateWindAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MAGNETIC + // --------------------------------------------------------------------- // + /// Global magnetic flux. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double3 mjc:option:magnetic = (0, -0.5, 0)` | + /// | C++ Type | GfVec3d | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double3 | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMagneticAttr() const; + + /// See GetMagneticAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMagneticAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // DENSITY + // --------------------------------------------------------------------- // + /// Density of medium. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:density = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetDensityAttr() const; + + /// See GetDensityAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateDensityAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // VISCOSITY + // --------------------------------------------------------------------- // + /// Viscosity of medium. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:viscosity = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetViscosityAttr() const; + + /// See GetViscosityAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateViscosityAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // OMARGIN + // --------------------------------------------------------------------- // + /// Replaces the margin parameter of all active contact pairs when + /// Contact override is enabled. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:o_margin = 0` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetOMarginAttr() const; + + /// See GetOMarginAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateOMarginAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // OSOLREF + // --------------------------------------------------------------------- // + /// Replaces the solref parameter of all active contact pairs when + /// Contact override is enabled. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:option:o_solref = [0.02, 1]` | + /// | C++ Type | VtArray | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->DoubleArray | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetOSolRefAttr() const; + + /// See GetOSolRefAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateOSolRefAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // OSOLIMP + // --------------------------------------------------------------------- // + /// Replaces the solimp parameter of all active contact pairs when + /// Contact override is enabled. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:option:o_solimp = [0.9, 0.95, 0.001, + /// 0.5, 2]` | | C++ Type | VtArray | | \ref Usd_Datatypes "Usd Type" + /// | SdfValueTypeNames->DoubleArray | | \ref SdfVariability "Variability" | + /// SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetOSolImpAttr() const; + + /// See GetOSolImpAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateOSolImpAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // OFRICTION + // --------------------------------------------------------------------- // + /// Replaces the friction parameter of all active contact pairs when + /// Contact override is enabled. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double[] mjc:option:o_friction = [1, 1, 0.005, + /// 0.0001, 0.0001]` | | C++ Type | VtArray | | \ref Usd_Datatypes + /// "Usd Type" | SdfValueTypeNames->DoubleArray | | \ref SdfVariability + /// "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetOFrictionAttr() const; + + /// See GetOFrictionAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateOFrictionAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // INTEGRATOR + // --------------------------------------------------------------------- // + /// Numerical integrator to be used. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:option:integrator = "euler"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | euler, rk4, implicit, + /// implicitfast | + MJCPHYSICS_API + UsdAttribute GetIntegratorAttr() const; + + /// See GetIntegratorAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateIntegratorAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CONE + // --------------------------------------------------------------------- // + /// The type of contact friction cone. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:option:cone = "pyramidal"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | pyramidal, elliptic | + MJCPHYSICS_API + UsdAttribute GetConeAttr() const; + + /// See GetConeAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateConeAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // JACOBIAN + // --------------------------------------------------------------------- // + /// The type of constraint Jacobian and matrices computed from it. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:option:jacobian = "auto"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | auto, dense, sparse | + MJCPHYSICS_API + UsdAttribute GetJacobianAttr() const; + + /// See GetJacobianAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateJacobianAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // SOLVER + // --------------------------------------------------------------------- // + /// Constraint solver algorithm to be used. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform token mjc:option:solver = "newton"` | + /// | C++ Type | TfToken | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Token | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + /// | \ref MjcPhysicsTokens "Allowed Values" | pgs, cg, newton | + MJCPHYSICS_API + UsdAttribute GetSolverAttr() const; + + /// See GetSolverAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateSolverAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // ITERATIONS + // --------------------------------------------------------------------- // + /// Maximum number of iterations of the constraint solver. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:iterations = 100` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetIterationsAttr() const; + + /// See GetIterationsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateIterationsAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // TOLERANCE + // --------------------------------------------------------------------- // + /// Tolerance threshold used for early termination of the iterative + /// solver. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:tolerance = 1e-8` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetToleranceAttr() const; + + /// See GetToleranceAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateToleranceAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // LSITERATIONS + // --------------------------------------------------------------------- // + /// Maximum number of linesearch iterations performed by CG/Newton + /// constraint solvers. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:ls_iterations = 50` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetLSIterationsAttr() const; + + /// See GetLSIterationsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateLSIterationsAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // LSTOLERANCE + // --------------------------------------------------------------------- // + /// Tolerance threshold used for early termination of the linesearch + /// algorithm. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:ls_tolerance = 0.01` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetLSToleranceAttr() const; + + /// See GetLSToleranceAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateLSToleranceAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // NOSLIPITERATIONS + // --------------------------------------------------------------------- // + /// Maximum number of iterations of the Noslip solver. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:noslip_iterations = 0` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetNoslipIterationsAttr() const; + + /// See GetNoslipIterationsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateNoslipIterationsAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // NOSLIPTOLERANCE + // --------------------------------------------------------------------- // + /// Tolerance threshold used for early termination of the Noslip solver. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:noslip_tolerance = 0.000001` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetNoslipToleranceAttr() const; + + /// See GetNoslipToleranceAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateNoslipToleranceAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CCDITERATIONS + // --------------------------------------------------------------------- // + /// Maximum number of iterations of the algorithm used for convex collisions. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:ccd_iterations = 50` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetCCDIterationsAttr() const; + + /// See GetCCDIterationsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateCCDIterationsAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CCDTOLERANCE + // --------------------------------------------------------------------- // + /// Tolerance threshold used for early termination of the convex + /// collision algorithm. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform double mjc:option:ccd_tolerance = 0.000001` | + /// | C++ Type | double | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Double | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetCCDToleranceAttr() const; + + /// See GetCCDToleranceAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateCCDToleranceAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // SDFITERATIONS + // --------------------------------------------------------------------- // + /// Number of iterations used for Signed Distance Field collisions + /// (per initial point). + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:sdf_iterations = 10` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetSDFIterationsAttr() const; + + /// See GetSDFIterationsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateSDFIterationsAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // SDFINITPOINTS + // --------------------------------------------------------------------- // + /// Number of starting points used for finding contacts with Signed + /// Distance Field collisions. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int mjc:option:sdf_initpoints = 40` | + /// | C++ Type | int | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Int | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetSDFInitPointsAttr() const; + + /// See GetSDFInitPointsAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateSDFInitPointsAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // ACTUATORGROUPDISABLE + // --------------------------------------------------------------------- // + /// List of actuator groups to disable. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform int[] mjc:option:actuatorgroupdisable` | + /// | C++ Type | VtArray | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->IntArray | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetActuatorGroupDisableAttr() const; + + /// See GetActuatorGroupDisableAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateActuatorGroupDisableAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CONSTRAINTFLAG + // --------------------------------------------------------------------- // + /// Enables constraint solver. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:constraint = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetConstraintFlagAttr() const; + + /// See GetConstraintFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateConstraintFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // EQUALITYFLAG + // --------------------------------------------------------------------- // + /// Enables all standard computations related to equality constraints. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:equality = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetEqualityFlagAttr() const; + + /// See GetEqualityFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateEqualityFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // FRICTIONLOSSFLAG + // --------------------------------------------------------------------- // + /// Enables all standard computations related to friction loss constraints. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:frictionloss = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetFrictionLossFlagAttr() const; + + /// See GetFrictionLossFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateFrictionLossFlagAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // LIMITFLAG + // --------------------------------------------------------------------- // + /// Enables all standard computations related to joint and tendon limit + /// constraints. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:limit = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetLimitFlagAttr() const; + + /// See GetLimitFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateLimitFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CONTACTFLAG + // --------------------------------------------------------------------- // + /// Enables collision detection and all standard computations related to + /// contact constraints. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:contact = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetContactFlagAttr() const; + + /// See GetContactFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateContactFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // PASSIVEFLAG + // --------------------------------------------------------------------- // + /// Enables the simulation of joint and tendon spring-dampers, fluid dynamics + /// forces, and custom passive forces. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:passive = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetPassiveFlagAttr() const; + + /// See GetPassiveFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreatePassiveFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // GRAVITYFLAG + // --------------------------------------------------------------------- // + /// Enables the application of gravitational acceleration as defined in + /// mjOption. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:gravity = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetGravityFlagAttr() const; + + /// See GetGravityFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateGravityFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // CLAMPCTRLFLAG + // --------------------------------------------------------------------- // + /// Enables the clamping of control inputs to all actuators, according to + /// actuator-specific attributes. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:clampctrl = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetClampCtrlFlagAttr() const; + + /// See GetClampCtrlFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateClampCtrlFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // WARMSTARTFLAG + // --------------------------------------------------------------------- // + /// Enables warm-starting of the constraint solver, using the solution from + /// the previous time step to initialize the iterative optimization. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:warmstart = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetWarmStartFlagAttr() const; + + /// See GetWarmStartFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateWarmStartFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // FILTERPARENTFLAG + // --------------------------------------------------------------------- // + /// Enables the filtering of contact pairs where the two geoms belong to a + /// parent and child body. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:filterparent = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetFilterParentFlagAttr() const; + + /// See GetFilterParentFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateFilterParentFlagAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // ACTUATIONFLAG + // --------------------------------------------------------------------- // + /// Enables all standard computations related to actuator forces, including + /// actuator dynamics. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:actuation = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetActuationFlagAttr() const; + + /// See GetActuationFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateActuationFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // REFSAFEFLAG + // --------------------------------------------------------------------- // + /// Enables a safety mechanism that prevents instabilities due to solref[0] + /// being too small compared to the simulation timestep. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:refsafe = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetRefSafeFlagAttr() const; + + /// See GetRefSafeFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateRefSafeFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // SENSORFLAG + // --------------------------------------------------------------------- // + /// Enables all computations related to sensors. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:sensor = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetSensorFlagAttr() const; + + /// See GetSensorFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateSensorFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MIDPHASEFLAG + // --------------------------------------------------------------------- // + /// Enables mid-phase collision filtering using a static AABB bounding volume + /// hierarchy (BVH). + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:midphase = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMidPhaseFlagAttr() const; + + /// See GetMidPhaseFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMidPhaseFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // NATIVECCDFLAG + // --------------------------------------------------------------------- // + /// Enables the native convex collision detection pipeline instead of using + /// the libccd library. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:nativeccd = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetNativeCCDFlagAttr() const; + + /// See GetNativeCCDFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateNativeCCDFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // EULERDAMPFLAG + // --------------------------------------------------------------------- // + /// Enables implicit integration with respect to joint damping in the Euler + /// integrator. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:eulerdamp = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetEulerDampFlagAttr() const; + + /// See GetEulerDampFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateEulerDampFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // AUTORESETFLAG + // --------------------------------------------------------------------- // + /// Enables the automatic resetting of the simulation state when numerical + /// issues are detected. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:autoreset = 1` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetAutoResetFlagAttr() const; + + /// See GetAutoResetFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateAutoResetFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // OVERRIDEFLAG + // --------------------------------------------------------------------- // + /// Enables the contact override mechanism. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:override = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetOverrideFlagAttr() const; + + /// See GetOverrideFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateOverrideFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // ENERGYFLAG + // --------------------------------------------------------------------- // + /// Enables the computation of potential and kinetic energy + /// (mjData.energy[0,1]). + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:energy = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetEnergyFlagAttr() const; + + /// See GetEnergyFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateEnergyFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // FWDINVFLAG + // --------------------------------------------------------------------- // + /// Enables the automatic comparison of forward and inverse dynamics. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:fwdinv = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetFwdinvFlagAttr() const; + + /// See GetFwdinvFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateFwdinvFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // INVDISCRETEFLAG + // --------------------------------------------------------------------- // + /// Enables discrete-time inverse dynamics with mj_inverse for integrators + /// other than RK4. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:invdiscrete = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetInvDiscreteFlagAttr() const; + + /// See GetInvDiscreteFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateInvDiscreteFlagAttr( + VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // MULTICCDFLAG + // --------------------------------------------------------------------- // + /// Enables multiple-contact collision detection for geom pairs using a + /// general-purpose convex-convex collider. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:multiccd = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetMultiCCDFlagAttr() const; + + /// See GetMultiCCDFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateMultiCCDFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // --------------------------------------------------------------------- // + // ISLANDFLAG + // --------------------------------------------------------------------- // + /// Enables the discovery of constraint islands. + /// + /// | || + /// | -- | -- | + /// | Declaration | `uniform bool mjc:flag:island = 0` | + /// | C++ Type | bool | + /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->Bool | + /// | \ref SdfVariability "Variability" | SdfVariabilityUniform | + MJCPHYSICS_API + UsdAttribute GetIslandFlagAttr() const; + + /// See GetIslandFlagAttr(), and also + /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create. + /// If specified, author \p defaultValue as the attribute's default, + /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true - + /// the default for \p writeSparsely is \c false. + MJCPHYSICS_API + UsdAttribute CreateIslandFlagAttr(VtValue const &defaultValue = VtValue(), + bool writeSparsely = false) const; + + public: + // ===================================================================== // + // Feel free to add custom code below this line, it will be preserved by + // the code generator. + // + // Just remember to: + // - Close the class declaration with }; + // - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE + // - Close the include guard with #endif + // ===================================================================== // + // --(BEGIN CUSTOM CODE)-- +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/mjcPhysics/schema.usda b/src/experimental/usd/mjcPhysics/schema.usda new file mode 100644 index 00000000..a59c6bb9 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/schema.usda @@ -0,0 +1,665 @@ +#usda 1.0 +( + subLayers = [ + @usd/schema.usda@, + ] +) + +over "GLOBAL" ( + customData = { + string libraryName = "mjcPhysics" + string libraryPath = "." + bool useLiteralIdentifier = 0 + dictionary libraryTokens = { + dictionary euler = { + string doc = """ + This token represents the Euler numerical integrator. + """ + } + + dictionary rk4 = { + string doc = """ + This token represents the RK4 numerical integrator. + """ + } + + dictionary implicit = { + string doc = """ + This token represents the implicit numerical integrator. + """ + } + + dictionary implicitfast = { + string doc = """ + This token represents the implicitfast numerical integrator. + """ + } + + dictionary pyramidal = { + string doc = """ + This token represents the pyramidal contact friction cone type. + """ + } + + dictionary elliptic = { + string doc = """ + This token represents the elliptic contact friction cone type. + """ + } + + dictionary dense = { + string doc = """ + This token represents the dense constraint Jacobian and matrices + computed from it. + """ + } + + dictionary sparse = { + string doc = """ + This token represents the sparse constraint Jacobian and matrices + computed from it. + """ + } + + dictionary auto = { + string doc = """ + This token represents the auto constraint Jacobian and matrices + computed from it. + """ + } + + dictionary pgs = { + string doc = """ + This token represents the PGS constraint solver algorithm. + """ + } + + dictionary cg = { + string doc = """ + This token represents the CG constraint solver algorithm. + """ + } + + dictionary newton = { + string doc = """ + This token represents the Newton constraint solver algorithm. + """ + } + } + } +) +{ + +} + +class "SceneAPI" +( + doc = """API providing global simulation options for Mujoco.""" + + inherits = +) +{ + uniform double mjc:option:timestep = 0.002 ( + customData = { + string apiName = "Timestep" + } + displayName = "Timestep" + doc = """Controls the timestep in seconds used by MuJoCo.""" + ) + + uniform double mjc:option:apirate = 100 ( + customData = { + string apiName = "ApiRate" + } + displayName = "ApiRate" + doc = """Determines the rate (in Hz) at which an external API allows + the update function to be executed.""" + ) + + uniform double mjc:option:impratio = 1.0 ( + customData = { + string apiName = "ImpRatio" + } + displayName = "Impedance Ratio" + doc = """Ratio of frictional-to-normal constraint impedance for elliptic + friction cones.""" + ) + + uniform double3 mjc:option:wind = (0.0, 0.0, 0.0) ( + customData = { + string apiName = "Wind" + } + displayName = "Wind Velocity" + doc = """Velocity vector of medium (i.e. wind).""" + ) + + uniform double3 mjc:option:magnetic = (0.0, -0.5, 0.0) ( + customData = { + string apiName = "Magnetic" + } + displayName = "Magnetic Flux" + doc = """Global magnetic flux.""" + ) + + uniform double mjc:option:density = 0.0 ( + customData = { + string apiName = "Density" + } + displayName = "Density" + doc = """Density of medium.""" + ) + + uniform double mjc:option:viscosity = 0.0 ( + customData = { + string apiName = "Viscosity" + } + displayName = "Viscosity" + doc = """Viscosity of medium.""" + ) + + uniform double mjc:option:o_margin = 0.0 ( + customData = { + string apiName = "OMargin" + } + displayName = "Contact Override Margin" + doc = """Replaces the margin parameter of all active contact pairs when + Contact override is enabled.""" + ) + + uniform double[] mjc:option:o_solref = [0.02, 1.0] ( + customData = { + string apiName = "OSolRef" + } + displayName = "Contact Override SolRef" + doc = """Replaces the solref parameter of all active contact pairs when + Contact override is enabled.""" + ) + + uniform double[] mjc:option:o_solimp = [0.9, 0.95, 0.001, 0.5, 2.0] ( + customData = { + string apiName = "OSolImp" + } + displayName = "Contact Override SolImp" + doc = """Replaces the solimp parameter of all active contact pairs when + Contact override is enabled.""" + ) + + uniform double[] mjc:option:o_friction = [1.0, 1.0, 0.005, 0.0001, 0.0001] ( + customData = { + string apiName = "OFriction" + } + displayName = "Contact Override Friction" + doc = """Replaces the friction parameter of all active contact pairs when + Contact override is enabled.""" + ) + + uniform token mjc:option:integrator = "euler" ( + allowedTokens = ["euler", "rk4", "implicit", "implicitfast"] + customData = { + string apiName = "Integrator" + } + displayName = "Integrator" + doc = """Numerical integrator to be used.""" + ) + + uniform token mjc:option:cone = "pyramidal" ( + allowedTokens = ["pyramidal", "elliptic"] + customData = { + string apiName = "Cone" + } + displayName = "Friction Cone Type" + doc = """The type of contact friction cone.""" + ) + + uniform token mjc:option:jacobian = "auto" ( + allowedTokens = ["auto", "dense", "sparse"] + customData = { + string apiName = "Jacobian" + } + displayName = "Jacobian Type" + doc = """The type of constraint Jacobian and matrices computed from it.""" + ) + + uniform token mjc:option:solver = "newton" ( + allowedTokens = ["pgs", "cg", "newton"] + customData = { + string apiName = "Solver" + } + displayName = "Solver" + doc = """Constraint solver algorithm to be used.""" + ) + + uniform int mjc:option:iterations = 100 ( + customData = { + string apiName = "Iterations" + } + displayName = "Solver Iterations" + doc = """Maximum number of iterations of the constraint solver.""" + ) + + uniform double mjc:option:tolerance = 1e-08 ( + customData = { + string apiName = "Tolerance" + } + displayName = "Solver Tolerance" + doc = """Tolerance threshold used for early termination of the iterative + solver.""" + ) + + uniform int mjc:option:ls_iterations = 50 ( + customData = { + string apiName = "LSIterations" + } + displayName = "Linesearch Iterations" + doc = """Maximum number of linesearch iterations performed by CG/Newton + constraint solvers.""" + ) + + uniform double mjc:option:ls_tolerance = 0.01 ( + customData = { + string apiName = "LSTolerance" + } + displayName = "Linesearch Tolerance" + doc = """Tolerance threshold used for early termination of the linesearch algorithm.""" + ) + + uniform int mjc:option:noslip_iterations = 0 ( + customData = { + string apiName = "NoslipIterations" + } + displayName = "Noslip Iterations" + doc = """Maximum number of iterations of the Noslip solver.""" + ) + + uniform double mjc:option:noslip_tolerance = 1e-06 ( + customData = { + string apiName = "NoslipTolerance" + } + displayName = "Noslip Tolerance" + doc = """Tolerance threshold used for early termination of the Noslip solver.""" + ) + + uniform int mjc:option:ccd_iterations = 50 ( + customData = { + string apiName = "CCDIterations" + } + displayName = "CCD Iterations" + doc = """Maximum number of iterations of the algorithm used for convex collisions.""" + ) + + uniform double mjc:option:ccd_tolerance = 1e-06 ( + customData = { + string apiName = "CCDTolerance" + } + displayName = "CCD Tolerance" + doc = """Tolerance threshold used for early termination of the convex + collision algorithm.""" + ) + + uniform int mjc:option:sdf_iterations = 10 ( + customData = { + string apiName = "SDFIterations" + } + displayName = "SDF Iterations" + doc = """Number of iterations used for Signed Distance Field collisions + (per initial point).""" + ) + + uniform int mjc:option:sdf_initpoints = 40 ( + customData = { + string apiName = "SDFInitPoints" + } + displayName = "SDF Initial Points" + doc = """Number of starting points used for finding contacts with Signed + Distance Field collisions.""" + ) + + uniform int[] mjc:option:actuatorgroupdisable ( + customData = { + string apiName = "ActuatorGroupDisable" + } + displayName = "Actuator Group Disable" + doc = """List of actuator groups to disable.""" + ) + + uniform bool mjc:flag:constraint = True ( + customData = { + string apiName = "ConstraintFlag" + } + displayName = "Constraint Solver Toggle" + doc = """Enables constraint solver.""" + ) + + uniform bool mjc:flag:equality = True ( + customData = { + string apiName = "EqualityFlag" + } + displayName = "Equality Constraints Toggle" + doc = """Enables all standard computations related to equality constraints.""" + ) + + uniform bool mjc:flag:frictionloss = True ( + customData = { + string apiName = "FrictionLossFlag" + } + displayName = "Friction Loss Constraints Toggle" + doc = """Enables all standard computations related to friction loss constraints.""" + ) + + uniform bool mjc:flag:limit = True ( + customData = { + string apiName = "LimitFlag" + } + displayName = "Joint and Tendon Limit Constraints Toggle" + doc = """Enables all standard computations related to joint and tendon limit constraints.""" + ) + + uniform bool mjc:flag:contact = True ( + customData = { + string apiName = "ContactFlag" + } + displayName = "Contact Constraints and Collision Detection Toggle" + doc = """Enables collision detection and all standard computations related to contact constraints.""" + ) + + uniform bool mjc:flag:passive = True ( + customData = { + string apiName = "PassiveFlag" + } + displayName = "Passive Forces Toggle" + doc = """Enables the simulation of joint and tendon spring-dampers, fluid dynamics forces, and custom passive forces.""" + ) + + uniform bool mjc:flag:gravity = True ( + customData = { + string apiName = "GravityFlag" + } + displayName = "Gravity Toggle" + doc = """Enables the application of gravitational acceleration as defined in mjOption.""" + ) + + uniform bool mjc:flag:clampctrl = True ( + customData = { + string apiName = "ClampCtrlFlag" + } + displayName = "Control Input Clamping Toggle" + doc = """Enables the clamping of control inputs to all actuators, according to actuator-specific attributes.""" + ) + + uniform bool mjc:flag:warmstart = True ( + customData = { + string apiName = "WarmStartFlag" + } + displayName = "Solver Warm-Starting Toggle" + doc = """Enables warm-starting of the constraint solver, using the solution from the previous time step to initialize the iterative optimization.""" + ) + + uniform bool mjc:flag:filterparent = True ( + customData = { + string apiName = "FilterParentFlag" + } + displayName = "Parent-Child Contact Filtering Toggle" + doc = """Enables the filtering of contact pairs where the two geoms belong to a parent and child body.""" + ) + + uniform bool mjc:flag:actuation = True ( + customData = { + string apiName = "ActuationFlag" + } + displayName = "Actuation Forces Toggle" + doc = """Enables all standard computations related to actuator forces, including actuator dynamics.""" + ) + + uniform bool mjc:flag:refsafe = True ( + customData = { + string apiName = "RefSafeFlag" + } + displayName = "Solver Reference Safety Mechanism Toggle" + doc = """Enables a safety mechanism that prevents instabilities due to solref[0] being too small compared to the simulation timestep.""" + ) + + uniform bool mjc:flag:sensor = True ( + customData = { + string apiName = "SensorFlag" + } + displayName = "Sensor Computations Toggle" + doc = """Enables all computations related to sensors.""" + ) + + uniform bool mjc:flag:midphase = True ( + customData = { + string apiName = "MidPhaseFlag" + } + displayName = "Mid-Phase Collision Filtering Toggle" + doc = """Enables mid-phase collision filtering using a static AABB bounding volume hierarchy (BVH).""" + ) + + uniform bool mjc:flag:nativeccd = True ( + customData = { + string apiName = "NativeCCDFlag" + } + displayName = "Native Convex Collision Detection Toggle" + doc = """Enables the native convex collision detection pipeline instead of using the libccd library.""" + ) + + uniform bool mjc:flag:eulerdamp = True ( + customData = { + string apiName = "EulerDampFlag" + } + displayName = "Euler Integrator Damping Toggle" + doc = """Enables implicit integration with respect to joint damping in the Euler integrator.""" + ) + + uniform bool mjc:flag:autoreset = True ( + customData = { + string apiName = "AutoResetFlag" + } + displayName = "Automatic Simulation Reset Toggle" + doc = """Enables the automatic resetting of the simulation state when numerical issues are detected.""" + ) + + uniform bool mjc:flag:override = False ( + customData = { + string apiName = "OverrideFlag" + } + displayName = "Contact Override Mechanism Toggle" + doc = """Enables the contact override mechanism.""" + ) + + uniform bool mjc:flag:energy = False ( + customData = { + string apiName = "EnergyFlag" + } + displayName = "Energy Computation Toggle" + doc = """Enables the computation of potential and kinetic energy (mjData.energy[0,1]).""" + ) + + uniform bool mjc:flag:fwdinv = False ( + customData = { + string apiName = "FwdinvFlag" + } + displayName = "Forward/Inverse Dynamics Comparison Toggle" + doc = """Enables the automatic comparison of forward and inverse dynamics.""" + ) + + uniform bool mjc:flag:invdiscrete = False ( + customData = { + string apiName = "InvDiscreteFlag" + } + displayName = "Discrete-Time Inverse Dynamics Toggle" + doc = """Enables discrete-time inverse dynamics with mj_inverse for integrators other than RK4.""" + ) + + uniform bool mjc:flag:multiccd = False ( + customData = { + string apiName = "MultiCCDFlag" + } + displayName = "Multiple Contact Collision Detection (CCD) Toggle" + doc = """Enables multiple-contact collision detection for geom pairs using a general-purpose convex-convex collider.""" + ) + + uniform bool mjc:flag:island = False ( + customData = { + string apiName = "IslandFlag" + } + displayName = "Constraint Island Discovery Toggle" + doc = """Enables the discovery of constraint islands.""" + ) +} + +class "SiteAPI" +( + doc = """API describing a Mujoco site.""" + + inherits = +) +{} + +class "CollisionAPI" +( + doc = """API describing a Mujoco collider.""" + + inherits = +) +{ + uniform bool mjc:shellinertia = False ( + customData = { + string apiName = "ShellInertia" + } + displayName = "Shell Inertia" + doc = """Enables handling of the inertia assuming mass is concentrated on the surface.""" + ) +} + +class "MeshCollisionAPI" +( + doc = """API describing a Mujoco collider.""" + + inherits = +) +{ + uniform token mjc:inertia = "legacy" ( + allowedTokens = ["legacy", "convex", "exact", "shell"] + customData = { + string apiName = "Inertia" + } + displayName = "Inertia" + doc = """Controls how a mesh is used when mass and inertia are inferred from geometry.""" + ) +} + +class "PhysicsActuatorAPI" +( + customData = { + string className = "ActuatorAPI" + } + doc = """API describing a Mujoco actuator.""" + + inherits = +) +{ + # Control/Force/Activation Limits + uniform token mjc:ctrlLimited = "auto" ( + doc = "If true, the control input to this actuator is automatically clamped to ctrlrange at runtime. If false, control input clamping is disabled. If 'auto' and autolimits is set in compiler, control clamping will automatically be set to true if ctrlrange is defined without explicitly setting this attribute to 'true'. Note that control input clamping can also be globally disabled with the clampctrl attribute of option/flag." + allowedTokens = ["false", "true", "auto"] + ) + uniform token mjc:forceLimited = "auto" ( + doc = "If true, the force output of this actuator is automatically clamped to forcerange at runtime. If false, force clamping is disabled. If 'auto' and autolimits is set in compiler, force clamping will automatically be set to true if forcerange is defined without explicitly setting this attribute to 'true'." + allowedTokens = ["false", "true", "auto"] + ) + uniform token mjc:actLimited = "auto" ( + doc = "If true, the internal state (activation) associated with this actuator is automatically clamped to actrange at runtime. If false, activation clamping is disabled. If 'auto' and autolimits is set in compiler, activation clamping will automatically be set to true if actrange is defined without explicitly setting this attribute to 'true'. See the Activation clamping section for more details." + allowedTokens = ["false", "true", "auto"] + ) + + uniform double mjc:ctrlRange:min = 0 ( + doc = "Minimum range for clamping the control input. The first value must be smaller than the second value. Setting this attribute without specifying ctrllimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:ctrlRange:max = 0 ( + doc = "Maximum range for clamping the control input. The first value must be smaller than the second value. Setting this attribute without specifying ctrllimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:forceRange:min = 0 ( + doc = "Minimum range for clamping the force output. The first value must be no greater than the second value. Setting this attribute without specifying forcelimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:forceRange:max = 0 ( + doc = "Maximum range for clamping the force output. The first value must be no greater than the second value. Setting this attribute without specifying forcelimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:actRange:min = 0 ( + doc = "Minimum range for clamping the activation state. The first value must be no greater than the second value. See the Activation clamping section for more details. Setting this attribute without specifying actlimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:actRange:max = 0 ( + doc = "Maximum range for clamping the activation state. The first value must be no greater than the second value. See the Activation clamping section for more details. Setting this attribute without specifying actlimited is an error if autolimits is 'false' in compiler." + ) + + uniform double mjc:lengthRange:min = 0 ( + doc = "Minimum range of feasible lengths of the actuator’s transmission." + ) + + uniform double mjc:lengthRange:max = 0 ( + doc = "Maximum range of feasible lengths of the actuator’s transmission." + ) + + # Transmission Properties + uniform double[] mjc:gear = [1, 0, 0, 0, 0, 0] ( + doc = "This attribute scales the length (and consequently moment arms, velocity and force) of the actuator, for all transmission types. It is different from the gain in the force generation mechanism, because the gain only scales the force output and does not affect the length, moment arms and velocity. For actuators with scalar transmission, only the first element of this vector is used. The remaining elements are needed for joint, jointinparent and site transmissions where this attribute is used to specify 3D force and torque axes." + ) + + uniform double mjc:crankLength = 0.0 ( + doc = "Used only for the slider-crank transmission type. Specifies the length of the connecting rod. The compiler expects this value to be positive when a slider-crank transmission is present." + ) + + uniform bool mjc:jointInParent = False ( + doc = "If true and applied to ball and free joints, the 3d rotation axis given by gear is defined in the parent frame (which is the world frame for free joints) rather than the child frame." + ) + + rel mjc:refSite ( + doc = "When applied to a site, measure the translation and rotation w.r.t the frame of the refsite. In this case the actuator does have length and position actuators can be used to directly control an end effector, see refsite.xml example model. As above, the length is the dot product of the gear vector and the frame difference. So gear='0 1 0 0 0 0' means 'Y-offset of site in the refsite frame', while gear='0 0 0 0 0 1' means rotation 'Z- rotation of site in the refsite frame'. It is recommended to use a normalized gear vector with nonzeros in only the first 3 or the last 3 elements of gear, so the actuator length will be in either length units or radians, respectively. As with ball joints (see joint above), for rotations which exceed a total angle of pi will wrap around, so tighter limits are recommended." + ) + + rel mjc:sliderSite ( + doc = "Used only for the slider-crank transmission type. The target site is the pin joining the slider and the connecting rod. The slider moves along the z-axis of the slidersite frame. Therefore the site should be oriented as needed when it is defined in the kinematic tree; its orientation cannot be changed in the actuator definition." + ) + + # Activation Dynamics and Force Generation + uniform int mjc:actDim = -1 ( + doc = "Dimension of the activation state. The default value of -1 instructs the compiler to set the dimension according to the dyntype. Values larger than 1 are only allowed for user-defined activation dynamics, as native types require dimensions of only 0 or 1. For activation dimensions bigger than 1, the last element is used to generate force." + ) + + uniform token mjc:dynType = "none" ( + doc = "Activation dynamics type for the actuator. The available dynamics types were already described in the Actuation model section." + allowedTokens = ["none", "integrator", "filter", "filterexact", "muscle", "user"] + ) + + uniform token mjc:gainType = "fixed" ( + doc = "The gain and bias together determine the output of the force generation mechanism, which is currently assumed to be affine." + allowedTokens = ["fixed", "affine", "muscle", "user"] + ) + + uniform token mjc:biasType = "none" ( + doc = "The gain and bias together determine the output of the force generation mechanism, which is currently assumed to be affine." + allowedTokens = ["none", "affine", "muscle", "user"] + ) + + uniform double[] mjc:dynPrm = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Activation dynamics parameters. The built-in activation types (except for muscle) use only the first parameter, but we provide additional parameters in case user callbacks implement a more elaborate model. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + + uniform double[] mjc:gainPrm = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Gain parameters. The built-in gain types (except for muscle) use only the first parameter, but we provide additional parameters in case user callbacks implement a more elaborate model. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + + uniform double[] mjc:biasPrm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] ( + doc = "Bias parameters. The affine bias type uses three parameters. The length of this array is not enforced by the parser, so the user can enter as many parameters as needed. These defaults are not compatible with muscle actuators." + ) + + uniform bool mjc:actEarly = False ( + doc = "If true, force computation will use the next value of the activation variable rather than the current one. Setting this flag reduces the delay between the control and accelerations by one time-step." + ) +} + + diff --git a/src/experimental/usd/mjcPhysics/siteAPI.cpp b/src/experimental/usd/mjcPhysics/siteAPI.cpp new file mode 100644 index 00000000..1693dd16 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/siteAPI.cpp @@ -0,0 +1,99 @@ +// Copyright 2025 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. + +#include "./siteAPI.h" + +#include "pxr/usd/sdf/assetPath.h" +#include "pxr/usd/sdf/types.h" +#include "pxr/usd/usd/schemaRegistry.h" +#include "pxr/usd/usd/typed.h" + +PXR_NAMESPACE_OPEN_SCOPE + +// Register the schema with the TfType system. +TF_REGISTRY_FUNCTION(TfType) { + TfType::Define >(); +} + +/* virtual */ +MjcPhysicsSiteAPI::~MjcPhysicsSiteAPI() {} + +/* static */ +MjcPhysicsSiteAPI MjcPhysicsSiteAPI::Get(const UsdStagePtr &stage, + const SdfPath &path) { + if (!stage) { + TF_CODING_ERROR("Invalid stage"); + return MjcPhysicsSiteAPI(); + } + return MjcPhysicsSiteAPI(stage->GetPrimAtPath(path)); +} + +/* virtual */ +UsdSchemaKind MjcPhysicsSiteAPI::_GetSchemaKind() const { + return MjcPhysicsSiteAPI::schemaKind; +} + +/* static */ +bool MjcPhysicsSiteAPI::CanApply(const UsdPrim &prim, std::string *whyNot) { + return prim.CanApplyAPI(whyNot); +} + +/* static */ +MjcPhysicsSiteAPI MjcPhysicsSiteAPI::Apply(const UsdPrim &prim) { + if (prim.ApplyAPI()) { + return MjcPhysicsSiteAPI(prim); + } + return MjcPhysicsSiteAPI(); +} + +/* static */ +const TfType &MjcPhysicsSiteAPI::_GetStaticTfType() { + static TfType tfType = TfType::Find(); + return tfType; +} + +/* static */ +bool MjcPhysicsSiteAPI::_IsTypedSchema() { + static bool isTyped = _GetStaticTfType().IsA(); + return isTyped; +} + +/* virtual */ +const TfType &MjcPhysicsSiteAPI::_GetTfType() const { + return _GetStaticTfType(); +} + +/*static*/ +const TfTokenVector &MjcPhysicsSiteAPI::GetSchemaAttributeNames( + bool includeInherited) { + static TfTokenVector localNames; + static TfTokenVector allNames = + UsdAPISchemaBase::GetSchemaAttributeNames(true); + + if (includeInherited) + return allNames; + else + return localNames; +} + +PXR_NAMESPACE_CLOSE_SCOPE + +// ===================================================================== // +// Feel free to add custom code below this line. It will be preserved by +// the code generator. +// +// Just remember to wrap code in the appropriate delimiters: +// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'. +// ===================================================================== // +// --(BEGIN CUSTOM CODE)-- diff --git a/src/experimental/usd/mjcPhysics/siteAPI.h b/src/experimental/usd/mjcPhysics/siteAPI.h new file mode 100644 index 00000000..537345ce --- /dev/null +++ b/src/experimental/usd/mjcPhysics/siteAPI.h @@ -0,0 +1,158 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_GENERATED_SITEAPI_H +#define MJCPHYSICS_GENERATED_SITEAPI_H + +/// \file mjcPhysics/siteAPI.h + +#include "./api.h" +#include "pxr/base/gf/matrix4d.h" +#include "pxr/base/gf/vec3d.h" +#include "pxr/base/gf/vec3f.h" +#include "pxr/base/tf/token.h" +#include "pxr/base/tf/type.h" +#include "pxr/base/vt/value.h" +#include "pxr/pxr.h" +#include "pxr/usd/usd/apiSchemaBase.h" +#include "pxr/usd/usd/prim.h" +#include "pxr/usd/usd/stage.h" + +PXR_NAMESPACE_OPEN_SCOPE + +class SdfAssetPath; + +// -------------------------------------------------------------------------- // +// SITEAPI // +// -------------------------------------------------------------------------- // + +/// \class MjcPhysicsSiteAPI +/// +/// API describing a Mujoco site. +/// +class MjcPhysicsSiteAPI : public UsdAPISchemaBase { + public: + /// Compile time constant representing what kind of schema this class is. + /// + /// \sa UsdSchemaKind + static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI; + + /// Construct a MjcPhysicsSiteAPI on UsdPrim \p prim . + /// Equivalent to MjcPhysicsSiteAPI::Get(prim.GetStage(), prim.GetPath()) + /// for a \em valid \p prim, but will not immediately throw an error for + /// an invalid \p prim + explicit MjcPhysicsSiteAPI(const UsdPrim &prim = UsdPrim()) + : UsdAPISchemaBase(prim) {} + + /// Construct a MjcPhysicsSiteAPI on the prim held by \p schemaObj . + /// Should be preferred over MjcPhysicsSiteAPI(schemaObj.GetPrim()), + /// as it preserves SchemaBase state. + explicit MjcPhysicsSiteAPI(const UsdSchemaBase &schemaObj) + : UsdAPISchemaBase(schemaObj) {} + + /// Destructor. + MJCPHYSICS_API + virtual ~MjcPhysicsSiteAPI(); + + /// Return a vector of names of all pre-declared attributes for this schema + /// class and all its ancestor classes. Does not include attributes that + /// may be authored by custom/extended methods of the schemas involved. + MJCPHYSICS_API + static const TfTokenVector &GetSchemaAttributeNames( + bool includeInherited = true); + + /// Return a MjcPhysicsSiteAPI holding the prim adhering to this + /// schema at \p path on \p stage. If no prim exists at \p path on + /// \p stage, or if the prim at that path does not adhere to this schema, + /// return an invalid schema object. This is shorthand for the following: + /// + /// \code + /// MjcPhysicsSiteAPI(stage->GetPrimAtPath(path)); + /// \endcode + /// + MJCPHYSICS_API + static MjcPhysicsSiteAPI Get(const UsdStagePtr &stage, const SdfPath &path); + + /// Returns true if this single-apply API schema can be applied to + /// the given \p prim. If this schema can not be a applied to the prim, + /// this returns false and, if provided, populates \p whyNot with the + /// reason it can not be applied. + /// + /// Note that if CanApply returns false, that does not necessarily imply + /// that calling Apply will fail. Callers are expected to call CanApply + /// before calling Apply if they want to ensure that it is valid to + /// apply a schema. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static bool CanApply(const UsdPrim &prim, std::string *whyNot = nullptr); + + /// Applies this single-apply API schema to the given \p prim. + /// This information is stored by adding "SiteAPI" to the + /// token-valued, listOp metadata \em apiSchemas on the prim. + /// + /// \return A valid MjcPhysicsSiteAPI object is returned upon success. + /// An invalid (or empty) MjcPhysicsSiteAPI object is returned upon + /// failure. See \ref UsdPrim::ApplyAPI() for conditions + /// resulting in failure. + /// + /// \sa UsdPrim::GetAppliedSchemas() + /// \sa UsdPrim::HasAPI() + /// \sa UsdPrim::CanApplyAPI() + /// \sa UsdPrim::ApplyAPI() + /// \sa UsdPrim::RemoveAPI() + /// + MJCPHYSICS_API + static MjcPhysicsSiteAPI Apply(const UsdPrim &prim); + + protected: + /// Returns the kind of schema this class belongs to. + /// + /// \sa UsdSchemaKind + MJCPHYSICS_API + UsdSchemaKind _GetSchemaKind() const override; + + private: + // needs to invoke _GetStaticTfType. + friend class UsdSchemaRegistry; + MJCPHYSICS_API + static const TfType &_GetStaticTfType(); + + static bool _IsTypedSchema(); + + // override SchemaBase virtuals. + MJCPHYSICS_API + const TfType &_GetTfType() const override; + + public: + // ===================================================================== // + // Feel free to add custom code below this line, it will be preserved by + // the code generator. + // + // Just remember to: + // - Close the class declaration with }; + // - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE + // - Close the include guard with #endif + // ===================================================================== // + // --(BEGIN CUSTOM CODE)-- +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/mjcPhysics/tokens.cpp b/src/experimental/usd/mjcPhysics/tokens.cpp new file mode 100644 index 00000000..2f33d763 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/tokens.cpp @@ -0,0 +1,240 @@ +// Copyright 2025 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. + +#include "./tokens.h" + +PXR_NAMESPACE_OPEN_SCOPE + +MjcPhysicsTokensType::MjcPhysicsTokensType() + : affine("affine", TfToken::Immortal), + auto_("auto", TfToken::Immortal), + cg("cg", TfToken::Immortal), + convex("convex", TfToken::Immortal), + dense("dense", TfToken::Immortal), + elliptic("elliptic", TfToken::Immortal), + euler("euler", TfToken::Immortal), + exact("exact", TfToken::Immortal), + false_("false", TfToken::Immortal), + filter("filter", TfToken::Immortal), + filterexact("filterexact", TfToken::Immortal), + fixed("fixed", TfToken::Immortal), + implicit("implicit", TfToken::Immortal), + implicitfast("implicitfast", TfToken::Immortal), + integrator("integrator", TfToken::Immortal), + legacy("legacy", TfToken::Immortal), + mjcActDim("mjc:actDim", TfToken::Immortal), + mjcActEarly("mjc:actEarly", TfToken::Immortal), + mjcActLimited("mjc:actLimited", TfToken::Immortal), + mjcActRangeMax("mjc:actRange:max", TfToken::Immortal), + mjcActRangeMin("mjc:actRange:min", TfToken::Immortal), + mjcBiasPrm("mjc:biasPrm", TfToken::Immortal), + mjcBiasType("mjc:biasType", TfToken::Immortal), + mjcCrankLength("mjc:crankLength", TfToken::Immortal), + mjcCrankSite("mjc:crankSite", TfToken::Immortal), + mjcCtrlLimited("mjc:ctrlLimited", TfToken::Immortal), + mjcCtrlRangeMax("mjc:ctrlRange:max", TfToken::Immortal), + mjcCtrlRangeMin("mjc:ctrlRange:min", TfToken::Immortal), + mjcDynPrm("mjc:dynPrm", TfToken::Immortal), + mjcDynType("mjc:dynType", TfToken::Immortal), + mjcFlagActuation("mjc:flag:actuation", TfToken::Immortal), + mjcFlagAutoreset("mjc:flag:autoreset", TfToken::Immortal), + mjcFlagClampctrl("mjc:flag:clampctrl", TfToken::Immortal), + mjcFlagConstraint("mjc:flag:constraint", TfToken::Immortal), + mjcFlagContact("mjc:flag:contact", TfToken::Immortal), + mjcFlagEnergy("mjc:flag:energy", TfToken::Immortal), + mjcFlagEquality("mjc:flag:equality", TfToken::Immortal), + mjcFlagEulerdamp("mjc:flag:eulerdamp", TfToken::Immortal), + mjcFlagFilterparent("mjc:flag:filterparent", TfToken::Immortal), + mjcFlagFrictionloss("mjc:flag:frictionloss", TfToken::Immortal), + mjcFlagFwdinv("mjc:flag:fwdinv", TfToken::Immortal), + mjcFlagGravity("mjc:flag:gravity", TfToken::Immortal), + mjcFlagInvdiscrete("mjc:flag:invdiscrete", TfToken::Immortal), + mjcFlagIsland("mjc:flag:island", TfToken::Immortal), + mjcFlagLimit("mjc:flag:limit", TfToken::Immortal), + mjcFlagMidphase("mjc:flag:midphase", TfToken::Immortal), + mjcFlagMulticcd("mjc:flag:multiccd", TfToken::Immortal), + mjcFlagNativeccd("mjc:flag:nativeccd", TfToken::Immortal), + mjcFlagOverride("mjc:flag:override", TfToken::Immortal), + mjcFlagPassive("mjc:flag:passive", TfToken::Immortal), + mjcFlagRefsafe("mjc:flag:refsafe", TfToken::Immortal), + mjcFlagSensor("mjc:flag:sensor", TfToken::Immortal), + mjcFlagWarmstart("mjc:flag:warmstart", TfToken::Immortal), + mjcForceLimited("mjc:forceLimited", TfToken::Immortal), + mjcForceRangeMax("mjc:forceRange:max", TfToken::Immortal), + mjcForceRangeMin("mjc:forceRange:min", TfToken::Immortal), + mjcGainPrm("mjc:gainPrm", TfToken::Immortal), + mjcGainType("mjc:gainType", TfToken::Immortal), + mjcGear("mjc:gear", TfToken::Immortal), + mjcInertia("mjc:inertia", TfToken::Immortal), + mjcJointInParent("mjc:jointInParent", TfToken::Immortal), + mjcLengthRangeMax("mjc:lengthRange:max", TfToken::Immortal), + mjcLengthRangeMin("mjc:lengthRange:min", TfToken::Immortal), + mjcOptionActuatorgroupdisable("mjc:option:actuatorgroupdisable", + TfToken::Immortal), + mjcOptionApirate("mjc:option:apirate", TfToken::Immortal), + mjcOptionCcd_iterations("mjc:option:ccd_iterations", TfToken::Immortal), + mjcOptionCcd_tolerance("mjc:option:ccd_tolerance", TfToken::Immortal), + mjcOptionCone("mjc:option:cone", TfToken::Immortal), + mjcOptionDensity("mjc:option:density", TfToken::Immortal), + mjcOptionImpratio("mjc:option:impratio", TfToken::Immortal), + mjcOptionIntegrator("mjc:option:integrator", TfToken::Immortal), + mjcOptionIterations("mjc:option:iterations", TfToken::Immortal), + mjcOptionJacobian("mjc:option:jacobian", TfToken::Immortal), + mjcOptionLs_iterations("mjc:option:ls_iterations", TfToken::Immortal), + mjcOptionLs_tolerance("mjc:option:ls_tolerance", TfToken::Immortal), + mjcOptionMagnetic("mjc:option:magnetic", TfToken::Immortal), + mjcOptionNoslip_iterations("mjc:option:noslip_iterations", + TfToken::Immortal), + mjcOptionNoslip_tolerance("mjc:option:noslip_tolerance", + TfToken::Immortal), + mjcOptionO_friction("mjc:option:o_friction", TfToken::Immortal), + mjcOptionO_margin("mjc:option:o_margin", TfToken::Immortal), + mjcOptionO_solimp("mjc:option:o_solimp", TfToken::Immortal), + mjcOptionO_solref("mjc:option:o_solref", TfToken::Immortal), + mjcOptionSdf_initpoints("mjc:option:sdf_initpoints", TfToken::Immortal), + mjcOptionSdf_iterations("mjc:option:sdf_iterations", TfToken::Immortal), + mjcOptionSolver("mjc:option:solver", TfToken::Immortal), + mjcOptionTimestep("mjc:option:timestep", TfToken::Immortal), + mjcOptionTolerance("mjc:option:tolerance", TfToken::Immortal), + mjcOptionViscosity("mjc:option:viscosity", TfToken::Immortal), + mjcOptionWind("mjc:option:wind", TfToken::Immortal), + mjcRefSite("mjc:refSite", TfToken::Immortal), + mjcShellinertia("mjc:shellinertia", TfToken::Immortal), + mjcSliderSite("mjc:sliderSite", TfToken::Immortal), + muscle("muscle", TfToken::Immortal), + newton("newton", TfToken::Immortal), + none("none", TfToken::Immortal), + pgs("pgs", TfToken::Immortal), + pyramidal("pyramidal", TfToken::Immortal), + rk4("rk4", TfToken::Immortal), + shell("shell", TfToken::Immortal), + sparse("sparse", TfToken::Immortal), + true_("true", TfToken::Immortal), + user("user", TfToken::Immortal), + CollisionAPI("CollisionAPI", TfToken::Immortal), + MeshCollisionAPI("MeshCollisionAPI", TfToken::Immortal), + PhysicsActuatorAPI("PhysicsActuatorAPI", TfToken::Immortal), + SceneAPI("SceneAPI", TfToken::Immortal), + SiteAPI("SiteAPI", TfToken::Immortal), + allTokens({affine, + auto_, + cg, + convex, + dense, + elliptic, + euler, + exact, + false_, + filter, + filterexact, + fixed, + implicit, + implicitfast, + integrator, + legacy, + mjcActDim, + mjcActEarly, + mjcActLimited, + mjcActRangeMax, + mjcActRangeMin, + mjcBiasPrm, + mjcBiasType, + mjcCrankLength, + mjcCrankSite, + mjcCtrlLimited, + mjcCtrlRangeMax, + mjcCtrlRangeMin, + mjcDynPrm, + mjcDynType, + mjcFlagActuation, + mjcFlagAutoreset, + mjcFlagClampctrl, + mjcFlagConstraint, + mjcFlagContact, + mjcFlagEnergy, + mjcFlagEquality, + mjcFlagEulerdamp, + mjcFlagFilterparent, + mjcFlagFrictionloss, + mjcFlagFwdinv, + mjcFlagGravity, + mjcFlagInvdiscrete, + mjcFlagIsland, + mjcFlagLimit, + mjcFlagMidphase, + mjcFlagMulticcd, + mjcFlagNativeccd, + mjcFlagOverride, + mjcFlagPassive, + mjcFlagRefsafe, + mjcFlagSensor, + mjcFlagWarmstart, + mjcForceLimited, + mjcForceRangeMax, + mjcForceRangeMin, + mjcGainPrm, + mjcGainType, + mjcGear, + mjcInertia, + mjcJointInParent, + mjcLengthRangeMax, + mjcLengthRangeMin, + mjcOptionActuatorgroupdisable, + mjcOptionApirate, + mjcOptionCcd_iterations, + mjcOptionCcd_tolerance, + mjcOptionCone, + mjcOptionDensity, + mjcOptionImpratio, + mjcOptionIntegrator, + mjcOptionIterations, + mjcOptionJacobian, + mjcOptionLs_iterations, + mjcOptionLs_tolerance, + mjcOptionMagnetic, + mjcOptionNoslip_iterations, + mjcOptionNoslip_tolerance, + mjcOptionO_friction, + mjcOptionO_margin, + mjcOptionO_solimp, + mjcOptionO_solref, + mjcOptionSdf_initpoints, + mjcOptionSdf_iterations, + mjcOptionSolver, + mjcOptionTimestep, + mjcOptionTolerance, + mjcOptionViscosity, + mjcOptionWind, + mjcRefSite, + mjcShellinertia, + mjcSliderSite, + muscle, + newton, + none, + pgs, + pyramidal, + rk4, + shell, + sparse, + true_, + user, + CollisionAPI, + MeshCollisionAPI, + PhysicsActuatorAPI, + SceneAPI, + SiteAPI}) {} + +TfStaticData MjcPhysicsTokens; + +PXR_NAMESPACE_CLOSE_SCOPE diff --git a/src/experimental/usd/mjcPhysics/tokens.h b/src/experimental/usd/mjcPhysics/tokens.h new file mode 100644 index 00000000..862d2ef8 --- /dev/null +++ b/src/experimental/usd/mjcPhysics/tokens.h @@ -0,0 +1,521 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJCPHYSICS_TOKENS_H +#define MJCPHYSICS_TOKENS_H + +/// \file mjcPhysics/tokens.h + +// XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX +// +// This is an automatically generated file (by usdGenSchema.py). +// Do not hand-edit! +// +// XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX + +#include + +#include "./api.h" +#include "pxr/base/tf/staticData.h" +#include "pxr/base/tf/token.h" +#include "pxr/pxr.h" + +PXR_NAMESPACE_OPEN_SCOPE + +/// \class MjcPhysicsTokensType +/// +/// \link MjcPhysicsTokens \endlink provides static, efficient +/// \link TfToken TfTokens\endlink for use in all public USD API. +/// +/// These tokens are auto-generated from the module's schema, representing +/// property names, for when you need to fetch an attribute or relationship +/// directly by name, e.g. UsdPrim::GetAttribute(), in the most efficient +/// manner, and allow the compiler to verify that you spelled the name +/// correctly. +/// +/// MjcPhysicsTokens also contains all of the \em allowedTokens values +/// declared for schema builtin attributes of 'token' scene description type. +/// Use MjcPhysicsTokens like so: +/// +/// \code +/// gprim.GetMyTokenValuedAttr().Set(MjcPhysicsTokens->affine); +/// \endcode +struct MjcPhysicsTokensType { + MJCPHYSICS_API MjcPhysicsTokensType(); + /// \brief "affine" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcBiasTypeAttr(), Possible + /// value for MjcPhysicsActuatorAPI::GetMjcGainTypeAttr() + const TfToken affine; + /// \brief "auto" + /// + /// Fallback value for MjcPhysicsSceneAPI::GetJacobianAttr(), Fallback value + /// for MjcPhysicsActuatorAPI::GetMjcActLimitedAttr(), Fallback value for + /// MjcPhysicsActuatorAPI::GetMjcCtrlLimitedAttr(), Fallback value for + /// MjcPhysicsActuatorAPI::GetMjcForceLimitedAttr(), This token represents + /// the auto constraint Jacobian and matrices computed from it. + const TfToken auto_; + /// \brief "cg" + /// + /// Possible value for MjcPhysicsSceneAPI::GetSolverAttr(), This token + /// represents the CG constraint solver algorithm. + const TfToken cg; + /// \brief "convex" + /// + /// Possible value for MjcPhysicsMeshCollisionAPI::GetInertiaAttr() + const TfToken convex; + /// \brief "dense" + /// + /// Possible value for MjcPhysicsSceneAPI::GetJacobianAttr(), This token + /// represents the dense constraint Jacobian and matrices computed from it. + const TfToken dense; + /// \brief "elliptic" + /// + /// Possible value for MjcPhysicsSceneAPI::GetConeAttr(), This token + /// represents the elliptic contact friction cone type. + const TfToken elliptic; + /// \brief "euler" + /// + /// Fallback value for MjcPhysicsSceneAPI::GetIntegratorAttr(), This token + /// represents the Euler numerical integrator. + const TfToken euler; + /// \brief "exact" + /// + /// Possible value for MjcPhysicsMeshCollisionAPI::GetInertiaAttr() + const TfToken exact; + /// \brief "false" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcActLimitedAttr(), Possible + /// value for MjcPhysicsActuatorAPI::GetMjcCtrlLimitedAttr(), Possible value + /// for MjcPhysicsActuatorAPI::GetMjcForceLimitedAttr() + const TfToken false_; + /// \brief "filter" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr() + const TfToken filter; + /// \brief "filterexact" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr() + const TfToken filterexact; + /// \brief "fixed" + /// + /// Fallback value for MjcPhysicsActuatorAPI::GetMjcGainTypeAttr() + const TfToken fixed; + /// \brief "implicit" + /// + /// Possible value for MjcPhysicsSceneAPI::GetIntegratorAttr(), This token + /// represents the implicit numerical integrator. + const TfToken implicit; + /// \brief "implicitfast" + /// + /// Possible value for MjcPhysicsSceneAPI::GetIntegratorAttr(), This token + /// represents the implicitfast numerical integrator. + const TfToken implicitfast; + /// \brief "integrator" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr() + const TfToken integrator; + /// \brief "legacy" + /// + /// Fallback value for MjcPhysicsMeshCollisionAPI::GetInertiaAttr() + const TfToken legacy; + /// \brief "mjc:actDim" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcActDim; + /// \brief "mjc:actEarly" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcActEarly; + /// \brief "mjc:actLimited" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcActLimited; + /// \brief "mjc:actRange:max" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcActRangeMax; + /// \brief "mjc:actRange:min" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcActRangeMin; + /// \brief "mjc:biasPrm" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcBiasPrm; + /// \brief "mjc:biasType" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcBiasType; + /// \brief "mjc:crankLength" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcCrankLength; + /// \brief "mjc:crankSite" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcCrankSite; + /// \brief "mjc:ctrlLimited" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcCtrlLimited; + /// \brief "mjc:ctrlRange:max" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcCtrlRangeMax; + /// \brief "mjc:ctrlRange:min" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcCtrlRangeMin; + /// \brief "mjc:dynPrm" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcDynPrm; + /// \brief "mjc:dynType" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcDynType; + /// \brief "mjc:flag:actuation" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagActuation; + /// \brief "mjc:flag:autoreset" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagAutoreset; + /// \brief "mjc:flag:clampctrl" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagClampctrl; + /// \brief "mjc:flag:constraint" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagConstraint; + /// \brief "mjc:flag:contact" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagContact; + /// \brief "mjc:flag:energy" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagEnergy; + /// \brief "mjc:flag:equality" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagEquality; + /// \brief "mjc:flag:eulerdamp" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagEulerdamp; + /// \brief "mjc:flag:filterparent" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagFilterparent; + /// \brief "mjc:flag:frictionloss" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagFrictionloss; + /// \brief "mjc:flag:fwdinv" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagFwdinv; + /// \brief "mjc:flag:gravity" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagGravity; + /// \brief "mjc:flag:invdiscrete" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagInvdiscrete; + /// \brief "mjc:flag:island" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagIsland; + /// \brief "mjc:flag:limit" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagLimit; + /// \brief "mjc:flag:midphase" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagMidphase; + /// \brief "mjc:flag:multiccd" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagMulticcd; + /// \brief "mjc:flag:nativeccd" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagNativeccd; + /// \brief "mjc:flag:override" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagOverride; + /// \brief "mjc:flag:passive" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagPassive; + /// \brief "mjc:flag:refsafe" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagRefsafe; + /// \brief "mjc:flag:sensor" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagSensor; + /// \brief "mjc:flag:warmstart" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcFlagWarmstart; + /// \brief "mjc:forceLimited" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcForceLimited; + /// \brief "mjc:forceRange:max" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcForceRangeMax; + /// \brief "mjc:forceRange:min" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcForceRangeMin; + /// \brief "mjc:gainPrm" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcGainPrm; + /// \brief "mjc:gainType" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcGainType; + /// \brief "mjc:gear" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcGear; + /// \brief "mjc:inertia" + /// + /// MjcPhysicsMeshCollisionAPI + const TfToken mjcInertia; + /// \brief "mjc:jointInParent" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcJointInParent; + /// \brief "mjc:lengthRange:max" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcLengthRangeMax; + /// \brief "mjc:lengthRange:min" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcLengthRangeMin; + /// \brief "mjc:option:actuatorgroupdisable" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionActuatorgroupdisable; + /// \brief "mjc:option:apirate" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionApirate; + /// \brief "mjc:option:ccd_iterations" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionCcd_iterations; + /// \brief "mjc:option:ccd_tolerance" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionCcd_tolerance; + /// \brief "mjc:option:cone" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionCone; + /// \brief "mjc:option:density" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionDensity; + /// \brief "mjc:option:impratio" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionImpratio; + /// \brief "mjc:option:integrator" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionIntegrator; + /// \brief "mjc:option:iterations" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionIterations; + /// \brief "mjc:option:jacobian" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionJacobian; + /// \brief "mjc:option:ls_iterations" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionLs_iterations; + /// \brief "mjc:option:ls_tolerance" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionLs_tolerance; + /// \brief "mjc:option:magnetic" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionMagnetic; + /// \brief "mjc:option:noslip_iterations" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionNoslip_iterations; + /// \brief "mjc:option:noslip_tolerance" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionNoslip_tolerance; + /// \brief "mjc:option:o_friction" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionO_friction; + /// \brief "mjc:option:o_margin" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionO_margin; + /// \brief "mjc:option:o_solimp" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionO_solimp; + /// \brief "mjc:option:o_solref" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionO_solref; + /// \brief "mjc:option:sdf_initpoints" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionSdf_initpoints; + /// \brief "mjc:option:sdf_iterations" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionSdf_iterations; + /// \brief "mjc:option:solver" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionSolver; + /// \brief "mjc:option:timestep" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionTimestep; + /// \brief "mjc:option:tolerance" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionTolerance; + /// \brief "mjc:option:viscosity" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionViscosity; + /// \brief "mjc:option:wind" + /// + /// MjcPhysicsSceneAPI + const TfToken mjcOptionWind; + /// \brief "mjc:refSite" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcRefSite; + /// \brief "mjc:shellinertia" + /// + /// MjcPhysicsCollisionAPI + const TfToken mjcShellinertia; + /// \brief "mjc:sliderSite" + /// + /// MjcPhysicsActuatorAPI + const TfToken mjcSliderSite; + /// \brief "muscle" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcBiasTypeAttr(), Possible + /// value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr(), Possible value for + /// MjcPhysicsActuatorAPI::GetMjcGainTypeAttr() + const TfToken muscle; + /// \brief "newton" + /// + /// Fallback value for MjcPhysicsSceneAPI::GetSolverAttr(), This token + /// represents the Newton constraint solver algorithm. + const TfToken newton; + /// \brief "none" + /// + /// Fallback value for MjcPhysicsActuatorAPI::GetMjcBiasTypeAttr(), Fallback + /// value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr() + const TfToken none; + /// \brief "pgs" + /// + /// Possible value for MjcPhysicsSceneAPI::GetSolverAttr(), This token + /// represents the PGS constraint solver algorithm. + const TfToken pgs; + /// \brief "pyramidal" + /// + /// Fallback value for MjcPhysicsSceneAPI::GetConeAttr(), This token + /// represents the pyramidal contact friction cone type. + const TfToken pyramidal; + /// \brief "rk4" + /// + /// Possible value for MjcPhysicsSceneAPI::GetIntegratorAttr(), This token + /// represents the RK4 numerical integrator. + const TfToken rk4; + /// \brief "shell" + /// + /// Possible value for MjcPhysicsMeshCollisionAPI::GetInertiaAttr() + const TfToken shell; + /// \brief "sparse" + /// + /// Possible value for MjcPhysicsSceneAPI::GetJacobianAttr(), This token + /// represents the sparse constraint Jacobian and matrices computed from it. + const TfToken sparse; + /// \brief "true" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcActLimitedAttr(), Possible + /// value for MjcPhysicsActuatorAPI::GetMjcCtrlLimitedAttr(), Possible value + /// for MjcPhysicsActuatorAPI::GetMjcForceLimitedAttr() + const TfToken true_; + /// \brief "user" + /// + /// Possible value for MjcPhysicsActuatorAPI::GetMjcBiasTypeAttr(), Possible + /// value for MjcPhysicsActuatorAPI::GetMjcDynTypeAttr(), Possible value for + /// MjcPhysicsActuatorAPI::GetMjcGainTypeAttr() + const TfToken user; + /// \brief "CollisionAPI" + /// + /// Schema identifer and family for MjcPhysicsCollisionAPI + const TfToken CollisionAPI; + /// \brief "MeshCollisionAPI" + /// + /// Schema identifer and family for MjcPhysicsMeshCollisionAPI + const TfToken MeshCollisionAPI; + /// \brief "PhysicsActuatorAPI" + /// + /// Schema identifer and family for MjcPhysicsActuatorAPI + const TfToken PhysicsActuatorAPI; + /// \brief "SceneAPI" + /// + /// Schema identifer and family for MjcPhysicsSceneAPI + const TfToken SceneAPI; + /// \brief "SiteAPI" + /// + /// Schema identifer and family for MjcPhysicsSiteAPI + const TfToken SiteAPI; + /// A vector of all of the tokens listed above. + const std::vector allTokens; +}; + +/// \var MjcPhysicsTokens +/// +/// A global variable with static, efficient \link TfToken TfTokens\endlink +/// for use in all public USD API. \sa MjcPhysicsTokensType +extern MJCPHYSICS_API TfStaticData MjcPhysicsTokens; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif diff --git a/src/experimental/usd/plugins/CMakeLists.txt b/src/experimental/usd/plugins/CMakeLists.txt new file mode 100644 index 00000000..6e506493 --- /dev/null +++ b/src/experimental/usd/plugins/CMakeLists.txt @@ -0,0 +1,137 @@ +# Copyright 2025 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 +# +# https://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. + + +# Plugin target name (used for library and plugInfo.json) +set(MJCF_PLUGIN_TARGET_NAME usdMjcf) + +add_library(${MJCF_PLUGIN_TARGET_NAME} SHARED + mjcf/mjcf_file_format.cc + mjcf/mjcf_file_format.h + mjcf/mujoco_to_usd.cc + mjcf/mujoco_to_usd.h + mjcf/utils.cc + mjcf/utils.h +) + +# We need to set the visibility to default until core type symbol visibility +# is resolved in OpenUSD https://github.com/PixarAnimationStudios/OpenUSD/issues/1475 +# Otherwise we will run into issues during composition on MacOS due to std::type_info +# comparisons failing for pxr::TfTokenVector and the like that we place in SdfAbstractData. +set_target_properties(${MJCF_PLUGIN_TARGET_NAME} PROPERTIES + OUTPUT_NAME ${MJCF_PLUGIN_TARGET_NAME} + CXX_VISIBILITY_PRESET default +) + +target_include_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}" +) + +if (MUJOCO_USD_TARGET STREQUAL "USD") + find_package(pxr REQUIRED) + + # --- Link Dependencies --- + # Link against the necessary OpenUSD components + target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE + usd + ar + kind + tf + gf + vt + usdShade + usdLux + usdGeom + usdImaging + usdPhysics + mujoco + tinyxml2 + ) +elseif (MUJOCO_USD_TARGET STREQUAL "Houdini") + + if (NOT DEFINED ENV{HFS}) + message(FATAL_ERROR "Environment variable 'HFS' is not defined: $ENV{HFS}. Please run houdini_setup.") + endif() + + # In Houdini, the Houdini package we would typically use via find_package + # does not have all the USD dependencies that we need (namely UsdPhysics) + # so we need to manually link all the required libraries. + + set(HFS_ENV "$ENV{HFS}") + set(HOUDINI_LIBS "${HFS_ENV}/../Libraries") + get_filename_component(HOUDINI_LIBS "${HOUDINI_LIBS}" ABSOLUTE) # Normalize the path + target_link_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE ${HOUDINI_LIBS}) + + target_include_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE + "${HFS_ENV}/toolkit/include" + "${HFS_ENV}/toolkit/include/python3.11" + ) + + # Assume everyone using Houdini on 3.11 for now. + set(USD_MJCF_PYTHON_LIB python3.11) + set(USD_MJCF_PYTHON_LIB_NUMBER python311) + set(PYTHON_LIB "${HFS_ENV}/Frameworks/Python.framework/Versions/3.11/Python") + set(PXR_LIB_PREFIX "pxr_") + + # --- Link Dependencies --- + # Link against the necessary OpenUSD components + target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE + ${PXR_LIB_PREFIX}usd + ${PXR_LIB_PREFIX}ar + ${PXR_LIB_PREFIX}kind + ${PXR_LIB_PREFIX}tf + ${PXR_LIB_PREFIX}gf + ${PXR_LIB_PREFIX}vt + ${PXR_LIB_PREFIX}sdf + ${PXR_LIB_PREFIX}usdShade + ${PXR_LIB_PREFIX}usdLux + ${PXR_LIB_PREFIX}usdGeom + ${PXR_LIB_PREFIX}usdImaging + ${PXR_LIB_PREFIX}usdPhysics + tbb + hboost_${USD_MJCF_PYTHON_LIB_NUMBER} + ${PYTHON_LIB} + mujoco + tinyxml2 + ) +endif() + + +# --- Generate plugInfo.json --- +if(CMAKE_SHARED_LIBRARY_PREFIX) + set(LIB_PREFIX ${CMAKE_SHARED_LIBRARY_PREFIX}) # Usually "lib" on Unix +else() + set(LIB_PREFIX "") +endif() +set(PLUG_INFO_LIBRARY_PATH "${LIB_PREFIX}${MJCF_PLUGIN_TARGET_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX}") + +# --- Installation --- + +set(USD_PLUGIN_INSTALL_DIR_LIB ${CMAKE_INSTALL_LIBDIR}/usdMjcf) + +message(STATUS "Copying plugInfo.json to ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json") +configure_file( + mjcf/plugInfo.json + ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json +) + +install(FILES ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json DESTINATION ${USD_PLUGIN_INSTALL_DIR_LIB}) + +# Install shared lib and plugInfo to same location for simplicity. +install(TARGETS ${MJCF_PLUGIN_TARGET_NAME} + LIBRARY DESTINATION ${USD_PLUGIN_INSTALL_DIR_LIB} +) + +message(STATUS "USD MJCF Plugin will be installed to: ${CMAKE_INSTALL_PREFIX}/${USD_PLUGIN_INSTALL_DIR_LIB}") +message(STATUS "Make sure PXR_PLUGINPATH_NAME includes: ${CMAKE_INSTALL_PREFIX}/${USD_PLUGIN_INSTALL_DIR_LIB}") diff --git a/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc b/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc new file mode 100644 index 00000000..904c84fd --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc @@ -0,0 +1,211 @@ +// Copyright 2025 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. + +#include "mjcf/mjcf_file_format.h" + +#include +#include +#include +#include + +#include +#include "mjcf/mujoco_to_usd.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "tinyxml2.h" + +PXR_NAMESPACE_OPEN_SCOPE + +TF_DEFINE_PUBLIC_TOKENS(UsdMjcfFileFormatTokens, USD_MJCF_FILE_FORMAT_TOKENS); + +TF_REGISTRY_FUNCTION(TfType) { + SDF_DEFINE_FILE_FORMAT(UsdMjcfFileFormat, SdfFileFormat); +} + +enum ErrorCodes { XmlParsingError }; +TF_REGISTRY_FUNCTION(TfEnum) { + TF_ADD_ENUM_NAME(XmlParsingError, "Error when parsing XML."); +}; + +namespace { + +void ResolveMjcfDependencies(const std::string &xml_string, + const std::string &resolved_path); + +void AccumulateFilesRecursive(std::unordered_set &files, + tinyxml2::XMLElement *elem, + const std::string &resolved_path) { + // get filename + const char *file = elem->Attribute("file"); + + if (file != nullptr) { + auto identifier = pxr::ArGetResolver().CreateIdentifier( + std::string(file), pxr::ArResolvedPath(resolved_path)); + if (!strcasecmp(elem->Value(), "include") || + !strcasecmp(elem->Value(), "model")) { + auto include_resolved_path = pxr::ArGetResolver().Resolve(identifier); + auto asset = pxr::ArGetResolver().OpenAsset(include_resolved_path); + ResolveMjcfDependencies(asset->GetBuffer().get(), include_resolved_path); + + // Neither of these elements should have children. + return; + } + + files.insert(identifier); + } + + if (!strcasecmp(elem->Value(), "texture")) { + static const char *attributes[] = {"fileright", "fileup", "fileleft", + "filedown", "filefront", "fileback"}; + for (const auto &attribute : attributes) { + const char *file = elem->Attribute(attribute); + if (file != nullptr) { + auto identifier = pxr::ArGetResolver().CreateIdentifier( + std::string(file), pxr::ArResolvedPath(resolved_path)); + files.insert(identifier); + } + } + } + + tinyxml2::XMLElement *child = elem->FirstChildElement(); + for (; child; child = child->NextSiblingElement()) { + AccumulateFilesRecursive(files, child, resolved_path); + } +} + +void ResolveMjcfDependencies(const std::string &xml_string, + const std::string &resolved_path) { + // load XML file or parse string + tinyxml2::XMLDocument doc; + doc.Parse(xml_string.c_str()); + + // error checking + if (doc.Error()) { + TF_ERROR(XmlParsingError, "%d:\n%s\n", doc.ErrorID(), doc.ErrorStr()); + return; + } + + // get top-level element + tinyxml2::XMLElement *root = doc.RootElement(); + if (!root) { + TF_ERROR(XmlParsingError, "XML root element not found"); + return; + } + + // Accumulate file dependencies. + std::unordered_set files = {}; + AccumulateFilesRecursive(files, root, resolved_path); + + auto open_asset = [resolved_path](const std::string &identifier) { + pxr::ArGetResolver().OpenAsset(pxr::ArGetResolver().Resolve(identifier)); + }; + // Open all assets in parallel. + pxr::WorkParallelForEach(files.begin(), files.end(), open_asset); +} +} // namespace + +UsdMjcfFileFormat::UsdMjcfFileFormat() + : SdfFileFormat( + UsdMjcfFileFormatTokens->Id, UsdMjcfFileFormatTokens->Version, + UsdMjcfFileFormatTokens->Target, UsdMjcfFileFormatTokens->Id) {} + +UsdMjcfFileFormat::~UsdMjcfFileFormat() {} + +bool UsdMjcfFileFormat::CanRead(const std::string &filePath) const { + auto extension = pxr::TfGetExtension(filePath); + if (extension.empty()) { + return false; + } + + return extension == this->GetFormatId(); +} + +bool UsdMjcfFileFormat::ReadImpl(pxr::SdfLayer *layer, mjSpec *spec) const { + auto args = layer->GetFileFormatArguments(); + + bool toggleUsdPhysics = false; + const auto it = + args.find(UsdMjcfFileFormatTokens->ToggleUsdPhysicsArg.GetString()); + if (it != args.end()) { + toggleUsdPhysics = pxr::TfUnstringify(it->second); + } + + auto data = InitData(args); + + auto success = mujoco::usd::WriteSpecToData(spec, data, toggleUsdPhysics); + mj_deleteSpec(spec); + if (!success) { + return false; + } + + _SetLayerData(layer, data); + + return true; +} + +bool UsdMjcfFileFormat::ReadFromString(pxr::SdfLayer *layer, + const std::string &str) const { + std::array error; + mjSpec *spec = + mj_parseXMLString(str.c_str(), nullptr, error.data(), error.size()); + if (spec == nullptr) { + TF_WARN(XmlParsingError, "%s", error.data()); + return false; + } + + return ReadImpl(layer, spec); +} + +bool UsdMjcfFileFormat::Read(pxr::SdfLayer *layer, + const std::string &resolved_path, + bool metadata_only) const { + // Resolved all dependencies so that they are accessible when parsing + // the XML. + std::shared_ptr asset = + pxr::ArGetResolver().OpenAsset(pxr::ArResolvedPath(resolved_path)); + auto buffer = asset->GetBuffer(); + ResolveMjcfDependencies(buffer.get(), resolved_path); + + // Parse to USD. + std::array error; + mjSpec *spec = + mj_parseXML(resolved_path.c_str(), nullptr, error.data(), error.size()); + if (spec == nullptr) { + TF_WARN(XmlParsingError, "%s", error.data()); + return false; + } + return ReadImpl(layer, spec); +} + +bool UsdMjcfFileFormat::WriteToString(const SdfLayer &layer, std::string *str, + const std::string &comment) const { + return SdfFileFormat::FindById(pxr::UsdUsdaFileFormatTokens->Id) + ->WriteToString(layer, str, comment); +} + +PXR_NAMESPACE_CLOSE_SCOPE diff --git a/src/experimental/usd/plugins/mjcf/mjcf_file_format.h b/src/experimental/usd/plugins/mjcf/mjcf_file_format.h new file mode 100644 index 00000000..5bc72337 --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/mjcf_file_format.h @@ -0,0 +1,90 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MJCF_FILE_FORMAT_H_ +#define MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MJCF_FILE_FORMAT_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include + +PXR_NAMESPACE_OPEN_SCOPE + +// clang-format off +// The Id should realistically be mjcf, but the id and extension need to match. +// So near term it just assumes the only .xml file we would import is MJCF. +#define USD_MJCF_FILE_FORMAT_TOKENS \ + ((Id, "xml")) \ + ((Version, "1.0")) \ + ((Target, "usd")) \ + ((ToggleUsdPhysicsArg, "usdMjcfToggleUsdPhysics")) +// clang-format on + +TF_DECLARE_PUBLIC_TOKENS(UsdMjcfFileFormatTokens, USD_MJCF_FILE_FORMAT_TOKENS); + +TF_DECLARE_WEAK_AND_REF_PTRS(UsdMjcfFileFormat); + +class UsdMjcfFileFormat : public SdfFileFormat { + public: + using SdfFileFormat::FileFormatArguments; + + // Returns true if 'file' can be read by this format plugin. + USD_API + bool CanRead(const std::string &file) const override; + + // Reads scene description from the asset specified by resolved_path into + // 'layer'. + // + // metadataOnly is a flag that asks for only the layer metadata to be read in, + // which can be much faster if that is all that is required but currently we + // ignore it. + // + // Returns true if the asset is successfully read into layer, false otherwise. + USD_API + bool Read(pxr::SdfLayer *layer, const std::string &resolved_path, + bool metadata_only) const override; + + // Reads data in the string 'str' into 'layer'. + // + // If the file is successfully read, this method returns true. Otherwise, + // false is returned and errors are posted. + USD_API + bool ReadFromString(SdfLayer *layer, const std::string &str) const override; + + // Writes the contents in 'layer' to 'str'. This just forwards to the usda + // implementation. + USD_API + bool WriteToString(const SdfLayer &layer, std::string *str, + const std::string &comment) const override; + + protected: + SDF_FILE_FORMAT_FACTORY_ACCESS; + + UsdMjcfFileFormat(); + virtual ~UsdMjcfFileFormat(); + + private: + // Function delegated to by Read and ReadFromString. + bool ReadImpl(SdfLayer *layer, mjSpec *spec) const; +}; + +PXR_NAMESPACE_CLOSE_SCOPE + +#endif // MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MJCF_FILE_FORMAT_H_ diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc new file mode 100644 index 00000000..aa5f7232 --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc @@ -0,0 +1,1586 @@ +// Copyright 2025 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. + +#include "mjcf/mujoco_to_usd.h" + +#include +#include +#include +#include +#include +#include + +#include +#include "third_party/mujoco/src/experimental/usd/mjcPhysics/tokens.h" +#include "mjcf/utils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// The ID of the World in mjModel and mjData. +static constexpr int kWorldIndex = 0; + +// Using to satisfy TF_DEFINE_PRIVATE_TOKENS macro below and avoid operating in +// PXR_NS. +using pxr::TfToken; +template +using TfStaticData = pxr::TfStaticData; + +// clang-format off +TF_DEFINE_PRIVATE_TOKENS(kTokens, + ((body, "Body")) + ((body_name, "mujoco:body_name")) + ((geom, "Geom")) + ((light, "Light")) + ((meshScope, "MeshSources")) + ((materialsScope, "Materials")) + ((surface, "PreviewSurface")) + ((world, "World")) + ((xformOpTransform, "xformOp:transform")) + ((xformOpScale, "xformOp:scale")) + (st) + ((primvarsSt, "primvars:st")) + ((outputsSt, "outputs:st")) + ((inputsSt, "inputs:st")) + ((inputsVarname, "inputs:varname")) + ((inputsFile, "inputs:file")) + ((inputsWrapS, "inputs:wrapS")) + ((inputsWrapT, "inputs:wrapT")) + ((inputsDiffuseColor, "inputs:diffuseColor")) + ((outputsRgb, "outputs:rgb")) + ((inputsMetallic, "inputs:metallic")) + (repeat) + ((sourceMesh, pxr::UsdGeomTokens->Mesh)) + ); + +// Using to satisfy TF_REGISTRY_FUNCTION macro below and avoid operating in PXR_NS. +using pxr::TfEnum; +using pxr::Tf_RegistryStaticInit; +using pxr::Tf_RegistryInit; +using pxr::TfEnum; +template +using Arch_PerLibInit = pxr::Arch_PerLibInit; +#if defined(ARCH_OS_DARWIN) +using Arch_ConstructorEntry = pxr::Arch_ConstructorEntry; +#endif +enum ErrorCodes { UnsupportedActuatorTypeError, UnsupportedGeomTypeError, MujocoCompilationError }; + +TF_REGISTRY_FUNCTION(pxr::TfEnum) { + TF_ADD_ENUM_NAME(UnsupportedGeomTypeError, "UsdGeom type is unsupported.") + TF_ADD_ENUM_NAME(MujocoCompilationError, "Mujoco spec failed to compile.") +} + +// Usings to satisfy TF_ERROR macro. +using pxr::TfCallContext; +using pxr::Tf_PostErrorHelper; +// clang-format on + +using pxr::MjcPhysicsTokens; + +using mujoco::usd::AddAttributeConnection; +using mujoco::usd::AddPrimInherit; +using mujoco::usd::AddPrimReference; +using mujoco::usd::ApplyApiSchema; +using mujoco::usd::CreateAttributeSpec; +using mujoco::usd::CreateClassSpec; +using mujoco::usd::CreatePrimSpec; +using mujoco::usd::CreateRelationshipSpec; +using mujoco::usd::SetAttributeDefault; +using mujoco::usd::SetAttributeMetadata; +using mujoco::usd::SetLayerMetadata; +using mujoco::usd::SetPrimKind; +using mujoco::usd::SetPrimMetadata; +using mujoco::usd::SetPrimPurpose; + +pxr::GfMatrix4d MujocoPosQuatToTransform(double *pos, double *quat) { + pxr::GfQuatd quaternion = pxr::GfQuatd::GetIdentity(); + quaternion.SetReal(quat[0]); + quaternion.SetImaginary(quat[1], quat[2], quat[3]); + pxr::GfRotation rotation(quaternion); + + pxr::GfVec3d translation(0.0, 0.0, 0.0); + translation.Set(pos[0], pos[1], pos[2]); + + pxr::GfMatrix4d transform; + transform.SetTransform(rotation, translation); + return transform; +} + +} // namespace + +class ModelWriter { + public: + ModelWriter(mjSpec *spec, mjModel *model, pxr::SdfAbstractDataRefPtr &data) + : spec_(spec), model_(model), data_(data), class_path_("/Bad_Path") { + body_paths_ = std::vector(model->nbody); + site_paths_ = std::vector(model->nsite); + } + ~ModelWriter() { mj_deleteModel(model_); } + + void Write(bool write_physics) { + // Create top level class holder. + class_path_ = CreateClassSpec(data_, pxr::SdfPath::AbsoluteRootPath(), + pxr::TfToken("__class__")); + + // Create the world body. + body_paths_[kWorldIndex] = WriteWorldBody(kWorldIndex); + + // Set working parameters. + write_physics_ = write_physics; + + SetLayerMetadata(data_, pxr::SdfFieldKeys->Documentation, + "Generated by mujoco model writer."); + // Mujoco is Z up by default. + SetLayerMetadata(data_, pxr::UsdGeomTokens->upAxis, pxr::UsdGeomTokens->z); + // Mujoco is authored in meters by default. + SetLayerMetadata(data_, pxr::UsdGeomTokens->metersPerUnit, + pxr::UsdGeomLinearUnits::meters); + + // Set the world body to be the default prim for referencing/payloads. + SetLayerMetadata(data_, pxr::SdfFieldKeys->DefaultPrim, + body_paths_[kWorldIndex].GetNameToken()); + + WritePhysicsScene(); + + // Author mesh scope + mesh prims to be referenced. + WriteMeshes(); + WriteMaterials(); + WriteBodies(); + if (write_physics_) { + WriteActuators(); + } + } + + private: + mjSpec *spec_; + mjModel *model_; + + // This is a handle to the Sdf data to be written into the generated USD + // layer. + pxr::SdfAbstractDataRefPtr &data_; + // Path to top level class spec that all classes should be children of. + pxr::SdfPath class_path_; + // Mapping from Mujoco body id to SdfPath. + std::vector body_paths_; + // Mapping from Mujoco site id to SdfPath. + std::vector site_paths_; + // Mapping from mesh names to Mesh prim path. + std::unordered_map mesh_paths_; + // Whether to write physics data. + bool write_physics_ = false; + + // Given a name index and a parent prim path this returns a + // token such that appending it to the parent prim path does not + // identify an existing prim spec. + // + // This is necessary since mujoco does not require names for elements + // so we must differentiate between elements of the same type. + // + // For example: + // + // + // + // + // + // + // + // We expect that the occurrence of this happens little enough that linear + // searching is plenty efficient. + pxr::TfToken GetAvailablePrimName(const std::string base_name, + const pxr::TfToken fallback_name, + const pxr::SdfPath &parent_path) { + const auto valid_base_name = pxr::TfMakeValidIdentifier( + base_name.empty() ? fallback_name : base_name); + std::string name = valid_base_name; + pxr::SdfPath test_path = parent_path.AppendChild(pxr::TfToken(name)); + int count = 1; + while (data_->HasSpec(test_path) && + data_->GetSpecType(test_path) == pxr::SdfSpecType::SdfSpecTypePrim) { + name = pxr::TfStringPrintf("%s_%d", valid_base_name.c_str(), count++); + test_path = parent_path.AppendChild(pxr::TfToken(name)); + } + return pxr::TfToken(name); + } + + // This function, conversely to GetAvailablePrimName will not handle + // collisions. This is useful when looking up a prim path that might exist or + // a path that you know must be unique. + pxr::TfToken GetValidPrimName(const std::string name) { + return pxr::TfToken(pxr::TfMakeValidIdentifier(name)); + } + + struct BodyPathComponents { + pxr::SdfPath parent_path; + pxr::TfToken body_name; + }; + + void WriteScaleXformOp(const pxr::SdfPath &prim_path, + const pxr::GfVec3f &scale) { + pxr::SdfPath scale_attr_path = + CreateAttributeSpec(data_, prim_path, kTokens->xformOpScale, + pxr::SdfValueTypeNames->Float3); + SetAttributeDefault(data_, scale_attr_path, scale); + } + + void WriteTransformXformOp(const pxr::SdfPath &prim_path, + const pxr::GfMatrix4d &transform) { + pxr::SdfPath transform_op_path = + CreateAttributeSpec(data_, prim_path, kTokens->xformOpTransform, + pxr::SdfValueTypeNames->Matrix4d); + SetAttributeDefault(data_, transform_op_path, transform); + } + + void WriteXformOpOrder(const pxr::SdfPath &prim_path, + const pxr::VtArray &order) { + pxr::SdfPath xform_op_order_path = + CreateAttributeSpec(data_, prim_path, pxr::UsdGeomTokens->xformOpOrder, + pxr::SdfValueTypeNames->TokenArray); + SetAttributeDefault(data_, xform_op_order_path, order); + } + + template + void WriteUniformAttribute(const pxr::SdfPath &prim_path, + const pxr::SdfValueTypeName &value_type_name, + const pxr::TfToken &token, const T &value) { + pxr::SdfPath attr_path = CreateAttributeSpec( + data_, prim_path, token, value_type_name, pxr::SdfVariabilityUniform); + SetAttributeDefault(data_, attr_path, value); + } + + void PrependToXformOpOrder(const pxr::SdfPath &prim_path, + const pxr::VtArray &order) { + auto xform_op_order_path = + prim_path.AppendProperty(pxr::UsdGeomTokens->xformOpOrder); + if (!data_->HasSpec(xform_op_order_path)) { + WriteXformOpOrder(prim_path, order); + return; + } + + auto existing_order = + data_->Get(xform_op_order_path, pxr::SdfFieldKeys->Default) + .UncheckedGet>(); + + pxr::VtArray new_order(order.size() + existing_order.size()); + std::copy(order.begin(), order.end(), new_order.begin()); + std::copy(existing_order.begin(), existing_order.end(), + new_order.begin() + order.size()); + + SetAttributeDefault(data_, xform_op_order_path, new_order); + } + + void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path) { + auto name = GetAvailablePrimName(*mesh->name, pxr::UsdGeomTokens->Mesh, + parent_path); + pxr::SdfPath subcomponent_path = + CreatePrimSpec(data_, parent_path, name, pxr::UsdGeomTokens->Xform); + pxr::SdfPath mesh_path = + CreatePrimSpec(data_, subcomponent_path, kTokens->sourceMesh, + pxr::UsdGeomTokens->Mesh); + mesh_paths_[*mesh->name] = subcomponent_path; + + if (write_physics_) { + ApplyApiSchema(data_, mesh_path, MjcPhysicsTokens->MeshCollisionAPI); + + pxr::TfToken inertia = MjcPhysicsTokens->legacy; + if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_EXACT) { + inertia = MjcPhysicsTokens->exact; + } else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_CONVEX) { + inertia = MjcPhysicsTokens->convex; + } else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_SHELL) { + inertia = MjcPhysicsTokens->shell; + } + + WriteUniformAttribute(mesh_path, pxr::SdfValueTypeNames->Token, + MjcPhysicsTokens->mjcInertia, inertia); + } + + // NOTE: The geometry data taken from the spec is the post-compilation + // data after it has been mjCMesh::Compile'd. So don't be surprised if + // things like user defined vertices have moved due to re-centering to + // CoM and other modifications (see mjCMesh::Process for other xforms). + int mesh_id = mjs_getId(mesh->element); + int vert_start_offset = model_->mesh_vertadr[mesh_id] * 3; + int nvert = model_->mesh_vertnum[mesh_id]; + pxr::VtArray points; + points.reserve(nvert); + for (int i = vert_start_offset; i < vert_start_offset + nvert * 3; i += 3) { + points.emplace_back(&model_->mesh_vert[i]); + } + + pxr::SdfPath points_attr_path = + CreateAttributeSpec(data_, mesh_path, pxr::UsdGeomTokens->points, + pxr::SdfValueTypeNames->Vector3fArray); + SetAttributeDefault(data_, points_attr_path, points); + + // NOTE: nface is never 0. + int nface = model_->mesh_facenum[mesh_id]; + pxr::VtArray faces; + faces.reserve(nface * 3); + int face_start_offset = model_->mesh_faceadr[mesh_id] * 3; + for (int i = face_start_offset; i < face_start_offset + nface * 3; i += 3) { + faces.push_back(model_->mesh_face[i]); + faces.push_back(model_->mesh_face[i + 1]); + faces.push_back(model_->mesh_face[i + 2]); + } + pxr::SdfPath face_vertex_idx_attr_path = CreateAttributeSpec( + data_, mesh_path, pxr::UsdGeomTokens->faceVertexIndices, + pxr::SdfValueTypeNames->IntArray); + SetAttributeDefault(data_, face_vertex_idx_attr_path, faces); + + pxr::VtArray vertex_counts; + for (int i = 0; i < nface; ++i) { + // Mujoco is always triangles. + vertex_counts.push_back(3); + } + pxr::SdfPath face_vertex_counts_attr_path = CreateAttributeSpec( + data_, mesh_path, pxr::UsdGeomTokens->faceVertexCounts, + pxr::SdfValueTypeNames->IntArray); + SetAttributeDefault(data_, face_vertex_counts_attr_path, vertex_counts); + + if (model_->mesh_normalnum[mesh_id]) { + // We have to convert from Mujoco's indexed normals to USD's faceVarying + // normals. + pxr::VtArray normals; + normals.reserve(nface * 3); + int normal_start_adr = model_->mesh_normaladr[mesh_id]; + int face_start_offset = model_->mesh_faceadr[mesh_id] * 3; + for (int i = face_start_offset; i < face_start_offset + nface * 3; ++i) { + int normal_adr = normal_start_adr + model_->mesh_facenormal[i]; + normals.emplace_back(&model_->mesh_normal[normal_adr * 3]); + } + pxr::SdfPath normals_attr_path = + CreateAttributeSpec(data_, mesh_path, pxr::UsdGeomTokens->normals, + pxr::SdfValueTypeNames->Vector3fArray); + SetAttributeDefault(data_, normals_attr_path, normals); + SetAttributeMetadata(data_, normals_attr_path, + pxr::UsdGeomTokens->interpolation, + pxr::UsdGeomTokens->faceVarying); + } + + if (model_->mesh_texcoordnum[mesh_id]) { + // We have to convert from Mujoco's indexed texcoords to USD's faceVarying + // texcoords. + pxr::VtArray texcoords; + texcoords.reserve(nface * 3); + int texcoord_start_adr = model_->mesh_texcoordadr[mesh_id]; + int face_start_offset = model_->mesh_faceadr[mesh_id] * 3; + for (int i = face_start_offset; i < face_start_offset + nface * 3; ++i) { + int texcoord_adr = texcoord_start_adr + model_->mesh_facetexcoord[i]; + // Invert the V coordinate, Mujoco assumes OpenGL 0,0 is top left. + // But USD UVs use image bottom left 0,0 convention. + pxr::GfVec2f uv(&model_->mesh_texcoord[texcoord_adr * 2]); + uv[1] = 1.0f - uv[1]; + texcoords.push_back(uv); + } + + pxr::SdfPath texcoords_attr_path = + CreateAttributeSpec(data_, mesh_path, kTokens->primvarsSt, + pxr::SdfValueTypeNames->TexCoord2fArray); + SetAttributeDefault(data_, texcoords_attr_path, texcoords); + SetAttributeMetadata(data_, texcoords_attr_path, + pxr::UsdGeomTokens->interpolation, + pxr::UsdGeomTokens->faceVarying); + } + + // Default subdivision scheme is catmull clark so explicitly set it + // to none here. + pxr::SdfPath subdivision_scheme_path = CreateAttributeSpec( + data_, mesh_path, pxr::UsdGeomTokens->subdivisionScheme, + pxr::SdfValueTypeNames->Token); + SetAttributeDefault(data_, subdivision_scheme_path, + pxr::UsdGeomTokens->none); + } + + void WritePhysicsScene() { + pxr::SdfPath physics_scene_path = CreatePrimSpec( + data_, body_paths_[kWorldIndex], pxr::UsdPhysicsTokens->PhysicsScene, + pxr::UsdPhysicsTokens->PhysicsScene); + + ApplyApiSchema(data_, physics_scene_path, MjcPhysicsTokens->SceneAPI); + + const std::vector> + option_double_attributes = { + {MjcPhysicsTokens->mjcOptionTimestep, spec_->option.timestep}, + {MjcPhysicsTokens->mjcOptionTolerance, spec_->option.tolerance}, + {MjcPhysicsTokens->mjcOptionLs_tolerance, + spec_->option.ls_tolerance}, + {MjcPhysicsTokens->mjcOptionNoslip_tolerance, + spec_->option.noslip_tolerance}, + {MjcPhysicsTokens->mjcOptionCcd_tolerance, + spec_->option.ccd_tolerance}, + {MjcPhysicsTokens->mjcOptionApirate, spec_->option.apirate}, + {MjcPhysicsTokens->mjcOptionImpratio, spec_->option.impratio}, + {MjcPhysicsTokens->mjcOptionDensity, spec_->option.density}, + {MjcPhysicsTokens->mjcOptionViscosity, spec_->option.viscosity}, + {MjcPhysicsTokens->mjcOptionO_margin, spec_->option.o_margin}, + }; + for (const auto &[token, value] : option_double_attributes) { + WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Double, + token, value); + } + + const std::vector> option_int_attributes = { + {MjcPhysicsTokens->mjcOptionIterations, spec_->option.iterations}, + {MjcPhysicsTokens->mjcOptionLs_iterations, spec_->option.ls_iterations}, + {MjcPhysicsTokens->mjcOptionNoslip_iterations, + spec_->option.noslip_iterations}, + {MjcPhysicsTokens->mjcOptionCcd_iterations, + spec_->option.ccd_iterations}, + {MjcPhysicsTokens->mjcOptionSdf_iterations, + spec_->option.sdf_iterations}, + {MjcPhysicsTokens->mjcOptionSdf_initpoints, + spec_->option.sdf_initpoints}, + }; + for (const auto &[token, value] : option_int_attributes) { + WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Int, + token, value); + } + + pxr::SdfPath cone_attr = CreateAttributeSpec( + data_, physics_scene_path, MjcPhysicsTokens->mjcOptionCone, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + + switch (spec_->option.cone) { + case mjCONE_PYRAMIDAL: + SetAttributeDefault(data_, cone_attr, MjcPhysicsTokens->pyramidal); + break; + case mjCONE_ELLIPTIC: + SetAttributeDefault(data_, cone_attr, MjcPhysicsTokens->elliptic); + break; + default: + break; + } + + pxr::SdfPath jacobian_attr = CreateAttributeSpec( + data_, physics_scene_path, MjcPhysicsTokens->mjcOptionJacobian, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + + switch (spec_->option.jacobian) { + case mjJAC_AUTO: + SetAttributeDefault(data_, jacobian_attr, MjcPhysicsTokens->auto_); + break; + case mjJAC_DENSE: + SetAttributeDefault(data_, jacobian_attr, MjcPhysicsTokens->dense); + break; + case mjJAC_SPARSE: + SetAttributeDefault(data_, jacobian_attr, MjcPhysicsTokens->sparse); + break; + default: + break; + } + + pxr::SdfPath solver_attr = CreateAttributeSpec( + data_, physics_scene_path, MjcPhysicsTokens->mjcOptionSolver, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + switch (spec_->option.solver) { + case mjSOL_NEWTON: + SetAttributeDefault(data_, solver_attr, MjcPhysicsTokens->newton); + break; + case mjSOL_PGS: + SetAttributeDefault(data_, solver_attr, MjcPhysicsTokens->pgs); + break; + case mjSOL_CG: + SetAttributeDefault(data_, solver_attr, MjcPhysicsTokens->cg); + break; + default: + break; + } + + pxr::GfVec3d wind(spec_->option.wind[0], spec_->option.wind[1], + spec_->option.wind[2]); + WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Double3, + MjcPhysicsTokens->mjcOptionWind, wind); + + pxr::GfVec3d magnetic(spec_->option.magnetic[0], spec_->option.magnetic[1], + spec_->option.magnetic[2]); + WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Double3, + MjcPhysicsTokens->mjcOptionMagnetic, magnetic); + + pxr::VtArray o_solref(spec_->option.o_solref, + spec_->option.o_solref + 2); + WriteUniformAttribute(physics_scene_path, + pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcOptionO_solref, o_solref); + + pxr::VtArray o_solimp(spec_->option.o_solimp, + spec_->option.o_solimp + 5); + WriteUniformAttribute(physics_scene_path, + pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcOptionO_solimp, o_solimp); + + pxr::VtArray o_friction(spec_->option.o_friction, + spec_->option.o_friction + 5); + WriteUniformAttribute(physics_scene_path, + pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcOptionO_friction, o_friction); + + pxr::SdfPath integrator_attr = CreateAttributeSpec( + data_, physics_scene_path, MjcPhysicsTokens->mjcOptionIntegrator, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + switch (spec_->option.integrator) { + case mjINT_EULER: + SetAttributeDefault(data_, integrator_attr, MjcPhysicsTokens->euler); + break; + case mjINT_RK4: + SetAttributeDefault(data_, integrator_attr, MjcPhysicsTokens->rk4); + break; + default: + break; + } + + auto create_flag_attr = [&](pxr::TfToken token, int flag, bool enable) { + int flags = + enable ? spec_->option.enableflags : spec_->option.disableflags; + bool value = enable ? (flags & flag) : !(flags & flag); + WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Bool, + token, value); + }; + + const std::vector> enable_flags = { + {MjcPhysicsTokens->mjcFlagMulticcd, mjENBL_MULTICCD}, + {MjcPhysicsTokens->mjcFlagIsland, mjENBL_ISLAND}, + {MjcPhysicsTokens->mjcFlagFwdinv, mjENBL_FWDINV}, + {MjcPhysicsTokens->mjcFlagEnergy, mjENBL_ENERGY}, + {MjcPhysicsTokens->mjcFlagOverride, mjENBL_OVERRIDE}, + {MjcPhysicsTokens->mjcFlagInvdiscrete, mjENBL_INVDISCRETE}}; + for (const auto &[token, flag] : enable_flags) { + create_flag_attr(token, flag, true); + } + + const std::vector> disable_flags = { + {MjcPhysicsTokens->mjcFlagConstraint, mjDSBL_CONSTRAINT}, + {MjcPhysicsTokens->mjcFlagEquality, mjDSBL_EQUALITY}, + {MjcPhysicsTokens->mjcFlagFrictionloss, mjDSBL_FRICTIONLOSS}, + {MjcPhysicsTokens->mjcFlagLimit, mjDSBL_LIMIT}, + {MjcPhysicsTokens->mjcFlagContact, mjDSBL_CONTACT}, + {MjcPhysicsTokens->mjcFlagPassive, mjDSBL_PASSIVE}, + {MjcPhysicsTokens->mjcFlagGravity, mjDSBL_GRAVITY}, + {MjcPhysicsTokens->mjcFlagClampctrl, mjDSBL_CLAMPCTRL}, + {MjcPhysicsTokens->mjcFlagWarmstart, mjDSBL_WARMSTART}, + {MjcPhysicsTokens->mjcFlagFilterparent, mjDSBL_FILTERPARENT}, + {MjcPhysicsTokens->mjcFlagActuation, mjDSBL_ACTUATION}, + {MjcPhysicsTokens->mjcFlagRefsafe, mjDSBL_REFSAFE}, + {MjcPhysicsTokens->mjcFlagSensor, mjDSBL_SENSOR}, + {MjcPhysicsTokens->mjcFlagMidphase, mjDSBL_MIDPHASE}, + {MjcPhysicsTokens->mjcFlagEulerdamp, mjDSBL_EULERDAMP}, + {MjcPhysicsTokens->mjcFlagAutoreset, mjDSBL_AUTORESET}, + {MjcPhysicsTokens->mjcFlagNativeccd, mjDSBL_NATIVECCD}}; + for (const auto &[token, flag] : disable_flags) { + create_flag_attr(token, flag, false); + } + } + + void WriteMeshes() { + // Create a scope for the meshes to keep things organized + pxr::SdfPath scope_path = + CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->meshScope, + pxr::UsdGeomTokens->Scope); + + // Make the mesh scope invisible since they will be referenced by the bits + // that should be visible. + SetPrimMetadata(data_, scope_path, pxr::SdfFieldKeys->Active, false); + + mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH)); + while (mesh) { + WriteMesh(mesh, scope_path); + mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element)); + } + } + + pxr::SdfPath AddTextureShader(const pxr::SdfPath &material_path, + const char *texture_file) { + // Shader "uvmap" + pxr::SdfPath uvmap_shader_path = + CreatePrimSpec(data_, material_path, pxr::TfToken("uvmap"), + pxr::UsdShadeTokens->Shader); + + pxr::SdfPath uvmap_info_id_attr = CreateAttributeSpec( + data_, uvmap_shader_path, pxr::UsdShadeTokens->infoId, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + SetAttributeDefault(data_, uvmap_info_id_attr, + pxr::UsdImagingTokens->UsdPrimvarReader_float2); + + pxr::SdfPath uvmap_varname_attr = + CreateAttributeSpec(data_, uvmap_shader_path, kTokens->inputsVarname, + pxr::SdfValueTypeNames->Token); + SetAttributeDefault(data_, uvmap_varname_attr, kTokens->st); + + pxr::SdfPath uvmap_st_output_attr = + CreateAttributeSpec(data_, uvmap_shader_path, kTokens->outputsSt, + pxr::SdfValueTypeNames->Float2); + + // Shader "texture" + pxr::SdfPath texture_shader_path = + CreatePrimSpec(data_, material_path, pxr::TfToken("texture"), + pxr::UsdShadeTokens->Shader); + pxr::SdfPath texture_info_id_attr = CreateAttributeSpec( + data_, texture_shader_path, pxr::UsdShadeTokens->infoId, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + SetAttributeDefault(data_, texture_info_id_attr, + pxr::UsdImagingTokens->UsdUVTexture); + + pxr::SdfPath texture_file_attr = + CreateAttributeSpec(data_, texture_shader_path, kTokens->inputsFile, + pxr::SdfValueTypeNames->Asset); + SetAttributeDefault(data_, texture_file_attr, + pxr::SdfAssetPath(texture_file)); + + pxr::SdfPath texture_st_input_attr = + CreateAttributeSpec(data_, texture_shader_path, kTokens->inputsSt, + pxr::SdfValueTypeNames->Float2); + AddAttributeConnection(data_, texture_st_input_attr, uvmap_st_output_attr); + + pxr::SdfPath texture_wrap_s_attr = + CreateAttributeSpec(data_, texture_shader_path, kTokens->inputsWrapS, + pxr::SdfValueTypeNames->Token); + SetAttributeDefault(data_, texture_wrap_s_attr, kTokens->repeat); + + pxr::SdfPath texture_wrap_t_attr = + CreateAttributeSpec(data_, texture_shader_path, kTokens->inputsWrapT, + pxr::SdfValueTypeNames->Token); + SetAttributeDefault(data_, texture_wrap_t_attr, kTokens->repeat); + + pxr::SdfPath texture_rgb_output_attr = + CreateAttributeSpec(data_, texture_shader_path, kTokens->outputsRgb, + pxr::SdfValueTypeNames->Float3); + + return texture_rgb_output_attr; + } + + void WriteMaterial(mjsMaterial *material, const pxr::SdfPath &parent_path) { + auto name = GetAvailablePrimName( + *material->name, pxr::UsdShadeTokens->Material, parent_path); + pxr::SdfPath material_path = + CreatePrimSpec(data_, parent_path, name, pxr::UsdShadeTokens->Material); + + // Shader "PreviewSurface" + pxr::SdfPath preview_surface_shader_path = CreatePrimSpec( + data_, material_path, kTokens->surface, pxr::UsdShadeTokens->Shader); + + pxr::SdfPath info_id_attr = CreateAttributeSpec( + data_, preview_surface_shader_path, pxr::UsdShadeTokens->infoId, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + SetAttributeDefault(data_, info_id_attr, + pxr::UsdImagingTokens->UsdPreviewSurface); + + pxr::SdfPath surface_output_attr = CreateAttributeSpec( + data_, preview_surface_shader_path, pxr::UsdShadeTokens->outputsSurface, + pxr::SdfValueTypeNames->Token); + + pxr::SdfPath displacement_output_attr = + CreateAttributeSpec(data_, preview_surface_shader_path, + pxr::UsdShadeTokens->outputsDisplacement, + pxr::SdfValueTypeNames->Token); + + pxr::SdfPath diffuse_color_attr = CreateAttributeSpec( + data_, preview_surface_shader_path, kTokens->inputsDiffuseColor, + pxr::SdfValueTypeNames->Color3f); + + // Find the main texture if specified. + std::string main_texture_name = (*material->textures)[mjTEXROLE_RGB]; + mjsTexture *main_texture = mjs_asTexture( + mjs_findElement(spec_, mjOBJ_TEXTURE, main_texture_name.c_str())); + if (main_texture) { + // Create the texture shader and connect it to the diffuse color + // attribute. + pxr::SdfPath texture_rgb_output_attr = + AddTextureShader(material_path, main_texture->file->c_str()); + AddAttributeConnection(data_, diffuse_color_attr, + texture_rgb_output_attr); + } else { + // If no texture is specified, use the rgba diffuse color. + SetAttributeDefault(data_, diffuse_color_attr, + pxr::GfVec3f(material->rgba[0], material->rgba[1], + material->rgba[2])); + } + + pxr::SdfPath metallic_attr = CreateAttributeSpec( + data_, preview_surface_shader_path, kTokens->inputsMetallic, + pxr::SdfValueTypeNames->Float); + SetAttributeDefault(data_, metallic_attr, material->metallic); + + pxr::SdfPath material_surface_output_attr = CreateAttributeSpec( + data_, material_path, pxr::UsdShadeTokens->outputsSurface, + pxr::SdfValueTypeNames->Token); + + AddAttributeConnection(data_, material_surface_output_attr, + surface_output_attr); + + pxr::SdfPath material_displacement_output_attr = CreateAttributeSpec( + data_, material_path, pxr::UsdShadeTokens->outputsDisplacement, + pxr::SdfValueTypeNames->Token); + + AddAttributeConnection(data_, material_displacement_output_attr, + displacement_output_attr); + } + + void WriteMaterials() { + // Create a scope for the meshes to keep things organized + pxr::SdfPath scope_path = + CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->materialsScope, + pxr::UsdGeomTokens->Scope); + + mjsMaterial *material = + mjs_asMaterial(mjs_firstElement(spec_, mjOBJ_MATERIAL)); + while (material) { + WriteMaterial(material, scope_path); + material = mjs_asMaterial(mjs_nextElement(spec_, material->element)); + } + } + + void WriteActuator(mjsActuator *actuator) { + pxr::SdfPath transmission_path; + if (actuator->trntype == mjtTrn::mjTRN_BODY) { + int body_id = mj_name2id(model_, mjOBJ_BODY, actuator->target->c_str()); + transmission_path = body_paths_[body_id]; + } else if (actuator->trntype == mjtTrn::mjTRN_SITE || + actuator->trntype == mjtTrn::mjTRN_SLIDERCRANK) { + int site_id = mj_name2id(model_, mjOBJ_SITE, actuator->target->c_str()); + transmission_path = site_paths_[site_id]; + } else { + TF_WARN(UnsupportedActuatorTypeError, + "Unsupported actuator type for actuator %d", + mjs_getId(actuator->element)); + return; + } + + ApplyApiSchema(data_, transmission_path, + MjcPhysicsTokens->PhysicsActuatorAPI); + + if (!actuator->refsite->empty()) { + int refsite_id = mj_name2id(model_, mjOBJ_SITE, actuator->refsite->c_str()); + pxr::SdfPath refsite_path = site_paths_[refsite_id]; + CreateRelationshipSpec(data_, transmission_path, + MjcPhysicsTokens->mjcRefSite, + refsite_path, pxr::SdfVariabilityUniform); + } + + if (!actuator->slidersite->empty()) { + int slidersite_id = mj_name2id(model_, mjOBJ_SITE, actuator->slidersite->c_str()); + pxr::SdfPath slidersite_path = site_paths_[slidersite_id]; + CreateRelationshipSpec(data_, transmission_path, + MjcPhysicsTokens->mjcSliderSite, + slidersite_path, pxr::SdfVariabilityUniform); + } + + + const std::vector> limited_attributes = { + {MjcPhysicsTokens->mjcCtrlLimited, actuator->ctrllimited}, + {MjcPhysicsTokens->mjcForceLimited, actuator->forcelimited}, + {MjcPhysicsTokens->mjcActLimited, actuator->actlimited}, + }; + for (const auto &[token, value] : limited_attributes) { + pxr::TfToken limited_token = pxr::MjcPhysicsTokens->auto_; + if (value == mjLIMITED_TRUE) { + limited_token = pxr::MjcPhysicsTokens->true_; + } else if (value == mjLIMITED_FALSE) { + limited_token = pxr::MjcPhysicsTokens->false_; + } + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Token, + token, limited_token); + } + + const std::vector> + actuator_double_attributes = { + {MjcPhysicsTokens->mjcCtrlRangeMin, actuator->ctrlrange[0]}, + {MjcPhysicsTokens->mjcCtrlRangeMax, actuator->ctrlrange[1]}, + {MjcPhysicsTokens->mjcForceRangeMin, actuator->forcerange[0]}, + {MjcPhysicsTokens->mjcForceRangeMax, actuator->forcerange[1]}, + {MjcPhysicsTokens->mjcActRangeMin, actuator->actrange[0]}, + {MjcPhysicsTokens->mjcActRangeMax, actuator->actrange[1]}, + {MjcPhysicsTokens->mjcLengthRangeMin, actuator->lengthrange[0]}, + {MjcPhysicsTokens->mjcLengthRangeMax, actuator->lengthrange[1]}, + {MjcPhysicsTokens->mjcCrankLength, actuator->cranklength}, + }; + for (const auto &[token, value] : actuator_double_attributes) { + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Double, + token, value); + } + + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Int, + MjcPhysicsTokens->mjcActDim, actuator->actdim); + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Bool, + MjcPhysicsTokens->mjcActEarly, + (bool)actuator->actearly); + + WriteUniformAttribute( + transmission_path, pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcGear, + pxr::VtDoubleArray(actuator->gear, actuator->gear + 6)); + + pxr::TfToken dyn_type; + if (actuator->dyntype == mjtDyn::mjDYN_NONE) { + dyn_type = MjcPhysicsTokens->none; + } else if (actuator->dyntype == mjtDyn::mjDYN_INTEGRATOR) { + dyn_type = MjcPhysicsTokens->integrator; + } else if (actuator->dyntype == mjtDyn::mjDYN_FILTER) { + dyn_type = MjcPhysicsTokens->filter; + } else if (actuator->dyntype == mjtDyn::mjDYN_FILTEREXACT) { + dyn_type = MjcPhysicsTokens->filterexact; + } else if (actuator->dyntype == mjtDyn::mjDYN_MUSCLE) { + dyn_type = MjcPhysicsTokens->muscle; + } else if (actuator->dyntype == mjtDyn::mjDYN_USER) { + dyn_type = MjcPhysicsTokens->user; + } + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Token, + MjcPhysicsTokens->mjcDynType, dyn_type); + WriteUniformAttribute( + transmission_path, pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcDynPrm, + pxr::VtDoubleArray(actuator->dynprm, actuator->dynprm + 10)); + + + pxr::TfToken gain_type; + if (actuator->gaintype == mjtGain::mjGAIN_FIXED) { + gain_type = MjcPhysicsTokens->fixed; + } else if (actuator->gaintype == mjtGain::mjGAIN_AFFINE) { + gain_type = MjcPhysicsTokens->affine; + } else if (actuator->gaintype == mjtGain::mjGAIN_MUSCLE) { + gain_type = MjcPhysicsTokens->muscle; + } else if (actuator->gaintype == mjtGain::mjGAIN_USER) { + gain_type = MjcPhysicsTokens->user; + } + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Token, + MjcPhysicsTokens->mjcGainType, gain_type); + WriteUniformAttribute( + transmission_path, pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcGainPrm, + pxr::VtDoubleArray(actuator->gainprm, actuator->gainprm + 10)); + + pxr::TfToken bias_type; + if (actuator->biastype == mjtBias::mjBIAS_NONE) { + bias_type = MjcPhysicsTokens->fixed; + } else if (actuator->biastype == mjtBias::mjBIAS_AFFINE) { + bias_type = MjcPhysicsTokens->affine; + } else if (actuator->biastype == mjtBias::mjBIAS_MUSCLE) { + bias_type = MjcPhysicsTokens->muscle; + } else if (actuator->biastype == mjtBias::mjBIAS_USER) { + bias_type = MjcPhysicsTokens->user; + } + WriteUniformAttribute(transmission_path, pxr::SdfValueTypeNames->Token, + MjcPhysicsTokens->mjcBiasType, bias_type); + WriteUniformAttribute( + transmission_path, pxr::SdfValueTypeNames->DoubleArray, + MjcPhysicsTokens->mjcBiasPrm, + pxr::VtDoubleArray(actuator->biasprm, actuator->biasprm + 10)); + } + + void WriteActuators() { + mjsActuator *actuator = + mjs_asActuator(mjs_firstElement(spec_, mjOBJ_ACTUATOR)); + while (actuator) { + WriteActuator(actuator); + actuator = mjs_asActuator(mjs_nextElement(spec_, actuator->element)); + } + } + + pxr::SdfPath WriteMeshGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + std::string mj_name = geom->name->empty() ? *geom->meshname : *geom->name; + auto name = + GetAvailablePrimName(mj_name, pxr::UsdGeomTokens->Mesh, body_path); + pxr::SdfPath subcomponent_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Xform); + + // Reference the mesh asset written in WriteMeshes. + AddPrimReference(data_, subcomponent_path, mesh_paths_[*geom->meshname]); + + // We want to use instancing with meshes, and it requires creating a parent + // scope to be referenced, with the Mesh prim as a child. + // To be able to actually manipulate the Mesh prim, we need to create and + // return the corresponding `over` prim as a child of the referencing prim. + pxr::SdfPath over_mesh_path = + CreatePrimSpec(data_, subcomponent_path, kTokens->sourceMesh, + pxr::UsdGeomTokens->Mesh, pxr::SdfSpecifierOver); + + return over_mesh_path; + } + + pxr::SdfPath WriteSiteGeom(const mjsSite *site, + const pxr::SdfPath &body_path) { + auto name = + GetAvailablePrimName(*site->name, pxr::UsdGeomTokens->Cube, body_path); + + int site_idx = mjs_getId(site->element); + const mjtNum *size = &model_->site_size[site_idx * 3]; + pxr::SdfPath site_path; + switch (site->type) { + case mjGEOM_BOX: + site_path = WriteBox(name, size, body_path); + break; + case mjGEOM_SPHERE: + site_path = WriteSphere(name, size, body_path); + break; + case mjGEOM_CAPSULE: + site_path = WriteCapsule(name, size, body_path); + break; + case mjGEOM_CYLINDER: + site_path = WriteCylinder(name, size, body_path); + break; + case mjGEOM_ELLIPSOID: + site_path = WriteEllipsoid(name, size, body_path); + break; + default: + break; + } + + return site_path; + } + + pxr::SdfPath WriteBox(const pxr::TfToken &name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath box_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Cube); + // MuJoCo uses half sizes. + pxr::SdfPath size_attr_path = + CreateAttributeSpec(data_, box_path, pxr::UsdGeomTokens->size, + pxr::SdfValueTypeNames->Double); + pxr::GfVec3f scale(static_cast(size[0]), static_cast(size[1]), + static_cast(size[2])); + SetAttributeDefault(data_, size_attr_path, 2.0); + + pxr::SdfPath extent_attr_path = + CreateAttributeSpec(data_, box_path, pxr::UsdGeomTokens->extent, + pxr::SdfValueTypeNames->Float3Array); + SetAttributeDefault(data_, extent_attr_path, + pxr::VtArray({ + pxr::GfVec3f(-size[0], -size[1], -size[2]), + pxr::GfVec3f(size[0], size[1], size[2]), + })); + + WriteScaleXformOp(box_path, scale); + WriteXformOpOrder(box_path, + pxr::VtArray{kTokens->xformOpScale}); + return box_path; + } + + pxr::SdfPath WriteBoxGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = + GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Cube, body_path); + + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + return WriteBox(name, geom_size, body_path); + } + + pxr::SdfPath WriteCapsule(const pxr::TfToken name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath capsule_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Capsule); + + pxr::SdfPath radius_attr_path = + CreateAttributeSpec(data_, capsule_path, pxr::UsdGeomTokens->radius, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, radius_attr_path, (double)size[0]); + + pxr::SdfPath height_attr_path = + CreateAttributeSpec(data_, capsule_path, pxr::UsdGeomTokens->height, + pxr::SdfValueTypeNames->Double); + // MuJoCo uses half sizes. + SetAttributeDefault(data_, height_attr_path, (double)(size[1] * 2)); + return capsule_path; + } + + pxr::SdfPath WriteCapsuleGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Capsule, + body_path); + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + + return WriteCapsule(name, geom_size, body_path); + } + + pxr::SdfPath WriteCylinder(const pxr::TfToken name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath cylinder_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Cylinder); + + pxr::SdfPath radius_attr_path = + CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->radius, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, radius_attr_path, (double)size[0]); + + pxr::SdfPath height_attr_path = + CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->height, + pxr::SdfValueTypeNames->Double); + // MuJoCo uses half sizes. + SetAttributeDefault(data_, height_attr_path, (double)(size[1] * 2)); + return cylinder_path; + } + + pxr::SdfPath WriteCylinderGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Cylinder, + body_path); + + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + return WriteCylinder(name, geom_size, body_path); + } + + pxr::SdfPath WriteEllipsoid(const pxr::TfToken name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath ellipsoid_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Sphere); + + pxr::GfVec3f scale = {static_cast(size[0]), + static_cast(size[1]), + static_cast(size[2])}; + + pxr::SdfPath radius_attr_path = + CreateAttributeSpec(data_, ellipsoid_path, pxr::UsdGeomTokens->radius, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, radius_attr_path, 1.0); + + WriteScaleXformOp(ellipsoid_path, scale); + WriteXformOpOrder(ellipsoid_path, + pxr::VtArray{kTokens->xformOpScale}); + return ellipsoid_path; + } + + pxr::SdfPath WriteEllipsoidGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Sphere, + body_path); + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + + return WriteEllipsoid(name, geom_size, body_path); + } + + pxr::SdfPath WriteSphere(const pxr::TfToken name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath sphere_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Sphere); + + pxr::SdfPath radius_attr_path = + CreateAttributeSpec(data_, sphere_path, pxr::UsdGeomTokens->radius, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, radius_attr_path, (double)size[0]); + return sphere_path; + } + + pxr::SdfPath WriteSphereGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Sphere, + body_path); + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + return WriteSphere(name, geom_size, body_path); + } + + pxr::SdfPath WritePlane(const pxr::TfToken &name, const mjtNum *size, + const pxr::SdfPath &body_path) { + pxr::SdfPath plane_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Plane); + + // MuJoCo uses half sizes. + // Note that UsdGeomPlane is infinite for simulation purposes but can have + // width/length for visualization, same as MuJoCo. + double width = size[0] * 2.0; + double length = size[1] * 2.0; + + pxr::SdfPath width_attr_path = + CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->width, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, width_attr_path, width); + + pxr::SdfPath length_attr_path = + CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->length, + pxr::SdfValueTypeNames->Double); + SetAttributeDefault(data_, length_attr_path, length); + + // MuJoCo plane is always a XY plane with +Z up. + // UsdGeomPlane is also a XY plane if axis is 'Z', which is default. + // So no need to set axis attribute explicitly. + + return plane_path; + } + + pxr::SdfPath WritePlaneGeom(const mjsGeom *geom, + const pxr::SdfPath &body_path) { + auto name = + GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Plane, body_path); + int geom_idx = mjs_getId(geom->element); + mjtNum *geom_size = &model_->geom_size[geom_idx * 3]; + return WritePlane(name, geom_size, body_path); + } + + void WriteSite(mjsSite *site, const mjsBody *body) { + const int body_id = mjs_getId(body->element); + const auto &body_path = body_paths_[body_id]; + auto name = + GetAvailablePrimName(*site->name, pxr::UsdGeomTokens->Xform, body_path); + + // Create a geom primitive and set its purpose to guide so it won't be + // rendered. + pxr::SdfPath site_path = WriteSiteGeom(site, body_path); + SetPrimPurpose(data_, site_path, pxr::UsdGeomTokens->guide); + + ApplyApiSchema(data_, site_path, MjcPhysicsTokens->SiteAPI); + + int site_id = mjs_getId(site->element); + auto transform = MujocoPosQuatToTransform(&model_->site_pos[3 * site_id], + &model_->site_quat[4 * site_id]); + WriteTransformXformOp(site_path, transform); + + PrependToXformOpOrder( + site_path, pxr::VtArray{kTokens->xformOpTransform}); + + site_paths_[site_id] = site_path; + } + + void WriteGeom(mjsGeom *geom, const mjsBody *body) { + const int body_id = mjs_getId(body->element); + const auto &body_path = body_paths_[body_id]; + + pxr::SdfPath geom_path; + int geom_id = mjs_getId(geom->element); + switch (geom->type) { + case mjGEOM_PLANE: + geom_path = WritePlaneGeom(geom, body_path); + break; + case mjGEOM_MESH: + geom_path = WriteMeshGeom(geom, body_path); + break; + case mjGEOM_BOX: + geom_path = WriteBoxGeom(geom, body_path); + break; + case mjGEOM_CAPSULE: + geom_path = WriteCapsuleGeom(geom, body_path); + break; + case mjGEOM_CYLINDER: + geom_path = WriteCylinderGeom(geom, body_path); + break; + case mjGEOM_ELLIPSOID: + geom_path = WriteEllipsoidGeom(geom, body_path); + break; + case mjGEOM_SPHERE: + geom_path = WriteSphereGeom(geom, body_path); + break; + default: + TF_WARN(UnsupportedGeomTypeError, "Unsupported geom type for geom %d", + geom_id); + return; + } + + // Apply the physics schemas if we are writing physics and the + // geom participates in collisions. + if (write_physics_ && (model_->geom_contype[geom_id] != 0 || + model_->geom_conaffinity[geom_id] != 0)) { + ApplyApiSchema(data_, geom_path, + pxr::UsdPhysicsTokens->PhysicsCollisionAPI); + ApplyApiSchema(data_, geom_path, MjcPhysicsTokens->CollisionAPI); + + WriteUniformAttribute( + geom_path, pxr::SdfValueTypeNames->Bool, + MjcPhysicsTokens->mjcShellinertia, + geom->typeinertia == mjtGeomInertia::mjINERTIA_SHELL); + + if (geom->mass >= mjMINVAL || geom->density >= mjMINVAL) { + ApplyApiSchema(data_, geom_path, pxr::UsdPhysicsTokens->PhysicsMassAPI); + } + + if (geom->mass >= mjMINVAL) { + pxr::SdfPath mass_attr = CreateAttributeSpec( + data_, geom_path, pxr::UsdPhysicsTokens->physicsMass, + pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform); + + // Make sure to cast to float here since mjtNum might be a double. + SetAttributeDefault(data_, mass_attr, (float)geom->mass); + } + + // Even though density is not used for mass computation when mass exists + // we want to retain the information anyways. + if (geom->density >= mjMINVAL) { + pxr::SdfPath density_attr = CreateAttributeSpec( + data_, geom_path, pxr::UsdPhysicsTokens->physicsDensity, + pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform); + + // Make sure to cast to float here since mjtNum might be a double. + SetAttributeDefault(data_, density_attr, (float)geom->density); + } + + // For meshes, also apply PhysicsMeshCollisionAPI and set the + // approximation attribute. + if (geom->type == mjGEOM_MESH) { + ApplyApiSchema(data_, geom_path, + pxr::UsdPhysicsTokens->PhysicsMeshCollisionAPI); + + // Note: MuJoCo documentation states that for collision purposes, meshes + // are always replaced with their convex hulls. Therefore, we set the + // approximation attribute to convexHull explicitly. + pxr::SdfPath approximation_attr = CreateAttributeSpec( + data_, geom_path, pxr::UsdPhysicsTokens->physicsApproximation, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + SetAttributeDefault(data_, approximation_attr, + pxr::UsdPhysicsTokens->convexHull); + } + } + + mjsDefault *spec_default = mjs_getDefault(geom->element); + pxr::TfToken valid_class_name = GetValidPrimName(*spec_default->name); + pxr::SdfPath geom_class_path = class_path_.AppendChild(valid_class_name); + if (!data_->HasSpec(geom_class_path)) { + pxr::SdfPath class_path = + CreateClassSpec(data_, class_path_, valid_class_name); + auto visibility_attr = + CreateAttributeSpec(data_, class_path, pxr::UsdGeomTokens->visibility, + pxr::SdfValueTypeNames->Token); + SetAttributeDefault(data_, visibility_attr, + pxr::UsdGeomTokens->inherited); + } + + // Bind material if it exists. + if (!geom->material->empty()) { + pxr::SdfPath material_path = + body_paths_[kWorldIndex] + .AppendChild(kTokens->materialsScope) + .AppendChild(GetValidPrimName(*geom->material)); + if (data_->HasSpec(material_path)) { + ApplyApiSchema(data_, geom_path, + pxr::UsdShadeTokens->MaterialBindingAPI); + // Bind the material to this geom. + CreateRelationshipSpec(data_, geom_path, + pxr::UsdShadeTokens->materialBinding, + material_path, pxr::SdfVariabilityUniform); + } + } + + // If geom rgba is not the default (0.5, 0.5, 0.5, 1), then set the + // displayColor attribute. + // No effort is made to properly handle the interaction between geom rgba + // and the material if both are specified. + if (geom->rgba[0] != 0.5f || geom->rgba[1] != 0.5f || + geom->rgba[2] != 0.5f || geom->rgba[3] != 1.0f) { + // Set the displayColor attribute. + pxr::SdfPath display_color_attr = CreateAttributeSpec( + data_, geom_path, pxr::UsdGeomTokens->primvarsDisplayColor, + pxr::SdfValueTypeNames->Color3fArray); + SetAttributeDefault(data_, display_color_attr, + pxr::VtArray{ + {geom->rgba[0], geom->rgba[1], geom->rgba[2]}}); + // Set the displayOpacity attribute, only if the opacity is not 1. + if (geom->rgba[3] != 1.0f) { + pxr::SdfPath display_opacity_attr = CreateAttributeSpec( + data_, geom_path, pxr::UsdGeomTokens->primvarsDisplayOpacity, + pxr::SdfValueTypeNames->FloatArray); + SetAttributeDefault(data_, display_opacity_attr, + pxr::VtArray{geom->rgba[3]}); + } + } + + if (body_id == kWorldIndex) { + SetPrimKind(data_, geom_path, pxr::KindTokens->component); + } + // Inherit from class. + AddPrimInherit(data_, geom_path, geom_class_path); + + auto transform = MujocoPosQuatToTransform(&model_->geom_pos[3 * geom_id], + &model_->geom_quat[4 * geom_id]); + WriteTransformXformOp(geom_path, transform); + + PrependToXformOpOrder( + geom_path, pxr::VtArray{kTokens->xformOpTransform}); + } + + void WriteSites(mjsBody *body) { + mjsSite *site = mjs_asSite(mjs_firstChild(body, mjOBJ_SITE, false)); + while (site) { + WriteSite(site, body); + site = mjs_asSite(mjs_nextChild(body, site->element, false)); + } + } + + void WriteGeoms(mjsBody *body) { + mjsGeom *geom = mjs_asGeom(mjs_firstChild(body, mjOBJ_GEOM, false)); + while (geom) { + WriteGeom(geom, body); + geom = mjs_asGeom(mjs_nextChild(body, geom->element, false)); + } + } + + void WriteCamera(mjsCamera *spec_cam, const mjsBody *body) { + const auto &body_path = body_paths_[mjs_getId(body->element)]; + auto name = GetAvailablePrimName(*spec_cam->name, + pxr::UsdGeomTokens->Camera, body_path); + // Create a root Xform for the world body with the model name if it exists + // otherwise called 'World'. + pxr::SdfPath camera_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Camera); + + int cam_id = mjs_getId(spec_cam->element); + auto transform = MujocoPosQuatToTransform(&model_->cam_pos[3 * cam_id], + &model_->cam_quat[4 * cam_id]); + WriteTransformXformOp(camera_path, transform); + WriteXformOpOrder(camera_path, + pxr::VtArray{kTokens->xformOpTransform}); + + // If the camera intrinsics are specified, then it is important that we + // reproduce the code in mujoco/src/engine/engine_vis_visualize.c + const float *cam_sensorsize = &model_->cam_sensorsize[cam_id * 2]; + bool use_intrinsic = cam_sensorsize[1] > 0.0f; + float znear = spec_->visual.map.znear * model_->stat.extent * 100; + float zfar = spec_->visual.map.zfar * model_->stat.extent * 100; + mjtNum fovy = model_->cam_fovy[cam_id]; + const float *cam_intrinsic = &model_->cam_intrinsic[cam_id * 4]; + + const float aspect_ratio = + use_intrinsic ? cam_sensorsize[0] / cam_sensorsize[1] : 4.0f / 3; + float vertical_apperture = + 2 * znear * + (use_intrinsic ? 1.0f / cam_intrinsic[1] * + (cam_sensorsize[1] / 2.f - cam_intrinsic[3]) + : mju_tan((fovy / 2) * (M_PI / 180.0))); + float horizontal_aperture = + use_intrinsic ? 2 * znear / cam_intrinsic[0] * + (cam_sensorsize[0] / 2.f - cam_intrinsic[2]) + : vertical_apperture * aspect_ratio; + + WriteUniformAttribute(camera_path, pxr::SdfValueTypeNames->Float2, + pxr::UsdGeomTokens->clippingRange, + pxr::GfVec2f(znear, zfar)); + WriteUniformAttribute(camera_path, pxr::SdfValueTypeNames->Float, + pxr::UsdGeomTokens->focalLength, znear); + WriteUniformAttribute(camera_path, pxr::SdfValueTypeNames->Float, + pxr::UsdGeomTokens->verticalAperture, + vertical_apperture); + WriteUniformAttribute(camera_path, pxr::SdfValueTypeNames->Float, + pxr::UsdGeomTokens->horizontalAperture, + horizontal_aperture); + } + + void WriteCameras(mjsBody *body) { + mjsCamera *cam = mjs_asCamera(mjs_firstChild(body, mjOBJ_CAMERA, false)); + while (cam) { + WriteCamera(cam, body); + cam = mjs_asCamera(mjs_nextChild(body, cam->element, false)); + } + } + + void WriteLight(mjsLight *light, const mjsBody *body) { + const auto &body_path = body_paths_[mjs_getId(body->element)]; + auto name = GetAvailablePrimName(*light->name, kTokens->light, body_path); + // Create a root Xform for the world body with the model name if it exists + // otherwise called 'World'. + pxr::SdfPath light_path = + CreatePrimSpec(data_, body_path, name, pxr::UsdLuxTokens->SphereLight); + + int light_id = mjs_getId(light->element); + auto transform = MujocoPosQuatToTransform(&model_->light_pos[3 * light_id], + &model_->light_dir[4 * light_id]); + WriteTransformXformOp(light_path, transform); + WriteXformOpOrder(light_path, + pxr::VtArray{kTokens->xformOpTransform}); + } + + void WriteLights(mjsBody *body) { + mjsLight *light = mjs_asLight(mjs_firstChild(body, mjOBJ_LIGHT, false)); + while (light) { + WriteLight(light, body); + light = mjs_asLight(mjs_nextChild(body, light->element, false)); + } + } + + void WriteBody(mjsBody *body) { + int body_id = mjs_getId(body->element); + // This should be safe as we process parent bodies before children. + mjsBody *parent = mjs_getParent(body->element); + int parent_id = mjs_getId(parent->element); + pxr::SdfPath parent_path = body_paths_[parent_id]; + pxr::TfToken body_name = GetValidPrimName(*body->name); + + // Create Xform prim for body. + pxr::SdfPath body_path = CreatePrimSpec(data_, parent_path, body_name, + pxr::UsdGeomTokens->Xform); + // The parent_path will be a component which makes the actual articulated + // bodies subcomponents. + auto kind = parent_id == kWorldIndex ? pxr::KindTokens->component + : pxr::KindTokens->subcomponent; + SetPrimKind(data_, body_path, kind); + + // Apply the PhysicsRigidBodyAPI schema if we are writing physics. + if (write_physics_) { + // If the body had a mass specified then it must have either inertia or + // fullinertia specified per inertia element XML documentation. + // Therefore it is sufficient to check if the mass is non-zero to see if + // we should set inertial attributes on the body. + // + // Note that if the user has NOT specified any inertial properties then + // we don't want to pull values from the compiled model since coming back + // into Mujoco would take those values instead of computing them + // automatically from the subtree. + if (body->mass > 0) { + // User might have specified the inertia via fullinertia and the + // compiler has extracted all values properly. So leverage those + // instead of doing the computation ourselves here. + ApplyApiSchema(data_, body_path, pxr::UsdPhysicsTokens->PhysicsMassAPI); + WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float, + pxr::UsdPhysicsTokens->physicsMass, + (float)model_->body_mass[body_id]); + + mjtNum *body_ipos = &model_->body_ipos[body_id * 3]; + pxr::GfVec3f inertial_pos(body_ipos[0], body_ipos[1], body_ipos[2]); + WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Point3f, + pxr::UsdPhysicsTokens->physicsCenterOfMass, + inertial_pos); + + mjtNum *body_iquat = &model_->body_iquat[body_id * 4]; + pxr::GfQuatf inertial_frame(body_iquat[0], body_iquat[1], body_iquat[2], + body_iquat[3]); + WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Quatf, + pxr::UsdPhysicsTokens->physicsPrincipalAxes, + inertial_frame); + + mjtNum *inertia = &model_->body_inertia[body_id * 3]; + pxr::GfVec3f diag_inertia(inertia[0], inertia[1], inertia[2]); + WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float3, + pxr::UsdPhysicsTokens->physicsDiagonalInertia, + diag_inertia); + } + + ApplyApiSchema(data_, body_path, + pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI); + + // If the parent is not the world body, but is child of the world body + // then we need to apply the articulation root API. + if (parent_id != kWorldIndex) { + int parent_parent_id = + mjs_getId(mjs_getParent(parent->element)->element); + if (parent_parent_id == kWorldIndex) { + ApplyApiSchema(data_, parent_path, + pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI); + } + } + } + + // Create classes if necessary + mjsDefault *spec_default = mjs_getDefault(body->element); + + pxr::TfToken body_class_name = GetValidPrimName(*spec_default->name); + pxr::SdfPath body_class_path = class_path_.AppendChild(body_class_name); + if (!data_->HasSpec(body_class_path)) { + CreateClassSpec(data_, class_path_, body_class_name); + } + + // Create XformOp attribute for body transform. + + pxr::SdfPath xform_op_path = + CreateAttributeSpec(data_, body_path, kTokens->xformOpTransform, + pxr::SdfValueTypeNames->Matrix4d); + // mjModel will have all frames already accounted for so no need to worry + // about them here. + auto body_xform = MujocoPosQuatToTransform(&model_->body_pos[body_id * 3], + &model_->body_quat[body_id * 4]); + SetAttributeDefault(data_, xform_op_path, body_xform); + + // Create XformOpOrder attribute for body transform order. + // For us this is simply the transform we authored above. + WriteXformOpOrder(body_path, + pxr::VtArray{kTokens->xformOpTransform}); + + pxr::VtDictionary customData; + customData[kTokens->body_name] = *body->name; + SetPrimMetadata(data_, body_path, pxr::SdfFieldKeys->CustomData, + customData); + + body_paths_[body_id] = body_path; + } + + void WriteBodies() { + mjsBody *body = mjs_asBody(mjs_firstElement(spec_, mjOBJ_BODY)); + while (body) { + // Only write a rigidbody if we are not the world body. + // We fall through since the world body might have static + // geom children. + if (mjs_getId(body->element) != kWorldIndex) { + WriteBody(body); + } + WriteSites(body); + WriteGeoms(body); + WriteCameras(body); + WriteLights(body); + body = mjs_asBody(mjs_nextElement(spec_, body->element)); + } + } + + pxr::SdfPath WriteWorldBody(const size_t body_index) { + // Create a root Xform for the world body with the model name if it exists + // otherwise called 'World'. + auto name = GetAvailablePrimName(*spec_->modelname, kTokens->world, + pxr::SdfPath::AbsoluteRootPath()); + pxr::SdfPath world_group_path = + CreatePrimSpec(data_, pxr::SdfPath::AbsoluteRootPath(), name, + pxr::UsdGeomTokens->Xform); + SetPrimKind(data_, world_group_path, pxr::KindTokens->group); + return world_group_path; + } +}; + +namespace mujoco { +namespace usd { + +bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data, + bool write_physics) { + // Create pseudo root first. + data->CreateSpec(pxr::SdfPath::AbsoluteRootPath(), + pxr::SdfSpecTypePseudoRoot); + + mjModel *model = mj_compile(spec, nullptr); + if (model == nullptr) { + TF_ERROR(MujocoCompilationError, "%s", mjs_getError(spec)); + return false; + } + + ModelWriter(spec, model, data).Write(write_physics); + + return true; +} + +} // namespace usd +} // namespace mujoco diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h new file mode 100644 index 00000000..9c0db6cf --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h @@ -0,0 +1,34 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MUJOCO_TO_USD_H_ +#define MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MUJOCO_TO_USD_H_ + +#include +#include + +namespace mujoco { +namespace usd { +// Given an mjSpec, write it to a SdfAbstractData. +// +// Args: +// spec: mjSpec built programmatically or via parsed XML. +// data: SdfAbstractDataRefPtr that will be written to. +// write_physics: Whether to write physics data. +bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data, + bool write_physics); +} // namespace usd +} // namespace mujoco + +#endif // MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MUJOCO_TO_USD_H_ diff --git a/src/experimental/usd/plugins/mjcf/plugInfo.json b/src/experimental/usd/plugins/mjcf/plugInfo.json new file mode 100644 index 00000000..da1f9ebd --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/plugInfo.json @@ -0,0 +1,25 @@ +{ + "Plugins": [ + { + "Info": { + "Types": { + "UsdMjcfFileFormat": { + "bases": ["SdfFileFormat"], + "displayName": "MJCF USD File Format", + "extensions": ["xml"], + "formatId": "xml", + "primary": true, + "supportsReading": true, + "supportsWriting": false, + "target": "usd" + } + } + }, + "LibraryPath": "@PLUG_INFO_LIBRARY_PATH@", + "Name": "usdMjcf", + "ResourcePath": "", + "Root": ".", + "Type": "library" + } + ] +} diff --git a/src/experimental/usd/plugins/mjcf/utils.cc b/src/experimental/usd/plugins/mjcf/utils.cc new file mode 100644 index 00000000..e919a672 --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/utils.cc @@ -0,0 +1,208 @@ +// Copyright 2025 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. + +#include "mjcf/utils.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace { + +template +void AppendChild(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& specPath, + const pxr::TfToken& childKey, const T& child) { + // Get existing children. + std::vector children; + pxr::SdfAbstractDataTypedValue> getter(&children); + data->Has(specPath, childKey, &getter); + + children.push_back(child); + data->Set(specPath, childKey, + pxr::SdfAbstractDataConstTypedValue>(&children)); +} + +template +void AppendListOp(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& spec_path, const pxr::TfToken& field, + const T& item) { + // Get existing list op. + pxr::SdfListOp list_op; + pxr::SdfAbstractDataTypedValue> getter(&list_op); + data->Has(spec_path, field, &getter); + + auto items = list_op.GetExplicitItems(); + items.push_back(item); + list_op.SetExplicitItems(items); + data->Set(spec_path, field, + pxr::SdfAbstractDataConstTypedValue>(&list_op)); +} + +template +void PrependListOp(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& spec_path, const pxr::TfToken& field, + const T& item) { + // Get existing list op. + pxr::SdfListOp listOp; + pxr::SdfAbstractDataTypedValue> getter(&listOp); + data->Has(spec_path, field, &getter); + + auto prependedItems = listOp.GetPrependedItems(); + prependedItems.insert(prependedItems.begin(), item); + listOp.SetPrependedItems(prependedItems); + data->Set(spec_path, field, + pxr::SdfAbstractDataConstTypedValue>(&listOp)); +} +} // namespace + +namespace mujoco { +namespace usd { + +pxr::SdfPath CreatePrimSpec(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& parent_path, + const pxr::TfToken& name, const pxr::TfToken& type, + pxr::SdfSpecifier specifier) { + const pxr::SdfPath prim_path = parent_path.AppendChild(name); + data->CreateSpec(prim_path, pxr::SdfSpecTypePrim); + data->Set(prim_path, pxr::SdfFieldKeys->Specifier, + pxr::SdfAbstractDataConstTypedValue(&specifier)); + if (!type.IsEmpty()) { + data->Set(prim_path, pxr::SdfFieldKeys->TypeName, + pxr::SdfAbstractDataConstTypedValue(&type)); + } + + AppendChild(data, parent_path, pxr::SdfChildrenKeys->PrimChildren, name); + + return prim_path; +} + +pxr::SdfPath CreateAttributeSpec(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& name, + const pxr::SdfValueTypeName& type_name, + pxr::SdfVariability variability) { + const pxr::SdfPath propertyPath = prim_path.AppendProperty(name); + data->CreateSpec(propertyPath, pxr::SdfSpecTypeAttribute); + + pxr::TfToken typeNameToken = type_name.GetAsToken(); + data->Set(propertyPath, pxr::SdfFieldKeys->TypeName, + pxr::SdfAbstractDataConstTypedValue(&typeNameToken)); + if (variability != pxr::SdfVariabilityVarying) { + data->Set( + propertyPath, pxr::SdfFieldKeys->Variability, + pxr::SdfAbstractDataConstTypedValue(&variability)); + } + + AppendChild(data, prim_path, pxr::SdfChildrenKeys->PropertyChildren, name); + + return propertyPath; +} + +pxr::SdfPath CreateRelationshipSpec(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& relationship_name, + const pxr::SdfPath& relationship_path, + pxr::SdfVariability variability) { + pxr::SdfPath prop_path = prim_path.AppendProperty(relationship_name); + data->CreateSpec(prop_path, pxr::SdfSpecTypeRelationship); + if (variability != pxr::SdfVariabilityVarying) { + data->Set( + prop_path, pxr::SdfFieldKeys->Variability, + pxr::SdfAbstractDataConstTypedValue(&variability)); + } + + AppendChild(data, prim_path, pxr::SdfChildrenKeys->PropertyChildren, + relationship_name); + + AppendChild(data, prop_path, pxr::SdfChildrenKeys->RelationshipTargetChildren, + relationship_path); + AppendListOp(data, prop_path, pxr::SdfFieldKeys->TargetPaths, + relationship_path); + + pxr::SdfPath target_path = prop_path.AppendTarget(relationship_path); + data->CreateSpec(target_path, pxr::SdfSpecTypeRelationshipTarget); + + return prop_path; +} + +pxr::SdfPath CreateClassSpec(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& class_name) { + pxr::SdfPath class_path = prim_path.AppendChild(class_name); + pxr::SdfSpecifier class_specifier = pxr::SdfSpecifier::SdfSpecifierClass; + data->CreateSpec(class_path, pxr::SdfSpecTypePrim); + data->Set( + class_path, pxr::SdfFieldKeys->Specifier, + pxr::SdfAbstractDataConstTypedValue(&class_specifier)); + + AppendChild(data, prim_path, pxr::SdfChildrenKeys->PrimChildren, class_name); + + return class_path; +} + +void AddAttributeConnection(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& attribute_path, + const pxr::SdfPath& target_attribute_path) { + AppendChild(data, attribute_path, pxr::SdfChildrenKeys->ConnectionChildren, + target_attribute_path); + AppendListOp(data, attribute_path, pxr::SdfFieldKeys->ConnectionPaths, + target_attribute_path); + + data->CreateSpec(attribute_path.AppendTarget(target_attribute_path), + pxr::SdfSpecTypeConnection); +} + +void AddPrimReference(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::SdfPath& referenced_prim_path) { + PrependListOp(data, prim_path, pxr::SdfFieldKeys->References, + pxr::SdfReference("", referenced_prim_path)); +} + +void AddPrimInherit(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::SdfPath& class_path) { + PrependListOp(data, prim_path, pxr::SdfFieldKeys->InheritPaths, class_path); +} + +void ApplyApiSchema(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& schema_name) { + PrependListOp(data, prim_path, pxr::UsdTokens->apiSchemas, schema_name); +} + +void SetPrimKind(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, pxr::TfToken kind) { + SetPrimMetadata(data, prim_path, pxr::TfToken("kind"), kind); +} + +void SetPrimPurpose(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, pxr::TfToken purpose) { + const pxr::SdfPath attr_path = CreateAttributeSpec( + data, prim_path, pxr::UsdGeomTokens->purpose, + pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform); + SetAttributeDefault(data, attr_path, purpose); +} + +} // namespace usd +} // namespace mujoco diff --git a/src/experimental/usd/plugins/mjcf/utils.h b/src/experimental/usd/plugins/mjcf/utils.h new file mode 100644 index 00000000..9f41a697 --- /dev/null +++ b/src/experimental/usd/plugins/mjcf/utils.h @@ -0,0 +1,130 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_UTILS_H_ +#define MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_UTILS_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace mujoco { +namespace usd { + +// Create a prim spec and append it as a child of parent_path. +pxr::SdfPath CreatePrimSpec( + pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& parent_path, + const pxr::TfToken& name, const pxr::TfToken& type = pxr::TfToken(), + pxr::SdfSpecifier specifier = pxr::SdfSpecifier::SdfSpecifierDef); + +// Create an attribute spec and append it as a child of parent_path. +// By default the attribute will be varying. +pxr::SdfPath CreateAttributeSpec( + pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& prim_path, + const pxr::TfToken& name, const pxr::SdfValueTypeName& type_name, + pxr::SdfVariability variability = pxr::SdfVariabilityVarying); + +// Create a relationship spec and append it as a child of prim_path. +pxr::SdfPath CreateRelationshipSpec( + pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& prim_path, + const pxr::TfToken& relationship_name, + const pxr::SdfPath& relationship_path, + pxr::SdfVariability variability = pxr::SdfVariabilityVarying); + +pxr::SdfPath CreateClassSpec(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& class_name); + +void AddAttributeConnection(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& attribute_path, + const pxr::SdfPath& target_attribute_path); + +void AddPrimReference(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::SdfPath& referenced_prim_path); + +void AddPrimInherit(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::SdfPath& class_path); + +void ApplyApiSchema(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, + const pxr::TfToken& schema_name); + +void SetPrimKind(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, pxr::TfToken kind); + +void SetPrimPurpose(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, pxr::TfToken purpose); + +// Set the value specified by key on any field at field_path. +template +void SetField(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& field_path, + const pxr::TfToken key, T&& value) { + using Deduced = typename std::remove_reference_t; + const auto typed_val = pxr::SdfAbstractDataConstTypedValue(&value); + const pxr::SdfAbstractDataConstValue& untyped_val = typed_val; + + data->Set(field_path, key, untyped_val); +} + +// Set the value specified by key on an attribute spec at attribute_path. +template +void SetAttribute(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& attribute_path, const pxr::TfToken key, + T&& value) { + SetField(data, attribute_path, key, value); +} + +// Set the value specified by key on a prim spec at prim_path. +template +void SetPrimMetadata(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& prim_path, const pxr::TfToken key, + T&& value) { + SetAttribute(data, prim_path, key, value); +} + +// Set the value specified by key on an attribute spec at attribute_path. +template +void SetAttributeMetadata(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& attribute_path, + const pxr::TfToken key, T&& value) { + SetAttribute(data, attribute_path, key, value); +} + +// Set the default value on an attribute spec at attribute_path. +template +void SetAttributeDefault(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& attribute_path, + T&& default_value) { + SetAttribute(data, attribute_path, pxr::SdfFieldKeys->Default, default_value); +} + +// Set the value specified by key on the root layer. +template +void SetLayerMetadata(pxr::SdfAbstractDataRefPtr& data, const pxr::TfToken& key, + T&& value) { + SetAttribute(data, pxr::SdfPath::AbsoluteRootPath(), key, value); +} + +} // namespace usd +} // namespace mujoco + +#endif // MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_UTILS_H_ diff --git a/src/experimental/usd/utils.cc b/src/experimental/usd/utils.cc new file mode 100644 index 00000000..220c4178 --- /dev/null +++ b/src/experimental/usd/utils.cc @@ -0,0 +1,44 @@ +// Copyright 2025 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. + +#include "third_party/mujoco/src/experimental/usd/utils.h" + +#include +#include + +namespace mujoco { +namespace usd { + +constexpr const char* kUsdPrimPathKey = "usd_primpath"; + +void SetUsdPrimPathUserValue(mjsElement* element, + const pxr::SdfPath& prim_path) { + // The value is a pointer to a newly allocated SdfPath, which will be deleted + // when the mjsElement is deleted. + const pxr::SdfPath* usd_primpath = new pxr::SdfPath(prim_path); + mjs_setUserValueWithCleanup( + element, kUsdPrimPathKey, usd_primpath, + [](const void* data) { delete static_cast(data); }); +} + +pxr::SdfPath GetUsdPrimPathUserValue(mjsElement* element) { + const void* user_data = mjs_getUserValue(element, kUsdPrimPathKey); + if (user_data) { + return *static_cast(user_data); + } + return pxr::SdfPath(); +} + +} // namespace usd +} // namespace mujoco diff --git a/src/experimental/usd/utils.h b/src/experimental/usd/utils.h new file mode 100644 index 00000000..f3dcd4b5 --- /dev/null +++ b/src/experimental/usd/utils.h @@ -0,0 +1,35 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_USD_UTILS_H_ +#define MUJOCO_SRC_EXPERIMENTAL_USD_UTILS_H_ + +#include +#include + +namespace mujoco { +namespace usd { + +// Sets a user value on an mjsElement with the key "usd_primpath". +void SetUsdPrimPathUserValue(mjsElement* element, + const pxr::SdfPath& prim_path); + +// Gets the user value associated with the key "usd_primpath" from an +// mjsElement. Returns empty pxr::SdfPath() if the value is not found. +pxr::SdfPath GetUsdPrimPathUserValue(mjsElement* element); + +} // namespace usd +} // namespace mujoco + +#endif // MUJOCO_SRC_EXPERIMENTAL_USD_UTILS_H_ diff --git a/src/render/glad/glad.c b/src/render/glad/glad.c index fa1642c4..72bd6f78 100644 --- a/src/render/glad/glad.c +++ b/src/render/glad/glad.c @@ -27,6 +27,7 @@ // GL_ARB_framebuffer_object, // GL_ARB_seamless_cube_map, // GL_ARB_vertex_buffer_object, +// GL_EXT_texture_sRGB, // GL_KHR_debug // Loader: True // Local files: False @@ -34,9 +35,9 @@ // Reproducible: False // // Commandline: -// --profile="compatibility" --api="gl=1.5" --generator="c" --spec="gl" --extensions="GL_ARB_clip_control,GL_ARB_depth_buffer_float,GL_ARB_framebuffer_object,GL_ARB_seamless_cube_map,GL_ARB_vertex_buffer_object,GL_KHR_debug" +// --profile="compatibility" --api="gl=1.5" --generator="c" --spec="gl" --extensions="GL_ARB_clip_control,GL_ARB_depth_buffer_float,GL_ARB_framebuffer_object,GL_ARB_seamless_cube_map,GL_ARB_vertex_buffer_object,GL_EXT_texture_sRGB,GL_KHR_debug" // Online: -// https://glad.dav1d.de/#profile=compatibility&language=c&specification=gl&loader=on&api=gl%3D1.5&extensions=GL_ARB_clip_control&extensions=GL_ARB_depth_buffer_float&extensions=GL_ARB_framebuffer_object&extensions=GL_ARB_seamless_cube_map&extensions=GL_ARB_vertex_buffer_object&extensions=GL_KHR_debug +// https://glad.dav1d.de/#profile=compatibility&language=c&specification=gl&loader=on&api=gl%3D1.5&extensions=GL_ARB_clip_control&extensions=GL_ARB_depth_buffer_float&extensions=GL_ARB_framebuffer_object&extensions=GL_ARB_seamless_cube_map&extensions=GL_ARB_vertex_buffer_object&extensions=GL_EXT_texture_sRGB&extensions=GL_KHR_debug #if defined(__GNUC__) && !defined(__clang__) #pragma GCC diagnostic push @@ -836,6 +837,7 @@ int mjGLAD_GL_ARB_depth_buffer_float = 0; int mjGLAD_GL_ARB_framebuffer_object = 0; int mjGLAD_GL_ARB_seamless_cube_map = 0; int mjGLAD_GL_ARB_vertex_buffer_object = 0; +int mjGLAD_GL_EXT_texture_sRGB = 0; int mjGLAD_GL_KHR_debug = 0; PFNGLCLIPCONTROLPROC mjGlad_glClipControl = NULL; PFNGLISRENDERBUFFERPROC mjGlad_glIsRenderbuffer = NULL; @@ -1433,6 +1435,7 @@ static int mjGlad_find_extensionsGL(void) { mjGLAD_GL_ARB_framebuffer_object = mjGlad_has_ext("GL_ARB_framebuffer_object"); mjGLAD_GL_ARB_seamless_cube_map = mjGlad_has_ext("GL_ARB_seamless_cube_map"); mjGLAD_GL_ARB_vertex_buffer_object = mjGlad_has_ext("GL_ARB_vertex_buffer_object"); + mjGLAD_GL_EXT_texture_sRGB = mjGlad_has_ext("GL_EXT_texture_sRGB"); mjGLAD_GL_KHR_debug = mjGlad_has_ext("GL_KHR_debug"); mjGlad_free_exts(); return 1; diff --git a/src/render/glad/glad.h b/src/render/glad/glad.h index 1b1fd449..620a7cbb 100644 --- a/src/render/glad/glad.h +++ b/src/render/glad/glad.h @@ -27,6 +27,7 @@ // GL_ARB_framebuffer_object, // GL_ARB_seamless_cube_map, // GL_ARB_vertex_buffer_object, +// GL_EXT_texture_sRGB, // GL_KHR_debug // Loader: True // Local files: False @@ -34,9 +35,9 @@ // Reproducible: False // // Commandline: -// --profile="compatibility" --api="gl=1.5" --generator="c" --spec="gl" --extensions="GL_ARB_clip_control,GL_ARB_depth_buffer_float,GL_ARB_framebuffer_object,GL_ARB_seamless_cube_map,GL_ARB_vertex_buffer_object,GL_KHR_debug" +// --profile="compatibility" --api="gl=1.5" --generator="c" --spec="gl" --extensions="GL_ARB_clip_control,GL_ARB_depth_buffer_float,GL_ARB_framebuffer_object,GL_ARB_seamless_cube_map,GL_ARB_vertex_buffer_object,GL_EXT_texture_sRGB,GL_KHR_debug" // Online: -// https://glad.dav1d.de/#profile=compatibility&language=c&specification=gl&loader=on&api=gl%3D1.5&extensions=GL_ARB_clip_control&extensions=GL_ARB_depth_buffer_float&extensions=GL_ARB_framebuffer_object&extensions=GL_ARB_seamless_cube_map&extensions=GL_ARB_vertex_buffer_object&extensions=GL_KHR_debug +// https://glad.dav1d.de/#profile=compatibility&language=c&specification=gl&loader=on&api=gl%3D1.5&extensions=GL_ARB_clip_control&extensions=GL_ARB_depth_buffer_float&extensions=GL_ARB_framebuffer_object&extensions=GL_ARB_seamless_cube_map&extensions=GL_ARB_vertex_buffer_object&extensions=GL_EXT_texture_sRGB&extensions=GL_KHR_debug #ifndef MUJOCO_SRC_RENDER_GLAD_GLAD_H_ #define MUJOCO_SRC_RENDER_GLAD_GLAD_H_ @@ -2420,6 +2421,22 @@ GLAPI PFNGLGETBUFFERPOINTERVPROC mjGlad_glGetBufferPointerv; #define GL_DYNAMIC_DRAW_ARB 0x88E8 #define GL_DYNAMIC_READ_ARB 0x88E9 #define GL_DYNAMIC_COPY_ARB 0x88EA +#define GL_SRGB_EXT 0x8C40 +#define GL_SRGB8_EXT 0x8C41 +#define GL_SRGB_ALPHA_EXT 0x8C42 +#define GL_SRGB8_ALPHA8_EXT 0x8C43 +#define GL_SLUMINANCE_ALPHA_EXT 0x8C44 +#define GL_SLUMINANCE8_ALPHA8_EXT 0x8C45 +#define GL_SLUMINANCE_EXT 0x8C46 +#define GL_SLUMINANCE8_EXT 0x8C47 +#define GL_COMPRESSED_SRGB_EXT 0x8C48 +#define GL_COMPRESSED_SRGB_ALPHA_EXT 0x8C49 +#define GL_COMPRESSED_SLUMINANCE_EXT 0x8C4A +#define GL_COMPRESSED_SLUMINANCE_ALPHA_EXT 0x8C4B +#define GL_COMPRESSED_SRGB_S3TC_DXT1_EXT 0x8C4C +#define GL_COMPRESSED_SRGB_ALPHA_S3TC_DXT1_EXT 0x8C4D +#define GL_COMPRESSED_SRGB_ALPHA_S3TC_DXT3_EXT 0x8C4E +#define GL_COMPRESSED_SRGB_ALPHA_S3TC_DXT5_EXT 0x8C4F #define GL_DEBUG_OUTPUT_SYNCHRONOUS 0x8242 #define GL_DEBUG_NEXT_LOGGED_MESSAGE_LENGTH 0x8243 #define GL_DEBUG_CALLBACK_FUNCTION 0x8244 @@ -2614,6 +2631,10 @@ typedef void (APIENTRYP PFNGLGETBUFFERPOINTERVARBPROC)(GLenum target, GLenum pna GLAPI PFNGLGETBUFFERPOINTERVARBPROC mjGlad_glGetBufferPointervARB; #define glGetBufferPointervARB mjGlad_glGetBufferPointervARB #endif +#ifndef GL_EXT_texture_sRGB +#define GL_EXT_texture_sRGB 1 +GLAPI int mjGLAD_GL_EXT_texture_sRGB; +#endif #ifndef GL_KHR_debug #define GL_KHR_debug 1 GLAPI int mjGLAD_GL_KHR_debug; diff --git a/src/render/render_context.c b/src/render/render_context.c index 6adb4f11..b1f279bb 100644 --- a/src/render/render_context.c +++ b/src/render/render_context.c @@ -1050,13 +1050,17 @@ static void makeShadow(const mjModel* m, mjrContext* con) { } glBindFramebuffer(GL_FRAMEBUFFER, con->shadowFBO); - // create shadow depth texture: in TEXTURE1 + // Create a shadow depth texture in TEXTURE1 and explicitly select an int24 + // depth buffer. A depth stencil format is used because that appears to be + // more widely supported (MacOS does not support GL_DEPTH_COMPONENT24). Using + // a fixed format makes it easier to choose glPolygonOffset parameters that + // result in reasonably consistent and artifact free shadows across platforms. glGenTextures(1, &con->shadowTex); glActiveTexture(GL_TEXTURE1); glEnable(GL_TEXTURE_2D); glBindTexture(GL_TEXTURE_2D, con->shadowTex); - glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, - con->shadowSize, con->shadowSize, 0, GL_DEPTH_COMPONENT, GL_FLOAT, NULL); + glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH24_STENCIL8, + con->shadowSize, con->shadowSize, 0, GL_DEPTH_STENCIL, GL_UNSIGNED_INT_24_8, NULL); glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST); glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST); glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE); @@ -1383,15 +1387,20 @@ void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { // assign data int type = 0; + int internaltype = 0; if (m->tex_nchannel[texid] == 3) { type = GL_RGB; + internaltype = (m->tex_colorspace[texid] == mjCOLORSPACE_SRGB) ? GL_SRGB8_EXT : GL_RGB; } else if (m->tex_nchannel[texid] == 4) { type = GL_RGBA; + internaltype = (m->tex_colorspace[texid] == mjCOLORSPACE_SRGB) ? GL_SRGB8_ALPHA8_EXT : GL_RGBA; } else { mju_error("Number of channels not supported: %d", m->tex_nchannel[texid]); } - glTexImage2D(GL_TEXTURE_2D, 0, type, m->tex_width[texid], m->tex_height[texid], 0, - type, GL_UNSIGNED_BYTE, m->tex_data + m->tex_adr[texid]); + + glTexImage2D(GL_TEXTURE_2D, 0, internaltype, m->tex_width[texid], + m->tex_height[texid], 0, type, GL_UNSIGNED_BYTE, + m->tex_data + m->tex_adr[texid]); // generate mipmaps glGenerateMipmap(GL_TEXTURE_2D); diff --git a/src/render/render_gl3.c b/src/render/render_gl3.c index dffa48d7..1660c67b 100644 --- a/src/render/render_gl3.c +++ b/src/render/render_gl3.c @@ -672,7 +672,7 @@ static void initLights(mjvScene* scn) { glLightfv(GL_LIGHT0+i, GL_SPECULAR, scn->lights[i].specular); // parameters for directional light - if (scn->lights[i].directional) { + if (scn->lights[i].type == mjLIGHT_DIRECTIONAL) { glLightf(GL_LIGHT0+i, GL_SPOT_EXPONENT, 0); glLightf(GL_LIGHT0+i, GL_SPOT_CUTOFF, 180); glLightf(GL_LIGHT0+i, GL_CONSTANT_ATTENUATION, 1); @@ -681,13 +681,17 @@ static void initLights(mjvScene* scn) { } // parameters for spot light - else { + else if (scn->lights[i].type == mjLIGHT_SPOT) { glLightf(GL_LIGHT0+i, GL_SPOT_EXPONENT, scn->lights[i].exponent); glLightf(GL_LIGHT0+i, GL_SPOT_CUTOFF, scn->lights[i].cutoff); glLightf(GL_LIGHT0+i, GL_CONSTANT_ATTENUATION, scn->lights[i].attenuation[0]); glLightf(GL_LIGHT0+i, GL_LINEAR_ATTENUATION, scn->lights[i].attenuation[1]); glLightf(GL_LIGHT0+i, GL_QUADRATIC_ATTENUATION, scn->lights[i].attenuation[2]); } + + else { + mju_error("Unsupported light type: %d", scn->lights[i].type); + } } // disable all lights (enable selectively in render) @@ -790,14 +794,16 @@ static void adjustLight(const mjvLight* thislight, int n) { float temp[4]; // set position and direction according to type - if (thislight->directional) { + if (thislight->type == mjLIGHT_DIRECTIONAL) { mjr_setf4(temp, -thislight->dir[0], -thislight->dir[1], -thislight->dir[2], 0); glLightfv(GL_LIGHT0+n, GL_POSITION, temp); - } else { + } else if (thislight->type == mjLIGHT_SPOT) { mjr_setf4(temp, thislight->dir[0], thislight->dir[1], thislight->dir[2], 0); glLightfv(GL_LIGHT0+n, GL_SPOT_DIRECTION, temp); mjr_setf4(temp, thislight->pos[0], thislight->pos[1], thislight->pos[2], 1); glLightfv(GL_LIGHT0+n, GL_POSITION, temp); + } else { + mju_error("Unsupported light type: %d", thislight->type); } } @@ -1183,13 +1189,15 @@ void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { // reverse Z rendering mapping without shift [znear, zfar] -> [1, -1] (ndc) glScalef(1.0f, 1.0f, -1.0f); } - if (thislight->directional) { + if (thislight->type == mjLIGHT_DIRECTIONAL) { glOrtho(-con->shadowClip, con->shadowClip, -con->shadowClip, con->shadowClip, cam.frustum_near, cam.frustum_far); - } else { + } else if (thislight->type == mjLIGHT_SPOT) { mjr_perspective(mju_min(2*thislight->cutoff*con->shadowScale, 160), 1, cam.frustum_near, cam.frustum_far); + } else { + mju_error("Unsupported light type: %d", thislight->type); } glGetFloatv(GL_PROJECTION_MATRIX, lightProject); @@ -1211,9 +1219,24 @@ void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { int cull_face = glIsEnabled(GL_CULL_FACE); glDisable(GL_CULL_FACE); // all faces cast shadows glEnable(GL_POLYGON_OFFSET_FILL); - float kOffsetFactor = -1.5f; - float kOffsetUnits = -4.0f; - glPolygonOffset(kOffsetFactor, kOffsetUnits); // prevents "shadow acne" + + // The limited resolution of the shadow maps means multiple fragments + // sample the same texel. When light and camera directions differ on + // surfaces that should be lit this causes "shadow acne" because some + // fragments will be lit while adjacent fragments are not. To mitigate + // this artifact, an offset is applied to the depth values in the + // shadow map. The offset must be large enough to ensure consistent + // depth comparison occurs within the limited precision of the depth + // buffer. The offset is computed by glPolygonOffset using parameters + // that are chosen empirically. We need different values when clip + // control is on/off because this setting changes the depth precision. + float kOffsetFactor = -16.0f; + float kOffsetUnits = -512.0f; + if (mjGLAD_GL_ARB_clip_control) { + kOffsetFactor = -1.5f; + kOffsetUnits = -4.0f; + } + glPolygonOffset(kOffsetFactor, kOffsetUnits); // render all geoms to depth texture for (int j=0; j < ngeom; j++) { diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 0e339780..c25049b1 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -73,7 +73,7 @@ mjSpec* mj_copySpec(const mjSpec* s) { try { modelC = new mjCModel(*static_cast(s->element)); } catch (mjCError& e) { - mju_error("Failed to copy spec: %s", e.message); + static_cast(s->element)->SetError(e); return nullptr; } return &modelC->spec; @@ -119,102 +119,170 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) { } +// set frame for all elements of a body +static void SetFrame(mjsBody* body, mjtObj objtype, mjsFrame* frame) { + mjsElement* el = mjs_firstChild(body, objtype, 0); + while (el) { + if (frame->element != el && mjs_getFrame(el) == nullptr) { + mjs_setFrame(el, frame); + } + el = mjs_nextChild(body, el, 0); + } +} + + // attach body to a frame of the parent -mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child, - const char* prefix, const char* suffix) { - if (!parent) { - mju_error("parent frame is null"); - return nullptr; - } - mjCFrame* frame_parent = static_cast(parent->element); - mjCBody* child_body = static_cast(child->element); +static mjsElement* attachBody(mjCFrame* parent, const mjCBody* child, + const char* prefix, const char* suffix) { try { - *frame_parent += std::string(prefix) + *child_body + std::string(suffix); + *parent += std::string(prefix) + *(mjCBody*)child + std::string(suffix); } catch (mjCError& e) { - frame_parent->model->SetError(e); + parent->model->SetError(e); return nullptr; } - mjsBody* attached_body = frame_parent->last_attached; - frame_parent->last_attached = nullptr; - return attached_body; + mjsBody* attached_body = parent->last_attached; + parent->last_attached = nullptr; + return attached_body->element; } // attach frame to a parent body -mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child, - const char* prefix, const char* suffix) { - if (!parent) { - mju_error("parent body is null"); - return nullptr; - } - mjCBody* body_parent = static_cast(parent->element); - mjCFrame* child_frame = static_cast(child->element); +static mjsElement* attachFrame(mjCBody* parent, const mjCFrame* child, + const char* prefix, const char* suffix) { try { - *body_parent += std::string(prefix) + *child_frame + std::string(suffix); + *parent += std::string(prefix) + *(mjCFrame*)child + std::string(suffix); } catch (mjCError& e) { - body_parent->model->SetError(e); + parent->model->SetError(e); return nullptr; } - mjsFrame* attached_frame = body_parent->last_attached; - body_parent->last_attached = nullptr; - return attached_frame; + mjsFrame* attached_frame = parent->last_attached; + parent->last_attached = nullptr; + return attached_frame->element; } // attach child body to a parent site -mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child, - const char* prefix, const char* suffix) { - if (!parent) { - mju_error("parent site is null"); - return nullptr; - } - mjSpec* spec = mjs_getSpec(parent->element); - mjCSite* site = static_cast(parent->element); - mjCBody* body = site->Body(); - mjCFrame* frame = body->AddFrame(site->frame); +static mjsElement* attachToSite(mjCSite* parent, const mjCBody* child, + const char* prefix, const char* suffix) { + mjSpec* spec = mjs_getSpec(parent->spec.element); + mjCBody* body = parent->Body(); + mjCFrame* frame = body->AddFrame(parent->frame); frame->SetParent(body); - frame->spec.pos[0] = site->spec.pos[0]; - frame->spec.pos[1] = site->spec.pos[1]; - frame->spec.pos[2] = site->spec.pos[2]; - frame->spec.quat[0] = site->spec.quat[0]; - frame->spec.quat[1] = site->spec.quat[1]; - frame->spec.quat[2] = site->spec.quat[2]; - frame->spec.quat[3] = site->spec.quat[3]; + frame->spec.pos[0] = parent->spec.pos[0]; + frame->spec.pos[1] = parent->spec.pos[1]; + frame->spec.pos[2] = parent->spec.pos[2]; + frame->spec.quat[0] = parent->spec.quat[0]; + frame->spec.quat[1] = parent->spec.quat[1]; + frame->spec.quat[2] = parent->spec.quat[2]; + frame->spec.quat[3] = parent->spec.quat[3]; mjs_resolveOrientation(frame->spec.quat, spec->compiler.degree, - spec->compiler.eulerseq, &site->spec.alt); - return mjs_attachBody(&frame->spec, child, prefix, suffix); + spec->compiler.eulerseq, &parent->spec.alt); + return attachBody(frame, child, prefix, suffix); } // attach child frame to a parent site -mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child, - const char* prefix, const char* suffix) { +static mjsElement* attachFrameToSite(mjCSite* parent, const mjCFrame* child, + const char* prefix, const char* suffix) { + mjSpec* spec = mjs_getSpec(parent->spec.element); + mjCBody* body = parent->Body(); + mjCFrame* frame = body->AddFrame(parent->frame); + frame->SetParent(body); + frame->spec.pos[0] = parent->spec.pos[0]; + frame->spec.pos[1] = parent->spec.pos[1]; + frame->spec.pos[2] = parent->spec.pos[2]; + frame->spec.quat[0] = parent->spec.quat[0]; + frame->spec.quat[1] = parent->spec.quat[1]; + frame->spec.quat[2] = parent->spec.quat[2]; + frame->spec.quat[3] = parent->spec.quat[3]; + mjs_resolveOrientation(frame->spec.quat, spec->compiler.degree, + spec->compiler.eulerseq, &parent->spec.alt); + + mjsElement* attached_frame = attachFrame(body, child, prefix, suffix); + mjs_setFrame(attached_frame, &frame->spec); + return attached_frame; +} + + +mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child, + const char* prefix, const char* suffix) { if (!parent) { - mju_error("parent site is null"); + mju_error("parent element is null"); return nullptr; } - mjSpec* spec = mjs_getSpec(parent->element); - mjCSite* site = static_cast(parent->element); - mjCBody* body = site->Body(); - mjCFrame* frame = body->AddFrame(site->frame); - frame->SetParent(body); - frame->spec.pos[0] = site->spec.pos[0]; - frame->spec.pos[1] = site->spec.pos[1]; - frame->spec.pos[2] = site->spec.pos[2]; - frame->spec.quat[0] = site->spec.quat[0]; - frame->spec.quat[1] = site->spec.quat[1]; - frame->spec.quat[2] = site->spec.quat[2]; - frame->spec.quat[3] = site->spec.quat[3]; - mjs_resolveOrientation(frame->spec.quat, spec->compiler.degree, - spec->compiler.eulerseq, &site->spec.alt); - - mjsFrame* attached_frame = mjs_attachFrame(&body->spec, child, prefix, suffix); - mjs_setFrame(attached_frame->element, &frame->spec); - return attached_frame; + if (!child) { + mju_error("child element is null"); + return nullptr; + } + mjCModel* model = static_cast(mjs_getSpec(parent)->element); + if (child->elemtype == mjOBJ_MODEL) { + mjCModel* child_model = static_cast((mjsElement*)child); + mjsBody* worldbody = mjs_findBody(&child_model->spec, "world"); + if (!worldbody) { + model->SetError(mjCError(0, "Child does not have a world body.")); + return nullptr; + } + mjsFrame* worldframe = mjs_addFrame(worldbody, nullptr); + SetFrame(worldbody, mjOBJ_BODY, worldframe); + SetFrame(worldbody, mjOBJ_SITE, worldframe); + SetFrame(worldbody, mjOBJ_FRAME, worldframe); + SetFrame(worldbody, mjOBJ_JOINT, worldframe); + SetFrame(worldbody, mjOBJ_GEOM, worldframe); + SetFrame(worldbody, mjOBJ_LIGHT, worldframe); + SetFrame(worldbody, mjOBJ_CAMERA, worldframe); + child = worldframe->element; + } + switch (parent->elemtype) { + case mjOBJ_FRAME: + if (child->elemtype == mjOBJ_BODY) { + return attachBody(static_cast(parent), + static_cast(child), prefix, suffix); + } else if (child->elemtype == mjOBJ_FRAME) { + mjsBody* parent_body = mjs_getParent(parent); + if (!parent_body) { + model->SetError(mjCError(0, "Frame does not have a parent body.")); + return nullptr; + } + mjCFrame* frame = static_cast(parent); + mjsElement* attached_frame = + attachFrame(static_cast(parent_body->element), + static_cast(child), prefix, suffix); + if (mjs_setFrame(attached_frame, &frame->spec)) { + return nullptr; + } + return attached_frame; + } else { + model->SetError(mjCError(0, "child element is not a body or frame")); + return nullptr; + } + case mjOBJ_BODY: + if (child->elemtype == mjOBJ_FRAME) { + return attachFrame(static_cast(parent), + static_cast(child), prefix, suffix); + } else { + model->SetError(mjCError(0, "child element is not a frame")); + return nullptr; + } + case mjOBJ_SITE: + if (child->elemtype == mjOBJ_BODY) { + return attachToSite(static_cast(parent), + static_cast(child), prefix, suffix); + } else if (child->elemtype == mjOBJ_FRAME) { + return attachFrameToSite(static_cast(parent), + static_cast(child), prefix, suffix); + } else { + model->SetError(mjCError(0, "child element is not a body or frame")); + return nullptr; + } + default: + model->SetError(mjCError(0, "parent element is not a frame, body or site")); + return nullptr; + } + return nullptr; } @@ -307,6 +375,14 @@ int mjs_setDeepCopy(mjSpec* s, int deepcopy) { +// copy real-valued arrays from model to spec, returns 1 on success +int mj_copyBack(mjSpec* s, const mjModel* m) { + mjCModel* model = static_cast(s->element); + return model->CopyBack(m); +} + + + // delete object, return 0 on success int mjs_delete(mjsElement* element) { mjCModel* model; @@ -758,13 +834,19 @@ mjsFrame* mjs_findFrame(mjSpec* s, const char* name) { // set frame -void mjs_setFrame(mjsElement* dest, mjsFrame* frame) { +int mjs_setFrame(mjsElement* dest, mjsFrame* frame) { if (!frame) { - return; + return -1; } mjCFrame* frameC = static_cast(frame->element); mjCBase* baseC = static_cast(dest); - baseC->SetFrame(frameC); + try { + baseC->SetFrame(frameC); + return 0; + } catch (mjCError& e) { + baseC->model->SetError(e); + return -1; + } } @@ -786,15 +868,17 @@ mjsFrame* mjs_bodyToFrame(mjsBody** body) { return &frameC->spec; } - - -// set user payload void mjs_setUserValue(mjsElement* element, const char* key, const void* data) { - mjCBase* baseC = static_cast(element); - baseC->SetUserValue(key, data); + mjs_setUserValueWithCleanup(element, key, data, nullptr); } - +// set user payload +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)) { + mjCBase* baseC = static_cast(element); + baseC->SetUserValue(key, data, cleanup); +} // return user payload or NULL if none found const void* mjs_getUserValue(mjsElement* element, const char* key) { @@ -1228,6 +1312,14 @@ void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes) { +// get plugin attributes +const void* mjs_getPluginAttributes(const mjsPlugin* plugin) { + mjCPlugin* pluginC = static_cast(plugin->element); + return &pluginC->config_attribs; +} + + + // -------------------------- GLOBAL ASSET CACHE ------------------------------- void mj_setCacheSize(mjCache cache, std::size_t size) { diff --git a/src/user/user_api.h b/src/user/user_api.h index 3c6c4420..04cdf13e 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -66,24 +66,15 @@ MJAPI int mjs_activatePlugin(mjSpec* s, const char* name); // Turn deep copy on or off attach. Returns 0 on success. MJAPI int mjs_setDeepCopy(mjSpec* s, int deepcopy); +// Copy real-valued arrays from model to spec, returns 1 on success. +MJAPI int mj_copyBack(mjSpec* s, const mjModel* m); + //---------------------------------- Attachment ---------------------------------------------------- -// Attach child body to a parent frame, return the attached body if success or NULL otherwise. -MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child, - const char* prefix, const char* suffix); - -// Attach child frame to a parent body, return the attached frame if success or NULL otherwise. -MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child, - const char* prefix, const char* suffix); - -// Attach child body to a parent site, return the attached body if success or NULL otherwise. -MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child, - const char* prefix, const char* suffix); - -// Attach child frame to a parent site, return the attached frame if success or NULL otherwise. -MJAPI mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child, - const char* prefix, const char* suffix); +// Attach child to a parent, return the attached element if success or NULL otherwise. +MJAPI mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child, + const char* prefix, const char* suffix); // Detach body from mjSpec, remove all references and delete the body, return 0 on success. MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b); @@ -364,14 +355,17 @@ MJAPI const char* mjs_getString(const mjString* source); // Get double array contents and optionally its size. MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size); +// Get plugin attributes. +MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin); + //---------------------------------- Other utilities ----------------------------------------------- // Set element's default. MJAPI void mjs_setDefault(mjsElement* element, const mjsDefault* def); -// Set element's enlcosing frame. -MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame); +// Set element's enclosing frame, return 0 on success. +MJAPI int mjs_setFrame(mjsElement* dest, mjsFrame* frame); // Resolve alternative orientations to quat, return error if any. MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence, @@ -383,6 +377,11 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/src/user/user_cache.cc b/src/user/user_cache.cc index 715e8e6a..ee22d7eb 100644 --- a/src/user/user_cache.cc +++ b/src/user/user_cache.cc @@ -102,7 +102,8 @@ bool mjCCache::Insert(const std::string& modelname, const mjResource *resource, -// populate data from the cache into the given function +// populate data from the cache into the given function, return true if data was +// copied bool mjCCache::PopulateData(const mjResource* resource, mjCDataFunc fn) { std::lock_guard lock(mutex_); auto it = lookup_.find(resource->name); @@ -121,8 +122,7 @@ bool mjCCache::PopulateData(const mjResource* resource, mjCDataFunc fn) { entries_.erase(asset); entries_.insert(asset); - asset->PopulateData(fn); - return true; + return asset->PopulateData(fn); } diff --git a/src/user/user_cache.h b/src/user/user_cache.h index a2d5577f..4b0e9288 100644 --- a/src/user/user_cache.h +++ b/src/user/user_cache.h @@ -28,7 +28,7 @@ #include -typedef std::function mjCDataFunc; +typedef std::function mjCDataFunc; typedef void (*mjCDeallocFunc)(const void*); // A class container for a thread-safe asset cache @@ -57,9 +57,9 @@ class mjCAsset { std::size_t InsertNum() const { return insert_num_; } std::size_t AccessCount() const { return access_count_; } - // pass data in the cache to the given function - void PopulateData(mjCDataFunc fn) const { - fn(data_.get()); + // pass data in the cache to the given function, return true if data was copied + bool PopulateData(mjCDataFunc fn) const { + return fn(data_.get()); } private: diff --git a/src/user/user_composite.cc b/src/user/user_composite.cc index 0c09164b..6cd21b51 100644 --- a/src/user/user_composite.cc +++ b/src/user/user_composite.cc @@ -12,25 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include #include "user/user_composite.h" -#include #include -#include -#include #include -#include -#include #include -#include #include -#include #include +#include #include #include "cc/array_safety.h" -#include "engine/engine_io.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" @@ -40,9 +32,9 @@ #include "user/user_util.h" namespace { + namespace mju = ::mujoco::util; -using mujoco::user::VectorToString; -using mujoco::user::StringToVector; + } // namespace // strncpy with 0, return false @@ -60,6 +52,8 @@ mjCComposite::mjCComposite(void) { type = mjCOMPTYPE_PARTICLE; count[0] = count[1] = count[2] = 1; mjuu_setvec(offset, 0, 0, 0); + mjuu_setvec(quat, 1, 0, 0, 0); + frame = nullptr; // plugin variables mjs_defaultPlugin(&plugin); @@ -267,19 +261,20 @@ bool mjCComposite::MakeCable(mjCModel* model, mjsBody* body, char* error, int er // populate uservert if not specified if (uservert.empty()) { for (int ix=0; ix < count[0]; ix++) { + double v[3]; for (int k=0; k < 3; k++) { switch (curve[k]) { case mjCOMPSHAPE_LINE: - uservert.push_back(ix*size[0]/(count[0]-1)); + v[k] = ix*size[0]/(count[0]-1); break; case mjCOMPSHAPE_COS: - uservert.push_back(size[1]*cos(mjPI*ix*size[2]/(count[0]-1))); + v[k] = size[1]*cos(mjPI*ix*size[2]/(count[0]-1)); break; case mjCOMPSHAPE_SIN: - uservert.push_back(size[1]*sin(mjPI*ix*size[2]/(count[0]-1))); + v[k] = size[1]*sin(mjPI*ix*size[2]/(count[0]-1)); break; case mjCOMPSHAPE_ZERO: - uservert.push_back(0); + v[k] = 0; break; default: // SHOULD NOT OCCUR @@ -287,6 +282,8 @@ bool mjCComposite::MakeCable(mjCModel* model, mjsBody* body, char* error, int er break; } } + mjuu_rotVecQuat(v, v, quat); + uservert.insert(uservert.end(), v, v+3); } } @@ -381,6 +378,9 @@ mjsBody* mjCComposite::AddCableBody(mjCModel* model, mjsBody* body, int ix, offset[1]+uservert[3*ix+1], offset[2]+uservert[3*ix+2]); mjuu_copyvec(body->quat, this_quat, 4); + if (frame) { + mjs_setFrame(body->element, frame); + } } else { mjuu_setvec(body->pos, length_prev, 0, 0); double negquat[4] = {prev_quat[0], -prev_quat[1], -prev_quat[2], -prev_quat[3]}; diff --git a/src/user/user_composite.h b/src/user/user_composite.h index 383dd3a2..17378004 100644 --- a/src/user/user_composite.h +++ b/src/user/user_composite.h @@ -73,13 +73,15 @@ class mjCComposite { std::string prefix; // name prefix mjtCompType type; // composite type int count[3]; // geom count in each dimension - double offset[3]; // position offset for particle and grid + double offset[3]; // position offset + double quat[4]; // quaternion offset // currently used only for cable std::string initial; // root boundary type std::vector uservert; // user-specified vertex positions double size[3]; // rope size (meaning depends on the shape) mjtCompShape curve[3]; // geometric shape + mjsFrame* frame; // frame where the composite is defined // body names used in the skin std::vector username; diff --git a/src/user/user_flexcomp.cc b/src/user/user_flexcomp.cc index 50c12d68..cb7124c1 100644 --- a/src/user/user_flexcomp.cc +++ b/src/user/user_flexcomp.cc @@ -33,6 +33,7 @@ #include "engine/engine_util_errmem.h" #include "user/user_flexcomp.h" #include +#include "user/user_api.h" #include "user/user_model.h" #include "user/user_objects.h" #include "user/user_resource.h" @@ -459,6 +460,14 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz) { // add new body at vertex coordinates mjsBody* pb = mjs_addBody(body, 0); + // add geom if vertcollide + if (dflex->vertcollide) { + mjsGeom* geom = mjs_addGeom(pb, 0); + geom->type = mjGEOM_SPHERE; + geom->size[0] = dflex->radius; + geom->group = 4; + } + // set frame and inertial pb->pos[0] = point[3*i]; pb->pos[1] = point[3*i+1]; @@ -544,6 +553,13 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz) { pb->inertia[2] = pb->mass*(2.0*inertiabox*inertiabox)/3.0; pb->explicitinertial = true; + // add geom if vertcollide + if (dflex->vertcollide) { + mjsGeom* geom = mjs_addGeom(pb, 0); + geom->type = mjGEOM_SPHERE; + geom->size[0] = dflex->radius; + } + for (int d=0; d < 3; d++) { mjsJoint* jnt = mjs_addJoint(pb, 0); jnt->type = mjJNT_SLIDE; @@ -1167,6 +1183,11 @@ void mjCFlexcomp::LoadGMSH41(char* buffer, int binary, int nodeend, throw mjCError(NULL, "All nodes must be in single block"); } + // require maximum number of nodes be equal to maximum number of nodes in a block + if (maxNodeTag != numNodesInBlock){ + throw mjCError(NULL, "Maximum number of nodes must be equal to number of nodes in a block"); + } + // check dimensionality and save if (entityDim < 1 || entityDim > 3) { throw mjCError(NULL, "Entity must be 1D, 2D or 3D"); diff --git a/src/user/user_init.c b/src/user/user_init.c index 76ea52ad..34199812 100644 --- a/src/user/user_init.c +++ b/src/user/user_init.c @@ -42,6 +42,7 @@ void mjs_defaultSpec(mjSpec* spec) { spec->compiler.usethread = 1; spec->compiler.inertiafromgeom = mjINERTIAFROMGEOM_AUTO; spec->compiler.inertiagrouprange[1] = mjNGROUP-1; + spec->compiler.saveinertial = 0; mj_defaultLROpt(&spec->compiler.LRopt); // engine data @@ -202,6 +203,8 @@ void mjs_defaultLight(mjsLight* light) { // intrinsics light->castshadow = 1; light->bulbradius = 0.02; + light->intensity = 0.0; + light->range = 10.0; light->active = 1; light->attenuation[0] = 1; light->cutoff = 45; @@ -270,6 +273,7 @@ void mjs_defaultSkin(mjsSkin* skin) { void mjs_defaultTexture(mjsTexture* texture) { memset(texture, 0, sizeof(mjsTexture)); texture->type = mjTEXTURE_CUBE; + texture->colorspace = mjCOLORSPACE_AUTO; texture->rgb1[0] = texture->rgb1[1] = texture->rgb1[2] = 0.8; texture->rgb2[0] = texture->rgb2[1] = texture->rgb2[2] = 0.5; texture->random = 0.01; @@ -403,4 +407,3 @@ void mjs_defaultKey(mjsKey* key) { void mjs_defaultPlugin(mjsPlugin* plugin) { memset(plugin, 0, sizeof(mjsPlugin)); } - diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index f999fc9a..03ad956c 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -1012,30 +1012,28 @@ void mjCMesh::LoadOBJ(mjResource* resource, bool remove_repeated) { // load mesh from cached asset, return true on success bool mjCMesh::LoadCachedMesh(mjCCache *cache, const mjResource* resource) { - // save previous mesh properties (in case different from cached mesh) - int maxhullvert = maxhullvert_; - mjtMeshInertia old_inertia = inertia; - double old_scale[3] = {scale[0], scale[1], scale[2]}; - auto process_mesh = [&](const void* data) { const mjCMesh* mesh = static_cast(data); // check if maxhullvert is different - maxhullvert_ = mesh->maxhullvert_; - if (maxhullvert != mesh->maxhullvert_) { - return; + if (maxhullvert_ != mesh->maxhullvert_) { + return false; } // check if inertia is different - inertia = mesh->inertia; - if (old_inertia != mesh->inertia) { - return; + if (inertia != mesh->inertia) { + return false; } // check if scale is different - memcpy(scale, mesh->scale, 3*sizeof(double)); - if (old_scale[0] != mesh->scale[0] || old_scale[1] != mesh->scale[1] || - old_scale[2] != mesh->scale[2]) { - return; + if (scale[0] != mesh->scale[0] || + scale[1] != mesh->scale[1] || + scale[2] != mesh->scale[2]) { + return false; + } + + // check if need hull + if (needhull_ && !mesh->szgraph_) { + return false; } processed_ = mesh->processed_; @@ -1047,11 +1045,14 @@ bool mjCMesh::LoadCachedMesh(mjCCache *cache, const mjResource* resource) { facetexcoord_ = mesh->facetexcoord_; halfedge_ = mesh->halfedge_; - szgraph_ = mesh->szgraph_; - graph_ = nullptr; - if (szgraph_) { - graph_ = (int*)mju_malloc(szgraph_*sizeof(int)); - std::copy(mesh->graph_, mesh->graph_ + szgraph_, graph_); + // only copy graph if needed + if (needhull_ || mesh->face_.empty()) { + szgraph_ = mesh->szgraph_; + graph_ = nullptr; + if (szgraph_) { + graph_ = (int*)mju_malloc(szgraph_*sizeof(int)); + std::copy(mesh->graph_, mesh->graph_ + szgraph_, graph_); + } } polygons_ = mesh->polygons_; @@ -1072,29 +1073,11 @@ bool mjCMesh::LoadCachedMesh(mjCCache *cache, const mjResource* resource) { } tree_ = mesh->tree_; face_aabb_ = mesh->face_aabb_; + return true; }; - // check that cached asset has all data, make sure no metadata has changed - if (!cache->PopulateData(resource, process_mesh)) { - return false; - } - - if (maxhullvert != maxhullvert_) { - maxhullvert_ = maxhullvert; - return false; - } - if (inertia != old_inertia) { - inertia = old_inertia; - return false; - } - if (scale[0] != old_scale[0] || scale[1] != old_scale[1] || - scale[2] != old_scale[2]) { - scale[0] = old_scale[0]; - scale[1] = old_scale[1]; - scale[2] = old_scale[2]; - return false; - } - return true; + // check that cached asset has all data + return cache->PopulateData(resource, process_mesh); } // load STL binary mesh @@ -3035,6 +3018,113 @@ void inline ComputeStiffness(std::vector& stiffness, MetricTensor(stiffness.data(), t, mu, la, basis); } +// local tetrahedron numbering +constexpr int kNumEdges = Stencil2D::kNumEdges; +constexpr int kNumVerts = Stencil2D::kNumVerts; +constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}}; + +// create map from triangles to vertices and edges and from edges to vertices +static void CreateFlapStencil(std::vector& flaps, + const std::vector& simplex, + const std::vector& edgeidx) { + // populate stencil + int ne = 0; + int nt = simplex.size() / kNumVerts; + std::vector elements(nt); + for (int t = 0; t < nt; t++) { + for (int v = 0; v < kNumVerts; v++) { + elements[t].vertices[v] = simplex[kNumVerts * t + v]; + } + } + + // map from edge vertices to their index in `edges` vector + std::unordered_map, int, PairHash> edge_indices; + + // loop over all triangles + for (int t = 0; t < nt; t++) { + int* v = elements[t].vertices; + + // compute edges to vertices map for fast computations + for (int e = 0; e < kNumEdges; e++) { + auto pair = std::pair(std::min(v[edge[e][0]], v[edge[e][1]]), + std::max(v[edge[e][0]], v[edge[e][1]])); + + // if edge is already present in the vector only store its index + auto [it, inserted] = edge_indices.insert({pair, ne}); + + if (inserted) { + StencilFlap flap; + flap.vertices[0] = v[edge[e][0]]; + flap.vertices[1] = v[edge[e][1]]; + flap.vertices[2] = v[(edge[e][1] + 1) % 3]; + flap.vertices[3] = -1; + flaps.push_back(flap); + elements[t].edges[e] = ne++; + } else { + elements[t].edges[e] = it->second; + flaps[it->second].vertices[3] = v[(edge[e][1] + 1) % 3]; + } + + // double check that the edge indices are consistent + if (!edgeidx.empty()) { + if (elements[t].edges[e] != edgeidx[kNumEdges * t + e]) { + mju_error("edge indices do not match in CreateFlapStencil"); + } + } + } + } +} + +// cotangent between two edges +double inline cot(double* x, int v0, int v1, int v2) { + double normal[3]; + double edge1[3] = {x[3*v1]-x[3*v0], x[3*v1+1]-x[3*v0+1], x[3*v1+2]-x[3*v0+2]}; + double edge2[3] = {x[3*v2]-x[3*v0], x[3*v2+1]-x[3*v0+1], x[3*v2+2]-x[3*v0+2]}; + + mjuu_crossvec(normal, edge1, edge2); + return mjuu_dot3(edge1, edge2) / sqrt(mjuu_dot3(normal, normal)); +} + +// area of a triangle +double inline ComputeVolume(const double* x, const int v[Stencil2D::kNumVerts]) { + double normal[3]; + double edge1[3] = {x[3*v[1]]-x[3*v[0]], x[3*v[1]+1]-x[3*v[0]+1], x[3*v[1]+2]-x[3*v[0]+2]}; + double edge2[3] = {x[3*v[2]]-x[3*v[0]], x[3*v[2]+1]-x[3*v[0]+1], x[3*v[2]+2]-x[3*v[0]+2]}; + + mjuu_crossvec(normal, edge1, edge2); + return sqrt(mjuu_dot3(normal, normal)) / 2; +} + +// compute bending stiffness for a single edge +template +void inline ComputeBending(double* bending, double* pos, const int v[4], double mu, + double thickness) { + int vadj[3] = {v[1], v[0], v[3]}; + + if (v[3]== -1) { + // skip boundary edges + return; + } + + // cotangent operator from Wardetzky at al., "Discrete Quadratic Curvature + // Energies", https://cims.nyu.edu/gcl/papers/wardetzky2007dqb.pdf + + mjtNum a01 = cot(pos, v[0], v[1], v[2]); + mjtNum a02 = cot(pos, v[0], v[3], v[1]); + mjtNum a03 = cot(pos, v[1], v[2], v[0]); + mjtNum a04 = cot(pos, v[1], v[0], v[3]); + mjtNum c[4] = {a03 + a04, a01 + a02, -(a01 + a03), -(a02 + a04)}; + mjtNum volume = ComputeVolume(pos, v) + + ComputeVolume(pos, vadj); + + for (int v1 = 0; v1 < T::kNumVerts; v1++) { + for (int v2 = 0; v2 < T::kNumVerts; v2++) { + bending[4 * v1 + v2] += + 1.5 * c[v1] * c[v2] / volume * mu * pow(thickness, 3) / 12; + } + } +} + //----------------------------- linear elasticity -------------------------------------------------- // Gauss Legendre quadrature points in 1 dimension on the interval [a, b] @@ -3292,9 +3382,13 @@ void mjCFlex::ResolveReferences(const mjCModel* m) { vertbodyid.clear(); nodebodyid.clear(); for (const auto& vertbody : vertbody_) { - mjCBase* pbody = m->FindObject(mjOBJ_BODY, vertbody); + mjCBody* pbody = static_cast(m->FindObject(mjOBJ_BODY, vertbody)); if (pbody) { vertbodyid.push_back(pbody->id); + if (pbody->joints.size() != 3 && dim == 2 && (elastic2d == 1 || elastic2d == 3)) { + // TODO(quaglino): add support for pins + throw mjCError(this, "pins are not supported for bending"); + } } else { throw mjCError(this, "unknown body '%s' in flex", vertbody.c_str()); } @@ -3530,11 +3624,18 @@ void mjCFlex::Compile(const mjVFS* vfs) { // set size nedge = (int)edge.size(); + // create flap stencil + if (dim == 2) { + CreateFlapStencil(flaps, elem_, edgeidx_); + } + // compute elasticity if (young > 0) { if (poisson < 0 || poisson >= 0.5) { throw mjCError(this, "Poisson ratio must be in [0, 0.5)"); } + + // linear elasticity stiffness.assign(21*nelem, 0); if (interpolated) { int min_size = ceil(nodexpos.size()*nodexpos.size() / 21); @@ -3543,11 +3644,13 @@ void mjCFlex::Compile(const mjVFS* vfs) { } ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson); } + + // geometrically nonlinear elasticity for (unsigned int t = 0; t < nelem; t++) { if (interpolated) { continue; } - if (dim == 2) { + if (dim == 2 && elastic2d >= 2 && thickness > 0) { ComputeStiffness(stiffness, vertxpos, elem_.data() + (dim + 1) * t, t, young, poisson, thickness); @@ -3557,13 +3660,22 @@ void mjCFlex::Compile(const mjVFS* vfs) { poisson); } } + + // bending stiffness (2D only) + if (dim == 2 && (elastic2d == 1 || elastic2d == 3)) { + if (thickness < 0) { + throw mjCError(this, "thickness must be positive for bending stiffness"); + } + bending.assign(nedge*16, 0); + + for (unsigned int e = 0; e < nedge; e++) { + ComputeBending(bending.data() + 16 * e, vertxpos.data(), flaps[e].vertices, + young / (2 * (1 + poisson)), thickness); + } + } } - // add plugins - std::string userface, useredge; - userface = VectorToString(elem_); - useredge = VectorToString(edgeidx_); - + // placeholder for setting plugins parameters, currently not used for (const auto& vbodyid : vertbodyid) { if (vbodyid < 0) { continue; @@ -3571,11 +3683,9 @@ void mjCFlex::Compile(const mjVFS* vfs) { if (model->Bodies()[vbodyid]->plugin.element) { mjCPlugin* plugin_instance = static_cast(model->Bodies()[vbodyid]->plugin.element); - if (damping > 0) { - plugin_instance->config_attribs["damping"] = std::to_string(damping); + if (!plugin_instance) { + throw mjCError(this, "plugin instance not found"); } - plugin_instance->config_attribs["face"] = userface; - plugin_instance->config_attribs["edge"] = useredge; } } @@ -3641,8 +3751,9 @@ void mjCFlex::CreateBVH() { elemaabb_[6*e+5] = 0.5*(xmax[2]-xmin[2]) + radius; // add bounding volume for this element + // contype and conaffinity are set to nonzero to force bvh generation const double* aabb = elemaabb_.data() + 6*e; - tree.AddBoundingVolume(e, contype, conaffinity, aabb, nullptr, aabb); + tree.AddBoundingVolume(e, 1, 1, aabb, nullptr, aabb); nbvh++; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index bc12a17e..8d4e2b5a 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -166,7 +166,7 @@ static void processlist(mjListKeyMap& ids, vector& list, // constructor mjCModel::mjCModel() { mjs_defaultSpec(&spec); - elemtype = mjOBJ_UNKNOWN; + elemtype = mjOBJ_MODEL; spec_comment_.clear(); spec_modelfiledir_.clear(); spec_meshdir_.clear(); @@ -211,15 +211,14 @@ mjCModel::mjCModel() { // create mjCBase lists from children lists CreateObjectLists(); - // the source spec is the model itself, overwritten in the copy constructor - source_spec_ = &spec; + // set the signature + spec.element->signature = 0; } mjCModel::mjCModel(const mjCModel& other) { CreateObjectLists(); - source_spec_ = (mjSpec*)&other.spec; *this = other; } @@ -235,12 +234,17 @@ mjCModel& mjCModel::operator=(const mjCModel& other) { // copy attached specs first so that we can resolve references to them for (const auto* s : other.specs_) { specs_.push_back(mj_copySpec(s)); + compiler2spec_[&s->compiler] = specs_.back(); } // the world copy constructor takes care of copying the tree mjCBody* world = new mjCBody(*other.bodies_[0], this); bodies_.push_back(world); + // update tree lists + ResetTreeLists(); + MakeTreeLists(); + // add everything else *this += other; @@ -255,6 +259,9 @@ mjCModel& mjCModel::operator=(const mjCModel& other) { for (int i=0; i < mjNOBJECT; i++) { ids[i] = other.ids[i]; } + + // update signature after we updated everything + spec.element->signature = Signature(); } deepcopy_ = other.deepcopy_; return *this; @@ -449,9 +456,8 @@ static bool IsPluginActive( mjCModel& mjCModel::operator+=(const mjCModel& other) { // create global lists - mjCBody *world = bodies_[0]; ResetTreeLists(); - MakeLists(world); + MakeTreeLists(); ProcessLists(/*checkrepeat=*/false); // copy all elements not in the tree @@ -500,12 +506,11 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) { nq = nv = na = nu = nmocap = 0; } - // restore to the original state - if (!compiled) { - ResetTreeLists(); - } - + // update pointers to local elements PointToLocal(); + + // update signature after we updated the tree lists and we updated the pointers + spec.element->signature = Signature(); return *this; } @@ -551,7 +556,7 @@ void mjCModel::RemoveFromList(std::vector& list, const mjCModel& other) { template <> void mjCModel::DeleteAll(std::vector& elements) { for (mjCKey* element : elements) { - delete element; + element->Release(); } elements.clear(); } @@ -609,13 +614,11 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) { // create global lists in the old model if not compiled if (!oldmodel.IsCompiled()) { - oldmodel.MakeLists(oldmodel.bodies_[0]); oldmodel.ProcessLists(/*checkrepeat=*/false); } // create global lists in this model if not compiled if (!IsCompiled()) { - MakeLists(bodies_[0]); ProcessLists(/*checkrepeat=*/false); } @@ -629,7 +632,7 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) { // update global lists ResetTreeLists(); - MakeLists(world); + MakeTreeLists(); ProcessLists(/*checkrepeat=*/false); // check if we have to remove anything else @@ -641,10 +644,8 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) { RemoveFromList(sensors_, oldmodel); RemovePlugins(); - // restore to the original state - if (!compiled) { - ResetTreeLists(); - } + // update signature before we reset the tree lists + spec.element->signature = Signature(); return *this; } @@ -742,18 +743,16 @@ void deletefromlist(std::vector* list, mjsElement* element) { // discard all invalid elements from all lists void mjCModel::DeleteElement(mjsElement* el) { - mjCBody *world = nullptr; - if (compiled) { - world = bodies_[0]; - ResetTreeLists(); - } + ResetTreeLists(); switch (el->elemtype) { case mjOBJ_BODY: + MakeTreeLists(); // rebuild lists that were reset at the beginning of the function throw mjCError(nullptr, "bodies cannot be deleted, use detach instead"); break; case mjOBJ_DEFAULT: + MakeTreeLists(); // rebuild lists that were reset at the beginning of the function throw mjCError(nullptr, "defaults cannot be deleted, use detach instead"); break; @@ -818,11 +817,12 @@ void mjCModel::DeleteElement(mjsElement* el) { break; } - if (compiled) { - ResetTreeLists(); // in case of a nested delete - MakeLists(world); - ProcessLists(/*checkrepeat=*/false); - } + ResetTreeLists(); // in case of a nested delete + MakeTreeLists(); + ProcessLists(/*checkrepeat=*/false); + + // update signature after we updated everything + spec.element->signature = Signature(); } @@ -1020,15 +1020,6 @@ void mjCModel::Clear() { nconmax = -1; nmocap = 0; - // pointer lists created by Compile - bodies_.clear(); - joints_.clear(); - geoms_.clear(); - sites_.clear(); - cameras_.clear(); - lights_.clear(); - frames_.clear(); - // internal variables hasImplicitPluginElem = false; compiled = false; @@ -1046,6 +1037,7 @@ T* mjCModel::AddObject(vector& list, string type) { T* obj = new T(this); obj->id = (int)list.size(); list.push_back(obj); + spec.element->signature = Signature(); return obj; } @@ -1057,6 +1049,7 @@ T* mjCModel::AddObjectDefault(vector& list, string type, mjCDef* def) { obj->id = (int)list.size(); obj->classname = def ? def->name : "main"; list.push_back(obj); + spec.element->signature = Signature(); return obj; } @@ -1165,9 +1158,13 @@ mjCPlugin* mjCModel::AddPlugin() { // append spec to spec -void mjCModel::AppendSpec(mjSpec* spec) { +void mjCModel::AppendSpec(mjSpec* spec, const mjsCompiler* compiler_) { // TODO: check if the spec is already in the list specs_.push_back(spec); + + if (compiler_) { + compiler2spec_[compiler_] = spec; + } } @@ -1461,10 +1458,15 @@ mjSpec* mjCModel::FindSpec(std::string name) const { // find spec by mjsCompiler pointer -mjSpec* mjCModel::FindSpec(const mjsCompiler* compiler_) const { - if (&GetSourceSpec()->compiler == compiler_) { - return (mjSpec*)&spec; +mjSpec* mjCModel::FindSpec(const mjsCompiler* compiler_) { + if (compiler_ == &spec.compiler) { + return &spec; } + + if (compiler2spec_.find(compiler_) != compiler2spec_.end()) { + return compiler2spec_[compiler_]; + } + for (auto s : specs_) { mjSpec* source = static_cast(s->element)->FindSpec(compiler_); if (source) { @@ -1476,17 +1478,14 @@ mjSpec* mjCModel::FindSpec(const mjsCompiler* compiler_) const { -// get the spec from which this model was created -mjSpec* mjCModel::GetSourceSpec() const { - return source_spec_; -} - - - //------------------------------- COMPILER PHASES -------------------------------------------------- // make lists of objects in tree: bodies, geoms, joints, sites, cameras, lights -void mjCModel::MakeLists(mjCBody* body) { +void mjCModel::MakeTreeLists(mjCBody* body) { + if (body == nullptr) { + body = bodies_[0]; + } + // add this body if not world if (body != bodies_[0]) { bodies_.push_back(body); @@ -1501,7 +1500,7 @@ void mjCModel::MakeLists(mjCBody* body) { for (mjCFrame *frame : body->frames) frames_.push_back(frame); // recursive call to all child bodies - for (mjCBody* body : body->bodies) MakeLists(body); + for (mjCBody* body : body->bodies) MakeTreeLists(body); } @@ -1550,7 +1549,7 @@ static void DeleteElements(std::vector& elements, int i = 0; for (int j=0; j < elements.size(); j++) { if (discard[j]) { - delete elements[j]; + elements[j]->Release(); } else { elements[i] = elements[j]; i++; @@ -1608,7 +1607,7 @@ void mjCModel::DeleteAll(std::vector& elements) { DeleteMaterial(sites_); DeleteMaterial(tendons_); for (mjCMaterial* element : elements) { - delete element; + element->Release(); } elements.clear(); } @@ -1618,7 +1617,7 @@ template <> void mjCModel::DeleteAll(std::vector& elements) { DeleteAllTextures(materials_); for (mjCTexture* element : elements) { - delete element; + element->Release(); } elements.clear(); } @@ -1778,7 +1777,6 @@ void mjCModel::IndexAssets(bool discard) { } } - // discard visual meshes and geoms if (discard) { std::vector discard_mesh(meshes_.size(), false); std::vector discard_geom(geoms_.size(), false); @@ -1792,6 +1790,28 @@ void mjCModel::IndexAssets(bool discard) { return geom->IsVisual(); }); + // update inertia in bodies + for (auto body : bodies_) { + if (body->spec.explicitinertial) { + continue; + } + for (auto geom : body->geoms) { + if (geom->IsVisual()) { + if (compiler.inertiafromgeom == mjINERTIAFROMGEOM_TRUE) { + compiler.inertiafromgeom = mjINERTIAFROMGEOM_AUTO; + } + body->explicitinertial = true; // for XML writer + body->spec.explicitinertial = true; + body->spec.mass = body->mass; + mjuu_copyvec(body->spec.ipos, body->ipos, 3); + mjuu_copyvec(body->spec.iquat, body->iquat, 4); + mjuu_copyvec(body->spec.inertia, body->inertia, 3); + break; + } + } + } + + // discard visual meshes and geoms Delete(meshes_, discard_mesh); Delete(geoms_, discard_geom); } @@ -2649,12 +2669,15 @@ void mjCModel::CopyTree(mjModel* m) { m->light_bodyid[lid] = pl->body->id; m->light_mode[lid] = (int)pl->mode; m->light_targetbodyid[lid] = pl->targetbodyid; - m->light_directional[lid] = (mjtByte)pl->directional; + m->light_type[lid] = pl->type; + m->light_texid[lid] = pl->texid; m->light_castshadow[lid] = (mjtByte)pl->castshadow; m->light_active[lid] = (mjtByte)pl->active; mjuu_copyvec(m->light_pos+3*lid, pl->pos, 3); mjuu_copyvec(m->light_dir+3*lid, pl->dir, 3); m->light_bulbradius[lid] = pl->bulbradius; + m->light_intensity[lid] = pl->intensity; + m->light_range[lid] = pl->range; mjuu_copyvec(m->light_attenuation+3*lid, pl->attenuation, 3); m->light_cutoff[lid] = pl->cutoff; m->light_exponent[lid] = pl->exponent; @@ -2745,31 +2768,6 @@ void mjCModel::CopyTree(mjModel* m) { } } m->nB = nB; - - // set dof_simplenum - int count = 0; - for (int i=nv-1; i >= 0; i--) { - if (m->body_simple[m->dof_bodyid[i]]) { - count++; // increment counter - } else { - count = 0; // reset - } - m->dof_simplenum[i] = count; - } - - // compute nC - int nOD = 0; // number of off-diagonal (non-simple) parent dofs - for (int i=0; i < nv; i++) { - // count ancestor (off-diagonal) dofs - if (!m->dof_simplenum[i]) { - int j = i; - while (j >= 0) { - if (j != i) nOD++; - j = m->dof_parentid[j]; - } - } - } - m->nC = nC = nOD + nv; } // copy plugin data @@ -3046,6 +3044,11 @@ void mjCModel::CopyObjects(mjModel* m) { } else { mjuu_zerovec(m->flex_stiffness + 21 * elem_adr, 21 * pfl->nelem); } + if (!pfl->bending.empty()) { + mjuu_copyvec(m->flex_bending + 16 * edge_adr, pfl->bending.data(), pfl->bending.size()); + } else { + mjuu_zerovec(m->flex_bending + 16 * edge_adr, 16 * pfl->nedge); + } m->flex_damping[i] = (mjtNum)pfl->damping; // set fields: mesh-like @@ -3163,6 +3166,13 @@ void mjCModel::CopyObjects(mjModel* m) { for (int k=0; k < pfl->nedge; k++) { m->flex_edge[2*(edge_adr+k)] = pfl->edge[k].first; m->flex_edge[2*(edge_adr+k)+1] = pfl->edge[k].second; + if (pfl->dim == 2 && (pfl->elastic2d == 1 || pfl->elastic2d == 3)) { + m->flex_edgeflap[2*(edge_adr+k)+0] = pfl->flaps[k].vertices[2]; + m->flex_edgeflap[2*(edge_adr+k)+1] = pfl->flaps[k].vertices[3]; + } else { + m->flex_edgeflap[2*(edge_adr+k)+0] = -1; + m->flex_edgeflap[2*(edge_adr+k)+1] = -1; + } if (pfl->rigid) { m->flexedge_rigid[edge_adr+k] = 1; @@ -3276,6 +3286,7 @@ void mjCModel::CopyObjects(mjModel* m) { // set fields m->tex_type[i] = ptex->type; + m->tex_colorspace[i] = ptex->colorspace; m->tex_height[i] = ptex->height; m->tex_width[i] = ptex->width; m->tex_nchannel[i] = ptex->nchannel; @@ -3356,6 +3367,7 @@ void mjCModel::CopyObjects(mjModel* m) { m->tendon_matid[i] = pte->matid; m->tendon_group[i] = pte->group; m->tendon_limited[i] = (mjtByte)pte->is_limited(); + m->tendon_actfrclimited[i] = (mjtByte)pte->is_actfrclimited(); m->tendon_width[i] = (mjtNum)pte->width; mjuu_copyvec(m->tendon_solref_lim+mjNREF*i, pte->solref_limit, mjNREF); mjuu_copyvec(m->tendon_solimp_lim+mjNIMP*i, pte->solimp_limit, mjNIMP); @@ -3363,9 +3375,12 @@ void mjCModel::CopyObjects(mjModel* m) { mjuu_copyvec(m->tendon_solimp_fri+mjNIMP*i, pte->solimp_friction, mjNIMP); m->tendon_range[2*i] = (mjtNum)pte->range[0]; m->tendon_range[2*i+1] = (mjtNum)pte->range[1]; + m->tendon_actfrcrange[2*i] = (mjtNum)pte->actfrcrange[0]; + m->tendon_actfrcrange[2*i+1] = (mjtNum)pte->actfrcrange[1]; m->tendon_margin[i] = (mjtNum)pte->margin; m->tendon_stiffness[i] = (mjtNum)pte->stiffness; m->tendon_damping[i] = (mjtNum)pte->damping; + m->tendon_armature[i] = (mjtNum)pte->armature; m->tendon_frictionloss[i] = (mjtNum)pte->frictionloss; m->tendon_lengthspring[2*i] = (mjtNum)pte->springlength[0]; m->tendon_lengthspring[2*i+1] = (mjtNum)pte->springlength[1]; @@ -3539,6 +3554,54 @@ void mjCModel::CopyObjects(mjModel* m) { +// finalize simple bodies/dofs including tendon information +void mjCModel::FinalizeSimple(mjModel* m) { + // demote bodies affected by inertia-bearing tendon to non-simple + for (int i=0; i < ntendon; i++) { + if (m->tendon_armature[i] == 0) { + continue; + } + int adr = m->tendon_adr[i]; + int num = m->tendon_num[i]; + for (int j=adr; j < adr+num; j++) { + int objid = m->wrap_objid[j]; + if (m->wrap_type[j] == mjWRAP_SITE) { + m->body_simple[m->site_bodyid[objid]] = 0; + } + if (m->wrap_type[j] == mjWRAP_CYLINDER || m->wrap_type[j] == mjWRAP_SPHERE) { + m->body_simple[m->geom_bodyid[objid]] = 0; + } + } + } + + // set dof_simplenum + int count = 0; + for (int i=nv-1; i >= 0; i--) { + if (m->body_simple[m->dof_bodyid[i]]) { + count++; // increment counter + } else { + count = 0; // reset + } + m->dof_simplenum[i] = count; + } + + // compute nC + int nOD = 0; // number of off-diagonal (non-simple) parent dofs + for (int i=0; i < nv; i++) { + // count ancestor (off-diagonal) dofs + if (!m->dof_simplenum[i]) { + int j = i; + while (j >= 0) { + if (j != i) nOD++; + j = m->dof_parentid[j]; + } + } + } + m->nC = nC = nOD + nv; +} + + + // save the current state template void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act, @@ -3567,7 +3630,7 @@ void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* } for (auto body : bodies_) { - if (!body->spec.mocap) { + if (!body->spec.mocap || body->mocapid == -1) { continue; } if (mpos) { @@ -3657,21 +3720,12 @@ template void mjCModel::RestoreState( // resolve keyframe references void mjCModel::StoreKeyframes(mjCModel* dest) { - bool resetlists = false; - if (this != dest && !key_pending_.empty()) { mju_warning( "Child model has pending keyframes. They will not be namespaced correctly. " "To prevent this, compile the child model before attaching it again."); } - // create tree lists if they are empty, occurs if an uncompiled model is attached - if (bodies_.size() == 1 && geoms_.empty() && sites_.empty() && joints_.empty() && - cameras_.empty() && lights_.empty() && frames_.empty()) { - MakeLists(bodies_[0]); - resetlists = true; - } - // do not change compilation quantities in case the user wants to recompile preserving the state if (!compiled) { SaveDofOffsets(/*computesize=*/true); @@ -3719,10 +3773,6 @@ void mjCModel::StoreKeyframes(mjCModel* dest) { key->spec_mpos_.data(), key->spec_mquat_.data()); } - if (resetlists) { - ResetTreeLists(); - } - if (!compiled) { nq = nv = na = nu = nmocap = 0; } @@ -3802,78 +3852,8 @@ void mjCModel::FuseStatic(void) { } //------------- add mass and inertia (if parent not world) - if (body->parent && body->parent->name != "world" && body->mass >= mjMINVAL) { - // body_ipose = body_pose * body_ipose - changeframe(body->ipos, body->iquat, body->pos, body->quat); - - // organize data - double mass[2] = { - par->mass, - body->mass - }; - double inertia[2][3] = { - {par->inertia[0], par->inertia[1], par->inertia[2]}, - {body->inertia[0], body->inertia[1], body->inertia[2]} - }; - double ipos[2][3] = { - {par->ipos[0], par->ipos[1], par->ipos[2]}, - {body->ipos[0], body->ipos[1], body->ipos[2]} - }; - double iquat[2][4] = { - {par->iquat[0], par->iquat[1], par->iquat[2], par->iquat[3]}, - {body->iquat[0], body->iquat[1], body->iquat[2], body->iquat[3]} - }; - - // compute total mass - par->mass = 0; - mjuu_setvec(par->ipos, 0, 0, 0); - for (int j=0; j < 2; j++) { - par->mass += mass[j]; - par->ipos[0] += mass[j]*ipos[j][0]; - par->ipos[1] += mass[j]*ipos[j][1]; - par->ipos[2] += mass[j]*ipos[j][2]; - } - - // small mass: allow for now, check for errors later - if (par->mass < mjMINVAL) { - par->mass = 0; - mjuu_setvec(par->inertia, 0, 0, 0); - mjuu_setvec(par->ipos, 0, 0, 0); - mjuu_setvec(par->iquat, 1, 0, 0, 0); - } - - // proceed with regular computation - else { - // locipos = center-of-mass - par->ipos[0] /= par->mass; - par->ipos[1] /= par->mass; - par->ipos[2] /= par->mass; - - // add inertias - double toti[6] = {0, 0, 0, 0, 0, 0}; - for (int j=0; j < 2; j++) { - double inertA[6], inertB[6]; - double dpos[3] = { - ipos[j][0] - par->ipos[0], - ipos[j][1] - par->ipos[1], - ipos[j][2] - par->ipos[2] - }; - - mjuu_globalinertia(inertA, inertia[j], iquat[j]); - mjuu_offcenter(inertB, mass[j], dpos); - for (int k=0; k < 6; k++) { - toti[k] += inertA[k] + inertB[k]; - } - } - - // compute principal axes of inertia - mjuu_copyvec(par->fullinertia, toti, 6); - const char* err1 = mjuu_fullInertia(par->iquat, par->inertia, par->fullinertia); - if (err1) { - throw mjCError(nullptr, "error '%s' in fusing static body inertias", err1); - } - } + par->AccumulateInertia(body); } //------------- replace body with its children in parent body list @@ -4054,11 +4034,7 @@ static void warninghandler(const char* msg) { // compiler mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) { if (compiled) { - // clear kinematic tree - mjCBody* world = bodies_[0]; - ResetTreeLists(); Clear(); - bodies_.push_back(world); } CopyFromSpec(); @@ -4105,9 +4081,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) { // deallocate everything allocated in Compile mj_deleteModel(model); mj_deleteData(data); - mjCBody* world = bodies_[0]; Clear(); - bodies_.push_back(world); // save error info errInfo = err; @@ -4354,14 +4328,14 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { AddKey(); } - // make lists of objects created in kinematic tree - MakeLists(bodies_[0]); - // clear subtreedofs for (int i=0; i < bodies_.size(); i++) { bodies_[i]->subtreedofs = 0; } + // initialize spec signature (needed if the user changed sensor or joint types) + spec.element->signature = Signature(); + // fill missing names and check that they are all filled for (const auto& asset : meshes_) asset->CopyFromSpec(); for (const auto& asset : skins_) asset->CopyFromSpec(); @@ -4506,6 +4480,9 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // copy objects outsite kinematic tree (including keyframes) CopyObjects(m); + // finalize simple bodies/dofs including tendon information + FinalizeSimple(m); + // compute non-zeros in actuator_moment m->nJmom = nJmom = CountNJmom(m); @@ -4616,6 +4593,96 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { mju::strcpy_arr(errInfo.message, warningtext); errInfo.warning = true; } + + // save signature + m->signature = Signature(); + + // special cases that are not caused by user edits + if (compiler.fusestatic || compiler.discardvisual || + !pairs_.empty() || !excludes_.empty()) { + spec.element->signature = m->signature; + } + + // check that the signature matches the spec + if (m->signature != spec.element->signature) { + throw mjCError(0, "signature mismatch"); // SHOULD NOT OCCUR + } +} + + + +std::string mjCModel::PrintTree(const mjCBody* body, std::string indent) { + std::string tree; + tree += indent + "\n"; + indent += " "; + for (const auto& joint : body->joints) { + tree += indent + "" + std::to_string(joint->nq()) + "\n"; + } + for (uint64_t i = 0; i < body->geoms.size(); ++i) { + tree += indent + "\n"; + } + for (uint64_t i = 0; i < body->sites.size(); ++i) { + tree += indent + "\n"; + } + for (uint64_t i = 0; i < body->cameras.size(); ++i) { + tree += indent + "\n"; + } + for (uint64_t i = 0; i < body->lights.size(); ++i) { + tree += indent + "\n"; + } + for (uint64_t i = 0; i < body->bodies.size(); ++i) { + tree += PrintTree(body->bodies[i], indent); + } + indent.pop_back(); + indent.pop_back(); + tree += indent + "\n"; + return tree; +} + + + +uint64_t mjCModel::Signature() { + std::string tree = "\n" + PrintTree(bodies_[0]); + for (unsigned int i = 0; i < flexes_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < meshes_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < skins_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < hfields_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < textures_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < materials_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < pairs_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < excludes_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 1; i < equalities_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < tendons_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < actuators_.size(); ++i) { + tree += "\n"; + } + for (unsigned int i = 0; i < sensors_.size(); ++i) { + tree += "" + std::to_string(sensors_[i]->spec.type) + "\n"; + } + for (unsigned int i = 0; i < keys_.size(); ++i) { + tree += "\n"; + } + return mj_hashString(tree.c_str(), UINT64_MAX); } @@ -4692,6 +4759,11 @@ bool mjCModel::CopyBack(const mjModel* m) { return false; } + if (spec.element->signature != m->signature) { + errInfo = mjCError(0, "incompatible signatures in CopyBack"); + return false; + } + // option and visual option = m->opt; visual = m->vis; @@ -4887,6 +4959,7 @@ bool mjCModel::CopyBack(const mjModel* m) { // tendons for (int i=0; i < ntendon; i++) { mjuu_copyvec(tendons_[i]->range, m->tendon_range+2*i, 2); + mjuu_copyvec(tendons_[i]->actfrcrange, m->tendon_actfrcrange+2*i, 2); mjuu_copyvec(tendons_[i]->solref_limit, m->tendon_solref_lim+mjNREF*i, mjNREF); mjuu_copyvec(tendons_[i]->solimp_limit, m->tendon_solimp_lim+mjNIMP*i, mjNIMP); mjuu_copyvec(tendons_[i]->solref_friction, m->tendon_solref_fri+mjNREF*i, mjNREF); @@ -4896,6 +4969,7 @@ bool mjCModel::CopyBack(const mjModel* m) { tendons_[i]->margin = (double)m->tendon_margin[i]; tendons_[i]->stiffness = (double)m->tendon_stiffness[i]; tendons_[i]->damping = (double)m->tendon_damping[i]; + tendons_[i]->armature = (double)m->tendon_armature[i]; tendons_[i]->frictionloss = (double)m->tendon_frictionloss[i]; if (nuser_tendon) { diff --git a/src/user/user_model.h b/src/user/user_model.h index 1078080f..cfd8082e 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -16,6 +16,7 @@ #define MUJOCO_SRC_USER_USER_MODEL_H_ #include +#include #include #include #include @@ -214,7 +215,9 @@ class mjCModel : public mjCModel_, private mjSpec { mjCTuple* AddTuple(); mjCKey* AddKey(); mjCPlugin* AddPlugin(); - void AppendSpec(mjSpec* spec); + + // append spec to this model, optionally map compiler options to the appended spec + void AppendSpec(mjSpec* spec, const mjsCompiler* compiler = nullptr); // delete elements marked as discard=true template void Delete(std::vector& elements, @@ -247,7 +250,7 @@ class mjCModel : public mjCModel_, private mjSpec { mjCBase* FindObject(mjtObj type, std::string name) const; // find object given type and name mjCBase* FindTree(mjCBody* body, mjtObj type, std::string name); // find tree object given name mjSpec* FindSpec(std::string name) const; // find spec given name - mjSpec* FindSpec(const mjsCompiler* compiler_) const; // find spec given mjsCompiler + mjSpec* FindSpec(const mjsCompiler* compiler_); // find spec given mjsCompiler void ActivatePlugin(const mjpPlugin* plugin, int slot); // activate plugin // accessors @@ -315,9 +318,6 @@ class mjCModel : public mjCModel_, private mjSpec { // map from default class name to default class pointer std::unordered_map def_map; - // get the spec from which this model was created - mjSpec* GetSourceSpec() const; - // set deepcopy flag void SetDeepCopy(bool deepcopy) { deepcopy_ = deepcopy; } @@ -328,15 +328,14 @@ class mjCModel : public mjCModel_, private mjSpec { // settings for each defaults class std::vector defaults_; - // spec from which this model was created in copy constructor - mjSpec* source_spec_; - // list of active plugins std::vector> active_plugins_; + // make lists of bodies and children + void MakeTreeLists(mjCBody* body = nullptr); + // compile phases void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs); - void MakeLists(mjCBody* body); // make lists of bodies, geoms, joints, sites void SetNuser(); // set nuser fields void IndexAssets(bool discard); // convert asset names into indices void CheckEmptyNames(); // check empty names @@ -347,6 +346,7 @@ class mjCModel : public mjCModel_, private mjSpec { void CopyPaths(mjModel*); // copy paths, compute path addresses void CopyObjects(mjModel*); // copy objects outside kinematic tree void CopyTree(mjModel*); // copy objects inside kinematic tree + void FinalizeSimple(mjModel* m); // finalize simple bodies/dofs including tendon information void CopyPlugins(mjModel*); // copy plugin data int CountNJmom(const mjModel* m); // compute number of non-zeros in actuator_moment matrix @@ -438,11 +438,17 @@ class mjCModel : public mjCModel_, private mjSpec { void MarkPluginInstance(std::unordered_map& instances, const std::vector& list); + // print the tree of a body + std::string PrintTree(const mjCBody* body, std::string indent = ""); + + // generate a signature for the model + uint64_t Signature(); mjListKeyMap ids; // map from object names to ids mjCError errInfo; // last error info std::vector key_pending_; // attached keyframes bool deepcopy_; // copy objects when attaching bool attached_ = false; // true if model is attached to a parent model + std::unordered_map compiler2spec_; // map from compiler to spec }; #endif // MUJOCO_SRC_USER_USER_MODEL_H_ diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 18c573d2..6a1b2451 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -55,6 +55,8 @@ class PNGImage { LodePNGColorType color_type); int Width() const { return width_; } int Height() const { return height_; } + bool IsSRGB() const { return is_srgb_; } + uint8_t operator[] (int i) const { return data_[i]; } std::vector& MoveData() { return data_; } @@ -65,6 +67,7 @@ class PNGImage { int width_; int height_; + bool is_srgb_; LodePNGColorType color_type_; std::vector data_; }; @@ -75,14 +78,19 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource, image.color_type_ = color_type; mjCCache *cache = reinterpret_cast(mj_globalCache()); + // cache callback + auto callback = [&image](const void* data) { + const PNGImage *cached_image = static_cast(data); + if (cached_image->color_type_ == image.color_type_) { + image = *cached_image; + return true; + } + return false; + }; + // try loading from cache - if (cache && cache->PopulateData(resource, [&image](const void* data) { - const PNGImage *cached_image = static_cast(data); - if (cached_image->color_type_ == image.color_type_) { - image = *cached_image; - } - })) { - if (!image.data_.empty()) return image; + if (cache && cache->PopulateData(resource, callback)) { + return image; } // open PNG resource @@ -99,8 +107,11 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource, // decode PNG from buffer unsigned int w, h; - unsigned err = lodepng::decode(image.data_, w, h, - buffer, nbuffer, image.color_type_, 8); + + lodepng::State state; + state.info_raw.colortype = image.color_type_; + state.info_raw.bitdepth = 8; + unsigned err = lodepng::decode(image.data_, w, h, state, buffer, nbuffer); // check for errors if (err) { @@ -111,6 +122,7 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource, image.width_ = w; image.height_ = h; + image.is_srgb_ = (state.info_png.srgb_defined == 1); if (image.width_ <= 0 || image.height_ < 0) { std::stringstream ss; @@ -773,18 +785,20 @@ void mjCBase::SetFrame(mjCFrame* _frame) { if (!_frame) { return; } + if (_frame->body && GetParent() != _frame->body) { + throw mjCError(this, "Frame and body '%s' have mismatched parents", name.c_str()); + } frame = _frame; } - -void mjCBase::SetUserValue(std::string_view key, const void* data) { - user_payload_[std::string(key)] = data; +void mjCBase::SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)) { + user_payload_[std::string(key)] = UserValue(data, cleanup); } - const void* mjCBase::GetUserValue(std::string_view key) { auto found = user_payload_.find(std::string(key)); - return found != user_payload_.end() ? found->second : nullptr; + return found != user_payload_.end() ? found->second.value : nullptr; } @@ -815,6 +829,7 @@ mjCBody::mjCBody(mjCModel* _model) { mjuu_zerovec(xpos0, 3); mjuu_setvec(xquat0, 1, 0, 0, 0); last_attached = nullptr; + mocapid = -1; // clear object lists bodies.clear(); @@ -890,8 +905,18 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) { for (int i=0; i < other.bodies.size(); i++) { bodies.push_back(new mjCBody(*other.bodies[i], model)); // triggers recursive call bodies.back()->parent = this; - bodies.back()->frame = - other.bodies[i]->frame ? frames[fmap[other.bodies[i]->frame]] : nullptr; + bodies.back()->frame = nullptr; + if (other.bodies[i]->frame) { + if (fmap.find(other.bodies[i]->frame) != fmap.end()) { + bodies.back()->frame = frames[fmap[other.bodies[i]->frame]]; + } else { + throw mjCError(this, "Frame '%s' not found in other body", + other.bodies[i]->frame->name.c_str()); + } + if (bodies.back()->frame && bodies.back()->frame->body != this) { + throw mjCError(this, "Frame and body '%s' have mismatched parents", name.c_str()); + } + } } return *this; @@ -902,8 +927,8 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) { // attach frame to body mjCBody& mjCBody::operator+=(const mjCFrame& other) { // append a copy of the attached spec - if (other.model != model && !model->FindSpec(mjs_getString(other.model->spec.modelname))) { - model->AppendSpec(mj_copySpec(&other.model->spec)); + if (other.model != model && !model->FindSpec(&other.model->spec.compiler)) { + model->AppendSpec(mj_copySpec(&other.model->spec), &other.model->spec.compiler); } // create a copy of the subtree that contains the frame @@ -1231,7 +1256,15 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) { obj->classname = _def ? _def->name : classname; bodies.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + obj->parent = this; + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1241,6 +1274,11 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) { mjCFrame* mjCBody::AddFrame(mjCFrame* _frame) { mjCFrame* obj = new mjCFrame(model, _frame ? _frame : NULL); frames.push_back(obj); + model->ResetTreeLists(); + model->MakeTreeLists(); + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1256,6 +1294,14 @@ mjCJoint* mjCBody::AddFreeJoint() { obj->body = this; joints.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1270,6 +1316,14 @@ mjCJoint* mjCBody::AddJoint(mjCDef* _def) { obj->body = this; joints.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1284,6 +1338,14 @@ mjCGeom* mjCBody::AddGeom(mjCDef* _def) { obj->body = this; geoms.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1298,6 +1360,14 @@ mjCSite* mjCBody::AddSite(mjCDef* _def) { obj->body = this; sites.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1312,6 +1382,14 @@ mjCCamera* mjCBody::AddCamera(mjCDef* _def) { obj->body = this; cameras.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1326,6 +1404,14 @@ mjCLight* mjCBody::AddLight(mjCDef* _def) { obj->body = this; lights.push_back(obj); + + // recompute lists + model->ResetTreeLists(); + model->MakeTreeLists(); + + + // update signature + model->spec.element->signature = model->Signature(); return obj; } @@ -1336,6 +1422,15 @@ mjCFrame* mjCBody::ToFrame() { mjCFrame* newframe = parent->AddFrame(frame); mjuu_copyvec(newframe->spec.pos, spec.pos, 3); mjuu_copyvec(newframe->spec.quat, spec.quat, 4); + if (parent->name != "world" && mass >= mjMINVAL) { + if (!parent->explicitinertial) { + parent->MakeInertialExplicit(); + mjuu_zerovec(parent->spec.ipos, 3); + mjuu_zerovec(parent->spec.iquat, 4); + mjuu_zerovec(parent->spec.inertia, 3); + } + parent->AccumulateInertia(&this->spec, &parent->spec); + } MapFrame(parent->bodies, bodies, newframe, parent); MapFrame(parent->geoms, geoms, newframe, parent); MapFrame(parent->joints, joints, newframe, parent); @@ -1347,11 +1442,9 @@ mjCFrame* mjCBody::ToFrame() { std::remove_if(parent->bodies.begin(), parent->bodies.end(), [this](mjCBody* body) { return body == this; }), parent->bodies.end()); - if (model->IsCompiled()) { - mjCBody *world = model->bodies_[0]; - model->ResetTreeLists(); - model->MakeLists(world); - } + model->ResetTreeLists(); + model->MakeTreeLists(); + model->spec.element->signature = model->Signature(); return newframe; } @@ -1637,6 +1730,90 @@ void mjCBody::MakeInertialExplicit() { +// accumulate inertia of another body into this body +void mjCBody::AccumulateInertia(const mjsBody* other, mjsBody* result) { + if (!result) { + result = this; // use the private mjsBody + } + + // body_ipose = body_pose * body_ipose + double other_ipos[3]; + double other_iquat[4]; + mjuu_copyvec(other_ipos, other->ipos, 3); + mjuu_copyvec(other_iquat, other->iquat, 4); + mjuu_frameaccum(other_ipos, other_iquat, other->pos, other->quat); + + // organize data + double mass[2] = { + result->mass, + other->mass + }; + double inertia[2][3] = { + {result->inertia[0], result->inertia[1], result->inertia[2]}, + {other->inertia[0], other->inertia[1], other->inertia[2]} + }; + double ipos[2][3] = { + {result->ipos[0], result->ipos[1], result->ipos[2]}, + {other_ipos[0], other_ipos[1], other_ipos[2]} + }; + double iquat[2][4] = { + {result->iquat[0], result->iquat[1], result->iquat[2], result->iquat[3]}, + {other->iquat[0], other->iquat[1], other->iquat[2], other->iquat[3]} + }; + + // compute total mass + result->mass = 0; + mjuu_setvec(result->ipos, 0, 0, 0); + for (int j=0; j < 2; j++) { + result->mass += mass[j]; + result->ipos[0] += mass[j]*ipos[j][0]; + result->ipos[1] += mass[j]*ipos[j][1]; + result->ipos[2] += mass[j]*ipos[j][2]; + } + + // small mass: allow for now, check for errors later + if (result->mass < mjMINVAL) { + result->mass = 0; + mjuu_setvec(result->inertia, 0, 0, 0); + mjuu_setvec(result->ipos, 0, 0, 0); + mjuu_setvec(result->iquat, 1, 0, 0, 0); + } + + // proceed with regular computation + else { + // locipos = center-of-mass + result->ipos[0] /= result->mass; + result->ipos[1] /= result->mass; + result->ipos[2] /= result->mass; + + // add inertias + double toti[6] = {0, 0, 0, 0, 0, 0}; + for (int j=0; j < 2; j++) { + double inertA[6], inertB[6]; + double dpos[3] = { + ipos[j][0] - result->ipos[0], + ipos[j][1] - result->ipos[1], + ipos[j][2] - result->ipos[2] + }; + + mjuu_globalinertia(inertA, inertia[j], iquat[j]); + mjuu_offcenter(inertB, mass[j], dpos); + for (int k=0; k < 6; k++) { + toti[k] += inertA[k] + inertB[k]; + } + } + + // compute principal axes of inertia + mjuu_copyvec(result->fullinertia, toti, 6); + const char* err1 = mjuu_fullInertia(result->iquat, result->inertia, result->fullinertia); + if (err1) { + throw mjCError(nullptr, "error '%s' in fusing static body inertias", err1); + } + } +} + + + // compute bounding volume hierarchy void mjCBody::ComputeBVH() { if (geoms.empty()) { @@ -1892,16 +2069,6 @@ void mjCBody::Compile(void) { } } - // if discarding visual geoms, use explicit inertias - if (compiler->discardvisual) { - for (int j=0; j < geoms.size(); j++) { - if (geoms[j]->IsVisual()) { - explicitinertial = true; - break; - } - } - } - // free joint alignment, phase 2 (transform sites, cameras and lights) if (align_free) { // frames have already been compiled and applied to children @@ -1966,8 +2133,8 @@ mjCFrame& mjCFrame::operator=(const mjCFrame& other) { // attach body to frame mjCFrame& mjCFrame::operator+=(const mjCBody& other) { // append a copy of the attached spec - if (other.model != model && !model->FindSpec(mjs_getString(other.model->spec.modelname))) { - model->AppendSpec(mj_copySpec(&other.model->spec)); + if (other.model != model && !model->FindSpec(&other.model->spec.compiler)) { + model->AppendSpec(mj_copySpec(&other.model->spec), &other.model->spec.compiler); } // apply namespace and store keyframes in the source model @@ -3479,7 +3646,10 @@ mjCLight::mjCLight(mjCModel* _model, mjCDef* _def) { // clear private variables body = 0; targetbodyid = -1; + texid = -1; spec_targetbody_.clear(); + spec_texture_.clear(); + // reset to default if given if (_def) { @@ -3519,6 +3689,7 @@ void mjCLight::PointToLocal() { spec.element = static_cast(this); spec.name = &name; spec.targetbody = &spec_targetbody_; + spec.texture = &spec_texture_; spec.info = &info; targetbody = nullptr; } @@ -3530,6 +3701,9 @@ void mjCLight::NameSpace(const mjCModel* m) { if (!spec_targetbody_.empty()) { spec_targetbody_ = m->prefix + spec_targetbody_ + m->suffix; } + if (!spec_texture_.empty()) { + spec_texture_ = m->prefix + spec_texture_ + m->suffix; + } } @@ -3537,6 +3711,7 @@ void mjCLight::NameSpace(const mjCModel* m) { void mjCLight::CopyFromSpec() { *static_cast(this) = spec; targetbody_ = spec_targetbody_; + texture_ = spec_texture_; } @@ -3569,6 +3744,16 @@ void mjCLight::Compile(void) { throw mjCError(this, "unknown target body in light"); } } + + // get texture + if (!texture_.empty()) { + mjCTexture* tex = (mjCTexture*)model->FindObject(mjOBJ_TEXTURE, texture_); + if (tex) { + texid = tex->id; + } else { + throw mjCError(this, "unknown target body in light"); + } + } } @@ -4190,7 +4375,7 @@ void mjCTexture::BuiltinCube(void) { // load PNG file void mjCTexture::LoadPNG(mjResource* resource, std::vector& image, - unsigned int& w, unsigned int& h) { + unsigned int& w, unsigned int& h, bool& is_srgb) { LodePNGColorType color_type; if (nchannel == 4) { color_type = LCT_RGBA; @@ -4205,13 +4390,36 @@ void mjCTexture::LoadPNG(mjResource* resource, PNGImage png_image = PNGImage::Load(this, resource, color_type); w = png_image.Width(); h = png_image.Height(); + is_srgb = png_image.IsSRGB(); image = png_image.MoveData(); } +// load KTX file +void mjCTexture::LoadKTX(mjResource* resource, + std::vector& image, unsigned int& w, + unsigned int& h, bool& is_srgb) { + const void* buffer = 0; + int buffer_sz = mju_readResource(resource, &buffer); + + // still not found + if (buffer_sz < 0) { + throw mjCError(this, "could not read texture file '%s'", resource->name); + } else if (!buffer_sz) { + throw mjCError(this, "texture file is empty: '%s'", resource->name); + } + + w = buffer_sz; + h = 1; + is_srgb = false; + + image.resize(buffer_sz); + memcpy(image.data(), buffer, buffer_sz); +} + // load custom file void mjCTexture::LoadCustom(mjResource* resource, std::vector& image, - unsigned int& w, unsigned int& h) { + unsigned int& w, unsigned int& h, bool& is_srgb) { const void* buffer = 0; int buffer_sz = mju_readResource(resource, &buffer); @@ -4228,6 +4436,9 @@ void mjCTexture::LoadCustom(mjResource* resource, w = pint[0]; h = pint[1]; + // assume linear color space + is_srgb = false; + // check dimensions if (w < 1 || h < 1) { throw mjCError(this, "Non-PNG texture, assuming custom binary file format,\n" @@ -4250,7 +4461,7 @@ void mjCTexture::LoadCustom(mjResource* resource, // load from PNG or custom file, flip if specified void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs, std::vector& image, - unsigned int& w, unsigned int& h) { + unsigned int& w, unsigned int& h, bool& is_srgb) { std::string asset_type = GetAssetContentType(filename, content_type_); // fallback to custom @@ -4258,7 +4469,7 @@ void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs, asset_type = "image/vnd.mujoco.texture"; } - if (asset_type != "image/png" && asset_type != "image/vnd.mujoco.texture") { + if (asset_type != "image/png" && asset_type != "image/ktx" && asset_type != "image/vnd.mujoco.texture") { throw mjCError(this, "unsupported content type: '%s'", asset_type.c_str()); } @@ -4266,9 +4477,14 @@ void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs, try { if (asset_type == "image/png") { - LoadPNG(resource, image, w, h); + LoadPNG(resource, image, w, h, is_srgb); + } else if (asset_type == "image/ktx") { + if (hflip || vflip) { + throw mjCError(this, "cannot flip KTX textures"); + } + LoadKTX(resource, image, w, h, is_srgb); } else { - LoadCustom(resource, image, w, h); + LoadCustom(resource, image, w, h, is_srgb); } mju_closeResource(resource); } catch(mjCError err) { @@ -4335,12 +4551,16 @@ void mjCTexture::LoadFlip(std::string filename, const mjVFS* vfs, void mjCTexture::Load2D(std::string filename, const mjVFS* vfs) { // load PNG or custom unsigned int w, h; + bool is_srgb; std::vector image; - LoadFlip(filename, vfs, image, w, h); + LoadFlip(filename, vfs, image, w, h, is_srgb); // assign size width = w; height = h; + if (colorspace == mjCOLORSPACE_AUTO) { + colorspace = is_srgb ? mjCOLORSPACE_SRGB : mjCOLORSPACE_LINEAR; + } // allocate and copy data std::int64_t size = static_cast(width)*height; @@ -4368,8 +4588,13 @@ void mjCTexture::LoadCubeSingle(std::string filename, const mjVFS* vfs) { // load PNG or custom unsigned int w, h; + bool is_srgb; std::vector image; - LoadFlip(filename, vfs, image, w, h); + LoadFlip(filename, vfs, image, w, h, is_srgb); + + if (colorspace == mjCOLORSPACE_AUTO) { + colorspace = is_srgb ? mjCOLORSPACE_SRGB : mjCOLORSPACE_LINEAR; + } // check gridsize for compatibility if (w/gridsize[1] != h/gridsize[0] || (w%gridsize[1]) || (h%gridsize[0])) { @@ -4482,8 +4707,14 @@ void mjCTexture::LoadCubeSeparate(const mjVFS* vfs) { // load PNG or custom unsigned int w, h; + bool is_srgb; std::vector image; - LoadFlip(filename.Str(), vfs, image, w, h); + LoadFlip(filename.Str(), vfs, image, w, h, is_srgb); + + // assume all faces have the same colorspace + if (colorspace == mjCOLORSPACE_AUTO) { + colorspace = is_srgb ? mjCOLORSPACE_SRGB : mjCOLORSPACE_LINEAR; + } // PNG must be square if (w != h) { @@ -5325,7 +5556,9 @@ mjCTendon& mjCTendon::operator=(const mjCTendon& other) { bool mjCTendon::is_limited() const { return islimited(limited, range); } - +bool mjCTendon::is_actfrclimited() const { + return islimited(actfrclimited, actfrcrange); +} void mjCTendon::PointToLocal() { spec.element = static_cast(this); @@ -5542,6 +5775,12 @@ void mjCTendon::Compile(void) { // spatial path else { + if (armature < 0) { + throw mjCError(this, + "tendon '%s' (id = %d): tendon armature cannot be negative", + name.c_str(), id); + } + switch (path[i]->type) { case mjWRAP_PULLEY: // pulley should not follow other pulley @@ -5583,6 +5822,12 @@ void mjCTendon::Compile(void) { name.c_str(), id, i); } + if (armature > 0) { + throw mjCError(this, + "tendon '%s' (id = %d): geom wrapping not supported by tendon armature", + name.c_str(), id); + } + // mark geoms as non visual model->Geoms()[path[i]->obj->id]->SetNotVisual(); break; @@ -5611,6 +5856,21 @@ void mjCTendon::Compile(void) { throw mjCError(this, "invalid limits in tendon"); } + // if limited is auto, set to 1 if range is specified, otherwise unlimited + if (actfrclimited == mjLIMITED_AUTO) { + bool hasactfrcrange = !(actfrcrange[0] == 0 && actfrcrange[1] == 0); + checklimited(this, compiler->autolimits, "tendon", "", actfrclimited, + hasactfrcrange); + } + + // check actfrclimits + if (actfrcrange[0] >= actfrcrange[1] && is_actfrclimited()) { + throw mjCError(this, "invalid actuatorfrcrange in tendon"); + } + if ((actfrcrange[0] > 0 || actfrcrange[1] < 0) && is_actfrclimited()) { + throw mjCError(this, "invalid actuatorfrcrange in tendon"); + } + // check springlength if (springlength[0] > springlength[1]) { throw mjCError(this, "invalid springlength in tendon"); @@ -6404,6 +6664,18 @@ void mjCSensor::Compile(void) { } break; + case mjSENS_TENDONACTFRC: + // must be attached to tendon + if (objtype != mjOBJ_TENDON) { + throw mjCError(this, "sensor must be attached to tendon"); + } + + // set + dim = 1; + datatype = mjDATATYPE_REAL; + needstage = mjSTAGE_ACC; + break; + case mjSENS_TENDONPOS: case mjSENS_TENDONVEL: // must be attached to tendon @@ -7167,6 +7439,8 @@ mjCPlugin::mjCPlugin(mjCModel* _model) { spec.plugin_name = &plugin_name; spec.name = &name; spec.info = &info; + + PointToLocal(); } @@ -7185,11 +7459,20 @@ mjCPlugin& mjCPlugin::operator=(const mjCPlugin& other) { parent = this; plugin_slot = other.plugin_slot; } + PointToLocal(); return *this; } +void mjCPlugin::PointToLocal() { + spec.element = static_cast(this); + spec.name = &name; + spec.info = &info; +} + + + // compiler void mjCPlugin::Compile(void) { mjCPlugin* plugin_instance = this; diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 47703182..18158c5c 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -253,6 +253,9 @@ class mjCBase : public mjCBase_ { // Copy plugins instantiated in this object virtual void CopyPlugin() {} + // Returns parent of this object + virtual mjCBase* GetParent() const { return nullptr; } + // Copy assignment mjCBase& operator=(const mjCBase& other); @@ -278,7 +281,8 @@ class mjCBase : public mjCBase_ { } // Set and get user payload - void SetUserValue(std::string_view key, const void* data); + void SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)); const void* GetUserValue(std::string_view key); void DeleteUserValue(std::string_view key); @@ -289,8 +293,44 @@ class mjCBase : public mjCBase_ { // reference count for allowing deleting an attached object int refcount = 1; + // Arbitrary user value that cleans up the data when destroyed. + struct UserValue { + const void* value = nullptr; + void (*cleanup)(const void*) = nullptr; + + UserValue() {} + UserValue(const void* value, void (*cleanup)(const void*)) + : value(value), cleanup(cleanup) {} + UserValue(const UserValue& other) = delete; + UserValue& operator=(const UserValue& other) = delete; + + UserValue(UserValue&& other) : value(other.value), cleanup(other.cleanup) { + other.value = nullptr; + other.cleanup = nullptr; + } + + UserValue& operator=(UserValue&& other) { + if (this != &other) { + if (cleanup && value) { + cleanup(value); + } + value = other.value; + cleanup = other.cleanup; + other.value = nullptr; + other.cleanup = nullptr; + } + return *this; + } + + ~UserValue() { + if (cleanup && value) { + cleanup(value); + } + } + }; + // user payload - std::unordered_map user_payload_; + std::unordered_map user_payload_; }; @@ -422,6 +462,10 @@ class mjCBody : public mjCBody_, private mjsBody { // getters std::vector Bodies() const { return bodies; } + // accumulate inertia of another body into this body, if `result` is not nullptr, the accumulated + // inertia will be stored in `result`, otherwise the body's private spec will be used. + void AccumulateInertia(const mjsBody* other, mjsBody* result = nullptr); + private: mjCBody(const mjCBody& other, mjCModel* _model); // copy constructor mjCBody& operator=(const mjCBody& other); // copy assignment @@ -753,6 +797,9 @@ class mjCLight_ : public mjCBase { protected: mjCBody* body; // light's body int targetbodyid; // id of target body; -1: none + int texid; // id of texture; -1: none + std::string texture_; + std::string spec_texture_; std::string targetbody_; std::string spec_targetbody_; }; @@ -774,6 +821,7 @@ class mjCLight : public mjCLight_, private mjsLight { // used by mjXWriter and mjCModel const std::string& get_targetbody() const { return targetbody_; } + const std::string& get_texture() const { return texture_; } void SetParent(mjCBody* _body) { body = _body; } mjCBody* GetParent() const { return body; } @@ -790,6 +838,11 @@ class mjCLight : public mjCLight_, private mjsLight { //------------------------- class mjCFlex ---------------------------------------------------------- // Describes a flex +struct StencilFlap { + static constexpr int kNumVerts = 4; + int vertices[kNumVerts]; +}; + class mjCFlex_ : public mjCBase { protected: int nvert; // number of vertices @@ -806,11 +859,13 @@ class mjCFlex_ : public mjCBase { std::vector shell; // shell fragment vertex ids (dim per fragment) std::vector elemlayer; // element layer (distance from border) std::vector evpair; // element-vertex pairs + std::vector flaps; // adjacent triangles std::vector vertxpos; // global vertex positions mjCBoundingVolumeHierarchy tree; // bounding volume hierarchy std::vector elemaabb_; // element bounding volume std::vector edgeidx_; // element edge ids std::vector stiffness; // elasticity stiffness matrix + std::vector bending; // bending stiffness matrix // variable-size data std::vector vertbody_; // vertex body names @@ -1281,14 +1336,17 @@ class mjCTexture : public mjCTexture_, private mjsTexture { void LoadFlip(std::string filename, const mjVFS* vfs, // load and flip std::vector& image, - unsigned int& w, unsigned int& h); + unsigned int& w, unsigned int& h, bool& is_srgb); void LoadPNG(mjResource* resource, std::vector& image, - unsigned int& w, unsigned int& h); + unsigned int& w, unsigned int& h, bool& is_srgb); + void LoadKTX(mjResource* resource, + std::vector& image, + unsigned int& w, unsigned int& h, bool& is_srgb); void LoadCustom(mjResource* resource, std::vector& image, - unsigned int& w, unsigned int& h); + unsigned int& w, unsigned int& h, bool& is_srgb); bool clear_data_; // if true, data_ is empty and should be filled by Compile }; @@ -1520,6 +1578,7 @@ class mjCTendon : public mjCTendon_, private mjsTendon { void SetModel(mjCModel* _model); bool is_limited() const; + bool is_actfrclimited() const; private: void Compile(void); // compiler @@ -1583,6 +1642,9 @@ class mjCPlugin : public mjCPlugin_ { mjCPlugin(mjCModel*); mjCPlugin(const mjCPlugin& other); mjCPlugin& operator=(const mjCPlugin& other); + + void PointToLocal(); + mjsPlugin spec; mjCBase* parent; // parent object (only used when generating error message) int plugin_slot; // global registered slot number of the plugin diff --git a/src/xml/xml.cc b/src/xml/xml.cc index 5277a534..57b1e123 100644 --- a/src/xml/xml.cc +++ b/src/xml/xml.cc @@ -395,7 +395,7 @@ mjSpec* ParseSpecFromString(std::string_view xml, const mjVFS* vfs, char* error, } // Main writer function - calls mjXWrite -std::string WriteXML(const mjModel* m, const mjSpec* spec, char* error, int nerror) { +std::string WriteXML(const mjModel* m, mjSpec* spec, char* error, int nerror) { LocaleOverride locale_override; // check for empty model diff --git a/src/xml/xml.h b/src/xml/xml.h index 44b9b23e..5f85d016 100644 --- a/src/xml/xml.h +++ b/src/xml/xml.h @@ -30,6 +30,6 @@ mjSpec* ParseSpecFromString(std::string_view xml, const mjVFS* vfs = nullptr, char* error = nullptr, int nerror = 0); // Main writer function -std::string WriteXML(const mjModel* m, const mjSpec* spec, char* error, int nerror); +std::string WriteXML(const mjModel* m, mjSpec* spec, char* error, int nerror); #endif // MUJOCO_SRC_XML_XML_H_ diff --git a/src/xml/xml_api.cc b/src/xml/xml_api.cc index bde10787..b289adb3 100644 --- a/src/xml/xml_api.cc +++ b/src/xml/xml_api.cc @@ -230,7 +230,8 @@ mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int er // save spec to XML file, return 0 on success, -1 otherwise int mj_saveXML(const mjSpec* s, const char* filename, char* error, int error_sz) { - std::string result = WriteXML(NULL, s, error, error_sz); + // cast to mjSpec since WriteXML can in principle perform mj_copyBack (not here) + std::string result = WriteXML(NULL, (mjSpec*)s, error, error_sz); if (result.empty()) { return -1; } @@ -247,7 +248,7 @@ int mj_saveXML(const mjSpec* s, const char* filename, char* error, int error_sz) // save spec to XML string, return 0 on success, -1 on failure // if length of the output buffer is too small, returns the required size int mj_saveXMLString(const mjSpec* s, char* xml, int xml_sz, char* error, int error_sz) { - std::string result = WriteXML(NULL, s, error, error_sz); + std::string result = WriteXML(NULL, (mjSpec*)s, error, error_sz); if (result.empty()) { return -1; } else if (result.size() >= xml_sz) { diff --git a/src/xml/xml_base.cc b/src/xml/xml_base.cc index 1a1dd1b9..914d99de 100644 --- a/src/xml/xml_base.cc +++ b/src/xml/xml_base.cc @@ -44,8 +44,8 @@ mjXBase::mjXBase() { // set model field -void mjXBase::SetModel(const mjSpec* _model, const mjModel* m) { - spec = (mjSpec*)_model; +void mjXBase::SetModel(mjSpec* _model, const mjModel* m) { + spec = _model; } diff --git a/src/xml/xml_base.h b/src/xml/xml_base.h index ba914229..9dc261b1 100644 --- a/src/xml/xml_base.h +++ b/src/xml/xml_base.h @@ -27,7 +27,7 @@ // keyword maps (defined in implementation files) extern const int joint_sz; extern const int camlight_sz; -extern const int light_sz; +extern const int lighttype_sz; extern const int integrator_sz; extern const int collision_sz; extern const int cone_sz; @@ -35,6 +35,7 @@ extern const int jac_sz; extern const int solver_sz; extern const int equality_sz; extern const int texture_sz; +extern const int colorspace_sz; extern const int builtin_sz; extern const int mark_sz; extern const int dyn_sz; @@ -50,7 +51,7 @@ extern const mjMap TFAuto_map[]; extern const mjMap joint_map[]; extern const mjMap geom_map[]; extern const mjMap camlight_map[]; -extern const mjMap light_map[]; +extern const mjMap lighttype_map[]; extern const mjMap integrator_map[]; extern const mjMap collision_map[]; extern const mjMap impedance_map[]; @@ -60,6 +61,7 @@ extern const mjMap jac_map[]; extern const mjMap solver_map[]; extern const mjMap equality_map[]; extern const mjMap texture_map[]; +extern const mjMap colorspace_map[]; extern const mjMap texrole_map[]; extern const mjMap builtin_map[]; extern const mjMap mark_map[]; @@ -71,6 +73,7 @@ extern const mjMap datatype_map[]; extern const mjMap meshtype_map[]; extern const mjMap meshinertia_map[]; extern const mjMap flexself_map[]; +extern const mjMap elastic2d_map[]; //---------------------------------- Base XML class ------------------------------------------------ @@ -89,7 +92,7 @@ class mjXBase : public mjXUtil { }; // set the model allocated externally - virtual void SetModel(const mjSpec*, const mjModel* = nullptr); + virtual void SetModel(mjSpec*, const mjModel* = nullptr); // read alternative orientation specification static int ReadAlternative(tinyxml2::XMLElement* elem, mjsOrientation& alt); diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 4f941c15..c5ad9c8f 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -97,10 +97,10 @@ static void UpdateString(string& psuffix, int count, int i) { const char* MJCF[nMJCF][mjXATTRNUM] = { {"mujoco", "!", "1", "model"}, {"<"}, - {"compiler", "*", "19", "autolimits", "boundmass", "boundinertia", "settotalmass", + {"compiler", "*", "20", "autolimits", "boundmass", "boundinertia", "settotalmass", "balanceinertia", "strippath", "coordinate", "angle", "fitaabb", "eulerseq", - "meshdir", "texturedir", "discardvisual", "usethread", - "fusestatic", "inertiafromgeom", "inertiagrouprange", "assetdir", "alignfree"}, + "meshdir", "texturedir", "discardvisual", "usethread", "fusestatic", "inertiafromgeom", + "inertiagrouprange", "saveinertial", "assetdir", "alignfree"}, {"<"}, {"lengthrange", "?", "10", "mode", "useexisting", "uselimit", "accel", "maxforce", "timeconst", "timestep", @@ -167,8 +167,9 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"camera", "?", "17", "orthographic", "fovy", "ipd", "resolution", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "mode", "focal", "focalpixel", "principal", "principalpixel", "sensorsize", "user"}, - {"light", "?", "13", "pos", "dir", "bulbradius", "directional", "castshadow", "active", - "attenuation", "cutoff", "exponent", "ambient", "diffuse", "specular", "mode"}, + {"light", "?", "16", "pos", "dir", "bulbradius", "intensity", "range", + "directional", "type", "castshadow", "active", "attenuation", "cutoff", "exponent", + "ambient", "diffuse", "specular", "mode"}, {"pair", "?", "7", "condim", "friction", "solref", "solreffriction", "solimp", "gap", "margin"}, {"equality", "?", "3", "active", "solref", "solimp"}, @@ -240,8 +241,8 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"<"}, {"bone", "*", "5", "body", "bindpos", "bindquat", "vertid", "vertweight"}, {">"}, - {"texture", "*", "23", "name", "type", "content_type", "file", "gridsize", "gridlayout", - "fileright", "fileleft", "fileup", "filedown", "filefront", "fileback", + {"texture", "*", "24", "name", "type", "colorspace", "content_type", "file", "gridsize", + "gridlayout", "fileright", "fileleft", "fileup", "filedown", "filefront", "fileback", "builtin", "rgb1", "rgb2", "mark", "markrgb", "random", "width", "height", "hflip", "vflip", "nchannel"}, {"material", "*", "12", "name", "class", "texture", "texrepeat", "texuniform", @@ -280,15 +281,15 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"camera", "*", "20", "name", "class", "orthographic", "fovy", "ipd", "resolution", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "mode", "target", "focal", "focalpixel", "principal", "principalpixel", "sensorsize", "user"}, - {"light", "*", "16", "name", "class", "directional", "castshadow", "active", - "pos", "dir", "bulbradius", "attenuation", "cutoff", "exponent", "ambient", "diffuse", - "specular", "mode", "target"}, + {"light", "*", "20", "name", "class", "directional", "type", "castshadow", "active", + "pos", "dir", "bulbradius", "intensity", "range", "attenuation", "cutoff", + "exponent", "ambient", "diffuse", "specular", "mode", "target", "texture"}, {"plugin", "*", "2", "plugin", "instance"}, {"<"}, {"config", "*", "2", "key", "value"}, {">"}, - {"composite", "*", "8", "prefix", "type", "count", "offset", - "vertex", "initial", "curve", "size"}, + {"composite", "*", "9", "prefix", "type", "count", "offset", + "vertex", "initial", "curve", "size", "quat"}, {"<"}, {"joint", "*", "17", "kind", "group", "stiffness", "damping", "armature", "solreffix", "solimpfix", "type", "axis", @@ -310,10 +311,10 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { "flatskin", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "origin"}, {"<"}, {"edge", "?", "5", "equality", "solref", "solimp", "stiffness", "damping"}, - {"elasticity", "?", "4", "young", "poisson", "damping", "thickness"}, - {"contact", "?", "13", "contype", "conaffinity", "condim", "priority", + {"elasticity", "?", "5", "young", "poisson", "damping", "thickness", "elastic2d"}, + {"contact", "?", "14", "contype", "conaffinity", "condim", "priority", "friction", "solmix", "solref", "solimp", "margin", "gap", - "internal", "selfcollide", "activelayers"}, + "internal", "selfcollide", "activelayers", "vertcollide"}, {"pin", "*", "4", "id", "range", "grid", "gridrange"}, {"plugin", "*", "2", "plugin", "instance"}, {"<"}, @@ -327,11 +328,11 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"flex", "*", "13", "name", "group", "dim", "radius", "material", "rgba", "flatskin", "body", "vertex", "element", "texcoord", "elemtexcoord", "node"}, {"<"}, - {"contact", "?", "13", "contype", "conaffinity", "condim", "priority", + {"contact", "?", "14", "contype", "conaffinity", "condim", "priority", "friction", "solmix", "solref", "solimp", "margin", "gap", - "internal", "selfcollide", "activelayers"}, + "internal", "selfcollide", "activelayers", "vertcollide"}, {"edge", "?", "2", "stiffness", "damping"}, - {"elasticity", "?", "4", "young", "poisson", "damping", "thickness"}, + {"elasticity", "?", "5", "young", "poisson", "damping", "thickness", "elastic2d"}, {">"}, {"skin", "*", "9", "name", "file", "material", "rgba", "inflate", "vertex", "texcoord", "face", "group"}, @@ -363,18 +364,18 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"tendon", "*", "0"}, {"<"}, - {"spatial", "*", "18", "name", "class", "group", "limited", "range", - "solreflimit", "solimplimit", "solreffriction", "solimpfriction", + {"spatial", "*", "21", "name", "class", "group", "limited", "actuatorfrclimited", "range", + "actuatorfrcrange", "solreflimit", "solimplimit", "solreffriction", "solimpfriction", "frictionloss", "springlength", "width", "material", - "margin", "stiffness", "damping", "rgba", "user"}, + "margin", "stiffness", "damping", "armature", "rgba", "user"}, {"<"}, {"site", "*", "1", "site"}, {"geom", "*", "2", "geom", "sidesite"}, {"pulley", "*", "1", "divisor"}, {">"}, - {"fixed", "*", "15", "name", "class", "group", "limited", "range", - "solreflimit", "solimplimit", "solreffriction", "solimpfriction", - "frictionloss", "springlength", "margin", "stiffness", "damping", "user"}, + {"fixed", "*", "18", "name", "class", "group", "limited", "actuatorfrclimited", "range", + "actuatorfrcrange","solreflimit", "solimplimit", "solreffriction", "solimpfriction", + "frictionloss", "springlength", "margin", "stiffness", "damping", "armature", "user"}, {"<"}, {"joint", "*", "2", "joint", "coef"}, {">"}, @@ -455,6 +456,7 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"actuatorvel", "*", "5", "name", "actuator", "cutoff", "noise", "user"}, {"actuatorfrc", "*", "5", "name", "actuator", "cutoff", "noise", "user"}, {"jointactuatorfrc", "*", "5", "name", "joint", "cutoff", "noise", "user"}, + {"tendonactuatorfrc", "*", "5", "name", "tendon", "cutoff", "noise", "user"}, {"ballquat", "*", "5", "name", "joint", "cutoff", "noise", "user"}, {"ballangvel", "*", "5", "name", "joint", "cutoff", "noise", "user"}, {"jointlimitpos", "*", "5", "name", "joint", "cutoff", "noise", "user"}, @@ -578,6 +580,17 @@ const mjMap camlight_map[camlight_sz] = { {"targetbodycom", mjCAMLIGHT_TARGETBODYCOM} }; + +// light type +const int lighttype_sz = 4; +const mjMap lighttype_map[lighttype_sz] = { + {"spot", mjLIGHT_SPOT}, + {"directional", mjLIGHT_DIRECTIONAL}, + {"point", mjLIGHT_POINT}, + {"image", mjLIGHT_IMAGE} +}; + + // texmat role type const int texrole_sz = mjNTEXROLE - 1; const mjMap texrole_map[texrole_sz] = { @@ -648,6 +661,15 @@ const mjMap texture_map[texture_sz] = { }; +// colorspace for texture +const int colorspace_sz = 3; +const mjMap colorspace_map[colorspace_sz] = { + {"auto", mjCOLORSPACE_AUTO}, + {"linear", mjCOLORSPACE_LINEAR}, + {"sRGB", mjCOLORSPACE_SRGB} +}; + + // builtin type for texture const int builtin_sz = 4; const mjMap builtin_map[builtin_sz] = { @@ -805,6 +827,15 @@ const mjMap flexself_map[5] = { }; +// flex elastic 2d type +const mjMap elastic2d_map[5] = { + {"none", 0}, + {"bend", 1}, + {"stretch", 2}, + {"both", 3}, +}; + + //---------------------------------- class mjXReader implementation -------------------------------- @@ -1015,6 +1046,9 @@ void mjXReader::Compiler(XMLElement* section, mjSpec* spec) { if (MapValue(section, "alignfree", &n, bool_map, 2)) { spec->compiler.alignfree = (n == 1); } + if (MapValue(section, "saveinertial", &n, bool_map, 2)) { + spec->compiler.saveinertial = (n == 1); + } // lengthrange subelement XMLElement* elem = FindSubElem(section, "lengthrange"); @@ -1376,6 +1410,9 @@ void mjXReader::OneFlex(XMLElement* elem, mjsFlex* flex) { flex->internal = (n == 1); } MapValue(cont, "selfcollide", &flex->selfcollide, flexself_map, 5); + if (MapValue(cont, "vertcollide", &flex->vertcollide, bool_map, 2)) { + flex->vertcollide = (n == 1); + } ReadAttrInt(cont, "activelayers", &flex->activelayers); } @@ -1393,6 +1430,7 @@ void mjXReader::OneFlex(XMLElement* elem, mjsFlex* flex) { ReadAttr(elasticity, "poisson", 1, &flex->poisson, text); ReadAttr(elasticity, "thickness", 1, &flex->thickness, text); ReadAttr(elasticity, "damping", 1, &flex->damping, text); + MapValue(elasticity, "elastic2d", &flex->elastic2d, elastic2d_map, 4); } // write error info @@ -1825,12 +1863,16 @@ void mjXReader::OneCamera(XMLElement* elem, mjsCamera* camera) { // light element parser void mjXReader::OneLight(XMLElement* elem, mjsLight* light) { int n; - string text, name, targetbody; + bool has_directional = false; + string text, name, texture, targetbody; // read attributes if (ReadAttrTxt(elem, "name", name)) { mjs_setString(light->name, name.c_str()); } + if (ReadAttrTxt(elem, "texture", texture)) { + mjs_setString(light->texture, texture.c_str()); + } if (ReadAttrTxt(elem, "target", targetbody)) { mjs_setString(light->targetbody, targetbody.c_str()); } @@ -1838,7 +1880,14 @@ void mjXReader::OneLight(XMLElement* elem, mjsLight* light) { light->mode = (mjtCamLight)n; } if (MapValue(elem, "directional", &n, bool_map, 2)) { - light->directional = (n == 1); + light->type = (n == 1) ? mjLIGHT_DIRECTIONAL : mjLIGHT_SPOT; + has_directional = true; + } + if (MapValue(elem, "type", &n, lighttype_map, lighttype_sz)) { + if (has_directional) { + throw mjXError(elem, "type and directional cannot both be defined"); + } + light->type = (mjtLightType)n; } if (MapValue(elem, "castshadow", &n, bool_map, 2)) { light->castshadow = (n == 1); @@ -1849,6 +1898,8 @@ void mjXReader::OneLight(XMLElement* elem, mjsLight* light) { ReadAttr(elem, "pos", 3, light->pos, text); ReadAttr(elem, "dir", 3, light->dir, text); ReadAttr(elem, "bulbradius", 1, &light->bulbradius, text); + ReadAttr(elem, "intensity", 1, &light->intensity, text); + ReadAttr(elem, "range", 1, &light->range, text); ReadAttr(elem, "attenuation", 3, light->attenuation, text); ReadAttr(elem, "cutoff", 1, &light->cutoff, text); ReadAttr(elem, "exponent", 1, &light->exponent, text); @@ -2049,15 +2100,18 @@ void mjXReader::OneTendon(XMLElement* elem, mjsTendon* tendon) { mjs_setString(tendon->material, material.c_str()); } MapValue(elem, "limited", &tendon->limited, TFAuto_map, 3); + MapValue(elem, "actuatorfrclimited", &tendon->actfrclimited, TFAuto_map, 3); ReadAttr(elem, "width", 1, &tendon->width, text); ReadAttr(elem, "solreflimit", mjNREF, tendon->solref_limit, text, false, false); ReadAttr(elem, "solimplimit", mjNIMP, tendon->solimp_limit, text, false, false); ReadAttr(elem, "solreffriction", mjNREF, tendon->solref_friction, text, false, false); ReadAttr(elem, "solimpfriction", mjNIMP, tendon->solimp_friction, text, false, false); ReadAttr(elem, "range", 2, tendon->range, text); + ReadAttr(elem, "actuatorfrcrange", 2, tendon->actfrcrange, text); ReadAttr(elem, "margin", 1, &tendon->margin, text); ReadAttr(elem, "stiffness", 1, &tendon->stiffness, text); ReadAttr(elem, "damping", 1, &tendon->damping, text); + ReadAttr(elem, "armature", 1, &tendon->armature, text); ReadAttr(elem, "frictionloss", 1, &tendon->frictionloss, text); // read springlength, either one or two values; if one, copy to second value if (ReadAttr(elem, "springlength", 2, tendon->springlength, text, false, false) == 1) { @@ -2388,7 +2442,7 @@ void mjXReader::OneActuator(XMLElement* elem, mjsActuator* actuator) { // make composite -void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, const mjsDefault* def) { +void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, mjsFrame* frame, const mjsDefault* def) { string text; int n; @@ -2402,6 +2456,8 @@ void mjXReader::OneComposite(XMLElement* elem, mjsBody* body, const mjsDefault* } ReadAttr(elem, "count", 3, comp.count, text, false, false); ReadAttr(elem, "offset", 3, comp.offset, text); + ReadAttr(elem, "quat", 4, comp.quat, text); + comp.frame = frame; // plugin XMLElement* eplugin = FirstChildElement(elem, "plugin"); @@ -2649,10 +2705,11 @@ void mjXReader::OneFlexcomp(XMLElement* elem, mjsBody* body, const mjVFS* vfs) { ReadAttr(elasticity, "poisson", 1, &dflex.poisson, text); ReadAttr(elasticity, "damping", 1, &dflex.damping, text); ReadAttr(elasticity, "thickness", 1, &dflex.thickness, text); + MapValue(elasticity, "elastic2d", &dflex.elastic2d, elastic2d_map, 4); } // check errors - if (elasticity && fcomp.equality) { + if (dflex.elastic2d >= 2 && fcomp.equality) { throw mjXError(elem, "elasticity and edge constraints cannot both be present"); } @@ -2673,6 +2730,9 @@ void mjXReader::OneFlexcomp(XMLElement* elem, mjsBody* body, const mjVFS* vfs) { dflex.internal = (n == 1); } MapValue(cont, "selfcollide", &dflex.selfcollide, flexself_map, 5); + if (MapValue(cont, "vertcollide", &n, bool_map, 2)) { + dflex.vertcollide = (n == 1); + } ReadAttrInt(cont, "activelayers", &dflex.activelayers); } @@ -3165,6 +3225,9 @@ void mjXReader::Asset(XMLElement* section, const mjVFS* vfs) { if (MapValue(elem, "type", &n, texture_map, texture_sz)) { texture->type = (mjtTexture)n; } + if (MapValue(elem, "colorspace", &n, colorspace_map, colorspace_sz)) { + texture->colorspace = (mjtColorSpace)n; + } if (ReadAttrTxt(elem, "name", texname)) { mjs_setString(texture->name, texname.c_str()); } @@ -3467,7 +3530,7 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame, // composite sub-element else if (name == "composite") { // parse composite - OneComposite(elem, body, def); + OneComposite(elem, body, frame, def); } // flexcomp sub-element @@ -3564,7 +3627,7 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame, UpdateString(suffix, count, i); // attach to parent - if (!mjs_attachFrame(body, pframe, /*prefix=*/"", suffix.c_str())) { + if (!mjs_attach(body->element, pframe->element, /*prefix=*/"", suffix.c_str())) { throw mjXError(elem, mjs_getError(spec)); } } @@ -3625,27 +3688,35 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame, else if (name == "attach") { string model_name, body_name, prefix; ReadAttrTxt(elem, "model", model_name, /*required=*/true); - ReadAttrTxt(elem, "body", body_name, /*required=*/true); + ReadAttrTxt(elem, "body", body_name, /*required=*/false); ReadAttrTxt(elem, "prefix", prefix, /*required=*/true); - mjsBody* child = mjs_findBody(spec, (prefix+body_name).c_str()); + mjsBody* child_body = mjs_findBody(spec, (prefix+body_name).c_str()); mjsFrame* pframe = frame ? frame : mjs_addFrame(body, nullptr); - if (!child) { + if (!child_body) { mjSpec* asset = mjs_findSpec(spec, model_name.c_str()); if (!asset) { throw mjXError(elem, "could not find model '%s'", model_name.c_str()); } - child = mjs_findBody(asset, body_name.c_str()); - if (!child) { - throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str()); + mjsElement* child; + if (body_name.empty()) { + child = asset->element; + } else { + child_body = mjs_findBody(asset, body_name.c_str()); + if (!child_body) { + throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str()); + } + child = child_body->element; } - if (!mjs_attachBody(pframe, child, prefix.c_str(), "")) { + if (!mjs_attach(pframe->element, child, prefix.c_str(), "")) { throw mjXError(elem, mjs_getError(spec)); } } else { // only set frame to existing body - mjs_setFrame(child->element, pframe); + if (mjs_setFrame(child_body->element, pframe)) { + throw mjXError(elem, mjs_getError(spec)); + } } } @@ -3787,7 +3858,7 @@ void mjXReader::Tendon(XMLElement* section) { def = mjs_getSpecDefault(spec); } - // create equality constraint and parse + // create tendon and parse mjsTendon* tendon = mjs_addTendon(spec, def); OneTendon(elem, tendon); @@ -3959,6 +4030,10 @@ void mjXReader::Sensor(XMLElement* section) { sensor->type = mjSENS_JOINTACTFRC; sensor->objtype = mjOBJ_JOINT; ReadAttrTxt(elem, "joint", objname, true); + } else if (type=="tendonactuatorfrc") { + sensor->type = mjSENS_TENDONACTFRC; + sensor->objtype = mjOBJ_TENDON; + ReadAttrTxt(elem, "tendon", objname, true); } // sensors related to ball joints diff --git a/src/xml/xml_native_reader.h b/src/xml/xml_native_reader.h index 41aacbbb..010ea2d5 100644 --- a/src/xml/xml_native_reader.h +++ b/src/xml/xml_native_reader.h @@ -80,7 +80,8 @@ class mjXReader : public mjXBase { void OneEquality(tinyxml2::XMLElement* elem, mjsEquality* pequality); void OneTendon(tinyxml2::XMLElement* elem, mjsTendon* ptendon); void OneActuator(tinyxml2::XMLElement* elem, mjsActuator* pactuator); - void OneComposite(tinyxml2::XMLElement* elem, mjsBody* pbody, const mjsDefault* def); + void OneComposite(tinyxml2::XMLElement* elem, mjsBody* pbody, mjsFrame* pframe, + const mjsDefault* def); void OneFlexcomp(tinyxml2::XMLElement* elem, mjsBody* pbody, const mjVFS* vfs); void OnePlugin(tinyxml2::XMLElement* elem, mjsPlugin* plugin); @@ -101,7 +102,7 @@ class mjXReader : public mjXBase { }; // MJCF schema -#define nMJCF 237 +#define nMJCF 238 extern const char* MJCF[nMJCF][mjXATTRNUM]; #endif // MUJOCO_SRC_XML_XML_NATIVE_READER_H_ diff --git a/src/xml/xml_native_writer.cc b/src/xml/xml_native_writer.cc index 77a7dc6a..de2aef96 100644 --- a/src/xml/xml_native_writer.cc +++ b/src/xml/xml_native_writer.cc @@ -194,6 +194,7 @@ void mjXWriter::OneFlex(XMLElement* elem, const mjCFlex* flex) { WriteAttr(elastic, "poisson", 1, &flex->poisson, &defflex.poisson); WriteAttr(elastic, "thickness", 1, &flex->thickness, &defflex.thickness); WriteAttr(elastic, "damping", 1, &flex->damping, &defflex.damping); + WriteAttrKey(elastic, "elastic2d", elastic2d_map, 2, flex->elastic2d, defflex.elastic2d); // edge subelement XMLElement* edge = InsertEnd(elem, "edge"); @@ -604,7 +605,10 @@ void mjXWriter::OneLight(XMLElement* elem, const mjCLight* light, mjCDef* def, // defaults and regular WriteAttr(elem, "bulbradius", 1, &light->bulbradius, &def->Light().bulbradius); - WriteAttrKey(elem, "directional", bool_map, 2, light->directional, def->Light().directional); + WriteAttr(elem, "intensity", 1, &light->intensity, &def->Light().intensity); + WriteAttr(elem, "range", 1, &light->range, &def->Light().range); + WriteAttrKey(elem, "type", lighttype_map, lighttype_sz, light->type, def->Light().type); + WriteAttrTxt(elem, "texture", light->get_texture()); WriteAttrKey(elem, "castshadow", bool_map, 2, light->castshadow, def->Light().castshadow); WriteAttrKey(elem, "active", bool_map, 2, light->active, def->Light().active); WriteAttr(elem, "attenuation", 3, light->attenuation, def->Light().attenuation); @@ -725,10 +729,13 @@ void mjXWriter::OneTendon(XMLElement* elem, const mjCTendon* tendon, mjCDef* def WriteAttr(elem, "solimpfriction", mjNIMP, tendon->solimp_friction, def->Tendon().solimp_friction, true); WriteAttrKey(elem, "limited", TFAuto_map, 3, tendon->limited, def->Tendon().limited); + WriteAttrKey(elem, "actuatorfrclimited", TFAuto_map, 3, tendon->actfrclimited, def->Tendon().actfrclimited); WriteAttr(elem, "range", 2, tendon->range, def->Tendon().range); + WriteAttr(elem, "actuatorfrcrange", 2, tendon->actfrcrange, def->Tendon().actfrcrange); WriteAttr(elem, "margin", 1, &tendon->margin, &def->Tendon().margin); WriteAttr(elem, "stiffness", 1, &tendon->stiffness, &def->Tendon().stiffness); WriteAttr(elem, "damping", 1, &tendon->damping, &def->Tendon().damping); + WriteAttr(elem, "armature", 1, &tendon->armature, &def->Tendon().armature); WriteAttr(elem, "frictionloss", 1, &tendon->frictionloss, &def->Tendon().frictionloss); if (tendon->springlength[0] != tendon->springlength[1] || def->Tendon().springlength[0] != def->Tendon().springlength[1]) { @@ -878,12 +885,12 @@ mjXWriter::mjXWriter(void) { // cast model -void mjXWriter::SetModel(const mjSpec* _spec, const mjModel* m) { +void mjXWriter::SetModel(mjSpec* _spec, const mjModel* m) { if (_spec) { model = static_cast(_spec->element); } if (m) { - model->CopyBack(m); + mj_copyBack(&model->spec, m); } } @@ -1482,6 +1489,7 @@ void mjXWriter::Asset(XMLElement* root) { // write common attributes WriteAttrKey(elem, "type", texture_map, texture_sz, texture->type); + WriteAttrKey(elem, "colorspace", colorspace_map, colorspace_sz, texture->colorspace); WriteAttrTxt(elem, "name", texture->name); // write builtin @@ -1548,6 +1556,7 @@ void mjXWriter::Asset(XMLElement* root) { if (mesh->Plugin().active) { elem = InsertEnd(section, "mesh"); WriteAttrTxt(elem, "name", mesh->name); + WriteAttrTxt(elem, "file", mesh->File()); OnePlugin(InsertEnd(elem, "plugin"), &mesh->Plugin()); } else{ elem = InsertEnd(section, "mesh"); @@ -1632,7 +1641,8 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi WriteVector(elem, "user", body->get_userdata()); // write inertial - if (body->explicitinertial && model->compiler.inertiafromgeom != mjINERTIAFROMGEOM_TRUE) { + if (model->compiler.saveinertial || + (body->explicitinertial && model->compiler.inertiafromgeom != mjINERTIAFROMGEOM_TRUE)) { XMLElement* inertial = InsertEnd(elem, "inertial"); WriteAttr(inertial, "pos", 3, body->ipos); WriteAttr(inertial, "quat", 4, body->iquat, unitq); @@ -2031,6 +2041,10 @@ void mjXWriter::Sensor(XMLElement* root) { elem = InsertEnd(section, "jointactuatorfrc"); WriteAttrTxt(elem, "joint", sensor->get_objname()); break; + case mjSENS_TENDONACTFRC: + elem = InsertEnd(section, "tendonactuatorfrc"); + WriteAttrTxt(elem, "tendon", sensor->get_objname()); + break; // sensors related to ball joints case mjSENS_BALLQUAT: diff --git a/src/xml/xml_native_writer.h b/src/xml/xml_native_writer.h index b4941429..09962522 100644 --- a/src/xml/xml_native_writer.h +++ b/src/xml/xml_native_writer.h @@ -29,7 +29,7 @@ class mjXWriter : public mjXBase { public: mjXWriter(); // constructor virtual ~mjXWriter() = default; // destructor - void SetModel(const mjSpec* _spec, const mjModel* m = nullptr); + void SetModel(mjSpec* _spec, const mjModel* m = nullptr); // write XML document to string std::string Write(char *error, std::size_t error_sz); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 0f92803d..a286a1c6 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -100,3 +100,4 @@ add_subdirectory(user) add_subdirectory(xml) add_subdirectory(plugin/elasticity) add_subdirectory(plugin/actuator) +add_subdirectory(experimental) diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index 0a23fafe..48a5a369 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -16,7 +16,7 @@ mujoco_test( ccd_benchmark_test MAIN_TARGET benchmark::benchmark_main - ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ccd ) mujoco_test( diff --git a/test/benchmark/ccd_benchmark_test.cc b/test/benchmark/ccd_benchmark_test.cc index 1ea8ce37..a7768dce 100644 --- a/test/benchmark/ccd_benchmark_test.cc +++ b/test/benchmark/ccd_benchmark_test.cc @@ -24,6 +24,9 @@ #include #include "test/fixture.h" +#include "src/engine/engine_collision_convex.h" +#include "src/engine/engine_collision_primitive.h" + namespace mujoco { namespace { @@ -116,11 +119,19 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BENCHMARK(BM_BoxMesh_LibCCD); void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox(benchmark::State& state) { - static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)"); + static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)", mjDSBL_NATIVECCD); harness.RunBenchmark(state); } BENCHMARK(BM_BoxBox); +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox_NativeCCD(benchmark::State& state) { + mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_Convex; + static TestHarness harness(kBoxBoxPath, "box.xml (NativeCCD)"); + harness.RunBenchmark(state); + mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_BoxBox; +} +BENCHMARK(BM_BoxBox_NativeCCD); + void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Ellipsoid_NativeCCD(benchmark::State& state) { static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (nativeccd)"); diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index 2c2ff6aa..fdeadb31 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -152,8 +152,8 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline( // transpose sparse matrix (uncompressed) void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline( mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz, - int* res_rowadr, int* res_colind, const int* rownnz, const int* rowadr, - const int* colind) { + int* res_rowadr, int* res_colind, int* res_rowsuper, + const int* rownnz, const int* rowadr, const int* colind) { memset(res_rownnz, 0, nc * sizeof(int)); for (int rt = 0; rt < nc; rt++) { res_rowadr[rt] = rt * nr; @@ -343,7 +343,7 @@ void ABSL_ATTRIBUTE_NOINLINE mulMatVecSparse_8(mjtNum* res, // ----------------------------- benchmark ------------------------------------ static void BM_MatVecSparse(benchmark::State& state, int unroll) { - static mjModel* m = LoadModelFromPath("plugin/elasticity/flag_flex.xml"); + static mjModel* m = LoadModelFromPath("flex/flag.xml"); mjData* d = mj_makeData(m); // warm-up rollout to get a typical state @@ -497,7 +497,14 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old( } BENCHMARK(BM_combineSparse_old); -static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) { +enum class Supernode { + None, + PostProcess, + Inline +}; + +static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func, + Supernode super) { static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml"); // force use of sparse matrices @@ -516,12 +523,19 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) { mjtNum* res = mj_stackAllocNum(d, m->nv * d->nefc); int* res_rownnz = mj_stackAllocInt(d, m->nv); int* res_rowadr = mj_stackAllocInt(d, m->nv); + int* res_rowsuper = mj_stackAllocInt(d, m->nv); int* res_colind = mj_stackAllocInt(d, m->nv * d->nefc); // time benchmark for (auto s : state) { - func(res, d->efc_J, d->nefc, m->nv, res_rownnz, res_rowadr, res_colind, + int* rowsuper = (super == Supernode::Inline) ? res_rowsuper : nullptr; + func(res, d->efc_J, d->nefc, m->nv, + res_rownnz, res_rowadr, res_colind, rowsuper, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); + if (super == Supernode::PostProcess) { + mju_superSparse(m->nv, res_rowsuper, + res_rownnz, res_rowadr, res_colind); + } } mj_freeStack(d); @@ -529,26 +543,44 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) { state.SetItemsProcessed(state.iterations()); } +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_transposeSparse_old(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_transposeSparse(state, &transposeSparse_baseline, Supernode::None); +} +BENCHMARK(BM_transposeSparse_old); + void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_transposeSparse_new(benchmark::State& state) { MujocoErrorTestGuard guard; - BM_transposeSparse(state, &mju_transposeSparse); + BM_transposeSparse(state, &mju_transposeSparse, Supernode::None); } BENCHMARK(BM_transposeSparse_new); void ABSL_ATTRIBUTE_NO_TAIL_CALL -BM_transposeSparse_old(benchmark::State& state) { +BM_transposeSparse_superpost(benchmark::State& state) { MujocoErrorTestGuard guard; - BM_transposeSparse(state, &transposeSparse_baseline); + BM_transposeSparse(state, &mju_transposeSparse, Supernode::PostProcess); } -BENCHMARK(BM_transposeSparse_old); +BENCHMARK(BM_transposeSparse_superpost); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_transposeSparse_superinline(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_transposeSparse(state, &mju_transposeSparse, Supernode::Inline); +} +BENCHMARK(BM_transposeSparse_superinline); static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) { - static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml"); - mjData* d = mj_makeData(m); + static mjModel* m = + LoadModelFromPath("../test/benchmark/testdata/2humanoid100.xml"); - // force use of sparse matrices + // force use of sparse matrices, Newton solver, no islands m->opt.jacobian = mjJAC_SPARSE; + m->opt.solver = mjSOL_NEWTON; + m->opt.enableflags &= ~mjENBL_ISLAND; + + mjData* d = mj_makeData(m); // warm-up rollout to get a typical state while (d->time < 2) { @@ -575,10 +607,13 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) { // time benchmark if (func) { - for (auto s : state) { - mju_sqrMatTDUncompressedInit(rowadr, m->nv); + mju_sqrMatTDSparseCount(rownnz, rowadr, m->nv, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + d->efc_JT_rownnz, d->efc_JT_rowadr, + d->efc_JT_colind, nullptr, d, 1); - // compute H = J'*D*J, uncompressed layout + for (auto s : state) { + // compute H = J'*D*J, compressed layout func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, @@ -606,18 +641,25 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) { } void ABSL_ATTRIBUTE_NO_TAIL_CALL -BM_sqrMatTDSparse_new(benchmark::State& state) { +BM_sqrMatTDSparse_col(benchmark::State& state) { MujocoErrorTestGuard guard; BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse); } -BENCHMARK(BM_sqrMatTDSparse_new); +BENCHMARK(BM_sqrMatTDSparse_col); void ABSL_ATTRIBUTE_NO_TAIL_CALL -BM_sqrMatTDSparse_old(benchmark::State& state) { +BM_sqrMatTDSparse_row(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse_row); +} +BENCHMARK(BM_sqrMatTDSparse_row); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_sqrMatTDSparse_uncompressed(benchmark::State& state) { MujocoErrorTestGuard guard; BM_sqrMatTDSparse(state, nullptr); } -BENCHMARK(BM_sqrMatTDSparse_old); +BENCHMARK(BM_sqrMatTDSparse_uncompressed); } // namespace } // namespace mujoco diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc index dd73d3da..62a42c6c 100644 --- a/test/benchmark/factorI_benchmark_test.cc +++ b/test/benchmark/factorI_benchmark_test.cc @@ -19,6 +19,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -44,10 +45,8 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { mj_markStack(d); // M: mass matrix in CSR format - mjtNum* M = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - M[i] = d->qM[d->mapM2M[i]]; - } + mjtNum* M = mj_stackAllocNum(d, m->nC); + mju_gather(M, d->qM, d->mapM2M, m->nC); // LDlegacy: legacy LD matrix (size nM) mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); @@ -60,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { } else { mju_copy(d->qLD, M, m->nC); mj_factorI(d->qLD, d->qLDiagInv, m->nv, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); } } } diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index e787fa24..457aa32a 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -20,6 +20,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -46,10 +47,8 @@ static void BM_solve(benchmark::State& state, SolveType type) { mj_markStack(d); // M: mass matrix in CSR format - mjtNum* M = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - M[i] = d->qM[d->mapM2M[i]]; - } + mjtNum* M = mj_stackAllocNum(d, m->nC); + mju_gather(M, d->qM, d->mapM2M, m->nC); // LDlegacy: legacy LD matrix (size nM) mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); @@ -74,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) { case SolveType::kCsr: mju_copy(d->qLD, M, m->nC); mj_factorI(d->qLD, d->qLDiagInv, m->nv, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); } } } diff --git a/test/benchmark/parse_benchmark_test.cc b/test/benchmark/parse_benchmark_test.cc index 059f0c58..95471c01 100644 --- a/test/benchmark/parse_benchmark_test.cc +++ b/test/benchmark/parse_benchmark_test.cc @@ -73,7 +73,7 @@ static void run_parse_benchmark(const std::string xml_path, // run_parse_benchmark). void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_ParseFlagPlugin(benchmark::State& state) { - run_parse_benchmark(GetModelPath("plugin/elasticity/flag_flex.xml"), state); + run_parse_benchmark(GetModelPath("flex/flag.xml"), state); } BENCHMARK(BM_ParseFlagPlugin); diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 64276ba9..65cf700c 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -19,6 +19,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -50,11 +51,10 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { vec[i] = 0.2 + 0.3*i; } - // make legacy matrix + // scatter into legacy matrix mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + mju_zero(LDlegacy, m->nM); + mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nC); // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { @@ -64,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv); } else { mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); } } } diff --git a/test/benchmark/step_benchmark_test.cc b/test/benchmark/step_benchmark_test.cc index 551e312c..1d34f2c9 100644 --- a/test/benchmark/step_benchmark_test.cc +++ b/test/benchmark/step_benchmark_test.cc @@ -73,7 +73,7 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) { void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepFlagPlugin(benchmark::State& state) { MujocoErrorTestGuard guard; - static mjModel* model = LoadModelFromPath("plugin/elasticity/flag_flex.xml"); + static mjModel* model = LoadModelFromPath("flex/flag.xml"); run_step_benchmark(model, state); } BENCHMARK(BM_StepFlagPlugin); diff --git a/test/benchmark/testdata/2humanoid100.xml b/test/benchmark/testdata/2humanoid100.xml new file mode 100644 index 00000000..424ed2c4 --- /dev/null +++ b/test/benchmark/testdata/2humanoid100.xml @@ -0,0 +1,123 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/benchmark/testdata/humanoid.xml b/test/benchmark/testdata/humanoid.xml new file mode 100644 index 00000000..7545e193 --- /dev/null +++ b/test/benchmark/testdata/humanoid.xml @@ -0,0 +1,252 @@ + + diff --git a/test/engine/CMakeLists.txt b/test/engine/CMakeLists.txt index 2acdc4b9..1884ae05 100644 --- a/test/engine/CMakeLists.txt +++ b/test/engine/CMakeLists.txt @@ -67,11 +67,4 @@ mujoco_test(engine_util_solve_test) mujoco_test(engine_util_spatial_test) -mujoco_test( - engine_vis_state_test - PROPERTIES - ENVIRONMENT - "MUJOCO_PLUGIN_DIR=$" -) - mujoco_test(engine_vis_visualize_test) diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc index 437d314d..9e5fa5e1 100644 --- a/test/engine/engine_collision_gjk_test.cc +++ b/test/engine/engine_collision_gjk_test.cc @@ -71,7 +71,7 @@ void CCDFree(void* data, void* buffer) { } mjtNum GeomDist(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3], - mjtNum x2[3], mjtNum cutoff = mjMAXVAL) { + mjtNum x2[3], mjtNum cutoff = mjMAX_LIMIT) { mjCCDConfig config; mjCCDStatus status; @@ -93,7 +93,7 @@ mjtNum GeomDist(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3], return dist; } -int Penetration(mjtNum& depth, std::vector& dir, +int Penetration(mjCCDStatus& status, mjtNum& depth, std::vector& dir, std::vector& pos, mjModel* model, mjData* data, int g1, int g2, mjtNum margin = 0, int max_contacts = 1) { mjCCDObj obj1, obj2; @@ -128,7 +128,6 @@ int Penetration(mjtNum& depth, std::vector& dir, #endif mjCCDConfig config; - mjCCDStatus status; // set config config.max_iterations = kMaxIterations; @@ -212,7 +211,7 @@ TEST_F(MjGjkTest, SphereSphereDistCutoff) { int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); mjtNum dist = GeomDist(model, data, geom1, geom2, nullptr, nullptr, .999999); - EXPECT_EQ(dist, mjMAXVAL); + EXPECT_EQ(dist, mjMAX_LIMIT); mj_deleteData(data); mj_deleteModel(model); } @@ -235,9 +234,11 @@ TEST_F(MjGjkTest, SphereSphereNoDist) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 0); mj_deleteData(data); @@ -262,9 +263,11 @@ TEST_F(MjGjkTest, SphereSphereIntersect) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 1); @@ -304,9 +307,10 @@ TEST_F(MjGjkTest, BoxBoxDepth) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 1); @@ -354,9 +358,11 @@ TEST_F(MjGjkTest, BoxBoxDepth2) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); if (ncons == 1) { EXPECT_NEAR(dist, -0.033401579411886845, kTolerance); @@ -421,9 +427,11 @@ TEST_F(MjGjkTest, BoxBoxDepth3) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 1); EXPECT_NEAR(dist, -0.003066, kTolerance); @@ -453,11 +461,14 @@ TEST_F(MjGjkTest, BoxBoxTouching) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 0); + EXPECT_EQ(status.epa_status, -1); mj_deleteData(data); mj_deleteModel(model); @@ -479,11 +490,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); EXPECT_NEAR(dist, -.1, kTolerance); @@ -517,11 +530,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD2) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); EXPECT_NEAR(dist, -.1, kTolerance); @@ -573,11 +588,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD3) { xpos[2] = 1.095456702630382306296041861060; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -634,11 +651,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD4) { xpos[1] = -0.023500601273213628239489025873; xpos[2] = -4.958782854594746325460619118530; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 8); EXPECT_NEAR(dist, -0.00060425119242707459, kTolerance); @@ -701,11 +720,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD5) { xpos[2] = -4.659108354876987156956147373421; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 8); EXPECT_NEAR(dist, -0.0001077858631973211, kTolerance); @@ -750,11 +771,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD6) { xpos[1] = 0.190777715293135141649827346555; xpos[2] = 0.100006658017411736993906856696; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 5); EXPECT_NEAR(dist, -0.00009843, kTolerance); @@ -817,11 +840,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD7) { xpos[2] = -4.958375812037025376355359185254; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 8); mj_deleteData(data); @@ -878,11 +903,13 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD8) { xpos[1] = -0.023505499999999998617106200527; xpos[2] = -4.958574289672835533338002278470; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -940,17 +967,323 @@ TEST_F(MjGjkTest, BoxBoxMultiCCD9) { xpos[2] = 0.2156259187793853615566774806211469694972; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); mj_deleteData(data); mj_deleteModel(model); } +TEST_F(MjGjkTest, BoxBoxMultiCCD10) { + static constexpr char xml[] = R"( + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + mjtNum* xpos = data->geom_xpos; + + xpos[0] = -0.1034859999999999946584949839234468527138; + xpos[1] = -0.0765140000000000264357424839545274153352; + xpos[2] = 0.1257628745456405572333835607423679903150; + + xpos = data->geom_xpos + 3; + + xpos[0] = -0.1034859999999999946584949839234468527138; + xpos[1] = -0.0765140000000000264357424839545274153352; + xpos[2] = 0.1751399999999999623767621415026951581240; + + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; + std::vector dir, pos; + mjtNum dist; + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 8); + + EXPECT_EQ(ncons, 4); + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(MjGjkTest, BoxBoxMultiCCD11) { + static constexpr char xml[] = R"( + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + mjtNum* xpos = data->geom_xpos; + mjtNum* xmat = data->geom_xmat; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = 0.0000000000000000000000000000000000000000; + xmat[2] = 0.0000000000000000000000000000000000000000; + xmat[3] = 0.0000000000000000000000000000000000000000; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = -0.0000000000000000013928437397151766790940; + xmat[6] = 0.0000000000000000000000000000000000000000; + xmat[7] = 0.0000000000000000013928437397151766790940; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = -0.1036549999999999971400654885655967518687; + xpos[1] = -0.1963450000000000195132798808117513544858; + xpos[2] = 0.1247685038468368534658736734854755923152; + + + xpos = data->geom_xpos + 3; + xmat = data->geom_xmat + 9; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = 0.0000000000000000000000000000000000000000; + xmat[2] = 0.0000000000000000000000000000000000000000; + xmat[3] = 0.0000000000000000000000000000000000000000; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = -0.0000000000000000018885268354605779111974; + xmat[6] = 0.0000000000000000000000000000000000000000; + xmat[7] = 0.0000000000000000018885268354605779111974; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = -0.1036549999999999971400654885655967518687; + xpos[1] = -0.1963450000000000195132798808117513544858; + xpos[2] = 0.1745248497897437800485676007156143896282; + + + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; + std::vector dir, pos; + mjtNum dist; + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 8); + + EXPECT_EQ(ncons, 4); + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(MjGjkTest, BoxBoxMultiCCD12) { + static constexpr char xml[] = R"( + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + mjtNum* xpos = data->geom_xpos; + mjtNum* xmat = data->geom_xmat; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = 0.0000000000000000000000000000000000000000; + xmat[2] = 0.0000000000000000000000000000000000000000; + xmat[3] = 0.0000000000000000000000000000000000000000; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = -0.0000000000000000032154383478277941584027; + xmat[6] = 0.0000000000000000000000000000000000000000; + xmat[7] = 0.0000000000000000032154383478277941584027; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = 0.0164299999999999862820843077315657865256; + xpos[1] = -0.0764300000000000256950016819246229715645; + xpos[2] = 0.1252706891962387103500731200256268493831; + + xpos = data->geom_xpos + 3; + xmat = data->geom_xmat + 9; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = 0.0000000000000000000000000000000000000000; + xmat[2] = 0.0000000000000000000000000000000000000000; + xmat[3] = 0.0000000000000000000000000000000000000000; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = -0.0000000000000000018997602302052549055743; + xmat[6] = 0.0000000000000000000000000000000000000000; + xmat[7] = 0.0000000000000000018997602302052549055743; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = 0.0164299999999999862820843077315657865256; + xpos[1] = -0.0764300000000000256950016819246229715645; + xpos[2] = 0.1748374248948718623353215662064030766487; + + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; + std::vector dir, pos; + mjtNum dist; + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 8); + + EXPECT_EQ(ncons, 4); + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(MjGjkTest, BoxBoxMultiCCD13) { + static constexpr char xml[] = R"( + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + mjtNum* xpos = data->geom_xpos; + mjtNum* xmat = data->geom_xmat; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = -0.0000000000000000000000000000001366192847; + xmat[2] = -0.0000000000000002451235402041944571528177; + xmat[3] = 0.0000000000000000000000000000001366707863; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = 0.0000000000000000002101047324941978855126; + xmat[6] = 0.0000000000000002451235402041944571528177; + xmat[7] = -0.0000000000000000002101047324941978855126; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = -0.1000000000000000055511151231257827021182; + xpos[1] = -0.2000000000000000111022302462515654042363; + xpos[2] = -0.0809921810760001470441693527391180396080; + + xpos = data->geom_xpos + 3; + xmat = data->geom_xmat + 9; + + xmat[0] = 1.0000000000000000000000000000000000000000; + xmat[1] = -0.0000000000000000000000000000000740327228; + xmat[2] = -0.0000000000000002557259745463766308177658; + xmat[3] = 0.0000000000000000000000000000000775428823; + xmat[4] = 1.0000000000000000000000000000000000000000; + xmat[5] = 0.0000000000000000137262533997081760161613; + xmat[6] = 0.0000000000000002557259745463766308177658; + xmat[7] = -0.0000000000000000137262533997081760161613; + xmat[8] = 1.0000000000000000000000000000000000000000; + + xpos[0] = -0.1000000000000000055511151231257827021182; + xpos[1] = -0.2000000000000000111022302462515654042363; + xpos[2] = -0.0418396695286432432348000531874276930466; + + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; + std::vector dir, pos; + mjtNum dist; + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 8); + + EXPECT_EQ(ncons, 4); + + EXPECT_NEAR(dir[0], 0, kTolerance); + EXPECT_NEAR(dir[1], 0, kTolerance); + EXPECT_NEAR(dir[2], 1, kTolerance); + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(MjGjkTest, BoxBoxMultiCCD14) { + static constexpr char xml[] = R"( + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + mjtNum* xpos = data->geom_xpos; + mjtNum* xmat = data->geom_xmat; + + xmat[0] = 0.9999999980312528347070610834634862840176; + xmat[1] = 0.0000179109150612445166097282805983681442; + xmat[2] = -0.0000601389470252842008382576644009986921; + xmat[3] = -0.0000179108686851742081238159781664265324; + xmat[4] = 0.9999999998393023226128661917755380272865; + xmat[5] = 0.0000007716871733595438989517368948145570; + xmat[6] = 0.0000601389608372434404702858157243383630; + xmat[7] = -0.0000007706100310572527002924239115932981; + xmat[8] = 0.9999999981913554325529958077822811901569; + + xpos[0] = 0.0002051257133161473724877743585182088282; + xpos[1] = 0.0000051793157380883478958571650152542531; + xpos[2] = -0.0800031938952457943869944756443146616220; + + xpos = data->geom_xpos + 3; + xmat = data->geom_xmat + 9; + + xmat[0] = 0.9999999606378873195922096783760935068130; + xmat[1] = -0.0000186818570733572177707156047876679850; + xmat[2] = -0.0002799557310143530259108346491814245383; + xmat[3] = 0.0000186853252997592718994551708178164517; + xmat[4] = 0.9999999997487241110150080203311517834663; + xmat[5] = 0.0000123858711158191162315369768243122905; + xmat[6] = 0.0002799554995529331168427344955773605761; + xmat[7] = -0.0000123911016921886008170612322731862776; + xmat[8] = 0.9999999607356884201436741932411678135395; + + xpos[0] = 0.0002145111032389043976328218965576866140; + xpos[1] = -0.0000051338999751368759734112059978095033; + xpos[2] = -0.0400059009625639144802633495601185131818; + + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; + std::vector dir, pos; + mjtNum dist; + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 8); + + EXPECT_EQ(ncons, 4); + + mj_deleteData(data); + mj_deleteModel(model); +} + TEST_F(MjGjkTest, SmallBoxMesh) { static constexpr char xml[] = R"( @@ -990,9 +1323,11 @@ TEST_F(MjGjkTest, SmallBoxMesh) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 1); EXPECT_NEAR(dist, 0, kTolerance); @@ -1032,11 +1367,13 @@ TEST_F(MjGjkTest, BoxMesh) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom2, geom1, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g2, g1, 0, 1000); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -1065,11 +1402,13 @@ TEST_F(MjGjkTest, BoxMesh2) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom2, geom1, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g2, g1, 0, 1000); EXPECT_EQ(ncons, 5); mj_deleteData(data); @@ -1098,11 +1437,13 @@ TEST_F(MjGjkTest, BoxMeshPrune) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom2, geom1, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g2, g1, 0, 4); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -1133,11 +1474,13 @@ TEST_F(MjGjkTest, MeshMesh) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 5); mj_deleteData(data); @@ -1168,11 +1511,13 @@ TEST_F(MjGjkTest, MeshMeshPrune) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 4); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -1206,11 +1551,13 @@ TEST_F(MjGjkTest, BoxEdge) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "box2"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "box2"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 4); EXPECT_EQ(ncons, 2); mj_deleteData(data); @@ -1279,11 +1626,13 @@ TEST_F(MjGjkTest, BoxEdge2) { xpos[1] = 0.9828851949225971829093850828940048813820; xpos[2] = 3.0930077345364814789263618877157568931580; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "box2"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "box2"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 4); EXPECT_EQ(ncons, 2); mj_deleteData(data); @@ -1351,11 +1700,13 @@ TEST_F(MjGjkTest, BoxEdgeEdge) { xpos[1] = -0.0000000000000000008679606505055748997840; xpos[2] = 2.8141526153588731773425024584867060184479; - int geom1 = mj_name2id(model, mjOBJ_GEOM, "box2"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "box2"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 4); EXPECT_EQ(ncons, 2); mj_deleteData(data); @@ -1399,11 +1750,13 @@ TEST_F(MjGjkTest, MeshEdge) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "box2"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "box2"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "box3"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 4); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 4); EXPECT_EQ(ncons, 2); mj_deleteData(data); @@ -1421,9 +1774,11 @@ TEST_F(MjGjkTest, EllipsoidEllipsoidPenetrating) { int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, geom1, geom2); EXPECT_EQ(ncons, 1); EXPECT_NEAR(dist, -0.00022548856248122027, kTolerance); @@ -1502,11 +1857,13 @@ static constexpr char xml[] = R"( mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2); EXPECT_EQ(ncons, 1); EXPECT_NEAR(dist, -0.01, kTolerance); @@ -1520,7 +1877,7 @@ static constexpr char xml[] = R"( EXPECT_NEAR(pos[2], -0.005, kTolerance); // multicontact - ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 0, 1000); + ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000); EXPECT_EQ(ncons, 4); mj_deleteData(data); @@ -1543,11 +1900,13 @@ TEST_F(MjGjkTest, EllipsoidEllipsoidIntersect) { mjData* data = mj_makeData(model); mj_forward(model, data); - int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); - int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1"); + int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2"); + + mjCCDStatus status; std::vector dir, pos; mjtNum dist; - int ncons = Penetration(dist, dir, pos, model, data, geom1, geom2, 15); + int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 15); EXPECT_EQ(ncons, 1); EXPECT_NEAR(dist, -14.245732934582151, kTolerance); diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index a0a8f149..9f970b18 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -25,12 +25,14 @@ #include #include "src/engine/engine_core_constraint.h" #include "src/engine/engine_support.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { namespace { using ::testing::DoubleNear; +using ::testing::NotNull; using ::testing::Pointwise; using CoreConstraintTest = MujocoTest; @@ -284,205 +286,17 @@ TEST_F(CoreConstraintTest, EqualityBodySite) { mj_deleteModel(model); } - static const char* const kIlslandEfcPath = "engine/testdata/island/island_efc.xml"; -TEST_F(CoreConstraintTest, MulJacVecIsland) { +// validate mj_constraintUpdate_impl +TEST_F(CoreConstraintTest, ConstraintUpdateImpl) { const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec_nv, fill with arbitrary values - mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - for (int i=0; i < model->nv; i++) { - vec_nv[i] = 0.2 + 0.3*i; - } - - // iterate through dense and sparse - for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { - model->opt.jacobian = sparsity; - - // simulate for 0.2 seconds - mj_resetData(model, data); - while (data->time < 0.2) { - mj_step(model, data); - } - mj_forward(model, data); - - // multiply by Jacobian: vec_nefc = J * vec_nv - mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - mj_mulJacVec(model, data, vec_nefc, vec_nv); - mjtNum* vec_nefc_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof and efc vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_nvi = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - int efcnum = data->island_efcnum[i]; - mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - - // get indices - int* dofind = data->island_dofind + data->island_dofadr[i]; - int* efcind = data->island_efcind + data->island_efcadr[i]; - - // copy values into vec_nvi - for (int j=0; j < dofnum; j++) { - vec_nvi[j] = vec_nv[dofind[j]]; - } - - // ===== both compressed - int flg_resunc = 0; - int flg_vecunc = 0; - mju_zero(vec_nefci, efcnum); // clear output - mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - // ===== input uncompressed: read from vec_nv - flg_resunc = 0; - flg_vecunc = 1; - mju_zero(vec_nefci, efcnum); // clear output - mj_mulJacVec_island(model, data, vec_nefci, vec_nv, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - // ===== output uncompressed: write to vec_nefc_tmp - flg_resunc = 1; - flg_vecunc = 0; - mju_zero(vec_nefc_tmp, data->nefc); // clear output - mj_mulJacVec_island(model, data, vec_nefc_tmp, vec_nvi, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefc_tmp[efcind[j]], - DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - mju_free(vec_nvi); - mju_free(vec_nefci); - } - - mju_free(vec_nefc_tmp); - mju_free(vec_nefc); - } - - mju_free(vec_nv); - mj_deleteData(data); - mj_deleteModel(model); -} - -TEST_F(CoreConstraintTest, MulJacTVecIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec_nv - mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - mjtNum* vec_nv_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - - // iterate through dense and sparse - for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { - model->opt.jacobian = sparsity; - - // simulate for 0.3 seconds - mj_resetData(model, data); - while (data->time < 0.3) { - mj_step(model, data); - } - mj_forward(model, data); - - // allocate vec_nefc, fill with arbitrary values - mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - for (int i=0; i < data->nefc; i++) { - vec_nefc[i] = 0.2 + 0.3*i; - } - - // multiply by Jacobian: vec_nv = J^T * vec_nefc - mj_mulJacTVec(model, data, vec_nv, vec_nefc); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof and efc vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_nvi = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - int efcnum = data->island_efcnum[i]; - mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - - // get indices - int* efcind = data->island_efcind + data->island_efcadr[i]; - int* dofind = data->island_dofind + data->island_dofadr[i]; - - // copy values into vec_nefci - for (int j=0; j < efcnum; j++) { - vec_nefci[j] = vec_nefc[efcind[j]]; - } - - // ==== both compressed - int flg_resunc = 0; - int flg_vecunc = 0; - mju_zero(vec_nvi, dofnum); // clear output - mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - // ===== input uncompressed: read from vec_nefc - flg_resunc = 0; - flg_vecunc = 1; - mju_zero(vec_nvi, dofnum); // clear output - mj_mulJacTVec_island(model, data, vec_nvi, vec_nefc, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - // ===== output uncompressed: write to vec_nv_tmp - flg_resunc = 1; - flg_vecunc = 0; - mju_zero(vec_nv_tmp, model->nv); // clear output - mj_mulJacTVec_island(model, data, vec_nv_tmp, vec_nefci, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nv_tmp[dofind[j]], - DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - mju_free(vec_nvi); - mju_free(vec_nefci); - } - mju_free(vec_nefc); - } - - mju_free(vec_nv_tmp); - mju_free(vec_nv); - mj_deleteData(data); - mj_deleteModel(model); -} - -// compare mj_constraintUpdate and mj_constraintUpdate_island -TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data1 = mj_makeData(model); - mjData* data2 = mj_makeData(model); + char err[1024]; + mjModel* model = mj_loadXML(xml_path.c_str(), 0, err, 1024); + ASSERT_THAT(model, NotNull()) << err; + mjData* d1 = mj_makeData(model); + mjData* d2 = mj_makeData(model); // iterate over sparsity and cone for (mjtJacobian sparsity : {mjJAC_SPARSE, mjJAC_DENSE}) { @@ -491,81 +305,84 @@ TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { model->opt.cone = cone; // simulate for 0.2 seconds - mj_resetData(model, data1); - mj_resetData(model, data2); - while (data1->time < 0.2) { - mj_step(model, data1); - mj_step(model, data2); + mj_resetData(model, d1); + mj_resetData(model, d2); + while (d1->time < 0.2) { + mj_step(model, d1); + mj_step(model, d2); } - mj_forward(model, data1); - mj_forward(model, data2); + mj_forward(model, d1); + mj_forward(model, d2); // get sizes - int nefc = data1->nefc; + int nefc = d1->nefc; int nv = model->nv; - int nisland = data1->nisland; + int nisland = d1->nisland; EXPECT_GT(nisland, 0); // get jar = J*a - aref mjtNum* jar = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc); - mj_mulJacVec(model, data1, jar, data1->qacc); - mju_subFrom(jar, data1->efc_aref, nefc); + mj_mulJacVec(model, d1, jar, d1->qacc); + mju_subFrom(jar, d1->efc_aref, nefc); // constraint update for data1 given jar mjtNum cost1; - mj_constraintUpdate(model, data1, jar, &cost1, /*flg_coneHessian=*/1); + mj_constraintUpdate(model, d1, jar, &cost1, /*flg_coneHessian=*/1); // iterate over islands, check match mjtNum cost2 = 0; for (int island=0; island < nisland; island++) { // clear outputs from data2 - for (int i=0; i < nefc; i++) data2->efc_state[i] = -1; - mju_zero(data2->efc_force, nefc); - mju_zero(data2->qfrc_constraint, nv); - for (int i=0; i < data2->ncon; i++) mju_zero(data2->contact[i].H, 36); + for (int i=0; i < nefc; i++) d2->efc_state[i] = -1; + mju_zero(d2->efc_force, nefc); + mju_zero(d2->qfrc_constraint, nv); + for (int i=0; i < d2->ncon; i++) mju_zero(d2->contact[i].H, 36); // sizes and indices, in this island - int dofnum = data2->island_dofnum[island]; - int efcnum = data2->island_efcnum[island]; - int* dofind = data2->island_dofind + data2->island_dofadr[island]; - int* efcind = data2->island_efcind + data2->island_efcadr[island]; + int efcnum = d2->island_nefc[island]; - // get jar restricted to island + // gather values into jari mjtNum* jari = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - for (int c=0; c < efcnum; c++) { - jari[c] = jar[efcind[c]]; - } + int* map2efc = d2->map_iefc2efc + d2->island_iefcadr[island]; + mju_gather(jari, jar, map2efc, efcnum); // update constraints for this island mjtNum cost2i; - mj_constraintUpdate_island(model, data2, jari, &cost2i, - /*flg_coneHessian=*/1, island); + int ne = d2->island_ne[island]; + int nf = d2->island_nf[island]; + int adr = d2->island_iefcadr[island]; + int* state = d2->iefc_state + adr; + mjtNum *force = d2->iefc_force + adr; + mj_constraintUpdate_impl(ne, nf, efcnum, + d2->iefc_D + adr, + d2->iefc_R + adr, + d2->iefc_frictionloss + adr, + jari, + d2->iefc_type + adr, + d2->iefc_id + adr, + d2->contact, + state, + force, + &cost2i, + /*flg_coneHessian=*/1); // compare nefc vectors for (int c=0; c < efcnum; c++) { - int i = efcind[c]; - EXPECT_EQ(data2->efc_island[i], island); - EXPECT_EQ(data2->efc_state[i], data1->efc_state[i]); - EXPECT_THAT(data2->efc_force[i], - DoubleNear(data1->efc_force[i], 1e-12)); - } - - // compare qfrc_constraint - for (int c=0; c < dofnum; c++) { - int i = dofind[c]; - EXPECT_THAT(data2->qfrc_constraint[i], - DoubleNear(data1->qfrc_constraint[i], 1e-12)); + int i = map2efc[c]; + EXPECT_EQ(d2->efc_island[i], island); + EXPECT_EQ(state[c], d1->efc_state[i]); + EXPECT_THAT(force[c], DoubleNear(d1->efc_force[i], 1e-12)); } // compare cone Hessians if (cone == mjCONE_ELLIPTIC) { - for (int c=0; c < data2->ncon; c++) { - int efcadr = data2->contact[c].efc_address; - if (data2->efc_island[efcadr] == island && - data2->efc_state[efcadr] == mjCNSTRSTATE_CONE) { + for (int c=0; c < d2->ncon; c++) { + int efcadr = d2->contact[c].efc_address; + if (d2->efc_island[efcadr] == island && + d2->efc_state[efcadr] == mjCNSTRSTATE_CONE) { for (int j=0; j < 36; j++) { - EXPECT_THAT(data2->contact[c].H[j], - DoubleNear(data1->contact[c].H[j], 1e-12)); + EXPECT_THAT(d2->contact[c].H[j], + DoubleNear(d1->contact[c].H[j], 1e-12)); } } } @@ -584,8 +401,8 @@ TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { } } - mj_deleteData(data2); - mj_deleteData(data1); + mj_deleteData(d2); + mj_deleteData(d1); mj_deleteModel(model); } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 3263acc5..880f3e7e 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -15,8 +15,10 @@ // Tests for engine/engine_core_smooth.c. #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "src/engine/engine_util_sparse.h" +#include #include #include #include @@ -159,6 +161,237 @@ TEST_F(CoreSmoothTest, FixedTendonSortedIndices) { mj_deleteModel(model); } +static const char* const kTen_J0 = "engine/testdata/core_smooth/ten_J0.xml"; +static const char* const kTen_J1 = "engine/testdata/core_smooth/ten_J1.xml"; +static const char* const kTen_J2 = "engine/testdata/core_smooth/ten_J2.xml"; +static const char* const kTen_J3 = "engine/testdata/core_smooth/ten_J3.xml"; + +TEST_F(CoreSmoothTest, TendonJdot) { + for (const char* local_path : {kTen_J0, kTen_J1, kTen_J2, kTen_J3}) { + const std::string xml_path = GetTestDataFilePath(local_path); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + int nv = m->nv; + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + EXPECT_EQ(m->ntendon, 1); + mjData* d = mj_makeData(m); + + for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = sparsity; + + if (m->nkey) { + mj_resetDataKeyframe(m, d, 0); + } else { + mj_resetData(m, d); + while (d->time < 1) { + mj_step(m, d); + } + } + + mj_forward(m, d); + + // get current J and Jdot for the tendon + vector ten_J(d->ten_J, d->ten_J + nv); + vector ten_Jdot(nv, 0); + mj_tendonDot(m, d, 0, ten_Jdot.data()); + + // compute finite-differenced Jdot + mjtNum h = 1e-7; + mj_integratePos(m, d->qpos, d->qvel, h); + mj_kinematics(m, d); + mj_comPos(m, d); + mj_tendon(m, d); + vector ten_Jh(d->ten_J, d->ten_J + nv); + mju_subFrom(ten_Jh.data(), ten_J.data(), nv); + mju_scl(ten_Jh.data(), ten_Jh.data(), 1.0 / h, nv); + + // expect analytic and FD derivatives to be similar to eps precision + mjtNum eps = 1e-6; + EXPECT_THAT(ten_Jdot, Pointwise(DoubleNear(eps), ten_Jh)); + } + + mj_deleteData(d); + mj_deleteModel(m); + } +} + +static const char* const kTen_offtree = + "engine/testdata/core_smooth/ten_armature_offtree.xml"; + +TEST_F(CoreSmoothTest, TendonArmature) { + const std::string xml_path = GetTestDataFilePath(kTen_offtree); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + int nv = m->nv; + mjData* d = mj_makeData(m); + + for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = sparsity; + + mj_forward(m, d); + + // get full M, includes both CRB and tendon inertia + vector M(nv*nv); + mj_fullM(m, M.data(), d->qM); + + // put only CRB inertia in M2 + mj_crb(m, d); + vector M2(nv*nv); + mj_fullM(m, M2.data(), d->qM); + + vector ten_J(nv); // tendon Jacobian + vector ten_M(nv*nv); // tendon inertia + + // add tendon inertias to M2 using outer product + for (int j=0; j < m->ntendon; j++) { + // get tendon Jacobian + if (mj_isSparse(m)) { + int rowadr = d->ten_J_rowadr[j]; + int* rownnz = d->ten_J_rownnz + j; + int zero = 0; + mju_sparse2dense(ten_J.data(), d->ten_J + rowadr, 1, nv, + rownnz, &zero, d->ten_J_colind + rowadr); + } else { + mju_copy(ten_J.data(), d->ten_J + j*nv, nv); + } + + // get tendon inertia only, using outer product + mju_mulMatMat(ten_M.data(), ten_J.data(), ten_J.data(), nv, 1, nv); + mju_scl(ten_M.data(), ten_M.data(), m->tendon_armature[j], nv * nv); + + // manually add values, at nonzeros only + for (int i=0; i < nv*nv; i++) { + if (M[i]) M2[i] += ten_M[i]; + } + } + + // expect matrices to match + EXPECT_THAT(M2, Pointwise(DoubleNear(1e-9), M)); + } + + mj_deleteData(d); + mj_deleteModel(m); +} + +static const char* const kTen_i0 = + "engine/testdata/core_smooth/ten_armature_0.xml"; +static const char* const kTen_i1 = + "engine/testdata/core_smooth/ten_armature_1.xml"; +static const char* const kTen_i2 = + "engine/testdata/core_smooth/ten_armature_2.xml"; +static const char* const kTen_i3 = + "engine/testdata/core_smooth/ten_armature_3.xml"; +static const char* const kTen_i4 = + "engine/testdata/core_smooth/ten_armature_4.xml"; + +TEST_F(CoreSmoothTest, TendonArmatureConservesEnergy) { + for (const char* local_path : {kTen_i0, kTen_i1, kTen_i2, kTen_i3, kTen_i4}) { + const std::string xml_path = GetTestDataFilePath(local_path); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + mjData* d = mj_makeData(m); + + for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = sparsity; + + mj_resetDataKeyframe(m, d, 0); + mj_forward(m, d); + + double energy_0 = d->energy[0] + d->energy[1]; + + double eps = std::max(energy_0, 1.0) * 1e-5; + while (d->time < 1) { + mj_step(m, d); + double energy_t = d->energy[0] + d->energy[1]; + EXPECT_THAT(energy_t, DoubleNear(energy_0, eps)); + } + } + mj_deleteData(d); + mj_deleteModel(m); + } +} + +TEST_F(CoreSmoothTest, TendonArmatureConservesMomentum) { + const std::string xml_path = GetTestDataFilePath(kTen_i4); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + mjData* d = mj_makeData(m); + + for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = sparsity; + + mj_resetData(m, d); + mj_forward(m, d); + + // this model contains subtreelinvel and subtreeangmom sensors + vector sdata_0 = AsVector(d->sensordata, m->nsensordata); + EXPECT_THAT(sdata_0, Each(Eq(0))); + + double eps = 1e-5; + while (d->time < 1) { + mj_step(m, d); + vector sdata_t = AsVector(d->sensordata, m->nsensordata); + EXPECT_THAT(sdata_t, Pointwise(DoubleNear(eps), sdata_0)); + } + + // momentum is conserved nontrivially (velocities are non-zero) + EXPECT_GT(d->energy[1], 0); + } + + mj_deleteData(d); + mj_deleteModel(m); +} + +static const char* const kTen_i0_equiv = + "engine/testdata/core_smooth/ten_armature_0_equiv.xml"; +static const char* const kTen_i1_equiv = + "engine/testdata/core_smooth/ten_armature_1_equiv.xml"; + +TEST_F(CoreSmoothTest, TendonInertiaEquivalent) { + for (const char* lpath : {kTen_i0, kTen_i1}) { + // load tendon model + const std::string path = GetTestDataFilePath(lpath); + char error[1024]; + mjModel* m = mj_loadXML(path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + int gid = mj_name2id(m, mjOBJ_GEOM, "query"); + mjData* d = mj_makeData(m); + + if (m->nkey) mj_resetDataKeyframe(m, d, 0); + + // load equivalent model + const char* lpath_e = lpath == kTen_i0 ? kTen_i0_equiv : kTen_i1_equiv; + const std::string path_e = GetTestDataFilePath(lpath_e); + mjModel* m_e = mj_loadXML(path_e.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + int gid_e = mj_name2id(m_e, mjOBJ_GEOM, "query"); + mjData* d_e = mj_makeData(m_e); + + if (m_e->nkey) mj_resetDataKeyframe(m_e, d_e, 0); + + // the equality constraint in kTen_i1_equiv reduces precision + double eps = lpath == kTen_i0 ? 1e-6 : 1e-3; + + while (d->time < 1) { + mj_step(m, d); + vector xpos = AsVector(d->geom_xpos + 3*gid, 3); + + mj_step(m_e, d_e); + vector xpos_e = AsVector(d_e->geom_xpos + 3*gid_e, 3); + + EXPECT_THAT(xpos, Pointwise(DoubleNear(eps), xpos_e)); + } + mj_deleteData(d); + mj_deleteModel(m); + mj_deleteData(d_e); + mj_deleteModel(m_e); + } +} + + // --------------------------- connect constraint ------------------------------ // test that bodies hanging on connects lead to expected force sensor readings @@ -401,66 +634,6 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) { mj_deleteModel(model); } -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; -static const char* const kModelPath = - "testdata/model.xml"; - -TEST_F(CoreSmoothTest, SolveMIsland) { - for (auto model_path : {kModelPath, kIlslandEfcPath}) { - const std::string xml_path = GetTestDataFilePath(model_path); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - int nv = model->nv; - - // allocate vec, fill with arbitrary values, copy to sol - mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv); - mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv); - for (int i=0; i < nv; i++) { - vec[i] = 0.2 + 0.3*i; - } - mju_copy(res, vec, nv); - - if (model->nkey > 0) mj_resetDataKeyframe(model, data, 0); - - for (int i=0; i < 6; i++) { - mj_step(model, data); - } - - mj_forward(model, data); - - // divide by mass matrix: sol = M^-1 * vec - mj_solveM(model, data, res, res, 1); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - - // copy values into sol_i - int* dofind = data->island_dofind + data->island_dofadr[i]; - for (int j=0; j < dofnum; j++) { - res_i[j] = vec[dofind[j]]; - } - - // divide by mass matrix, for this island - mj_solveM_island(model, data, res_i, i); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12)); - } - mju_free(res_i); - } - - mju_free(res); - mju_free(vec); - mj_deleteData(data); - mj_deleteModel(model); - } -} - static const char* const kInertiaPath = "engine/testdata/inertia.xml"; TEST_F(CoreSmoothTest, FactorI) { @@ -521,12 +694,11 @@ TEST_F(CoreSmoothTest, SolveLDs) { int nv = m->nv; int nM = m->nM; + int nC = m->nC; // copy M into LD: Legacy format - vector LDlegacy(nM); - for (int i=0; i < nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + vector LDlegacy(nM, 0); + mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nC); // compare LD and LDs densified matrices vector LDdense(nv*nv); @@ -550,7 +722,7 @@ TEST_F(CoreSmoothTest, SolveLDs) { mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); // expect vectors to match up to floating point precision for (int i=0; i < nv; i++) { @@ -571,13 +743,10 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { mj_forward(m, d); int nv = m->nv; - int nM = m->nM; // copy LD into LDlegacy: Legacy format - vector LDlegacy(nM); - for (int i=0; i < nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + vector LDlegacy(m->nM, 0); + mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, m->nC); // compare n LD and LDs vector solve int n = 3; @@ -588,7 +757,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); // expect vectors to match up to floating point precision for (int i=0; i < nv*n; i++) { @@ -626,7 +795,7 @@ TEST_F(CoreSmoothTest, SolveM2) { mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v) for (int i=0; i < n; i++) { @@ -647,29 +816,25 @@ TEST_F(CoreSmoothTest, FactorIs) { mjData* d = mj_makeData(m); mj_forward(m, d); - int nM = m->nM, nv = m->nv; + int nC = m->nC, nM = m->nM, nv = m->nv; // copy qM into into qLDlegacy and factorize vector qLDlegacy(nM); mj_factorI_legacy(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv); // copy qLDlegacy into qLDexpected: CSR format - vector qLDexpected(nM); - for (int i=0; i < nM; i++) { - qLDexpected[i] = qLDlegacy[d->mapM2M[i]]; - } + vector qLDexpected(nC); + mju_gather(qLDexpected.data(), qLDlegacy.data(), d->mapM2M, nC); // copy qM into qLD: CSR format - vector qLD(nM); - for (int i=0; i < nM; i++) { - qLD[i] = d->qM[d->mapM2M[i]]; // mj_factorI is in-place - } + vector qLD(nC); + mju_gather(qLD.data(), d->qM, d->mapM2M, nC); vector qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); vector qLDiagInv(nv, 0); mj_factorI(qLD.data(), qLDiagInv.data(), nv, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); // expect outputs to match to floating point precision EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected)); @@ -679,12 +844,12 @@ TEST_F(CoreSmoothTest, FactorIs) { vector LDdense(nv*nv); mju_sparse2dense(LDdense.data(), qLDexpected.data(), nv, nv, - d->M_rownnz, d->M_rowadr, d->M_colind); - PrintMatrix(LDdense.data(), nv, nv, 2, "qLDexpected"); + d->C_rownnz, d->C_rowadr, d->C_colind); + PrintMatrix(LDdense.data(), nv, nv, 2); - mju_sparse2dense(LDdense.data(), qLD.data(), nv, nv, - d->M_rownnz, d->M_rowadr, d->M_colind); - PrintMatrix(LDdense.data(), nv, nv, 2, "qLD"); + mju_sparse2dense(LDdense.data(), qLDs.data(), nv, nv, + d->C_rownnz, d->C_rowadr, d->C_colind); + PrintMatrix(LDdense.data(), nv, nv, 2); */ mj_deleteData(d); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 027578e7..3b0a8d92 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -436,7 +436,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; } mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); // A = [dt*Ac; Ac] mju_transpose(A, Ac, 2*nv, nv); @@ -464,7 +464,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz, d->moment_rowadr, d->moment_colind); mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); + d->M_rownnz, d->M_rowadr, d->M_colind); mju_transpose(BcT, Bc, nu, nv); mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv); diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index 16a52f43..d345b065 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -46,6 +46,8 @@ static const char* const kDampedActuatorsPath = "engine/testdata/derivative/damped_actuators.xml"; static const char* const kJointForceClamp = "engine/testdata/actuation/joint_force_clamp.xml"; +static const char* const kTendonForceClamp = + "engine/testdata/actuation/tendon_force_clamp.xml"; using ::testing::Pointwise; using ::testing::DoubleNear; @@ -1393,5 +1395,57 @@ TEST_F(ActuatorTest, DisableActuatorOutOfRange) { mj_deleteModel(model); } +TEST_F(ActuatorTest, TendonActuatorForceRange) { + const std::string xml_path = GetTestDataFilePath(kTendonForceClamp); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + + EXPECT_EQ(model->tendon_actfrclimited[0], 0); + EXPECT_EQ(model->tendon_actfrcrange[0], 0); + EXPECT_EQ(model->tendon_actfrcrange[1], 0); + + EXPECT_EQ(model->tendon_actfrclimited[1], 1); + EXPECT_EQ(model->tendon_actfrcrange[2], -1); + EXPECT_EQ(model->tendon_actfrcrange[3], 1); + + EXPECT_EQ(model->tendon_actfrclimited[2], 1); + EXPECT_EQ(model->tendon_actfrcrange[4], -10); + EXPECT_EQ(model->tendon_actfrcrange[5], 10); + + EXPECT_EQ(model->tendon_actfrclimited[3], 1); + EXPECT_EQ(model->tendon_actfrcrange[6], 0); + EXPECT_EQ(model->tendon_actfrcrange[7], 1); + + data->ctrl[0] = 1; + data->ctrl[1] = 1; + data->ctrl[2] = 1; + + data->ctrl[3] = -1; + data->ctrl[4] = 1; + + data->ctrl[5] = -20; + data->ctrl[6] = 5; + data->ctrl[7] = -5; + + mj_forward(model, data); + + EXPECT_NEAR(data->actuator_force[0], 1, 1e-6); + EXPECT_NEAR(data->actuator_force[1], 1, 1e-6); + EXPECT_NEAR(data->actuator_force[2], 1, 1e-6); + EXPECT_NEAR(data->actuator_force[3], -1, 1e-6); + EXPECT_NEAR(data->actuator_force[4], 1, 1e-6); + EXPECT_NEAR(data->actuator_force[5], -10, 1e-6); + EXPECT_NEAR(data->actuator_force[6], 5, 1e-6); + EXPECT_NEAR(data->actuator_force[7], -5, 1e-6); + + EXPECT_EQ(data->sensordata[0], 3); + EXPECT_EQ(data->sensordata[1], 0); + EXPECT_EQ(data->sensordata[2], -10); + EXPECT_EQ(data->sensordata[3], 0); + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index 470f6f02..72e942d1 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -226,11 +226,42 @@ TEST_F(EngineIoTest, MjvCopyModel) { EXPECT_FLOAT_EQ(model2->mesh_vert[0], 0.1); // unchanged EXPECT_FLOAT_EQ(model2->geom_rgba[0], 0.4); - // mj_deleteData(data); mj_deleteModel(model2); mj_deleteModel(model1); } +TEST_F(EngineIoTest, MjvCopyData) { + static constexpr char xml[] = R"( + + + + + + + + + + )"; + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + + mjData* data1 = mj_makeData(model); + mj_forward(model, data1); + EXPECT_THAT(data1->efc_J, NotNull()); + + mjData* data2 = mj_copyData(nullptr, model, data1); + EXPECT_THAT(data2->efc_J, NotNull()); + + mj_deleteData(data2); + data2 = mjv_copyData(nullptr, model, data1); + EXPECT_THAT(data2->efc_J, IsNull()); + + mj_deleteData(data2); + mj_deleteData(data1); + mj_deleteModel(model); +} + using ValidateReferencesTest = MujocoTest; TEST_F(ValidateReferencesTest, BodyReferences) { diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3688c15b..e59ac3b4 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -208,17 +208,19 @@ TEST_F(IslandTest, Abacus) { int nv = model->nv; int nefc = data->nefc; int nisland = data->nisland; + int nidof = data->nidof; // 4 dofs, 12 constraints, 2 islands EXPECT_EQ(nv, 4); + EXPECT_EQ(nidof, 3); EXPECT_EQ(nefc, 12); // 3 pyramidal contacts EXPECT_EQ(nisland, 2); // the islands begin at dofs 0 and 1 - EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1)); // number of dofs in the 2 islands - EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2)); // dof 0 in island 0 // dof 1 in no island @@ -228,19 +230,19 @@ TEST_F(IslandTest, Abacus) { // dof 0 constitutes first island // dofs 2, 3 are the second island // last index is unassigned since dof 1 is unconstrained - EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 2, 3, -1)); + EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 2, 3, 1)); // dof 0 constitutes first island // dofs 1 is unassigned // dofs 2, 3 are second island - EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, -1, 0, 1)); + EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 3, 1, 2)); // island 0 starts at constraint 0 // island 1 starts at constraint 4 - EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4)); + EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4)); // number of constraints in the 2 islands - EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 8)); + EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 8)); // first contact (4 constraints) is in island 0 // second contact (8 constraints) is in island 1 @@ -248,7 +250,7 @@ TEST_F(IslandTest, Abacus) { ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1)); // index lists for islands 0 and 1 - EXPECT_THAT(AsVector(data->island_efcind, nefc), + EXPECT_THAT(AsVector(data->map_iefc2efc, nefc), ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)); // reset, push 0 to the left, 3 to the right, 1,2 to the middle @@ -266,18 +268,20 @@ TEST_F(IslandTest, Abacus) { // local variables nefc = data->nefc; nisland = data->nisland; + nidof = data->nidof; EXPECT_EQ(nisland, 3); - EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1, 3)); - EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2, 1)); + EXPECT_EQ(nidof, 4); + EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1, 3)); + EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2, 1)); EXPECT_THAT(AsVector(data->dof_island, nv), ElementsAre(0, 1, 1, 2)); - EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 1, 2, 3)); - EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, 0, 1, 0)); - EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4, 8)); - EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 4, 4)); + EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 1, 2, 3)); + EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 1, 2, 3)); + EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4, 8)); + EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 4, 4)); EXPECT_THAT(AsVector(data->efc_island, nefc), ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2)); - EXPECT_THAT(AsVector(data->island_efcind, nefc), + EXPECT_THAT(AsVector(data->map_iefc2efc, nefc), ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)); mj_deleteData(data); @@ -311,27 +315,30 @@ TEST_F(IslandTest, DenseSparse) { int nisland = data1->nisland; // expect sparse and dense to be identical + EXPECT_EQ(data1->nidof, data2->nidof); EXPECT_EQ(data1->nefc, data2->nefc); EXPECT_EQ(data1->nisland, data2->nisland); EXPECT_EQ(data1->nefc, data2->nefc); - EXPECT_EQ(AsVector(data1->island_dofadr, nisland), - AsVector(data2->island_dofadr, nisland)); - EXPECT_EQ(AsVector(data1->island_dofnum, nisland), - AsVector(data2->island_dofnum, nisland)); + EXPECT_EQ(AsVector(data1->island_idofadr, nisland), + AsVector(data2->island_idofadr, nisland)); + EXPECT_EQ(AsVector(data1->island_nv, nisland), + AsVector(data2->island_nv, nisland)); EXPECT_EQ(AsVector(data1->dof_island, nv), AsVector(data2->dof_island, nv)); - EXPECT_EQ(AsVector(data1->island_dofind, nv), - AsVector(data2->island_dofind, nv)); - EXPECT_EQ(AsVector(data1->dof_islandind, nv), - AsVector(data2->dof_islandind, nv)); - EXPECT_EQ(AsVector(data1->island_efcadr, nisland), - AsVector(data2->island_efcadr, nisland)); - EXPECT_EQ(AsVector(data1->island_efcnum, nisland), - AsVector(data2->island_efcnum, nisland)); + EXPECT_EQ(AsVector(data1->map_idof2dof, nv), + AsVector(data2->map_idof2dof, nv)); + EXPECT_EQ(AsVector(data1->map_dof2idof, nv), + AsVector(data2->map_dof2idof, nv)); + EXPECT_EQ(AsVector(data1->island_iefcadr, nisland), + AsVector(data2->island_iefcadr, nisland)); + EXPECT_EQ(AsVector(data1->island_nefc, nisland), + AsVector(data2->island_nefc, nisland)); EXPECT_EQ(AsVector(data1->efc_island, nefc), AsVector(data2->efc_island, nefc)); - EXPECT_EQ(AsVector(data1->island_efcind, nefc), - AsVector(data2->island_efcind, nefc)); + EXPECT_EQ(AsVector(data1->map_iefc2efc, nefc), + AsVector(data2->map_iefc2efc, nefc)); + EXPECT_EQ(AsVector(data1->map_efc2iefc, nefc), + AsVector(data2->map_efc2iefc, nefc)); mj_deleteData(data2); mj_deleteData(data1); @@ -361,6 +368,156 @@ TEST_F(IslandTest, IslandEfc) { mj_deleteModel(model); } +static const char* const k2H100Path = "engine/testdata/island/2humanoid100.xml"; + +TEST_F(IslandTest, IslandJacobian) { + for (const char* local_path : {kIlslandEfcPath, k2H100Path}) { + const std::string xml_path = GetTestDataFilePath(local_path); + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int jac0 = m->opt.jacobian; + mjData* d = mj_makeData(m); + + for (mjtNum t_stop : {0.0, 0.2, 2.0}) { + while (d->time < t_stop) { + mj_step(m, d); + } + + for (mjtJacobian jac : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = jac; + mj_forward(m, d); + + int nv = m->nv; + int nefc = d->nefc; + int nisland = d->nisland; + int nidof = d->nidof; + + mjtNum* J = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nv); + mjtNum* iJ = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nidof); + + // get local dense Jacobian + if (jac == mjJAC_DENSE) { + mju_copy(J, d->efc_J, nefc * nv); + mju_copy(iJ, d->iefc_J, nefc * nidof); + } else { + mju_sparse2dense(J, d->efc_J, nefc, nv, d->efc_J_rownnz, + d->efc_J_rowadr, d->efc_J_colind); + } + + // compare random access in efc_J to contiguous memory in iefc_J + for (int island=0; island < nisland; island++) { + int idof = d->island_idofadr[island]; + int iefc = d->island_iefcadr[island]; + int nefc_island = d->island_nefc[island]; + int nv_island = d->island_nv[island]; + + // === test J + + // get pointer to J_island, dense (nefc_island x nv_island) submatrix + mjtNum* J_island; + if (jac == mjJAC_DENSE) { + // point to starting address of island in efc_J + J_island = iJ + iefc * nidof; + } else { + // dense copy of island in iJ (here used as scratch) + mju_sparse2dense(iJ, d->iefc_J, nefc_island, nv_island, + d->iefc_J_rownnz + iefc, + d->iefc_J_rowadr + iefc, + d->iefc_J_colind); + J_island = iJ; + } + + // sequential memory in J_island equals random access memory in J + for (int i=0; i < nefc_island; i++) { + for (int j=0; j < nv_island; j++) { + int efc = d->map_iefc2efc[iefc + i]; + int dof = d->map_idof2dof[idof + j]; + EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]); + } + } + + // === test JT (if sparse) + + // get pointer to J_island, dense (nefc_island x nv_island) submatrix + if (jac == mjJAC_SPARSE) { + // dense copy of island in iJ (here used as scratch) + mju_sparse2dense(iJ, d->iefc_JT, nv_island, nefc_island, + d->iefc_JT_rownnz + idof, + d->iefc_JT_rowadr + idof, + d->iefc_JT_colind); + J_island = iJ; + + // sequential memory in J_island equals random access memory in J + for (int i=0; i < nv_island; i++) { + for (int j=0; j < nefc_island; j++) { + int dof = d->map_idof2dof[idof + i]; + int efc = d->map_iefc2efc[iefc + j]; + EXPECT_EQ(J_island[i * nefc_island + j], J[efc * nv + dof]); + } + } + } + } + + mju_free(iJ); + mju_free(J); + } + + // reset opt.jacobian to initial value + m->opt.jacobian = jac0; + } + + mj_deleteData(d); + mj_deleteModel(m); + } +} + +TEST_F(IslandTest, IslandInertia) { + for (const char* local_path : {kIlslandEfcPath, k2H100Path}) { + const std::string xml_path = GetTestDataFilePath(local_path); + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int nv = m->nv; + mjData* d = mj_makeData(m); + mjtNum* M = (mjtNum*)mju_malloc(sizeof(mjtNum) * nv * nv); + + for (mjtNum t_stop : {0.0, 0.2, 2.0}) { + while (d->time < t_stop) { + mj_step(m, d); + } + mj_forward(m, d); + + int nisland = d->nisland; + + // get dense inertia (lower only) + mj_fullM(m, M, d->qM); + + // compare iM sub-matrix to full M + for (int island=0; island < nisland; island++) { + int nvi = d->island_nv[island]; + mjtNum* Mi = (mjtNum*)mju_malloc(sizeof(mjtNum) * nvi * nvi); + + int adr = d->island_idofadr[island]; + mju_sparse2dense(Mi, d->iM, nvi, nvi, + d->iM_rownnz + adr, + d->iM_rowadr + adr, + d->iM_colind); + + // compare Mi to M (lower triangle only) + for (int i=0; i < nvi; i++) { + for (int j=0; j <= i; j++) { + int dofi = d->map_idof2dof[adr + j]; + int dofj = d->map_idof2dof[adr + i]; + EXPECT_EQ(Mi[i * nvi + j], M[dofi * nv + dofj]); + } + } + mju_free(Mi); + } + } + + mju_free(M); + mj_deleteData(d); + mj_deleteModel(m); + } +} + TEST_F(IslandTest, IslandEfcElliptic) { const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); diff --git a/test/engine/engine_sensor_test.cc b/test/engine/engine_sensor_test.cc index 5aa2aa5b..b5f0125a 100644 --- a/test/engine/engine_sensor_test.cc +++ b/test/engine/engine_sensor_test.cc @@ -454,6 +454,33 @@ TEST_F(SensorTest, PotentialEnergy) { mj_deleteModel(model); } +TEST_F(SensorTest, PotentialEnergyFreeJointSpring) { + constexpr char xml[] = R"( + + + )"; + mjModel* model = LoadModelFromString(xml); + mjData* data = mj_makeData(model); + data->qpos[0] = 1; + data->qpos[1] = 2; + data->qpos[2] = 3; + mj_forward(model, data); + EXPECT_EQ(data->sensordata[0], 0.5*2*14); + + mj_deleteData(data); + mj_deleteModel(model); +} + TEST_F(SensorTest, KineticEnergy) { constexpr char xml[] = R"( diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index 6343ce24..75b42cf7 100644 --- a/test/engine/engine_solver_test.cc +++ b/test/engine/engine_solver_test.cc @@ -17,7 +17,6 @@ #include #include #include -#include #include #include @@ -29,158 +28,135 @@ namespace { using ::testing::DoubleNear; using ::testing::NotNull; -using ::testing::Pointwise; -using ::std::vector; using ::std::abs; using ::std::max; -// compare two vectors, relative error (reduces size of large vector elements) -inline void ExpectEqRel(vector v1, vector v2, mjtNum rtol) { - ASSERT_TRUE(v1.size() == v2.size()); - - // make scale vector - int n = v1.size(); - vector scale(n); - for (int i = 0; i < n; i++) { - scale[i] = max(1.0, abs(v1[i]) + abs(v2[i])); - } - - // scale and compare - for (int i = 0; i < n; i++) { - v1[i] /= scale[i]; - v2[i] /= scale[i]; - } - EXPECT_THAT(v1, Pointwise(DoubleNear(rtol), v2)); -} - using SolverTest = MujocoTest; -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; +static const char* const kModelPath = + "testdata/model.xml"; // compare accelerations produced by CG solver with and without islands TEST_F(SolverTest, IslandsEquivalent) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); + const std::string xml_path = GetTestDataFilePath(kModelPath); + char error[1024]; + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + model->opt.solver = mjSOL_CG; // use CG solver + model->opt.jacobian = mjJAC_SPARSE; // use sparse + model->opt.tolerance = 0; // set tolerance to 0 + model->opt.ls_tolerance = 0; // set ls_tolerance to 0 + + int nv = model->nv; + + int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); + mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size); + + mjData* data_island = mj_makeData(model); + mjData* data_noisland = mj_makeData(model); + + // Below are 3 tolerances associated with 3 different iteration counts, + // they are only moderately tight, 2x higher than x86-64 failure on Linux, + // i.e. in that case the test fails with rtol smaller than {5e-3, 5e-4, 5e-5}. + // The point of this test is to show that CG convergence is actually not very + // precise, simply changing whether islands are used changes the solution by + // quite a lot, even at high iteration count and zero {ls_}tolerance. + // Increasing the iteration count higher than 60 does not improve convergence. + constexpr int kNumTol = 3; + mjtNum maxiter[kNumTol] = {30, 40, 60}; + mjtNum rtol[kNumTol] = {1e-2, 1e-3, 1e-4}; + + for (int i = 0; i < kNumTol; ++i) { + model->opt.iterations = maxiter[i]; + model->opt.ls_iterations = maxiter[i]; + + for (bool coldstart : {true, false}) { + mj_resetDataKeyframe(model, data_noisland, 0); + + if (coldstart) { + model->opt.disableflags |= mjDSBL_WARMSTART; + } else { + model->opt.disableflags &= ~mjDSBL_WARMSTART; + } + + while (data_noisland->time < .1) { + mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION); + mj_setState(model, data_island, state, mjSTATE_INTEGRATION); + + model->opt.enableflags |= mjENBL_ISLAND; // enable islands + mj_forward(model, data_island); + + model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands + mj_forward(model, data_noisland); + + auto time = std::to_string(data_noisland->time); + for (int j = 0; j < nv; j++) { + // increase tolerance for large elements + mjtNum scale = 0.5 * max(2.0, abs(data_noisland->qacc[j]) + + abs(data_island->qacc[j])); + EXPECT_THAT(data_noisland->qacc[j], + DoubleNear(data_island->qacc[j], scale * rtol[i])) + << "time: " << time << '\n' + << "dof: " << j << '\n' + << "maxiter: " << maxiter[i] << '\n' + << "rtol: " << scale * rtol[i]; + } + + mj_step(model, data_noisland); + } + } + } + + mj_deleteData(data_noisland); + mj_deleteData(data_island); + mju_free(state); + mj_deleteModel(model); +} + +// compare accelerations produced by CG solver with and without islands +TEST_F(SolverTest, IslandsEquivalentForward) { + const std::string xml_path = GetTestDataFilePath(kModelPath); char error[1024]; mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); ASSERT_THAT(model, NotNull()) << error; model->opt.solver = mjSOL_CG; // use CG solver model->opt.tolerance = 0; // set tolerance to 0 - model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands + model->opt.ls_tolerance = 0; // set ls_tolerance to 0 - int nv = model->nv; - - int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); - mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size); - mjtNum* qacc_diff = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv); + mjtNum rtol = 2e-6; mjData* data_island = mj_makeData(model); mjData* data_noisland = mj_makeData(model); - mjtNum rtol = 1e-5; + for (bool coldstart : {true, false}) { + mj_resetDataKeyframe(model, data_island, 0); + mj_resetDataKeyframe(model, data_noisland, 0); - for (bool warmstart : {true, false}) { - if (warmstart) { + if (coldstart) { model->opt.disableflags |= mjDSBL_WARMSTART; } else { model->opt.disableflags &= ~mjDSBL_WARMSTART; } - mj_resetData(model, data_noisland); - while (data_noisland->time < .3) { - mj_step(model, data_noisland); - - mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION); - mj_setState(model, data_island, state, mjSTATE_INTEGRATION); - - mj_forward(model, data_noisland); - - model->opt.enableflags |= mjENBL_ISLAND; // enable islands - mj_forward(model, data_island); - model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands - - ExpectEqRel(AsVector(data_noisland->qacc, nv), - AsVector(data_island->qacc, nv), rtol); - } - } - - mj_deleteData(data_noisland); - mj_deleteData(data_island); - mju_free(qacc_diff); - mju_free(state); - mj_deleteModel(model); -} - -// compare qacc from 1 iteration of monolithic CG solver and one big island -TEST_F(SolverTest, OneBigIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - ASSERT_THAT(model, NotNull()); - model->opt.solver = mjSOL_CG; // use CG solver - model->opt.disableflags |= mjDSBL_WARMSTART; // disable warmstart - model->opt.tolerance = 0; // set tolerance to 0 - model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands - - int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); - mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size); - - mjData* data_island = mj_makeData(model); - mjData* data_noisland = mj_makeData(model); - - int nv = model->nv; - mjtNum rtol = 1e-7; - - // save current (default) iterations - int iterations_default = model->opt.iterations; - - while (data_noisland->time < .2) { - // step and copy the state to data_island - mj_step(model, data_noisland); - mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION); - mj_setState(model, data_island, state, mjSTATE_INTEGRATION); - - // set small number of iterations - model->opt.iterations = 1; - - // call forward on data_noisland + model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands mj_forward(model, data_noisland); - // enable islands - model->opt.enableflags |= mjENBL_ISLAND; - - // call forward (just for smooth dynamics and to allocate islands) + model->opt.enableflags |= mjENBL_ISLAND; // enable islands mj_forward(model, data_island); - - // overwrite island structure with one big island - data_island->nisland = 1; - data_island->island_dofnum[0] = nv; - data_island->island_dofadr[0] = 0; - for (int i = 0; i < nv; i++) { - data_island->island_dofind[i] = data_island->dof_islandind[i] = i; + for (int j = 0; j < model->nv; j++) { + mjtNum scale = 0.5 * max(2.0, abs(data_noisland->qacc[j]) + + abs(data_island->qacc[j])); + EXPECT_THAT(data_noisland->qacc[j], + DoubleNear(data_island->qacc[j], scale * rtol)) + << "dof: " << j << '\n' + << "rtol: " << scale * rtol; } - int nefc = data_island->nefc; - data_island->island_efcnum[0] = nefc; - data_island->island_efcadr[0] = 0; - for (int i = 0; i < nefc; i++) data_island->island_efcind[i] = i; - - // solve using using one big island - mj_fwdConstraint(model, data_island); - - // re-disable islands and reset iterations - model->opt.enableflags &= ~mjENBL_ISLAND; - model->opt.iterations = iterations_default; - - // compare accelerations (relative error) - ExpectEqRel(AsVector(data_noisland->qacc, nv), - AsVector(data_island->qacc, nv), rtol); } mj_deleteData(data_noisland); mj_deleteData(data_island); - mju_free(state); mj_deleteModel(model); } - } // namespace } // namespace mujoco diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index c7d4f018..71b8e820 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -830,77 +830,6 @@ TEST_F(InertiaTest, mulM2) { mj_deleteModel(model); } -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; - -TEST_F(SupportTest, MulMIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec, fill with arbitrary values - mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - for (int i=0; i < model->nv; i++) { - vec[i] = 0.2 + 0.3*i; - } - - // simulate for 0.2 seconds - mj_resetData(model, data); - while (data->time < 0.2) { - mj_step(model, data); - } - mj_forward(model, data); - - // multiply by Mass matrix: Mvec = M * vec - mjtNum* Mvec = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - mj_mulM(model, data, Mvec, vec); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - mjtNum* Mvec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - - // copy values into vec_i - int* dofind = data->island_dofind + data->island_dofadr[i]; - for (int j=0; j < dofnum; j++) { - vec_i[j] = vec[dofind[j]]; - } - - // === compressed: use vec_i - - // multiply by Jacobian, for this island - int flg_vecunc = 0; - mj_mulM_island(model, data, Mvec_i, vec_i, i, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12)); - } - - // === uncompressed: use vec - mju_zero(Mvec_i, dofnum); // clear output - - // multiply by Jacobian, for this island - flg_vecunc = 1; - mj_mulM_island(model, data, Mvec_i, vec, i, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12)); - } - - mju_free(vec_i); - mju_free(Mvec_i); - } - - mju_free(Mvec); - mju_free(vec); - mj_deleteData(data); - mj_deleteModel(model); -} - static constexpr char GeomDistanceTestingModel[] = R"( - + - + diff --git a/test/engine/testdata/collision_convex/dice.xml b/test/engine/testdata/collision_convex/dice.xml index 2a3b6880..5f077f4c 100644 --- a/test/engine/testdata/collision_convex/dice.xml +++ b/test/engine/testdata/collision_convex/dice.xml @@ -25,7 +25,7 @@ - > + @@ -37,7 +37,7 @@ - + diff --git a/test/engine/testdata/collision_convex/dice_boxmesh.xml b/test/engine/testdata/collision_convex/dice_boxmesh.xml new file mode 100644 index 00000000..cfd3ae86 --- /dev/null +++ b/test/engine/testdata/collision_convex/dice_boxmesh.xml @@ -0,0 +1,47 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_J0.xml b/test/engine/testdata/core_smooth/ten_J0.xml new file mode 100644 index 00000000..b4fec9ba --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_J0.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_J1.xml b/test/engine/testdata/core_smooth/ten_J1.xml new file mode 100644 index 00000000..d5c3cf93 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_J1.xml @@ -0,0 +1,29 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_J2.xml b/test/engine/testdata/core_smooth/ten_J2.xml new file mode 100644 index 00000000..e5777029 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_J2.xml @@ -0,0 +1,34 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_J3.xml b/test/engine/testdata/core_smooth/ten_J3.xml new file mode 100644 index 00000000..520969d2 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_J3.xml @@ -0,0 +1,66 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_0.xml b/test/engine/testdata/core_smooth/ten_armature_0.xml new file mode 100644 index 00000000..b27e2438 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_0.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_0_compare.xml b/test/engine/testdata/core_smooth/ten_armature_0_compare.xml new file mode 100644 index 00000000..d2875a6e --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_0_compare.xml @@ -0,0 +1,17 @@ + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_0_equiv.xml b/test/engine/testdata/core_smooth/ten_armature_0_equiv.xml new file mode 100644 index 00000000..0da57dfb --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_0_equiv.xml @@ -0,0 +1,20 @@ + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_1.xml b/test/engine/testdata/core_smooth/ten_armature_1.xml new file mode 100644 index 00000000..35b48b38 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_1.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_1_compare.xml b/test/engine/testdata/core_smooth/ten_armature_1_compare.xml new file mode 100644 index 00000000..31f2a02f --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_1_compare.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_1_equiv.xml b/test/engine/testdata/core_smooth/ten_armature_1_equiv.xml new file mode 100644 index 00000000..dcf62b56 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_1_equiv.xml @@ -0,0 +1,34 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_2.xml b/test/engine/testdata/core_smooth/ten_armature_2.xml new file mode 100644 index 00000000..9313d875 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_2.xml @@ -0,0 +1,29 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_3.xml b/test/engine/testdata/core_smooth/ten_armature_3.xml new file mode 100644 index 00000000..afaec4e9 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_3.xml @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_4.xml b/test/engine/testdata/core_smooth/ten_armature_4.xml new file mode 100644 index 00000000..232c4094 --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_4.xml @@ -0,0 +1,47 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_smooth/ten_armature_offtree.xml b/test/engine/testdata/core_smooth/ten_armature_offtree.xml new file mode 100644 index 00000000..46865eac --- /dev/null +++ b/test/engine/testdata/core_smooth/ten_armature_offtree.xml @@ -0,0 +1,77 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/island/2humanoid100.xml b/test/engine/testdata/island/2humanoid100.xml new file mode 100644 index 00000000..424ed2c4 --- /dev/null +++ b/test/engine/testdata/island/2humanoid100.xml @@ -0,0 +1,123 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/island/humanoid.xml b/test/engine/testdata/island/humanoid.xml new file mode 100644 index 00000000..7545e193 --- /dev/null +++ b/test/engine/testdata/island/humanoid.xml @@ -0,0 +1,252 @@ + + diff --git a/test/experimental/CMakeLists.txt b/test/experimental/CMakeLists.txt new file mode 100644 index 00000000..9f5100db --- /dev/null +++ b/test/experimental/CMakeLists.txt @@ -0,0 +1,19 @@ +# 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 +# +# https://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. + +if(MUJOCO_BUILD_USD_PLUGINS AND MUJOCO_USD_TARGET STREQUAL "USD") + add_subdirectory(usd/plugins/mjcf) +endif() + + diff --git a/test/experimental/usd/mjcPhysics/mjc_physics_scene_test.cc b/test/experimental/usd/mjcPhysics/mjc_physics_scene_test.cc new file mode 100644 index 00000000..534dcffe --- /dev/null +++ b/test/experimental/usd/mjcPhysics/mjc_physics_scene_test.cc @@ -0,0 +1,190 @@ +// Copyright 2025 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. + +#include +#include +#include +#include "src/experimental/usd/mjcPhysics/sceneAPI.h" +#include "test/fixture.h" +#include +#include +#include +#include +#include +#include + +namespace mujoco { +namespace { + +using pxr::SdfPath; +using MjcPhysicsSceneTest = MujocoTest; +using testing::NotNull; + +// clang-format off +#define EXPECT_TYPE_USD_FALLBACK_EQ_MODEL_DEFAULT(type, usd_attr, mjc_attr) \ + { \ + type value; \ + mjc_phys_scene.Get##usd_attr##Attr().Get(&value); \ + EXPECT_EQ((mjtNum)value, default_model->opt.mjc_attr); \ + } + + +#define EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_attr) \ + EXPECT_TYPE_USD_FALLBACK_EQ_MODEL_DEFAULT(double, usd_attr, mjc_attr) + +#define EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_flag) \ + { \ + bool flag; \ + mjc_phys_scene.Get##usd_attr##Attr().Get(&flag); \ + EXPECT_NE(flag, default_model->opt.disableflags & (mjc_flag)); \ + } + +#define EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_flag) \ + { \ + bool flag; \ + mjc_phys_scene.Get##usd_attr##Attr().Get(&flag); \ + EXPECT_EQ(flag, default_model->opt.enableflags & (mjc_flag)); \ + } + + +#define EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_attr) \ + EXPECT_TYPE_USD_FALLBACK_EQ_MODEL_DEFAULT(int, usd_attr, mjc_attr) + +#define EXPECT_VEC3_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_attr) \ + { \ + pxr::GfVec3d value; \ + mjc_phys_scene.Get##usd_attr##Attr().Get(&value); \ + EXPECT_EQ(value[0], default_model->opt.mjc_attr[0]); \ + EXPECT_EQ(value[1], default_model->opt.mjc_attr[1]); \ + EXPECT_EQ(value[2], default_model->opt.mjc_attr[2]); \ + } + +#define EXPECT_TYPE_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(type, usd_attr, mjc_attr)\ + { \ + pxr::Vt##type##Array mjc_attr; \ + mjc_phys_scene.Get##usd_attr##Attr().Get(&mjc_attr); \ + EXPECT_THAT( \ + mjc_attr, \ + testing::ElementsAreArray(default_model->opt.mjc_attr) \ + ); \ + } + +#define EXPECT_REAL_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_attr)\ + EXPECT_TYPE_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(Double, usd_attr, mjc_attr) + +#define EXPECT_INT_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(usd_attr, mjc_attr)\ + EXPECT_TYPE_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(Int, usd_attr, mjc_attr) +// clang-format on + +TEST_F(MjcPhysicsSceneTest, TestDefaults) { + auto stage = pxr::UsdStage::CreateInMemory(); + + auto physics_scene = + pxr::UsdPhysicsScene::Define(stage, SdfPath("/World/PhysicsScene")); + + auto mjc_phys_scene = pxr::MjcPhysicsSceneAPI::Apply(physics_scene.GetPrim()); + + mjSpec* empty_spec = mj_makeSpec(); + mjModel* default_model = mj_compile(empty_spec, nullptr); + EXPECT_THAT(default_model, NotNull()); + + // Check that all the USD schema fallback values are the same + // as the model defaults. + // If this test is failing due to an update of defaults in Mujoco you need to + // update mjcPhysics/schema.usda. + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(Timestep, timestep); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(ApiRate, apirate); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(ImpRatio, impratio); + + EXPECT_VEC3_USD_FALLBACK_EQ_MODEL_DEFAULT(Wind, wind); + EXPECT_VEC3_USD_FALLBACK_EQ_MODEL_DEFAULT(Magnetic, magnetic); + + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(Density, density); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(Viscosity, viscosity); + + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(OMargin, o_margin); + EXPECT_REAL_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(OSolRef, o_solref); + EXPECT_REAL_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(OSolImp, o_solimp); + EXPECT_REAL_ARR_USD_FALLBACK_EQ_MODEL_DEFAULT(OFriction, o_friction); + + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(Iterations, iterations); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(Tolerance, tolerance); + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(LSIterations, ls_iterations); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(LSTolerance, ls_tolerance); + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(NoslipIterations, noslip_iterations); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(NoslipTolerance, noslip_tolerance); + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(CCDIterations, ccd_iterations); + EXPECT_REAL_USD_FALLBACK_EQ_MODEL_DEFAULT(CCDTolerance, ccd_tolerance); + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(SDFIterations, sdf_iterations); + EXPECT_INT_USD_FALLBACK_EQ_MODEL_DEFAULT(SDFInitPoints, sdf_initpoints); + + // We store the actuator disable groups as an array of integers, but the + // model stores it as a bitmask. + pxr::VtIntArray actuator_group_disable; + mjc_phys_scene.GetActuatorGroupDisableAttr().Get(&actuator_group_disable); + int bitmask = 0; + for (int ind : actuator_group_disable) { + bitmask |= 1 << ind; + } + EXPECT_EQ(default_model->opt.disableactuator, bitmask); + + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(ConstraintFlag, + mjDSBL_CONSTRAINT); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(EqualityFlag, + mjDSBL_EQUALITY); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(FrictionLossFlag, + mjDSBL_FRICTIONLOSS); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(LimitFlag, mjDSBL_LIMIT); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(ContactFlag, + mjDSBL_CONTACT); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(PassiveFlag, + mjDSBL_PASSIVE); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(GravityFlag, + mjDSBL_GRAVITY); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(ClampCtrlFlag, + mjDSBL_CLAMPCTRL); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(WarmStartFlag, + mjDSBL_WARMSTART); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(FilterParentFlag, + mjDSBL_FILTERPARENT); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(ActuationFlag, + mjDSBL_ACTUATION); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(RefSafeFlag, + mjDSBL_REFSAFE); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(SensorFlag, mjDSBL_SENSOR); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(MidPhaseFlag, + mjDSBL_MIDPHASE); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(NativeCCDFlag, + mjDSBL_NATIVECCD); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(EulerDampFlag, + mjDSBL_EULERDAMP); + EXPECT_DISABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(AutoResetFlag, + mjDSBL_AUTORESET); + + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(OverrideFlag, + mjENBL_OVERRIDE); + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(EnergyFlag, mjENBL_ENERGY); + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(FwdinvFlag, mjENBL_FWDINV); + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(InvDiscreteFlag, + mjENBL_INVDISCRETE); + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(MultiCCDFlag, + mjENBL_MULTICCD); + EXPECT_ENABLE_FLAG_USD_FALLBACK_EQ_MODEL_DEFAULT(IslandFlag, mjENBL_ISLAND); + + mj_deleteModel(default_model); + mj_deleteSpec(empty_spec); +} + +} // namespace +} // namespace mujoco diff --git a/test/experimental/usd/mjcPhysics/mjc_site_api_test.cc b/test/experimental/usd/mjcPhysics/mjc_site_api_test.cc new file mode 100644 index 00000000..a193dbb1 --- /dev/null +++ b/test/experimental/usd/mjcPhysics/mjc_site_api_test.cc @@ -0,0 +1,96 @@ +// Copyright 2025 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. + +#include +#include +#include +#include +#include "src/experimental/usd/mjcPhysics/siteAPI.h" +#include "test/fixture.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define EXPECT_SITE_TYPE(spec, site_path, expected_type) \ + { \ + mjsElement* site_element = \ + mjs_findElement(spec, mjOBJ_SITE, site_path.GetString().c_str()); \ + EXPECT_THAT(site_element, NotNull()); \ + mjsSite* site = mjs_asSite(site_element); \ + EXPECT_EQ(site->type, expected_type); \ + } + +namespace mujoco { +namespace { + +using pxr::MjcPhysicsSiteAPI; +using pxr::SdfPath; +using MjcSiteApiTest = MujocoTest; +using testing::NotNull; + +TEST_F(MjcSiteApiTest, TestApply) { + auto stage = pxr::UsdStage::CreateInMemory(); + + auto test_body_path = SdfPath("/World/TestBody"); + auto body = pxr::UsdGeomXform::Define(stage, test_body_path); + pxr::UsdPhysicsRigidBodyAPI::Apply(body.GetPrim()); + + auto test_collider_path = + test_body_path.AppendChild(pxr::TfToken("Collider")); + auto collider = pxr::UsdGeomSphere::Define(stage, test_collider_path); + pxr::UsdPhysicsCollisionAPI::Apply(collider.GetPrim()); + + auto test_sphere_site_path = + test_body_path.AppendChild(pxr::TfToken("SphereSite")); + auto test_cylinder_site_path = + test_body_path.AppendChild(pxr::TfToken("CylinderSite")); + auto test_capsule_site_path = + test_body_path.AppendChild(pxr::TfToken("CapsuleSite")); + auto test_box_site_path = test_body_path.AppendChild(pxr::TfToken("BoxSite")); + + auto sphere = pxr::UsdGeomSphere::Define(stage, test_sphere_site_path); + MjcPhysicsSiteAPI::Apply(sphere.GetPrim()); + + auto cylinder = pxr::UsdGeomCylinder::Define(stage, test_cylinder_site_path); + MjcPhysicsSiteAPI::Apply(cylinder.GetPrim()); + + auto capsule = pxr::UsdGeomCapsule::Define(stage, test_capsule_site_path); + MjcPhysicsSiteAPI::Apply(capsule.GetPrim()); + + auto box = pxr::UsdGeomCube::Define(stage, test_box_site_path); + MjcPhysicsSiteAPI::Apply(box.GetPrim()); + + mjSpec* spec = mj_parseUSDStage(stage); + mjModel* default_model = mj_compile(spec, nullptr); + EXPECT_THAT(default_model, NotNull()) << mjs_getError(spec); + + EXPECT_SITE_TYPE(spec, test_sphere_site_path, mjGEOM_SPHERE); + EXPECT_SITE_TYPE(spec, test_cylinder_site_path, mjGEOM_CYLINDER); + EXPECT_SITE_TYPE(spec, test_capsule_site_path, mjGEOM_CAPSULE); + EXPECT_SITE_TYPE(spec, test_box_site_path, mjGEOM_BOX); + + mj_deleteModel(default_model); + mj_deleteSpec(spec); +} + +} // namespace +} // namespace mujoco diff --git a/test/experimental/usd/plugins/mjcf/CMakeLists.txt b/test/experimental/usd/plugins/mjcf/CMakeLists.txt new file mode 100644 index 00000000..1b747aee --- /dev/null +++ b/test/experimental/usd/plugins/mjcf/CMakeLists.txt @@ -0,0 +1,30 @@ +# Copyright 2025 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 +# +# https://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. + +find_package(pxr REQUIRED) + +add_library(usd_fixture STATIC fixture.h fixture.cc) +target_include_directories(usd_fixture PUBLIC ${MUJOCO_TEST_INCLUDE}) +target_compile_definitions(usd_fixture PUBLIC MJSTATIC) + +target_link_libraries( + usd_fixture + PUBLIC usd + tf + gtest + gmock + mujoco +) + +mujoco_test(mjcf_file_format_test ADDITIONAL_LINK_LIBRARIES usd tf usdGeom usdImaging usdPhysics usdShade usd_fixture) diff --git a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc new file mode 100644 index 00000000..b8f57bca --- /dev/null +++ b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc @@ -0,0 +1,1448 @@ +// Copyright 2025 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. + +#include +#include + +#include +#include +#include "src/experimental/usd/mjcPhysics/actuatorAPI.h" +#include "src/experimental/usd/mjcPhysics/collisionAPI.h" +#include "src/experimental/usd/mjcPhysics/meshCollisionAPI.h" +#include "src/experimental/usd/mjcPhysics/sceneAPI.h" +#include "src/experimental/usd/mjcPhysics/siteAPI.h" +#include "src/experimental/usd/mjcPhysics/tokens.h" +#include "test/experimental/usd/test_utils.h" +#include "test/fixture.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // IWYU pragma: keep, used for TraverseAll +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +PXR_NAMESPACE_OPEN_SCOPE +// clang-format off +TF_DEFINE_PRIVATE_TOKENS(_tokens, + (st) + ); +// clang-format on +PXR_NAMESPACE_CLOSE_SCOPE + +namespace mujoco { +namespace usd { +namespace { + +using pxr::MjcPhysicsSiteAPI; +using pxr::MjcPhysicsTokens; +using pxr::SdfPath; +using MjcfSdfFileFormatPluginTest = MujocoTest; + +static const char* kMaterialsPath = + "experimental/usd/plugins/mjcf/testdata/materials.xml"; +static const char* kMeshObjPath = + "experimental/usd/plugins/mjcf/testdata/mesh_obj.xml"; + +TEST_F(MjcfSdfFileFormatPluginTest, TestClassAuthored) { + static constexpr char kXml[] = R"( + + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_VALID(stage, "/__class__"); + EXPECT_PRIM_VALID(stage, "/__class__/test"); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestBasicMeshSources) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_VALID(stage, "/mesh_test"); + EXPECT_PRIM_VALID(stage, "/mesh_test/test_body/tetrahedron"); + EXPECT_PRIM_VALID(stage, "/mesh_test/test_body/tetrahedron/Mesh"); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) { + const std::string xml_path = GetTestDataFilePath(kMaterialsPath); + + auto stage = pxr::UsdStage::Open(xml_path); + EXPECT_THAT(stage, testing::NotNull()); + + EXPECT_PRIM_VALID(stage, "/mesh_test"); + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials"); + + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_red"); + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_red/PreviewSurface"); + ExpectAttributeEqual( + stage, + "/mesh_test/Materials/material_red/PreviewSurface.inputs:diffuseColor", + pxr::GfVec3f(0.8, 0, 0)); + + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_texture"); + EXPECT_PRIM_VALID(stage, + "/mesh_test/Materials/material_texture/PreviewSurface"); + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_texture/uvmap"); + EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_texture/texture"); + ExpectAttributeHasConnection( + stage, + "/mesh_test/Materials/material_texture/" + "PreviewSurface.inputs:diffuseColor", + "/mesh_test/Materials/material_texture/texture.outputs:rgb"); + ExpectAttributeEqual( + stage, "/mesh_test/Materials/material_texture/texture.inputs:file", + pxr::SdfAssetPath("textures/cube.png")); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestGeomRgba) { + static constexpr char kXml[] = R"( + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + auto stage = pxr::UsdStage::Open(layer); + + EXPECT_PRIM_VALID(stage, "/test/sphere_red"); + ExpectAttributeEqual(stage, "/test/sphere_red.primvars:displayColor", + pxr::VtArray{{1, 0, 0}}); + EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage, + "/test/sphere_red.primvars:displayOpacity"); + + // There's no mechanism in Mujoco to specify whether an attribute was set + // explicitly or not. We do the same as Mujoco does, which is to compare with + // the default value. + // Which explains why not setting rgba is the same as setting it to the + // default value of (0.5, 0.5, 0.5, 1). + EXPECT_PRIM_VALID(stage, "/test/sphere_default"); + EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage, + "/test/sphere_default.primvars:displayColor"); + EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage, + "/test/sphere_default.primvars:displayOpacity"); + + EXPECT_PRIM_VALID(stage, "/test/sphere_also_default"); + EXPECT_ATTRIBUTE_HAS_NO_VALUE( + stage, "/test/sphere_also_default.primvars:displayColor"); + EXPECT_ATTRIBUTE_HAS_NO_VALUE( + stage, "/test/sphere_also_default.primvars:displayOpacity"); + + EXPECT_PRIM_VALID(stage, "/test/sphere_almost_default"); + ExpectAttributeEqual(stage, + "/test/sphere_almost_default.primvars:displayColor", + pxr::VtArray{{0.5, 0.5, 0.5}}); + ExpectAttributeEqual(stage, + "/test/sphere_almost_default.primvars:displayOpacity", + pxr::VtArray{0.9}); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestFaceVaryingMeshSourcesSimpleMjcfMesh) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + + auto mesh = pxr::UsdGeomMesh::Get( + stage, SdfPath("/mesh_test/test_body/tetrahedron/Mesh")); + ASSERT_TRUE(mesh); + pxr::VtArray face_vertex_counts; + mesh.GetFaceVertexCountsAttr().Get(&face_vertex_counts); + EXPECT_EQ(face_vertex_counts.size(), 4); + EXPECT_EQ(face_vertex_counts, pxr::VtArray({3, 3, 3, 3})); + + pxr::VtArray face_vertex_indices; + mesh.GetFaceVertexIndicesAttr().Get(&face_vertex_indices); + EXPECT_EQ(face_vertex_indices.size(), 12); + EXPECT_EQ(face_vertex_indices, + pxr::VtArray({0, 3, 2, 0, 1, 3, 0, 2, 1, 1, 2, 3})); + + pxr::VtArray normals; + mesh.GetNormalsAttr().Get(&normals); + EXPECT_EQ(normals.size(), face_vertex_indices.size()); + // We can't directly check the normals values because they are altered by + // Mujoco's compiling step. So we at least check that normals with the same + // original index are the same. + for (int i = 0; i < face_vertex_indices.size(); ++i) { + for (int j = i + 1; j < face_vertex_indices.size(); ++j) { + if (face_vertex_indices[i] == face_vertex_indices[j]) { + EXPECT_EQ(normals[i], normals[j]); + } + } + } + + auto primvars_api = pxr::UsdGeomPrimvarsAPI(mesh.GetPrim()); + + pxr::VtArray texcoords; + EXPECT_TRUE(primvars_api.HasPrimvar(pxr::_tokens->st)); + auto primvar_st = primvars_api.GetPrimvar(pxr::_tokens->st); + primvar_st.Get(&texcoords); + EXPECT_EQ(texcoords.size(), face_vertex_indices.size()); + + // Check the faceVarying texcoords against the manually indexed source + // texcoords. + pxr::VtArray source_texcoords{ + {0.5, 0.5}, {0, 0.5}, {1, 0}, {1, 1}}; + for (int i = 0; i < face_vertex_indices.size(); ++i) { + EXPECT_EQ(texcoords[i], source_texcoords[face_vertex_indices[i]]); + } +} + +TEST_F(MjcfSdfFileFormatPluginTest, + TestFaceVaryingMeshSourcesObjWithIndexedNormals) { + const std::string xml_path = GetTestDataFilePath(kMeshObjPath); + + auto stage = pxr::UsdStage::Open(xml_path); + EXPECT_THAT(stage, testing::NotNull()); + + auto mesh = + pxr::UsdGeomMesh::Get(stage, SdfPath("/mesh_test/test_body/mesh/Mesh")); + ASSERT_TRUE(mesh); + pxr::VtArray face_vertex_counts; + mesh.GetFaceVertexCountsAttr().Get(&face_vertex_counts); + EXPECT_EQ(face_vertex_counts.size(), 4); + EXPECT_EQ(face_vertex_counts, pxr::VtArray({3, 3, 3, 3})); + + pxr::VtArray face_vertex_indices; + mesh.GetFaceVertexIndicesAttr().Get(&face_vertex_indices); + EXPECT_EQ(face_vertex_indices.size(), 12); + EXPECT_EQ(face_vertex_indices, + pxr::VtArray({0, 3, 2, 0, 1, 3, 0, 2, 1, 1, 2, 3})); + + pxr::VtArray normals; + mesh.GetNormalsAttr().Get(&normals); + EXPECT_EQ(normals.size(), face_vertex_indices.size()); + // We can't directly check the normals values because they are altered by + // Mujoco's compiling step. + // We also can't access the normals indexing data, and can't use the vertex + // indexing data here because they are separate. + // So we check that the first half of the normals are the same, then the + // second half, as set in the OBJ file. + pxr::GfVec3f first_half_normal = normals[0]; + pxr::GfVec3f second_half_normal = normals[face_vertex_indices.size() / 2]; + EXPECT_NE(first_half_normal, second_half_normal); + int i = 0; + for (; i < face_vertex_indices.size() / 2; ++i) { + EXPECT_EQ(normals[i], first_half_normal); + } + for (; i < face_vertex_indices.size(); ++i) { + EXPECT_EQ(normals[i], second_half_normal); + } + + auto primvars_api = pxr::UsdGeomPrimvarsAPI(mesh.GetPrim()); + + pxr::VtArray texcoords; + EXPECT_TRUE(primvars_api.HasPrimvar(pxr::_tokens->st)); + auto primvar_st = primvars_api.GetPrimvar(pxr::_tokens->st); + primvar_st.Get(&texcoords); + EXPECT_EQ(texcoords.size(), face_vertex_indices.size()); + + // Check the faceVarying texcoords against the manually indexed source + // texcoords. + // NOTE: For OBJ we must use different indices for the texcoords than for the + // vertices! + std::vector source_face_texcoord_indices{0, 1, 2, 1, 2, 3, + 2, 3, 0, 3, 0, 1}; + pxr::VtArray source_texcoords{ + {0.5, 0.5}, {0, 0.5}, {1, 0}, {1, 1}}; + // NOTE: The v component of the texcoords is flipped when Mujoco loads the + // OBJ. + for (auto& uv : source_texcoords) { + uv[1] = 1 - uv[1]; + } + for (int i = 0; i < source_face_texcoord_indices.size(); ++i) { + EXPECT_EQ(texcoords[i], source_texcoords[source_face_texcoord_indices[i]]); + } +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestBody) { + static constexpr char kXml[] = R"( + + + + + + + + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + + EXPECT_PRIM_VALID(stage, "/body_test"); + EXPECT_PRIM_VALID(stage, "/body_test/test_body"); + EXPECT_PRIM_VALID(stage, "/body_test/test_body/test_body_2"); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestBasicParenting) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_VALID(stage, "/test/root"); + EXPECT_PRIM_VALID(stage, "/test/root/root_body_1"); + EXPECT_PRIM_VALID(stage, "/test/root/root_body_2"); + EXPECT_PRIM_VALID(stage, "/test/root/root_body_2/root_body_3"); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestJointsDoNotAffectParenting) { + static constexpr char kXml[] = R"( + + + + + + + + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_VALID(stage, "/test/root"); + EXPECT_PRIM_VALID(stage, "/test/root/middle"); + EXPECT_PRIM_VALID(stage, "/test/root/middle/tet"); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestKindAuthoring) { + static constexpr char kXml[] = R"( + + + + + + + + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_KIND(stage, "/test", pxr::KindTokens->group); + EXPECT_PRIM_KIND(stage, "/test/root", pxr::KindTokens->component); + EXPECT_PRIM_KIND(stage, "/test/root/middle", pxr::KindTokens->subcomponent); + EXPECT_PRIM_KIND(stage, "/test/root/middle/tet", + pxr::KindTokens->subcomponent); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestAttributesMatchSchemaTypes) { + // TODO(robinalazard): Make the scene much more comprehensive. We ideally want + // to test all the prims that the plugin can generate. + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + auto stage = pxr::UsdStage::Open(layer); + + for (const auto& prim : stage->TraverseAll()) { + ExpectAllAuthoredAttributesMatchSchemaTypes(prim); + } +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestGeomsPrims) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfLayerRefPtr layer = LoadLayer(kXml); + auto stage = pxr::UsdStage::Open(layer); + + // Note that all sizes are multiplied by 2 because Mujoco uses half sizes. + + // Plane + EXPECT_PRIM_VALID(stage, "/test/plane_geom"); + EXPECT_PRIM_IS_A(stage, "/test/plane_geom", pxr::UsdGeomPlane); + ExpectAttributeEqual(stage, "/test/plane_geom.width", 2 * 10.0); + ExpectAttributeEqual(stage, "/test/plane_geom.length", 2 * 20.0); + // Box + EXPECT_PRIM_VALID(stage, "/test/box_geom"); + EXPECT_PRIM_IS_A(stage, "/test/box_geom", pxr::UsdGeomCube); + // Box is a special case, it uses a UsdGeomCube and scales it with + // xformOp:scale. The radius is always set to 2. + ExpectAttributeEqual(stage, "/test/box_geom.size", 2.0); + ExpectAttributeEqual(stage, "/test/box_geom.xformOp:scale", + pxr::GfVec3f(10.0, 20.0, 30.0)); + // Sphere + EXPECT_PRIM_VALID(stage, "/test/sphere_geom"); + EXPECT_PRIM_IS_A(stage, "/test/sphere_geom", pxr::UsdGeomSphere); + ExpectAttributeEqual(stage, "/test/sphere_geom.radius", 10.0); + // Capsule + EXPECT_PRIM_VALID(stage, "/test/capsule_geom"); + EXPECT_PRIM_IS_A(stage, "/test/capsule_geom", pxr::UsdGeomCapsule); + ExpectAttributeEqual(stage, "/test/capsule_geom.radius", 10.0); + ExpectAttributeEqual(stage, "/test/capsule_geom.height", 2 * 20.0); + // Cylinder + EXPECT_PRIM_VALID(stage, "/test/cylinder_geom"); + EXPECT_PRIM_IS_A(stage, "/test/cylinder_geom", pxr::UsdGeomCylinder); + ExpectAttributeEqual(stage, "/test/cylinder_geom.radius", 10.0); + ExpectAttributeEqual(stage, "/test/cylinder_geom.height", 2 * 20.0); + // Ellipsoid + EXPECT_PRIM_VALID(stage, "/test/ellipsoid_geom"); + // Ellipsoid is a special case, it uses a UsdGeomSphere and scales it with + // xformOp:scale. The radius is always set to 1. + EXPECT_PRIM_IS_A(stage, "/test/ellipsoid_geom", pxr::UsdGeomSphere); + ExpectAttributeEqual(stage, "/test/ellipsoid_geom.radius", 1.0); + ExpectAttributeEqual(stage, "/test/ellipsoid_geom.xformOp:scale", + pxr::GfVec3f(10.0, 20.0, 30.0)); +} + +static const pxr::SdfPath kPhysicsScenePrimPath("/test/PhysicsScene"); + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimTimestep) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionTimestep), + 0.005); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimCone) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionCone), + MjcPhysicsTokens->elliptic); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimWind) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionWind), + pxr::GfVec3d(1, 2, 3)); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimApirate) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionApirate), + 1.2); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimImpratio) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionImpratio), + 0.8); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimMagnetic) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionMagnetic), + pxr::GfVec3d(1, 2, 3)); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimDensity) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionDensity), + 1.2); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimViscosity) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionViscosity), + 0.8); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimO_margin) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionO_margin), + 0.001); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimO_solref) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionO_solref), + pxr::VtArray({0.1, 0.2})); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimO_solimp) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionO_solimp), + pxr::VtArray({0.1, 0.2, 0.3, 0.4, 0.5})); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimTolerance) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionTolerance), + 0.0012); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimLSTolerance) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionLs_tolerance), + 0.0034); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimNoslipTolerance) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionNoslip_tolerance), + 0.0056); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimCCDTolerance) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionCcd_tolerance), + 0.0078); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimOFriction) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionO_friction), + pxr::VtArray({0.1, 0.2, 0.3, 0.4, 0.5})); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimIntegrator) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionIntegrator), + MjcPhysicsTokens->rk4); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimJacobian) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionJacobian), + MjcPhysicsTokens->sparse); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimSolver) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual( + stage, + kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionSolver), + MjcPhysicsTokens->cg); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimIterations) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionIterations), + 10); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimLSIterations) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionLs_iterations), + 20); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimNoslipIterations) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionNoslip_iterations), + 30); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimCCDIterations) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionCcd_iterations), + 40); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimSDFInitPoints) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionSdf_initpoints), + 50); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimSDFIterations) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + ExpectAttributeEqual(stage, + kPhysicsScenePrimPath.AppendProperty( + MjcPhysicsTokens->mjcOptionSdf_iterations), + 60); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimDisableFlags) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + const std::vector kFlags = { + MjcPhysicsTokens->mjcFlagConstraint, + MjcPhysicsTokens->mjcFlagEquality, + MjcPhysicsTokens->mjcFlagFrictionloss, + MjcPhysicsTokens->mjcFlagLimit, + MjcPhysicsTokens->mjcFlagContact, + MjcPhysicsTokens->mjcFlagPassive, + MjcPhysicsTokens->mjcFlagGravity, + MjcPhysicsTokens->mjcFlagClampctrl, + MjcPhysicsTokens->mjcFlagWarmstart, + MjcPhysicsTokens->mjcFlagFilterparent, + MjcPhysicsTokens->mjcFlagActuation, + MjcPhysicsTokens->mjcFlagRefsafe, + MjcPhysicsTokens->mjcFlagSensor, + MjcPhysicsTokens->mjcFlagMidphase, + MjcPhysicsTokens->mjcFlagNativeccd, + MjcPhysicsTokens->mjcFlagEulerdamp, + MjcPhysicsTokens->mjcFlagAutoreset, + }; + for (const auto& flag : kFlags) { + ExpectAttributeEqual(stage, kPhysicsScenePrimPath.AppendProperty(flag), + false); + } +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimEnableFlags) { + auto stage = pxr::UsdStage::Open(LoadLayer(R"( + + + + )")); + + // clang-format off + const std::vector kFlags = { + MjcPhysicsTokens->mjcFlagOverride, + MjcPhysicsTokens->mjcFlagEnergy, + MjcPhysicsTokens->mjcFlagFwdinv, + MjcPhysicsTokens->mjcFlagInvdiscrete, + MjcPhysicsTokens->mjcFlagMulticcd, + MjcPhysicsTokens->mjcFlagIsland, + }; + // clang-format on + for (const auto& flag : kFlags) { + ExpectAttributeEqual(stage, kPhysicsScenePrimPath.AppendProperty(flag), + true); + } +} + +static constexpr char kSiteXml[] = R"( + + + + + + + + + + + + + )"; + +TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsAuthored) { + pxr::SdfLayerRefPtr layer = LoadLayer(kSiteXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_VALID(stage, "/test/box_site"); + EXPECT_PRIM_IS_A(stage, "/test/box_site", pxr::UsdGeomCube); + EXPECT_PRIM_API_APPLIED(stage, "/test/box_site", pxr::MjcPhysicsSiteAPI); + EXPECT_PRIM_VALID(stage, "/test/ball/sphere_site"); + EXPECT_PRIM_IS_A(stage, "/test/ball/sphere_site", pxr::UsdGeomSphere); + EXPECT_PRIM_API_APPLIED(stage, "/test/ball/sphere_site", + pxr::MjcPhysicsSiteAPI); + EXPECT_PRIM_VALID(stage, "/test/ball/capsule_site"); + EXPECT_PRIM_IS_A(stage, "/test/ball/capsule_site", pxr::UsdGeomCapsule); + EXPECT_PRIM_API_APPLIED(stage, "/test/ball/capsule_site", + pxr::MjcPhysicsSiteAPI); + EXPECT_PRIM_VALID(stage, "/test/ball/cylinder_site"); + EXPECT_PRIM_IS_A(stage, "/test/ball/cylinder_site", pxr::UsdGeomCylinder); + EXPECT_PRIM_API_APPLIED(stage, "/test/ball/cylinder_site", + pxr::MjcPhysicsSiteAPI); + EXPECT_PRIM_VALID(stage, "/test/ball/ellipsoid_site"); + EXPECT_PRIM_IS_A(stage, "/test/ball/ellipsoid_site", pxr::UsdGeomSphere); + EXPECT_PRIM_API_APPLIED(stage, "/test/ball/ellipsoid_site", + pxr::MjcPhysicsSiteAPI); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsPurpose) { + pxr::SdfLayerRefPtr layer = LoadLayer(kSiteXml); + + auto stage = pxr::UsdStage::Open(layer); + EXPECT_PRIM_PURPOSE(stage, "/test/box_site", pxr::UsdGeomTokens->guide); + EXPECT_PRIM_PURPOSE(stage, "/test/ball/sphere_site", + pxr::UsdGeomTokens->guide); + EXPECT_PRIM_PURPOSE(stage, "/test/ball/capsule_site", + pxr::UsdGeomTokens->guide); + EXPECT_PRIM_PURPOSE(stage, "/test/ball/cylinder_site", + pxr::UsdGeomTokens->guide); + EXPECT_PRIM_PURPOSE(stage, "/test/ball/ellipsoid_site", + pxr::UsdGeomTokens->guide); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsToggleSdfFormatArg) { + std::string xml_path = GetTestDataFilePath(kMeshObjPath); + + // Test that the default is no physics. + auto stage_no_physics = pxr::UsdStage::Open(xml_path); + EXPECT_THAT(stage_no_physics, testing::NotNull()); + EXPECT_PRIM_VALID(stage_no_physics, "/mesh_test/test_body"); + EXPECT_PRIM_API_NOT_APPLIED(stage_no_physics, "/mesh_test/test_body", + pxr::UsdPhysicsRigidBodyAPI); + + // Then test that the physics flag enables physics. + std::string xml_path_physics_flag = + xml_path + ":SDF_FORMAT_ARGS:usdMjcfToggleUsdPhysics=true"; + auto stage_with_physics = pxr::UsdStage::Open(xml_path_physics_flag); + EXPECT_THAT(stage_with_physics, testing::NotNull()); + + EXPECT_PRIM_VALID(stage_with_physics, "/mesh_test/test_body"); + EXPECT_PRIM_API_APPLIED(stage_with_physics, "/mesh_test/test_body", + pxr::UsdPhysicsRigidBodyAPI); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsRigidBody) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfFileFormat::FileFormatArguments args; + args["usdMjcfToggleUsdPhysics"] = "true"; + pxr::SdfLayerRefPtr layer = LoadLayer(kXml, args); + auto stage = pxr::UsdStage::Open(layer); + + EXPECT_THAT(stage, testing::NotNull()); + EXPECT_PRIM_VALID(stage, "/physics_test"); + EXPECT_PRIM_VALID(stage, "/physics_test/test_body"); + EXPECT_PRIM_VALID(stage, "/physics_test/test_body/test_body_2"); + + EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body", + pxr::UsdPhysicsArticulationRootAPI); + EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body/test_body_2", + pxr::UsdPhysicsRigidBodyAPI); + + // Only the root body should have the articulation API applied. + EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body/test_body_2", + pxr::UsdPhysicsArticulationRootAPI); + + // Geoms should not have RigidBodyAPI applied either. + EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body/test_geom", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, + "/physics_test/test_body/test_body_2/test_geom_2", + pxr::UsdPhysicsRigidBodyAPI); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) { + static constexpr char kXml[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + )"; + + pxr::SdfFileFormat::FileFormatArguments args; + args["usdMjcfToggleUsdPhysics"] = "true"; + pxr::SdfLayerRefPtr layer = LoadLayer(kXml, args); + auto stage = pxr::UsdStage::Open(layer); + + EXPECT_THAT(stage, testing::NotNull()); + EXPECT_PRIM_VALID(stage, "/test"); + + // Expected hierarchy under /test: + // + // ground [collider] + // + // body_0 [rigidbody] + // body_0/body_0_col [collider] + // + // body_0/body_0_0 [rigidbody] + // body_0/body_0_0/body_0_0_col [collider] + // + // body_1 [rigidbody] + // body_1/body_1_col_0 [collider] + // body_1/body_1_col_1 [collider] + // + // body_2 [rigidbody] + // body_2/body_2_nocol [] + // + // body_3 [rigidbody] + // body_3/body_3_col [] <-- Intermediate prim for mesh instancing + // body_3/body_3_col/Mesh [collider, mesh collider] + + // ground [collider] (Static collider) + EXPECT_PRIM_VALID(stage, "/test/ground"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/ground", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::MjcPhysicsCollisionAPI); + + // body_0/body_0_0 [rigidbody] (Nested body - reparented) + EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0", + pxr::MjcPhysicsCollisionAPI); + // body_0/body_0_0/body_0_0_col [collider] + EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0/body_0_0_col"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col", + pxr::MjcPhysicsCollisionAPI); + + // body_1 [rigidbody] + EXPECT_PRIM_VALID(stage, "/test/body_1"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_1", pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1", + pxr::MjcPhysicsCollisionAPI); + // body_1/body_1_col_0 [collider] + EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_0"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_0", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0", + pxr::MjcPhysicsCollisionAPI); + // body_1/body_1_col_1 [collider] + EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_1"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_1", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1", + pxr::MjcPhysicsCollisionAPI); + + // body_2 [rigidbody] + EXPECT_PRIM_VALID(stage, "/test/body_2"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_2", pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_2", + pxr::UsdPhysicsCollisionAPI); + // body_2/body_2_nocol [] (No physics APIs applied) + EXPECT_PRIM_VALID(stage, "/test/body_2/body_2_nocol"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_2/body_2_nocol", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_2/body_2_nocol", + pxr::UsdPhysicsCollisionAPI); + + // body_3 [rigidbody] + EXPECT_PRIM_VALID(stage, "/test/body_3"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_3", pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_3", + pxr::UsdPhysicsCollisionAPI); + // body_3/body_3_col [] (Intermediate prim for mesh instancing) + EXPECT_PRIM_VALID(stage, "/test/body_3/body_3_col"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_3/body_3_col", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_3/body_3_col", + pxr::UsdPhysicsCollisionAPI); + // body_3/body_3_col/Mesh [collider, mesh collider] + EXPECT_PRIM_VALID(stage, "/test/body_3/body_3_col/Mesh"); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_3/body_3_col/Mesh", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh", + pxr::UsdPhysicsCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh", + pxr::UsdPhysicsMeshCollisionAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh", + pxr::MjcPhysicsMeshCollisionAPI); + ExpectAttributeEqual(stage, + "/test/body_3/body_3_col/Mesh.physics:approximation", + pxr::UsdPhysicsTokens->convexHull); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsCollisionAPI) { + static constexpr char xml[] = R"( + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + ExpectAttributeEqual(stage, "/test/body/box.mjc:shellinertia", true); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsMeshCollisionAPI) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + ExpectAttributeEqual(stage, "/test/body/tet_legacy/Mesh.mjc:inertia", + MjcPhysicsTokens->legacy); + ExpectAttributeEqual(stage, "/test/body/tet_exact/Mesh.mjc:inertia", + MjcPhysicsTokens->exact); + ExpectAttributeEqual(stage, "/test/body/tet_convex/Mesh.mjc:inertia", + MjcPhysicsTokens->convex); + ExpectAttributeEqual(stage, "/test/body/tet_shell/Mesh.mjc:inertia", + MjcPhysicsTokens->shell); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIApplied) { + static constexpr char xml[] = R"( + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + EXPECT_PRIM_VALID(stage, "/test/body"); + EXPECT_PRIM_VALID(stage, "/test/body/box"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI); + ExpectAttributeEqual(stage, "/test/body/box.physics:mass", 0.1f); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIAppliedToBody) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + EXPECT_PRIM_VALID(stage, "/test/body"); + EXPECT_PRIM_VALID(stage, "/test/body/box"); + EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI); + EXPECT_PRIM_API_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI); + // Make sure that body gets it's inertial elements from the inertial element + // and not from the subtree. + ExpectAttributeEqual(stage, "/test/body.physics:mass", 3.0f); + ExpectAttributeEqual(stage, "/test/body.physics:centerOfMass", + pxr::GfVec3f(1, 2, 3)); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIDensity) { + static constexpr char xml[] = R"( + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + ExpectAttributeEqual(stage, "/test/body/box.physics:density", 1234.0f); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsActuatorGeneral) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + EXPECT_PRIM_API_APPLIED(stage, "/test/body/site", pxr::MjcPhysicsActuatorAPI); + EXPECT_REL_HAS_TARGET(stage, "/test/body/site.mjc:refSite", "/test/body/ref"); + ExpectAttributeEqual(stage, "/test/body/site.mjc:ctrlLimited", + pxr::MjcPhysicsTokens->true_); + ExpectAttributeEqual(stage, "/test/body/site.mjc:ctrlRange:min", 0.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:ctrlRange:max", 1.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:forceLimited", + pxr::MjcPhysicsTokens->true_); + ExpectAttributeEqual(stage, "/test/body/site.mjc:forceRange:min", 2.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:forceRange:max", 3.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:actLimited", + pxr::MjcPhysicsTokens->false_); + ExpectAttributeEqual(stage, "/test/body/site.mjc:actRange:min", 4.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:actRange:max", 5.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:lengthRange:min", 6.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:lengthRange:max", 7.0); + ExpectAttributeEqual(stage, "/test/body/site.mjc:actDim", 1); + ExpectAttributeEqual(stage, "/test/body/site.mjc:dynType", + MjcPhysicsTokens->filter); + ExpectAttributeEqual(stage, "/test/body/site.mjc:gainType", + MjcPhysicsTokens->user); + ExpectAttributeEqual(stage, "/test/body/site.mjc:biasType", + MjcPhysicsTokens->user); + ExpectAttributeEqual(stage, "/test/body/site.mjc:actEarly", true); + ExpectAttributeEqual(stage, "/test/body/site.mjc:gear", + pxr::VtDoubleArray{{1, 2, 3, 4, 5, 6}}); + ExpectAttributeEqual(stage, "/test/body/site.mjc:dynPrm", + pxr::VtDoubleArray{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}}); + ExpectAttributeEqual(stage, "/test/body/site.mjc:gainPrm", + pxr::VtDoubleArray{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}}); + ExpectAttributeEqual(stage, "/test/body/site.mjc:biasPrm", + pxr::VtDoubleArray{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}}); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsBodyActuator) { + static constexpr char xml[] = R"( + + + + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + EXPECT_PRIM_API_APPLIED(stage, "/test/body", pxr::MjcPhysicsActuatorAPI); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsSliderCrankActuator) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + )"; + auto stage = OpenStageWithPhysics(xml); + + EXPECT_PRIM_API_APPLIED(stage, "/test/body/crank", + pxr::MjcPhysicsActuatorAPI); + EXPECT_REL_HAS_TARGET(stage, "/test/body/crank.mjc:sliderSite", + "/test/body/slider"); + ExpectAttributeEqual(stage, "/test/body/crank.mjc:crankLength", 1.23); +} + +} // namespace +} // namespace usd +} // namespace mujoco diff --git a/test/experimental/usd/plugins/mjcf/testdata/materials.xml b/test/experimental/usd/plugins/mjcf/testdata/materials.xml new file mode 100644 index 00000000..735a789d --- /dev/null +++ b/test/experimental/usd/plugins/mjcf/testdata/materials.xml @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + + diff --git a/test/experimental/usd/plugins/mjcf/testdata/mesh_obj.xml b/test/experimental/usd/plugins/mjcf/testdata/mesh_obj.xml new file mode 100644 index 00000000..11465088 --- /dev/null +++ b/test/experimental/usd/plugins/mjcf/testdata/mesh_obj.xml @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/test/experimental/usd/plugins/mjcf/testdata/meshes/tetrahedron.obj b/test/experimental/usd/plugins/mjcf/testdata/meshes/tetrahedron.obj new file mode 100644 index 00000000..98ab2e51 --- /dev/null +++ b/test/experimental/usd/plugins/mjcf/testdata/meshes/tetrahedron.obj @@ -0,0 +1,16 @@ +v 0 1 0 +v 0 0 0 +v 1 0 1 +v 1 0 -1 +vn 1 0 0 +vn -1 0 0 +vt 0.5 0.5 +vt 0 0.5 +vt 1 1 +vt 1 0 + +f 1/1/1 4/2/1 3/3/1 +f 1/2/1 2/3/1 4/4/1 +f 1/3/2 3/4/2 2/1/2 +f 2/4/2 3/1/2 4/2/2 + diff --git a/test/experimental/usd/plugins/mjcf/testdata/textures/cube.png b/test/experimental/usd/plugins/mjcf/testdata/textures/cube.png new file mode 100644 index 00000000..8c244eca Binary files /dev/null and b/test/experimental/usd/plugins/mjcf/testdata/textures/cube.png differ diff --git a/test/experimental/usd/test_utils.cc b/test/experimental/usd/test_utils.cc new file mode 100644 index 00000000..aedb13b7 --- /dev/null +++ b/test/experimental/usd/test_utils.cc @@ -0,0 +1,123 @@ +// Copyright 2025 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. + +#include "test/experimental/usd/test_utils.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace mujoco { +namespace usd { + +using pxr::SdfPath; + +pxr::SdfLayerRefPtr LoadLayer( + const std::string& xml, + const pxr::SdfFileFormat::FileFormatArguments& args) { + auto layer = pxr::SdfLayer::CreateAnonymous( + "test_layer", pxr::SdfFileFormat::FindByExtension("xml"), args); + layer->ImportFromString(xml); + EXPECT_THAT(layer, testing::NotNull()); + return layer; +} + +pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml) { + pxr::SdfFileFormat::FileFormatArguments args; + args["usdMjcfToggleUsdPhysics"] = "true"; + pxr::SdfLayerRefPtr layer = LoadLayer(xml, args); + auto stage = pxr::UsdStage::Open(layer); + EXPECT_THAT(stage, testing::NotNull()); + return stage; +} + +template <> +void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, + pxr::SdfPath path, + const pxr::SdfAssetPath& value) { + auto attr = stage->GetAttributeAtPath(path); + EXPECT_TRUE(attr.IsValid()); + pxr::SdfAssetPath attr_value; + attr.Get(&attr_value); + EXPECT_EQ(attr_value.GetAssetPath(), value.GetAssetPath()); +} + +void ExpectAttributeHasConnection(pxr::UsdStageRefPtr stage, const char* path, + const char* connection_path) { + auto attr = stage->GetAttributeAtPath(SdfPath(path)); + EXPECT_TRUE(attr.IsValid()); + pxr::SdfPathVector sources; + attr.GetConnections(&sources); + EXPECT_EQ(sources.size(), 1); + EXPECT_EQ(sources[0], SdfPath(connection_path)); +} + +void ExpectAllAuthoredAttributesMatchSchemaTypes(const pxr::UsdPrim& prim) { + // Get all properties on the prim that have authored opinions. + for (const pxr::UsdProperty& prop : prim.GetAuthoredProperties()) { + // We only care about attributes, as they are the ones with a typeName. + if (pxr::UsdAttribute attr = prop.As()) { + // 1. Get the official, composed schema type name for the attribute. + const pxr::TfToken schemaTypeName = attr.GetTypeName().GetAsToken(); + + // An empty schema type name means the attribute is not defined by + // a schema, or is of a dynamically-determined type. We can't + // check for a mismatch in this case. + if (schemaTypeName.IsEmpty()) { + continue; + } + + // 2. Get the property stack to check for authored opinions. + // The stack is ordered from strongest to weakest. + const pxr::SdfPropertySpecHandleVector propStack = + attr.GetPropertyStack(); + + for (const pxr::SdfPropertySpecHandle& spec : propStack) { + // We only care about attribute specs. + if (auto attrSpec = TfDynamic_cast(spec)) { + // 3. Check if this spec has an authored `typeName`. + if (attrSpec->HasField(pxr::SdfFieldKeys->TypeName)) { + const pxr::TfToken authoredTypeName = + attrSpec->GetTypeName().GetAsToken(); + + EXPECT_EQ(authoredTypeName, schemaTypeName) + << "Type mismatch for attribute <" << attr.GetPath() + << ">: expected schema-defined type '" + << schemaTypeName.GetString() << "', got authored type '" + << authoredTypeName.GetString() << "' in layer @" + << attrSpec->GetLayer()->GetIdentifier() << "@"; + + // We've found the strongest authored opinion for `typeName`, + // so we can stop checking the stack for this attribute. + break; + } + } + } + } + } +} +} // namespace usd +} // namespace mujoco diff --git a/test/experimental/usd/test_utils.h b/test/experimental/usd/test_utils.h new file mode 100644 index 00000000..2aa62a6c --- /dev/null +++ b/test/experimental/usd/test_utils.h @@ -0,0 +1,120 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_TEST_EXPERIMENTAL_USD_PLUGINS_MJCF_FIXTURE_H_ +#define MUJOCO_TEST_EXPERIMENTAL_USD_PLUGINS_MJCF_FIXTURE_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#define EXPECT_PRIM_VALID(stage, path) \ + EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsValid()); + +#define EXPECT_PRIM_IS_A(stage, path, type) \ + { \ + EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsA()); \ + } + +#define EXPECT_PRIM_API_APPLIED(stage, path, api) \ + { \ + EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI()); \ + } + +#define EXPECT_PRIM_API_NOT_APPLIED(stage, path, api) \ + { \ + EXPECT_FALSE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI()); \ + } + +#define EXPECT_PRIM_KIND(stage, path, kind) \ + { \ + pxr::TfToken prim_kind; \ + pxr::UsdModelAPI::Get(stage, SdfPath(path)).GetKind(&prim_kind); \ + EXPECT_EQ(kind, prim_kind); \ + } + +#define EXPECT_PRIM_PURPOSE(stage, path, purpose) \ + { \ + pxr::TfToken prim_purpose; \ + pxr::UsdGeomImageable::Get(stage, SdfPath(path)) \ + .GetPurposeAttr() \ + .Get(&prim_purpose); \ + EXPECT_EQ(prim_purpose, purpose); \ + } + +#define EXPECT_ATTRIBUTE_HAS_VALUE(stage, path) \ + EXPECT_TRUE((stage)->GetAttributeAtPath(SdfPath(path)).HasValue()); + +#define EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage, path) \ + EXPECT_FALSE((stage)->GetAttributeAtPath(SdfPath(path)).HasValue()); + +#define EXPECT_REL_HAS_TARGET(stage, path, target_path) \ + { \ + pxr::SdfPathVector targets; \ + (stage)->GetRelationshipAtPath(SdfPath(path)).GetTargets(&targets); \ + EXPECT_TRUE(std::find(targets.begin(), targets.end(), \ + SdfPath(target_path)) != targets.end()); \ + } + +namespace mujoco { +namespace usd { + +pxr::SdfLayerRefPtr LoadLayer( + const std::string& xml, + const pxr::SdfFileFormat::FileFormatArguments& args = {}); + +pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml); + +template +void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, pxr::SdfPath path, + const T& value) { + auto attr = stage->GetAttributeAtPath(pxr::SdfPath(path)); + EXPECT_TRUE(attr.IsValid()) << "Attribute " << path << " is not valid"; + T attr_value; + attr.Get(&attr_value); + EXPECT_EQ(attr_value, value) << "Attribute " << path << " has value " + << attr_value << ". Expected: " << value; +} + +template +void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path, + const T& value) { + ExpectAttributeEqual(stage, pxr::SdfPath(path), value); +} + +// Specialization for SdfAssetPath, so that we can compare only the asset path +// and not care about whatever the resolved path is. +// Otherwise the default operator== would fail because it tests for equality of +// the asset path AND the resolved path. +template <> +void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, + pxr::SdfPath, + const pxr::SdfAssetPath& value); + +void ExpectAttributeHasConnection(pxr::UsdStageRefPtr stage, const char* path, + const char* connection_path); + +// Checks that all authored attributes on the given prim have types that match +// the schema types. +void ExpectAllAuthoredAttributesMatchSchemaTypes(const pxr::UsdPrim& prim); +} // namespace usd +} // namespace mujoco +#endif // MUJOCO_TEST_EXPERIMENTAL_USD_PLUGINS_MJCF_FIXTURE_H_ diff --git a/test/fixture.cc b/test/fixture.cc index 545691b4..cdf714b4 100644 --- a/test/fixture.cc +++ b/test/fixture.cc @@ -130,9 +130,7 @@ std::string GetFileContents(const char* path) { return sstream.str(); } -std::string SaveAndReadXml(const mjModel* model) { - EXPECT_THAT(model, testing::NotNull()); - +std::string SaveAndReadXmlImpl(const mjModel* model, const mjSpec* spec) { constexpr int kMaxPathLen = 1024; std::string path_template = std::filesystem::temp_directory_path().append("tmp.XXXXXX").string(); @@ -148,7 +146,11 @@ std::string SaveAndReadXml(const mjModel* model) { EXPECT_NE(_mktemp_s(filepath), EINVAL); #endif - mj_saveLastXML(filepath, model, nullptr, 0); + if (spec) { + mj_saveXML(spec, filepath, nullptr, 0); + } else if (model) { + mj_saveLastXML(filepath, model, nullptr, 0); + } std::string contents = GetFileContents(filepath); #if defined(_POSIX_VERSION) && _POSIX_VERSION >= 200112L @@ -159,6 +161,16 @@ std::string SaveAndReadXml(const mjModel* model) { return contents; } +std::string SaveAndReadXml(const mjModel* model) { + EXPECT_THAT(model, testing::NotNull()); + return SaveAndReadXmlImpl(model, nullptr); +} + +std::string SaveAndReadXml(const mjSpec* spec) { + EXPECT_THAT(spec, testing::NotNull()); + return SaveAndReadXmlImpl(nullptr, spec); +} + std::vector GetCtrlNoise(const mjModel* m, int nsteps, mjtNum ctrlnoise) { std::vector ctrl; diff --git a/test/fixture.h b/test/fixture.h index df568943..f9d0cb1f 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -53,6 +53,9 @@ class MujocoErrorTestGuard { // By default, any MuJoCo operation which triggers a warning or error will // trigger a test failure. class MujocoTest : public ::testing::Test { + public: + ~MujocoTest() { mj_freeLastXML(); } + private: MujocoErrorTestGuard error_guard; }; @@ -99,6 +102,9 @@ mjModel* LoadModelFromPath(const char* model_path); // Returns a string loaded from first saving the model given an input. std::string SaveAndReadXml(const mjModel* model); +// Returns a string loaded from first saving the spec given an input. +std::string SaveAndReadXml(const mjSpec* spec); + // Adds control noise. std::vector GetCtrlNoise(const mjModel* m, int nsteps, mjtNum ctrlnoise = 0.01); @@ -108,8 +114,9 @@ std::vector GetCtrlNoise(const mjModel* m, int nsteps, mjtNum CompareModel(const mjModel* m1, const mjModel* m2, std::string& field); // Returns a vector containing the elements of the array. -inline std::vector AsVector(const mjtNum* array, int n) { - return std::vector(array, array + n); +template +std::vector AsVector(const T* array, int n) { + return std::vector(array, array + n); } // Prints a matrix to stderr, useful for debugging. diff --git a/test/pipeline_test.cc b/test/pipeline_test.cc index 6f1eb737..c602da54 100644 --- a/test/pipeline_test.cc +++ b/test/pipeline_test.cc @@ -45,24 +45,35 @@ TEST_F(PipelineTest, SparseDenseEquivalent) { mjtNum tol = 1e-11; - for (mjtSolver solver : {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG}) { - model->opt.solver = solver; + const char* sname[4] = {"NEWTON", "PGS", "CG", "NOSLIP"}; + mjtSolver solver[4] = {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG, mjSOL_NEWTON}; - // set dense jacobian, call mj_forward, save accelerations + for (int i : {0, 1, 2, 3}) { + model->opt.solver = solver[i]; + if (i == 3) { + model->opt.noslip_iterations = 2; + } + + // set dense jacobian, call mj_step, save qacc and new qpos model->opt.jacobian = mjJAC_DENSE; mj_resetDataKeyframe(model, data, 0); - mj_forward(model, data); + mj_step(model, data); std::vector qacc_dense = AsVector(data->qacc, model->nv); + std::vector qpos_dense = AsVector(data->qpos, model->nq); - // set sparse jacobian, call mj_forward, save accelerations + // set sparse jacobian, call mj_step, save qacc and new qpos model->opt.jacobian = mjJAC_SPARSE; mj_resetDataKeyframe(model, data, 0); - mj_forward(model, data); + mj_step(model, data); std::vector qacc_sparse = AsVector(data->qacc, model->nv); + std::vector qpos_sparse = AsVector(data->qpos, model->nq); // expect accelerations to be insignificantly different EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse)) - << "failed equivalence for solver=" << solver; + << "failed qacc equivalence for solver=" << sname[i]; + // expect positions to be insignificantly different + EXPECT_THAT(qpos_dense, Pointwise(DoubleNear(tol), qpos_sparse)) + << "failed qpos equivalence for solver=" << sname[i]; } mj_deleteData(data); diff --git a/test/plugin/elasticity/elasticity_test.cc b/test/plugin/elasticity/elasticity_test.cc index f38559b5..ea4df2e3 100644 --- a/test/plugin/elasticity/elasticity_test.cc +++ b/test/plugin/elasticity/elasticity_test.cc @@ -22,7 +22,6 @@ #include #include #include "test/fixture.h" -#include "plugin/elasticity/shell.h" namespace mujoco { namespace { @@ -58,18 +57,10 @@ TEST_F(ElasticityTest, FlexCompatibility) { TEST_F(ElasticityTest, ElasticEnergyShell) { static constexpr char cantilever_xml[] = R"( - - - - - - - - - + @@ -79,12 +70,15 @@ TEST_F(ElasticityTest, ElasticEnergyShell) { mjModel* m = LoadModelFromString(cantilever_xml, error, sizeof(error)); ASSERT_THAT(m, testing::NotNull()) << error; mjData* d = mj_makeData(m); - auto* shell = reinterpret_cast(d->plugin_data[0]); + mj_kinematics(m, d); + mj_flex(m, d); // check that a plane is in the kernel of the energy for (mjtNum scale = 1; scale < 4; scale++) { - for (int e = 0; e < shell->ne; e++) { - int* v = shell->flaps[e].vertices; + for (int e = 0; e < m->flex_edgenum[0]; e++) { + int* edge = m->flex_edge + 2*(m->flex_edgeadr[0] + e); + int* flap = m->flex_edgeflap + 2*(m->flex_edgeadr[0] + e); + int v[4] = {edge[0], edge[1], flap[0], flap[1]}; if (v[3]== -1) { continue; } @@ -93,9 +87,9 @@ TEST_F(ElasticityTest, ElasticEnergyShell) { for (int i = 0; i < 4; i++) { for (int j = 0; j < 4; j++) { for (int x = 0; x < 3; x++) { - mjtNum elongation1 = scale * shell->position[3*v[i]+x]; - mjtNum elongation2 = scale * shell->position[3*v[j]+x]; - energy += shell->bending[16*e+4*i+j] * elongation1 * elongation2; + mjtNum elongation1 = scale * d->flexvert_xpos[3*v[i]+x]; + mjtNum elongation2 = scale * d->flexvert_xpos[3*v[j]+x]; + energy += m->flex_bending[16*e+4*i+j] * elongation1 * elongation2; } } } @@ -115,7 +109,7 @@ TEST_F(PluginTest, ElasticEnergyMembrane) { - + @@ -159,31 +153,6 @@ TEST_F(PluginTest, ElasticEnergyMembrane) { mj_deleteModel(m); } -TEST_F(ElasticityTest, InvalidThickness) { - static constexpr char xml[] = R"( - - - - - - - - - - - - - - - )"; - - char error[1024] = {0}; - mjModel* m = LoadModelFromString(xml, error, sizeof(error)); - ASSERT_THAT(m, testing::IsNull()); - EXPECT_THAT(error, ::testing::HasSubstr("Invalid parameter")); -} - // -------------------------------- solid ----------------------------------- TEST_F(ElasticityTest, ElasticEnergySolid) { static constexpr char cantilever_xml[] = R"( diff --git a/test/testdata/model.xml b/test/testdata/model.xml index e122f927..5a1f9c81 100644 --- a/test/testdata/model.xml +++ b/test/testdata/model.xml @@ -168,7 +168,6 @@ diff --git a/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.msh b/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.msh new file mode 100644 index 00000000..7f88379f --- /dev/null +++ b/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.msh @@ -0,0 +1,67 @@ +$MeshFormat +4.1 0 8 +$EndMeshFormat +$Entities +0 0 0 1 +1 -0.5 -0.5 0 0.5 0.5 1 0 0 +$EndEntities +$Nodes +1 14 1 100 +3 1 0 14 +1 +2 +3 +4 +5 +6 +7 +8 +9 +10 +11 +12 +13 +14 +-0.5 -0.5 0 +0.5 -0.5 0 +0.5 0.5 0 +-0.5 0.5 0 +-0.5 -0.5 1 +0.5 -0.5 1 +0.5 0.5 1 +-0.5 0.5 1 +0 0 0 +0 -0.5 0.5 +0.5 0 0.5 +0 0.5 0.5 +-0.5 0 0.5 +0 0 1 +$EndNodes +$Elements +1 24 1 24 +3 1 4 24 +1 10 9 12 11 +2 10 12 13 14 +3 9 10 12 13 +4 10 12 14 11 +5 9 10 1 2 +6 7 11 12 3 +7 4 1 9 13 +8 13 12 4 8 +9 1 10 13 5 +10 13 8 5 14 +11 12 8 14 7 +12 14 5 10 6 +13 14 6 11 7 +14 9 4 12 3 +15 3 11 9 2 +16 11 6 10 2 +17 6 10 14 11 +18 1 10 9 13 +19 12 4 9 13 +20 5 10 13 14 +21 8 13 12 14 +22 12 14 11 7 +23 9 12 11 3 +24 10 9 11 2 +$EndElements diff --git a/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.xml b/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.xml new file mode 100644 index 00000000..1e179a81 --- /dev/null +++ b/test/user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.xml @@ -0,0 +1,27 @@ + + diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 84febf8e..a415d740 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -279,20 +279,23 @@ TEST_F(PluginTest, AttachPlugin) { EXPECT_THAT(body_1, NotNull()); mjsFrame* attachment_frame = mjs_addFrame(body_1, 0); EXPECT_THAT(attachment_frame, NotNull()); - mjs_attachBody(attachment_frame, mjs_findBody(spec_1, "body"), "child-", ""); + mjs_attach(attachment_frame->element, mjs_findBody(spec_1, "body")->element, + "child-", ""); mjModel* model_1 = mj_compile(parent, nullptr); EXPECT_THAT(model_1, NotNull()); EXPECT_THAT(model_1->nbody, 3); // attach it a second time to test namespacing and compile ASSERT_THAT(spec_2, NotNull()) << err.data(); - mjs_attachBody(attachment_frame, mjs_findBody(spec_2, "body"), "copy-", ""); + mjs_attach(attachment_frame->element, mjs_findBody(spec_2, "body")->element, + "copy-", ""); mjModel* model_2 = mj_compile(parent, nullptr); EXPECT_THAT(model_2, NotNull()); EXPECT_THAT(model_2->nbody, 4); // attach a body not referencing the plugin and compile - mjs_attachBody(attachment_frame, mjs_findBody(spec_3, "empty"), "empty-", ""); + mjs_attach(attachment_frame->element, mjs_findBody(spec_3, "empty")->element, + "empty-", ""); mjModel* model_3 = mj_compile(parent, nullptr); EXPECT_THAT(model_3, NotNull()); EXPECT_THAT(model_3->nbody, 5); @@ -314,9 +317,9 @@ TEST_F(PluginTest, DetachPlugin) { ASSERT_THAT(child, NotNull()) << err.data(); // attach a body referencing the plugin to the frame - mjsFrame* frame = mjs_addFrame(mjs_findBody(parent, "world"), 0); - mjsBody* body = mjs_findBody(child, "body"); - EXPECT_THAT(mjs_attachBody(frame, body, "child-", ""), NotNull()); + mjsElement* frame = mjs_addFrame(mjs_findBody(parent, "world"), 0)->element; + mjsElement* body = mjs_findBody(child, "body")->element; + EXPECT_THAT(mjs_attach(frame, body, "child-", ""), NotNull()); // detach the body and compile mjsBody* body_to_detach = mjs_findBody(parent, "child-body"); @@ -372,7 +375,8 @@ TEST_F(PluginTest, AttachExplicitPlugin) { mjsFrame* attachment_frame = mjs_addFrame(body_parent, 0); EXPECT_THAT(attachment_frame, NotNull()); - mjs_attachBody(attachment_frame, mjs_findBody(child, "body"), "child-", ""); + mjs_attach(attachment_frame->element, mjs_findBody(child, "body")->element, + "child-", ""); mjModel* model = mj_compile(parent, nullptr); EXPECT_THAT(model, NotNull()); EXPECT_THAT(model->nplugin, 1); @@ -520,8 +524,7 @@ TEST_F(PluginTest, RecompileCompare) { // if file is meant to fail, skip it if (absl::StrContains(p.path().string(), "malformed_") || absl::StrContains(p.path().string(), "touch_grid") || - absl::StrContains(p.path().string(), "cow") || - absl::StrContains(p.path().string(), "discardvisual")) { + absl::StrContains(p.path().string(), "cow")) { continue; } @@ -535,6 +538,9 @@ TEST_F(PluginTest, RecompileCompare) { // copy spec mjSpec* s_copy = mj_copySpec(s); + // compare signature + EXPECT_EQ(s->element->signature, s_copy->element->signature) << xml; + // compile twice and compare mjModel* m_old = mj_compile(s, nullptr); @@ -544,6 +550,10 @@ TEST_F(PluginTest, RecompileCompare) { mjModel* m_new = mj_compile(s, nullptr); mjModel* m_copy = mj_compile(s_copy, nullptr); + // compare signature + EXPECT_EQ(m_old->signature, m_new->signature) << xml; + EXPECT_EQ(m_old->signature, m_copy->signature) << xml; + ASSERT_THAT(m_new, NotNull()) << "Failed to recompile " << xml << ": " << mjs_getError(s); ASSERT_THAT(m_copy, NotNull()) @@ -938,7 +948,8 @@ TEST_F(MujocoTest, AttachSame) { EXPECT_THAT(body, NotNull()); // attach child to parent frame - mjsBody* attached = mjs_attachBody(frame, body, "attached-", "-1"); + mjsBody* attached = + mjs_asBody(mjs_attach(frame->element, body->element, "attached-", "-1")); EXPECT_THAT(attached, mjs_findBody(parent, "attached-body-1")); // check that the spec was not copied @@ -1074,7 +1085,8 @@ TEST_F(MujocoTest, AttachDifferent) { EXPECT_THAT(body, NotNull()); // attach child to parent frame - mjsBody* attached = mjs_attachBody(frame, body, "attached-", "-1"); + mjsBody* attached = + mjs_asBody(mjs_attach(frame->element, body->element, "attached-", "-1")); EXPECT_THAT(attached, mjs_findBody(parent, "attached-body-1")); // check that the spec was copied @@ -1213,7 +1225,8 @@ TEST_F(MujocoTest, AttachFrame) { EXPECT_THAT(frame, NotNull()); // attach child frame to parent body - mjsFrame* attached = mjs_attachFrame(body, frame, "attached-", "-1"); + mjsFrame* attached = + mjs_asFrame(mjs_attach(body->element, frame->element, "attached-", "-1")); EXPECT_THAT(attached, mjs_findFrame(parent, "attached-pframe-1")); // check that the spec was copied @@ -1287,7 +1300,7 @@ TEST_F(MujocoTest, AttachCompiled) { // attach child body to the frame mjsBody* to_attach = mjs_findBody(child, "base"); EXPECT_THAT(to_attach, NotNull()) << mjs_getError(child); - mjs_attachBody(frame, to_attach, "", ""); + mjs_attach(frame->element, to_attach->element, "", ""); // check that attached model can be compiled mjModel* m_attached = mj_compile(parent, 0); @@ -1434,7 +1447,8 @@ TEST_F(MujocoTest, AttachToSite) { EXPECT_THAT(site, NotNull()); mjsBody* body = mjs_findBody(child, "sphere"); EXPECT_THAT(body, NotNull()); - mjsBody* attached = mjs_attachToSite(site, body, "attached-", "-1"); + mjsBody* attached = + mjs_asBody(mjs_attach(site->element, body->element, "attached-", "-1")); EXPECT_THAT(attached, NotNull()); mjModel* model = mj_compile(parent, 0); @@ -1499,7 +1513,8 @@ TEST_F(MujocoTest, AttachFrameToSite) { EXPECT_THAT(site, NotNull()); mjsFrame* frame = mjs_findFrame(child, "frame"); EXPECT_THAT(frame, NotNull()); - mjsFrame* attached = mjs_attachFrameToSite(site, frame, "attached-", "-1"); + mjsFrame* attached = + mjs_asFrame(mjs_attach(site->element, frame->element, "attached-", "-1")); EXPECT_THAT(attached, NotNull()); mjModel* model = mj_compile(parent, 0); @@ -1574,7 +1589,8 @@ TEST_F(MujocoTest, BodyToFrame) { EXPECT_THAT(frame, NotNull()); mjsBody* body = mjs_findBody(child1, "sphere"); EXPECT_THAT(body, NotNull()); - mjsBody* attached = mjs_attachBody(frame, body, "attached-", "-1"); + mjsBody* attached = + mjs_asBody(mjs_attach(frame->element, body->element, "attached-", "-1")); EXPECT_THAT(attached, NotNull()); mjModel* model1 = mj_compile(parent, 0); EXPECT_THAT(model1, NotNull()); @@ -1582,7 +1598,8 @@ TEST_F(MujocoTest, BodyToFrame) { // attach the world to the same frame and convert it to a frame mjsBody* world = mjs_findBody(child2, "world"); EXPECT_THAT(world, NotNull()); - mjsBody* child_world = mjs_attachBody(frame, world, "attached-", "-2"); + mjsBody* child_world = + mjs_asBody(mjs_attach(frame->element, world->element, "attached-", "-2")); EXPECT_THAT(child_world, NotNull()); mjsFrame* frame_world = mjs_bodyToFrame(&child_world); EXPECT_THAT(frame_world, NotNull()); @@ -1605,6 +1622,36 @@ TEST_F(MujocoTest, BodyToFrame) { mj_deleteModel(expected); } +TEST_F(MujocoTest, BodyToFrameWithInertial) { + static constexpr char xml_child[] = R"( + + + + + + + + + )"; + + std::array er; + mjSpec* spec = mj_parseXMLString(xml_child, 0, er.data(), er.size()); + EXPECT_THAT(spec, NotNull()) << er.data(); + mjModel* model = mj_compile(spec, 0); + EXPECT_THAT(model, NotNull()); + mjsBody* parent = mjs_findBody(spec, "parent"); + EXPECT_THAT(parent, NotNull()); + mjsBody* child = mjs_findBody(spec, "child"); + EXPECT_THAT(child, NotNull()); + mjs_bodyToFrame(&child); + EXPECT_THAT(parent->mass, 1); + EXPECT_THAT(parent->fullinertia[0], 1); + EXPECT_THAT(parent->fullinertia[1], 2); + EXPECT_THAT(parent->fullinertia[2], 3); + mj_deleteSpec(spec); + mj_deleteModel(model); +} + TEST_F(MujocoTest, AttachSpecToSite) { std::array er; mjtNum tol = 0; @@ -1663,7 +1710,8 @@ TEST_F(MujocoTest, AttachSpecToSite) { mjs_setFrame(mjs_firstChild(world, mjOBJ_CAMERA, 0), frame); // attach the entire spec to the site - mjsFrame* worldframe = mjs_attachFrameToSite(site, frame, "attached-", "-1"); + mjsFrame* worldframe = + mjs_asFrame(mjs_attach(site->element, frame->element, "attached-", "-1")); EXPECT_THAT(worldframe, NotNull()); // compile and compare @@ -1740,7 +1788,8 @@ TEST_F(MujocoTest, AttachSpecToBody) { mjs_setFrame(mjs_firstChild(world, mjOBJ_CAMERA, 0), frame); // attach the entire spec to the site - mjsFrame* worldframe = mjs_attachFrame(body, frame, "attached-", "-1"); + mjsFrame* worldframe = + mjs_asFrame(mjs_attach(body->element, frame->element, "attached-", "-1")); EXPECT_THAT(worldframe, NotNull()); worldframe->pos[0] = 1; worldframe->pos[1] = 2; @@ -1938,8 +1987,8 @@ TEST_F(MujocoTest, RecompileAttach) { mjSpec* child2 = mj_parseXMLString(xml, 0, er.data(), er.size()); EXPECT_THAT(child2, NotNull()); - mjsFrame* frame1 = mjs_addFrame(mjs_findBody(parent, "world"), 0); - mjs_attachBody(frame1, mjs_findBody(child1, "body"), "child-", "-1"); + mjsElement* frame1 = mjs_addFrame(mjs_findBody(parent, "world"), 0)->element; + mjs_attach(frame1, mjs_findBody(child1, "body")->element, "child-", "-1"); mjModel* model = mj_compile(parent, 0); EXPECT_THAT(model, NotNull()); @@ -1951,8 +2000,8 @@ TEST_F(MujocoTest, RecompileAttach) { mj_step(model, data); } - mjsFrame* frame2 = mjs_addFrame(mjs_findBody(parent, "world"), 0); - mjs_attachBody(frame2, mjs_findBody(child2, "body"), "child-", "-2"); + mjsElement* frame2 = mjs_addFrame(mjs_findBody(parent, "world"), 0)->element; + mjs_attach(frame2, mjs_findBody(child2, "body")->element, "child-", "-2"); EXPECT_EQ(mj_recompile(parent, 0, model, data), 0); EXPECT_THAT(model, NotNull()); @@ -2004,8 +2053,8 @@ TEST_F(MujocoTest, AttachMocap) { mjsBody* world = mjs_findBody(spec, "world"); EXPECT_THAT(world, NotNull()); - mjsFrame* frame = mjs_addFrame(world, NULL); - mjs_attachBody(frame, body, "attached-", "-1"); + mjsElement* frame = mjs_addFrame(world, NULL)->element; + mjs_attach(frame, body->element, "attached-", "-1"); mjsBody* attached_body = mjs_findBody(spec, "attached-mocap-1"); EXPECT_THAT(attached_body, NotNull()); @@ -2084,7 +2133,7 @@ TEST_F(MujocoTest, AttachUnnamedAssets) { geom->type = mjGEOM_MESH; mjSpec* spec = mj_makeSpec(); - mjs_attachFrame(mjs_findBody(spec, "world"), frame, "_", ""); + mjs_attach(mjs_findBody(spec, "world")->element, frame->element, "_", ""); mjModel* model = mj_compile(spec, vfs.get()); EXPECT_THAT(model, NotNull()); @@ -2230,8 +2279,8 @@ void AttachNestedKeyframe(bool compile) { mjs_setDeepCopy(child, true); // attach gchild to child - mjs_attachBody(mjs_findFrame(child, "frame"), - mjs_findBody(gchild, "body"), "gchild-", ""); + mjs_attach(mjs_findFrame(child, "frame")->element, + mjs_findBody(gchild, "body")->element, "gchild-", ""); // compile required before further attachment mjModel* m_child = compile ? mj_compile(child, 0) : nullptr; @@ -2244,8 +2293,8 @@ void AttachNestedKeyframe(bool compile) { }; // attach child to parent - mjs_attachBody(mjs_findFrame(parent, "frame"), - mjs_findBody(child, "body"), "child-", ""); + mjs_attach(mjs_findFrame(parent, "frame")->element, + mjs_findBody(child, "body")->element, "child-", ""); EXPECT_THAT(warning, HasSubstr(compile ? "" : "model has pending keyframes")); @@ -2306,11 +2355,11 @@ TEST_F(MujocoTest, RepeatedAttachKeyframe) { EXPECT_THAT(child, NotNull()) << er.data(); mjsBody* body_1 = mjs_findBody(parent, "body"); - mjsFrame* attachment_frame = mjs_addFrame(body_1, 0); - mjs_attachBody(attachment_frame, mjs_findBody(child, "b1"), "b1-", ""); + mjsElement* attachment_frame = mjs_addFrame(body_1, 0)->element; + mjs_attach(attachment_frame, mjs_findBody(child, "b1")->element, "b1-", ""); mjModel* model_1 = mj_compile(parent, 0); EXPECT_THAT(model_1, NotNull()); - mjs_attachBody(attachment_frame, mjs_findBody(child, "b2"), "b2-", ""); + mjs_attach(attachment_frame, mjs_findBody(child, "b2")->element, "b2-", ""); mjModel* model_2 = mj_compile(parent, 0); EXPECT_THAT(model_2, NotNull()); @@ -2375,8 +2424,8 @@ TEST_F(MujocoTest, ResizeParentKeyframe) { mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size()); EXPECT_THAT(child, NotNull()) << er.data(); - mjs_attachBody(mjs_findFrame(parent, "frame"), mjs_findBody(child, "body"), - "child-", ""); + mjs_attach(mjs_findFrame(parent, "frame")->element, + mjs_findBody(child, "body")->element, "child-", ""); mjModel* model = mj_compile(parent, 0); EXPECT_THAT(model, NotNull()); @@ -2467,12 +2516,10 @@ TEST_F(MujocoTest, DifferentUnitsAllowed) { mjSpec* child = mj_parseXMLString(child_xml, 0, error.data(), error.size()); mjSpec* spec = mj_parseXMLString(parent_xml, 0, error.data(), error.size()); ASSERT_THAT(spec, NotNull()) << error.data(); - mjs_attachBody(mjs_findFrame(child, "frame"), - mjs_findBody(gchild, "gchild"), - "gchild_", ""); - mjs_attachBody(mjs_findFrame(spec, "frame"), - mjs_findBody(child, "child"), - "child_", ""); + mjs_attach(mjs_findFrame(child, "frame")->element, + mjs_findBody(gchild, "gchild")->element, "gchild_", ""); + mjs_attach(mjs_findFrame(spec, "frame")->element, + mjs_findBody(child, "child")->element, "child_", ""); mjModel* model = mj_compile(spec, 0); EXPECT_THAT(model, NotNull()); @@ -2514,6 +2561,70 @@ TEST_F(MujocoTest, DifferentUnitsAllowed) { mj_deleteModel(copied_model); } +TEST_F(MujocoTest, DifferentOptionsInAttachedFrame) { + static constexpr char xml_parent[] = R"( + + + + )"; + + static constexpr char xml_child[] = R"( + + + + + + + + + )"; + + // load specs and compile child + mjSpec* parent = mj_parseXMLString(xml_parent, 0, nullptr, 0); + EXPECT_THAT(parent, NotNull()); + mjSpec* child1 = mj_parseXMLString(xml_child, 0, nullptr, 0); + EXPECT_THAT(child1, NotNull()); + mjModel* m_child1 = mj_compile(child1, 0); + EXPECT_THAT(m_child1, NotNull()); + mjSpec* child2 = mj_parseXMLString(xml_child, 0, nullptr, 0); + EXPECT_THAT(child2, NotNull()); + mjModel* m_child2 = mj_compile(child1, 0); + EXPECT_THAT(m_child2, NotNull()); + + // attach child frame to parent worldbody + mjsBody* world = mjs_findBody(parent, "world"); + EXPECT_THAT(world, NotNull()); + mjsFrame* child1_frame = mjs_findFrame(child1, "child"); + EXPECT_THAT(child1_frame, NotNull()); + mjsFrame* child2_frame = mjs_findFrame(child2, "child"); + EXPECT_THAT(child2_frame, NotNull()); + mjsElement* attached_frame1 = + mjs_attach(world->element, child1_frame->element, "child-", "-1"); + EXPECT_THAT(attached_frame1, NotNull()); + mjsElement* attached_frame2 = + mjs_attach(world->element, child2_frame->element, "child-", "-2"); + EXPECT_THAT(attached_frame2, NotNull()); + + // wrap the child frame in the parent frame and compile + mjModel* m_attached = mj_compile(parent, 0); + EXPECT_THAT(m_attached, NotNull()); + EXPECT_NEAR(m_attached->site_quat[0], m_child1->site_quat[0], 1e-6); + EXPECT_NEAR(m_attached->site_quat[1], m_child1->site_quat[1], 1e-6); + EXPECT_NEAR(m_attached->site_quat[2], m_child1->site_quat[2], 1e-6); + EXPECT_NEAR(m_attached->site_quat[3], m_child1->site_quat[3], 1e-6); + EXPECT_NEAR(m_attached->site_quat[4], m_child2->site_quat[0], 1e-6); + EXPECT_NEAR(m_attached->site_quat[5], m_child2->site_quat[1], 1e-6); + EXPECT_NEAR(m_attached->site_quat[6], m_child2->site_quat[2], 1e-6); + EXPECT_NEAR(m_attached->site_quat[7], m_child2->site_quat[3], 1e-6); + + mj_deleteSpec(parent); + mj_deleteSpec(child1); + mj_deleteModel(m_child1); + mj_deleteSpec(child2); + mj_deleteModel(m_child2); + mj_deleteModel(m_attached); +} + TEST_F(MujocoTest, CopyAttachedSpec) { static constexpr char xml_parent[] = R"( @@ -2612,8 +2723,9 @@ TEST_F(MujocoTest, ApplyNameSpaceToDefaults) { mjSpec* parent = mj_parseXMLString(xml_p, 0, err.data(), err.size()); EXPECT_THAT(parent, NotNull()) << err.data(); - mjsBody* attached = mjs_attachBody(mjs_findFrame(parent, "parent"), - mjs_findBody(child, "body"), "child-", ""); + mjsElement* attached = + mjs_attach(mjs_findFrame(parent, "parent")->element, + mjs_findBody(child, "body")->element, "child-", ""); EXPECT_THAT(attached, NotNull()); mjModel* model = mj_compile(parent, vfs.get()); @@ -2727,7 +2839,7 @@ TEST_F(MujocoTest, ErrorWhenCompilingOrphanedSpec) { EXPECT_THAT(body, NotNull()); mjsFrame* frame = mjs_addFrame(mjs_findBody(parent, "world"), nullptr); EXPECT_THAT(frame, NotNull()); - mjs_attachBody(frame, body, "child-", ""); + mjs_attach(frame->element, body->element, "child-", ""); mj_deleteSpec(parent); mjModel* model = mj_compile(child, 0); EXPECT_THAT(model, IsNull()); @@ -2763,6 +2875,14 @@ TEST_F(MujocoTest, UserValue) { EXPECT_STREQ(static_cast(payload), data.c_str()); mjs_deleteUserValue(body->element, "key"); EXPECT_THAT(mjs_getUserValue(body->element, "key"), IsNull()); + + std::string* heap_data = new std::string("heap_data"); + mjs_setUserValueWithCleanup( + body->element, "key", heap_data, + [](const void* data) { delete static_cast(data); }); + payload = mjs_getUserValue(body->element, "key"); + EXPECT_STREQ(static_cast(payload)->c_str(), + heap_data->c_str()); mj_deleteSpec(spec); } diff --git a/test/user/user_cache_test.cc b/test/user/user_cache_test.cc index f226827b..82d0f57b 100644 --- a/test/user/user_cache_test.cc +++ b/test/user/user_cache_test.cc @@ -65,6 +65,7 @@ GetCachedText(mjCCache& cache, const std::string& model, bool inserted = cache.PopulateData(resource, [&cached_text](const void* data) { cached_text = *(static_cast(data)); + return true; }); mju_closeResource(resource); mj_deleteVFS(&vfs); diff --git a/test/user/user_flex_test.cc b/test/user/user_flex_test.cc index 67dfb8f6..492c4170 100644 --- a/test/user/user_flex_test.cc +++ b/test/user/user_flex_test.cc @@ -243,6 +243,27 @@ TEST_F(UserFlexTest, RigidFlex) { mj_deleteModel(m); mj_deleteData(d); } + +TEST_F(UserFlexTest, FlexNotCollide) { + static constexpr char xml[] = R"( + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(m, NotNull()) << error.data(); + mjData* d = mj_makeData(m); + mj_step(m, d); + mj_deleteModel(m); + mj_deleteData(d); +} + TEST_F(UserFlexTest, BoundingBoxCoordinates) { static constexpr char xml[] = R"( @@ -697,6 +718,18 @@ TEST_F(UserFlexTest, LoadMSHASCII_41_MissingElement_Fail) { mj_deleteModel(m); } +TEST_F(UserFlexTest, + LoadMSHASCII_41_MismatchBetweenMaxNodesAndNodesInBlock_Fail) { + const std::string xml_path = + GetTestDataFilePath( + "user/testdata/malformed_cube_41_ascii_mismatch_between_max_nodes_and_nodes_in_block.xml"); + std::array error; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error.data(), error.size()); + EXPECT_THAT(error.data(), HasSubstr( + "XML Error: Error: Maximum number of nodes must be equal to number of nodes in a block\nElement 'flexcomp', line 22\n")); + mj_deleteModel(m); +} + TEST_F(UserFlexTest, LoadMSHASCII_22_MissingNumNodes_Fail) { const std::string xml_path = GetTestDataFilePath( diff --git a/test/user/user_mesh_test.cc b/test/user/user_mesh_test.cc index b9c1f3af..759561b2 100644 --- a/test/user/user_mesh_test.cc +++ b/test/user/user_mesh_test.cc @@ -42,14 +42,14 @@ static const char* const kDuplicateVerticesPath = "user/testdata/duplicate_vertices.xml"; static const char* const kCubePath = "user/testdata/cube.xml"; +static const char* const kCubeCompletePath = + "user/testdata/cube_complete.obj"; static const char* const kTorusPath = "user/testdata/torus.xml"; static const char* const kTorusMaxhullVertPath = "user/testdata/torus_maxhullvert.xml"; static const char* const kTorusDefaultMaxhullVertPath = "user/testdata/torus_maxhullvert_default.xml"; -static const char* const kTorusShellPath = - "user/testdata/torus_shell.xml"; static const char* const kCompareInertiaPath = "user/testdata/inertia_compare.xml"; static const char* const kConvexInertiaPath = @@ -1207,6 +1207,48 @@ TEST_F(MjCMeshTest, InvalidIndexInFace) { mj_deleteModel(model); } +TEST_F(MjCMeshTest, QhullCache) { + static constexpr char xml1[] = R"( + + + + + + + + + )"; + + static constexpr char xml2[] = R"( + + + + + + + + + )"; + + mjVFS vfs; + mj_defaultVFS(&vfs); + mj_addFileVFS(&vfs, "", GetTestDataFilePath(kCubeCompletePath).c_str()); + + std::array error; + mjModel* model = LoadModelFromString(xml1, error.data(), error.size(), &vfs); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + EXPECT_THAT(model->mesh_graphadr[0], -1); + + mj_deleteModel(model); + + model = LoadModelFromString(xml2, error.data(), error.size(), &vfs); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + EXPECT_GT(model->mesh_graphadr[0], -1); + + mj_deleteModel(model); + mj_deleteVFS(&vfs); +} + TEST_F(MjCMeshTest, LoadSkin) { const std::string xml_path = GetTestDataFilePath(kCubeSkinPath); std::array error; diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 9f85669e..8b00753d 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -623,7 +623,7 @@ TEST_F(MujocoTest, Modeldir) { mjSpec* spec = mj_makeSpec(); mjs_setDeepCopy(spec, true); mjs_setString(spec->meshdir, "asset"); - mjs_attachFrame(mjs_findBody(spec, "world"), frame, "_", ""); + mjs_attach(mjs_findBody(spec, "world")->element, frame->element, "_", ""); mjModel* model = mj_compile(spec, vfs.get()); EXPECT_THAT(model, NotNull()); diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc index efd51a99..396b2ea0 100644 --- a/test/user/user_objects_test.cc +++ b/test/user/user_objects_test.cc @@ -2070,6 +2070,51 @@ TEST_F(TendonTest, SiteBetweenPulleyNotAllowed) { EXPECT_THAT(error.data(), HasSubstr("line 9")); } +TEST_F(TendonTest, ActuatorForceRangeNotAllowed) { + std::string xml = R"( + + + + + + + + + + + + + + + + )"; + + std::array error; + std::string str_replace = "{}"; + size_t rng_ind = xml.find(str_replace); + + std::string xml0 = xml; + std::string range0 = "-2 -1"; + xml0.replace(rng_ind, str_replace.length(), range0); + mjModel* m0 = LoadModelFromString(xml0.c_str(), error.data(), error.size()); + EXPECT_THAT(m0, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("invalid actuatorfrcrange in tendon")); + + std::string xml1 = xml; + std::string range1 = "1 2"; + xml1.replace(rng_ind, str_replace.length(), range1); + mjModel* m1 = LoadModelFromString(xml1.c_str(), error.data(), error.size()); + EXPECT_THAT(m1, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("invalid actuatorfrcrange in tendon")); + + std::string xml2 = xml; + std::string range2 = "1 0"; + xml2.replace(rng_ind, str_replace.length(), range2); + mjModel* m2 = LoadModelFromString(xml2.c_str(), error.data(), error.size()); + EXPECT_THAT(m2, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("invalid actuatorfrcrange in tendon")); +} + // ------------- tests for tendon springrange ---------------------------------- using SpringrangeTest = MujocoTest; @@ -2249,6 +2294,29 @@ TEST_F(UserObjectsTest, FrameTransformsLight) { mj_deleteModel(m); } +TEST_F(ContentTypeTest, ImageLightsReferenceTexture) { + static constexpr char xml[] = R"( + + + + + + + + + + )"; + + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, NotNull()); + EXPECT_EQ(m->ntex, 1); + EXPECT_EQ(m->nlight, 1); + EXPECT_THAT(m->light_texid[0], 0); + mj_deleteModel(m); +} + // ------------- test bvh ------------------------------------------------------ TEST_F(UserObjectsTest, RobustBVH) { @@ -2508,7 +2576,8 @@ TEST_F(UserObjectsTest, Inertial) { mjtNum quat[4]; const mjtNum euler[3] = {3, 4, 5}; mju_euler2Quat(quat, euler, "xyz"); - EXPECT_EQ(AsVector(m->body_iquat+4, 4), AsVector(quat, 4)); + EXPECT_THAT(AsVector(m->body_iquat+4, 4), + Pointwise(DoubleNear(1e-8), AsVector(quat, 4))); EXPECT_EQ(m->body_mass[2], 2); EXPECT_THAT(AsVector(m->body_ipos+6, 3), ElementsAre(1, 2, 3)); diff --git a/test/xml/testdata/lights.xml b/test/xml/testdata/lights.xml new file mode 100644 index 00000000..7937e1e7 --- /dev/null +++ b/test/xml/testdata/lights.xml @@ -0,0 +1,12 @@ + + + + + + + + + + + + diff --git a/test/xml/xml_api_test.cc b/test/xml/xml_api_test.cc index fdd1421b..d8b7a58c 100644 --- a/test/xml/xml_api_test.cc +++ b/test/xml/xml_api_test.cc @@ -199,5 +199,15 @@ TEST_F(MujocoTest, SaveXmlWithDefaultMesh) { mj_deleteModel(saved_model); } +TEST_F(MujocoTest, FreeLastXml) { + static constexpr char xml[] = ""; + mjModel* model = LoadModelFromString(xml, 0, 0); + ASSERT_THAT(model, NotNull()); + mj_deleteModel(model); + ASSERT_NE(mj_saveLastXML(nullptr, nullptr, nullptr, 0), 0); + mj_freeLastXML(); + ASSERT_EQ(mj_saveLastXML(nullptr, nullptr, nullptr, 0), 0); +} + } // namespace } // namespace mujoco diff --git a/test/xml/xml_native_reader_test.cc b/test/xml/xml_native_reader_test.cc index 52650330..05a6c61c 100644 --- a/test/xml/xml_native_reader_test.cc +++ b/test/xml/xml_native_reader_test.cc @@ -15,6 +15,7 @@ // Tests for xml/xml_native_reader.cc. #include +#include #include #include #include @@ -1148,6 +1149,88 @@ TEST_F(XMLReaderTest, ParsePolycoef) { mj_deleteModel(m); } +TEST_F(XMLReaderTest, TendonArmature) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, NotNull()) << error.data(); + EXPECT_EQ(m->ntendon, 3); + EXPECT_FLOAT_EQ(m->tendon_armature[0], 1.5); + EXPECT_FLOAT_EQ(m->tendon_armature[1], 2.5); + EXPECT_FLOAT_EQ(m->tendon_armature[2], 0); + mj_deleteModel(m); +} + +TEST_F(XMLReaderTest, TendonArmatureNegative) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("tendon armature cannot be negative")); +} + +TEST_F(XMLReaderTest, TendonArmatureGeomWrap) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("geom wrapping not supported")); +} + // ------------------------ test frame parsing --------------------------------- TEST_F(XMLReaderTest, ParseFrame) { static constexpr char xml[] = R"( @@ -1449,6 +1532,67 @@ TEST_F(XMLReaderTest, ParseReplicateRepeatedName) { EXPECT_THAT(error.data(), HasSubstr("Element 'replicate'")); } +TEST_F(XMLReaderTest, RepeatedPrefix) { + static constexpr char parent[] = R"( + + + + + + + + + + + + + + )"; + + static constexpr char child_1[] = R"( + + + + + + + + + + + + + + )"; + + static constexpr char child_2[] = R"( + + + + + + )"; + + auto vfs = std::make_unique(); + mj_defaultVFS(vfs.get()); + mj_addBufferVFS(vfs.get(), "child_1.xml", child_1, sizeof(child_1)); + mj_addBufferVFS(vfs.get(), "child_2.xml", child_2, sizeof(child_2)); + + std::array err; + mjSpec* c2 = mj_parseXMLString(child_2, 0, err.data(), err.size()); + EXPECT_THAT(c2, NotNull()) << err.data(); + mjSpec* c1 = mj_parseXMLString(child_1, vfs.get(), err.data(), err.size()); + EXPECT_THAT(c1, NotNull()) << err.data(); + mj_deleteSpec(c1); + mj_deleteSpec(c2); + + mjSpec* spec = mj_parseXMLString(parent, vfs.get(), err.data(), err.size()); + EXPECT_THAT(spec, IsNull()); + EXPECT_THAT(err.data(), HasSubstr("mismatched parents")); + mj_deleteSpec(spec); + mj_deleteVFS(vfs.get()); +} + TEST_F(XMLReaderTest, ParseReplicateExcludeTendon) { static constexpr char xml[] = R"( @@ -1801,7 +1945,7 @@ TEST_F(XMLReaderTest, CameraInvalidFovyAndSensorsize) { EXPECT_THAT(error.data(), HasSubstr("line 6")); } -TEST_F(XMLReaderTest, CameraPricipalRequiresSensorsize) { +TEST_F(XMLReaderTest, CameraPrincipalRequiresSensorsize) { static constexpr char xml[] = R"( diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index 4de0a7fd..2ad13bc7 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -22,7 +22,7 @@ #include #include #include -#include +#include // NOLINT(build/c++17) #include #include @@ -32,7 +32,6 @@ #include #include #include -#include "src/cc/array_safety.h" #include "src/xml/xml_numeric_format.h" #include "test/fixture.h" @@ -135,6 +134,23 @@ TEST_F(XMLWriterTest, SavesDisableSensor) { mj_deleteModel(model); } +TEST_F(XMLWriterTest, SavesInertial) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + mjModel* model = LoadModelFromString(xml); + std::string saved_xml = SaveAndReadXml(model); + EXPECT_THAT(saved_xml, HasSubstr("mass=\"1\"")); + mj_deleteModel(model); +} + TEST_F(XMLWriterTest, EmptyUserSensor) { static constexpr char xml[] = R"( @@ -957,8 +973,10 @@ TEST_F(XMLWriterTest, WritesSkin) { ASSERT_THAT(model, NotNull()); EXPECT_THAT(model->nskin, 1); - mjModel* mtemp = LoadModelFromString(SaveAndReadXml(model)); - ASSERT_THAT(mtemp, NotNull()); + char error[1024]; + mjModel* mtemp = LoadModelFromString(SaveAndReadXml(model), + error, sizeof(error)); + ASSERT_THAT(mtemp, NotNull()) << error; EXPECT_THAT(mtemp->nskin, 1); mj_deleteModel(model); @@ -1363,20 +1381,30 @@ TEST_F(XMLWriterTest, WriteReadCompare) { std::string xml = p.path().string(); // if file is meant to fail, skip it - if (absl::StrContains(p.path().string(), "100_humanoids") || - absl::StrContains(p.path().string(), "malformed_") || - absl::StrContains(p.path().string(), "touch_grid") || + if (absl::StrContains(p.path().string(), "malformed_") || + // exclude files that are too slow to load + absl::StrContains(p.path().string(), "cow") || absl::StrContains(p.path().string(), "gmsh_") || absl::StrContains(p.path().string(), "shark_") || - absl::StrContains(p.path().string(), "cow") || - absl::StrContains(p.path().string(), "frameless_contact_hfield") || - absl::StrContains(p.path().string(), "spheremesh")) { + absl::StrContains(p.path().string(), "spheremesh") || + // exclude files that fail the comparison test + absl::StrContains(p.path().string(), "usd") || + absl::StrContains(p.path().string(), "torus_maxhull") || + absl::StrContains(p.path().string(), "fitmesh_") || + absl::StrContains(p.path().string(), "lengthrange") || + absl::StrContains(p.path().string(), "hfield_xml") || + absl::StrContains(p.path().string(), "fromto_convex") || + absl::StrContains(p.path().string(), "cube_skin") || + absl::StrContains(p.path().string(), "cube_3x3x3")) { continue; } // load model std::array error; - mjModel* m = mj_loadXML( - xml.c_str(), nullptr, error.data(), error.size()); + mjSpec* s = + mj_parseXML(xml.c_str(), nullptr, error.data(), error.size()); + ASSERT_THAT(s, NotNull()) + << "Failed to load " << xml.c_str() << ": " << error.data(); + mjModel* m = mj_compile(s, nullptr); ASSERT_THAT(m, NotNull()) << "Failed to load " << xml.c_str() << ": " << error.data(); @@ -1385,32 +1413,33 @@ TEST_F(XMLWriterTest, WriteReadCompare) { ASSERT_THAT(d, testing::NotNull()) << "Failed to create data\n"; // save and load back - mjModel* mtemp = - LoadModelFromString(SaveAndReadXml(m), error.data(), error.size()); + auto abs_path = p.path(); + mjSpec* stemp = mj_parseXMLString(SaveAndReadXml(s).c_str(), 0, + error.data(), error.size()); + mjs_setString(stemp->modelfiledir, + abs_path.remove_filename().string().c_str()); + mjModel* mtemp = mj_compile(stemp, nullptr); - if (!mtemp) { - // if failing because assets are missing, accept the test - ASSERT_THAT(error.data(), HasSubstr("file")) - << error.data() << " from " << xml.c_str(); - } else { - mjtNum tol = 0; + ASSERT_THAT(mtemp, NotNull()) + << error.data() << " from " << xml.c_str(); - // for particularly sensitive models, relax the tolerance - if (absl::StrContains(p.path().string(), "belt.xml") || - absl::StrContains(p.path().string(), "cable.xml")) { - tol = 1e-13; - } + mjtNum tol = 0; - // compare and delete - std::string field = ""; - mjtNum result = CompareModel(m, mtemp, field); - EXPECT_LE(result, tol) - << "Loaded and saved models are different!\n" - << "Affected file " << p.path().string() << '\n' - << "Different field: " << field << '\n'; - mj_deleteModel(mtemp); + // for particularly sensitive models, relax the tolerance + if (absl::StrContains(p.path().string(), "belt.xml") || + absl::StrContains(p.path().string(), "cable.xml")) { + tol = 1e-13; } + // compare and delete + std::string field = ""; + mjtNum result = CompareModel(m, mtemp, field); + EXPECT_LE(result, tol) + << "Loaded and saved models are different!\n" + << "Affected file " << p.path().string() << '\n' + << "Different field: " << field << '\n'; + mj_deleteModel(mtemp); + // check for stack memory leak mj_step(m, d); EXPECT_EQ(d->pstack, 0) << "mjData stack memory leak detected in " << @@ -1435,21 +1464,21 @@ TEST_F(XMLWriterTest, WriteReadCompare) { ASSERT_THAT(mtemp, NotNull()); // compare with 0 tolerance - std::string field = ""; - mjtNum result = CompareModel(m, mtemp, field); + field = ""; + result = CompareModel(m, mtemp, field); EXPECT_EQ(result, 0) << "Loaded and saved binary models are different!\n" << "Affected file " << p.path().string() << '\n' << "Different field: " << field << '\n'; // clean up + mj_deleteSpec(s); + mj_deleteSpec(stemp); + mj_deleteModel(m); mj_deleteModel(mtemp); mj_deleteVFS(vfs); mju_free(vfs); mju_free(buffer); - - // delete model - mj_deleteModel(m); } } } @@ -1481,6 +1510,65 @@ TEST_F(DecompilerTest, SavesStatistics) { mj_deleteModel(model); } +TEST_F(DecompilerTest, SaveAndReadXml) { + static constexpr char xml1[] = R"( + + + + + + + )"; + static constexpr char xml2[] = R"( + + + + + + + + )"; + std::array error; + mjModel* m1 = LoadModelFromString(xml1, error.data(), error.size()); + ASSERT_THAT(m1, NotNull()) << error.data(); + m1->geom_size[0] = 10; + m1->geom_size[3] = 20; + std::string saved_xml = SaveAndReadXml(m1); + EXPECT_THAT(saved_xml, HasSubstr("geom size=\"10\"")); + EXPECT_THAT(saved_xml, HasSubstr("geom size=\"20\"")); + + // parse the mjSpec, save it and read it back + mjSpec* spec = mj_parseXMLString(xml2, nullptr, error.data(), error.size()); + EXPECT_THAT(spec, NotNull()) << error.data(); + mjModel* m2 = mj_compile(spec, nullptr); + std::string saved_xml1 = SaveAndReadXml(spec); + EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"1\"")); + EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"2\"")); + EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"3\"")); + + // modify the mjModel, save it and read it back + m2->geom_size[0] = .1; + m2->geom_size[3] = .2; + m2->geom_size[6] = .3; + EXPECT_EQ(mj_copyBack(spec, m1), 0); + EXPECT_THAT(mjs_getError(spec), HasSubstr("CopyBack")); + EXPECT_EQ(mj_copyBack(spec, m2), 1); + std::string saved_xml2 = SaveAndReadXml(spec); + EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.1\"")); + EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.2\"")); + EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.3\"")); + + // check that using mjModel as argument writes in the wrong mjSpec + std::string saved_xml3 = SaveAndReadXml(m2); + EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.1\""))); + EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.2\""))); + EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.3\""))); + + mj_deleteSpec(spec); + mj_deleteModel(m1); + mj_deleteModel(m2); +} + TEST_F(DecompilerTest, DoesntSaveInferredStatistics) { static constexpr char xml[] = R"( diff --git a/unity/Editor/Bindings/MujocoBinaryRetriever.cs b/unity/Editor/Bindings/MujocoBinaryRetriever.cs index a1721d2f..3ff96ab5 100644 --- a/unity/Editor/Bindings/MujocoBinaryRetriever.cs +++ b/unity/Editor/Bindings/MujocoBinaryRetriever.cs @@ -37,7 +37,7 @@ public class MujocoBinaryRetriever { if (AssetDatabase.LoadMainAssetAtPath(mujocoPath + "/mujoco.dylib") == null) { File.Copy( "/Applications/MuJoCo.app/Contents/Frameworks" + - "/mujoco.framework/Versions/Current/libmujoco.3.3.1.dylib", + "/mujoco.framework/Versions/Current/libmujoco.3.3.3.dylib", mujocoPath + "/mujoco.dylib"); AssetDatabase.Refresh(); } @@ -45,7 +45,7 @@ public class MujocoBinaryRetriever { if (AssetDatabase.LoadMainAssetAtPath(mujocoPath + "/libmujoco.so") == null) { File.Copy( Environment.GetFolderPath(Environment.SpecialFolder.UserProfile) + - "/.mujoco/mujoco-3.3.1/lib/libmujoco.so.3.3.1", + "/.mujoco/mujoco-3.3.3/lib/libmujoco.so.3.3.3", mujocoPath + "/libmujoco.so"); AssetDatabase.Refresh(); } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 845c5c91..11698ed6 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -110,7 +110,7 @@ public const int mjMAXLINEPNT = 1000; public const int mjMAXPLANEGRID = 200; public const bool THIRD_PARTY_MUJOCO_MJXMACRO_H_ = true; public const bool THIRD_PARTY_MUJOCO_MUJOCO_H_ = true; -public const int mjVERSION_HEADER = 331; +public const int mjVERSION_HEADER = 333; // ------------------------------------Enums------------------------------------ @@ -207,6 +207,12 @@ public enum mjtCamLight : int{ mjCAMLIGHT_TARGETBODY = 3, mjCAMLIGHT_TARGETBODYCOM = 4, } +public enum mjtLightType : int{ + mjLIGHT_SPOT = 0, + mjLIGHT_DIRECTIONAL = 1, + mjLIGHT_POINT = 2, + mjLIGHT_IMAGE = 3, +} public enum mjtTexture : int{ mjTEXTURE_2D = 0, mjTEXTURE_CUBE = 1, @@ -225,6 +231,11 @@ public enum mjtTextureRole : int{ mjTEXROLE_ORM = 9, mjNTEXROLE = 10, } +public enum mjtColorSpace : int{ + mjCOLORSPACE_AUTO = 0, + mjCOLORSPACE_LINEAR = 1, + mjCOLORSPACE_SRGB = 2, +} public enum mjtIntegrator : int{ mjINT_EULER = 0, mjINT_RK4 = 1, @@ -320,6 +331,7 @@ public enum mjtObj : int{ mjNOBJECT = 26, mjOBJ_FRAME = 100, mjOBJ_DEFAULT = 101, + mjOBJ_MODEL = 102, } public enum mjtConstraint : int{ mjCNSTR_EQUALITY = 0, @@ -356,34 +368,35 @@ public enum mjtSensor : int{ mjSENS_ACTUATORVEL = 14, mjSENS_ACTUATORFRC = 15, mjSENS_JOINTACTFRC = 16, - mjSENS_BALLQUAT = 17, - mjSENS_BALLANGVEL = 18, - mjSENS_JOINTLIMITPOS = 19, - mjSENS_JOINTLIMITVEL = 20, - mjSENS_JOINTLIMITFRC = 21, - mjSENS_TENDONLIMITPOS = 22, - mjSENS_TENDONLIMITVEL = 23, - mjSENS_TENDONLIMITFRC = 24, - mjSENS_FRAMEPOS = 25, - mjSENS_FRAMEQUAT = 26, - mjSENS_FRAMEXAXIS = 27, - mjSENS_FRAMEYAXIS = 28, - mjSENS_FRAMEZAXIS = 29, - mjSENS_FRAMELINVEL = 30, - mjSENS_FRAMEANGVEL = 31, - mjSENS_FRAMELINACC = 32, - mjSENS_FRAMEANGACC = 33, - mjSENS_SUBTREECOM = 34, - mjSENS_SUBTREELINVEL = 35, - mjSENS_SUBTREEANGMOM = 36, - mjSENS_GEOMDIST = 37, - mjSENS_GEOMNORMAL = 38, - mjSENS_GEOMFROMTO = 39, - mjSENS_E_POTENTIAL = 40, - mjSENS_E_KINETIC = 41, - mjSENS_CLOCK = 42, - mjSENS_PLUGIN = 43, - mjSENS_USER = 44, + mjSENS_TENDONACTFRC = 17, + mjSENS_BALLQUAT = 18, + mjSENS_BALLANGVEL = 19, + mjSENS_JOINTLIMITPOS = 20, + mjSENS_JOINTLIMITVEL = 21, + mjSENS_JOINTLIMITFRC = 22, + mjSENS_TENDONLIMITPOS = 23, + mjSENS_TENDONLIMITVEL = 24, + mjSENS_TENDONLIMITFRC = 25, + mjSENS_FRAMEPOS = 26, + mjSENS_FRAMEQUAT = 27, + mjSENS_FRAMEXAXIS = 28, + mjSENS_FRAMEYAXIS = 29, + mjSENS_FRAMEZAXIS = 30, + mjSENS_FRAMELINVEL = 31, + mjSENS_FRAMEANGVEL = 32, + mjSENS_FRAMELINACC = 33, + mjSENS_FRAMEANGACC = 34, + mjSENS_SUBTREECOM = 35, + mjSENS_SUBTREELINVEL = 36, + mjSENS_SUBTREEANGMOM = 37, + mjSENS_GEOMDIST = 38, + mjSENS_GEOMNORMAL = 39, + mjSENS_GEOMFROMTO = 40, + mjSENS_E_POTENTIAL = 41, + mjSENS_E_KINETIC = 42, + mjSENS_CLOCK = 43, + mjSENS_PLUGIN = 44, + mjSENS_USER = 45, } public enum mjtStage : int{ mjSTAGE_NONE = 0, @@ -417,6 +430,12 @@ public enum mjtFlexSelf : int{ mjFLEXSELF_SAP = 3, mjFLEXSELF_AUTO = 4, } +public enum mjtSDFType : int{ + mjSDFTYPE_SINGLE = 0, + mjSDFTYPE_INTERSECTION = 1, + mjSDFTYPE_MIDSURFACE = 2, + mjSDFTYPE_COLLISION = 3, +} public enum mjtPluginCapabilityBit : int{ mjPLUGIN_ACTUATOR = 1, mjPLUGIN_SENSOR = 2, @@ -4824,7 +4843,6 @@ public unsafe struct mjData_ { public mjSolverStat_ solver3997; public mjSolverStat_ solver3998; public mjSolverStat_ solver3999; - public int solver_nisland; public fixed int solver_niter[20]; public fixed int solver_nnz[20]; public fixed double solver_fwdinv[2]; @@ -4859,6 +4877,7 @@ public unsafe struct mjData_ { public int nJ; public int nA; public int nisland; + public int nidof; public double time; public fixed double energy[2]; public void* buffer; @@ -4921,6 +4940,7 @@ public unsafe struct mjData_ { public double* actuator_moment; public double* crb; public double* qM; + public double* M; public double* qLD; public double* qLDiagInv; public double* bvh_aabb_dyn; @@ -4947,10 +4967,6 @@ public unsafe struct mjData_ { public int* M_rowadr; public int* M_colind; public int* mapM2M; - public int* C_rownnz; - public int* C_rowadr; - public int* C_colind; - public int* mapM2C; public int* D_rownnz; public int* D_rowadr; public int* D_diag; @@ -4990,14 +5006,42 @@ public unsafe struct mjData_ { public double* efc_R; public int* tendon_efcadr; public int* dof_island; - public int* island_dofnum; + public int* island_nv; + public int* island_idofadr; public int* island_dofadr; - public int* island_dofind; - public int* dof_islandind; + public int* map_dof2idof; + public int* map_idof2dof; + public double* ifrc_smooth; + public double* iacc_smooth; + public int* iM_rownnz; + public int* iM_rowadr; + public int* iM_colind; + public double* iM; + public double* iLD; + public double* iLDiagInv; + public double* iacc; public int* efc_island; - public int* island_efcnum; - public int* island_efcadr; - public int* island_efcind; + public int* island_ne; + public int* island_nf; + public int* island_nefc; + public int* island_iefcadr; + public int* map_efc2iefc; + public int* map_iefc2efc; + public int* iefc_type; + public int* iefc_id; + public int* iefc_J_rownnz; + public int* iefc_J_rowadr; + public int* iefc_J_rowsuper; + public int* iefc_J_colind; + public int* iefc_JT_rownnz; + public int* iefc_JT_rowadr; + public int* iefc_JT_rowsuper; + public int* iefc_JT_colind; + public double* iefc_J; + public double* iefc_JT; + public double* iefc_frictionloss; + public double* iefc_D; + public double* iefc_R; public int* efc_AR_rownnz; public int* efc_AR_rowadr; public int* efc_AR_colind; @@ -5005,9 +5049,14 @@ public unsafe struct mjData_ { public double* efc_vel; public double* efc_aref; public double* efc_b; - public double* efc_force; + public double* iefc_aref; + public int* iefc_state; + public double* iefc_force; public int* efc_state; + public double* efc_force; + public double* ifrc_constraint; public UIntPtr threadpool; + public UInt64 signature; } [StructLayout(LayoutKind.Sequential)] @@ -5391,9 +5440,12 @@ public unsafe struct mjModel_ { public int* light_mode; public int* light_bodyid; public int* light_targetbodyid; - public byte* light_directional; + public int* light_type; + public int* light_texid; public byte* light_castshadow; public float* light_bulbradius; + public float* light_intensity; + public float* light_range; public byte* light_active; public double* light_pos; public double* light_dir; @@ -5441,6 +5493,7 @@ public unsafe struct mjModel_ { public int* flex_nodebodyid; public int* flex_vertbodyid; public int* flex_edge; + public int* flex_edgeflap; public int* flex_elem; public int* flex_elemtexcoord; public int* flex_elemedge; @@ -5455,6 +5508,7 @@ public unsafe struct mjModel_ { public double* flexedge_invweight0; public double* flex_radius; public double* flex_stiffness; + public double* flex_bending; public double* flex_damping; public double* flex_edgestiffness; public double* flex_edgedamping; @@ -5527,6 +5581,7 @@ public unsafe struct mjModel_ { public float* hfield_data; public int* hfield_pathadr; public int* tex_type; + public int* tex_colorspace; public int* tex_height; public int* tex_width; public int* tex_nchannel; @@ -5567,15 +5622,18 @@ public unsafe struct mjModel_ { public int* tendon_matid; public int* tendon_group; public byte* tendon_limited; + public byte* tendon_actfrclimited; public double* tendon_width; public double* tendon_solref_lim; public double* tendon_solimp_lim; public double* tendon_solref_fri; public double* tendon_solimp_fri; public double* tendon_range; + public double* tendon_actfrcrange; public double* tendon_margin; public double* tendon_stiffness; public double* tendon_damping; + public double* tendon_armature; public double* tendon_frictionloss; public double* tendon_lengthspring; public double* tendon_length0; @@ -5672,6 +5730,7 @@ public unsafe struct mjModel_ { public char* names; public int* names_map; public char* paths; + public UInt64 signature; } [StructLayout(LayoutKind.Sequential)] @@ -5762,6 +5821,7 @@ public unsafe struct mjsCompiler_ { public byte fusestatic; public int inertiafromgeom; public fixed int inertiagrouprange[2]; + public byte saveinertial; public int alignfree; public mjLROpt_ LRopt; } @@ -6022,6 +6082,8 @@ public unsafe struct mjvGeom_ { public unsafe struct mjvLight_ { public fixed float pos[3]; public fixed float dir[3]; + public int type; + public int texid; public fixed float attenuation[3]; public float cutoff; public float exponent; @@ -6029,9 +6091,10 @@ public unsafe struct mjvLight_ { public fixed float diffuse[3]; public fixed float specular[3]; public byte headlight; - public byte directional; public byte castshadow; public float bulbradius; + public float intensity; + public float range; } [StructLayout(LayoutKind.Sequential)] @@ -6227,267 +6290,6 @@ public unsafe struct mjvFigure_ { public fixed int yaxispixel[2]; public fixed float xaxisdata[2]; public fixed float yaxisdata[2]; -} - -[StructLayout(LayoutKind.Sequential)] -public unsafe struct model { - public int nv; - public int nu; - public int na; - public int nbody; - public int nbvh; - public int nbvhstatic; - public int njnt; - public int ngeom; - public int nsite; - public int ncam; - public int nlight; - public int nmesh; - public int nskin; - public int nflex; - public int nflexvert; - public int nflextexcoord; - public int nskinvert; - public int nskinface; - public int nskinbone; - public int nskinbonevert; - public int nmat; - public int neq; - public int ntendon; - public int ntree; - public int nwrap; - public int nsensor; - public int nnames; - public int npaths; - public int nsensordata; - public int narena; - public mjOption_ opt; - public mjVisual_ vis; - public mjStatistic_ stat; - public int* body_parentid; - public int* body_rootid; - public int* body_weldid; - public int* body_mocapid; - public int* body_jntnum; - public int* body_jntadr; - public int* body_dofnum; - public int* body_dofadr; - public int* body_geomnum; - public int* body_geomadr; - public double* body_iquat; - public double* body_mass; - public double* body_inertia; - public int* body_bvhadr; - public int* body_bvhnum; - public int* bvh_depth; - public int* bvh_child; - public int* bvh_nodeid; - public double* bvh_aabb; - public int* jnt_type; - public int* jnt_bodyid; - public int* jnt_group; - public int* geom_type; - public int* geom_bodyid; - public int* geom_contype; - public int* geom_conaffinity; - public int* geom_dataid; - public int* geom_matid; - public int* geom_group; - public double* geom_size; - public double* geom_aabb; - public double* geom_rbound; - public float* geom_rgba; - public int* site_type; - public int* site_bodyid; - public int* site_matid; - public int* site_group; - public double* site_size; - public float* site_rgba; - public int* cam_orthographic; - public double* cam_fovy; - public double* cam_ipd; - public int* cam_resolution; - public float* cam_sensorsize; - public float* cam_intrinsic; - public byte* light_directional; - public byte* light_castshadow; - public float* light_bulbradius; - public byte* light_active; - public float* light_attenuation; - public float* light_cutoff; - public float* light_exponent; - public float* light_ambient; - public float* light_diffuse; - public float* light_specular; - public byte* flex_flatskin; - public int* flex_dim; - public int* flex_matid; - public int* flex_group; - public int* flex_interp; - public int* flex_nodeadr; - public int* flex_nodenum; - public int* flex_nodebodyid; - public int* flex_vertadr; - public int* flex_vertnum; - public int* flex_elem; - public int* flex_elemtexcoord; - public int* flex_elemlayer; - public int* flex_elemadr; - public int* flex_elemnum; - public int* flex_elemdataadr; - public int* flex_shell; - public int* flex_shellnum; - public int* flex_shelldataadr; - public int* flex_texcoordadr; - public int* flex_bvhadr; - public int* flex_bvhnum; - public byte* flex_centered; - public double* flex_node; - public double* flex_radius; - public float* flex_rgba; - public float* flex_texcoord; - public int* hfield_pathadr; - public int* mesh_bvhadr; - public int* mesh_bvhnum; - public int* mesh_texcoordadr; - public int* mesh_graphadr; - public int* mesh_pathadr; - public int* skin_matid; - public int* skin_group; - public float* skin_rgba; - public float* skin_inflate; - public int* skin_vertadr; - public int* skin_vertnum; - public int* skin_texcoordadr; - public int* skin_faceadr; - public int* skin_facenum; - public int* skin_boneadr; - public int* skin_bonenum; - public float* skin_vert; - public int* skin_face; - public int* skin_bonevertadr; - public int* skin_bonevertnum; - public float* skin_bonebindpos; - public float* skin_bonebindquat; - public int* skin_bonebodyid; - public int* skin_bonevertid; - public float* skin_bonevertweight; - public int* skin_pathadr; - public int* tex_pathadr; - public int* mat_texid; - public byte* mat_texuniform; - public float* mat_texrepeat; - public float* mat_emission; - public float* mat_specular; - public float* mat_shininess; - public float* mat_reflectance; - public float* mat_metallic; - public float* mat_roughness; - public float* mat_rgba; - public int* eq_type; - public int* eq_obj1id; - public int* eq_obj2id; - public int* eq_objtype; - public double* eq_data; - public int* tendon_num; - public int* tendon_matid; - public int* tendon_group; - public byte* tendon_limited; - public double* tendon_width; - public double* tendon_range; - public double* tendon_stiffness; - public double* tendon_damping; - public double* tendon_frictionloss; - public double* tendon_lengthspring; - public float* tendon_rgba; - public int* actuator_trntype; - public int* actuator_dyntype; - public int* actuator_trnid; - public int* actuator_actadr; - public int* actuator_actnum; - public int* actuator_group; - public byte* actuator_ctrllimited; - public byte* actuator_actlimited; - public double* actuator_ctrlrange; - public double* actuator_actrange; - public double* actuator_cranklength; - public int* sensor_type; - public int* sensor_objid; - public int* sensor_adr; - public int* name_bodyadr; - public int* name_jntadr; - public int* name_geomadr; - public int* name_siteadr; - public int* name_camadr; - public int* name_lightadr; - public int* name_eqadr; - public int* name_tendonadr; - public int* name_actuatoradr; - public char* names; - public char* paths; -} - -[StructLayout(LayoutKind.Sequential)] -public unsafe struct data { - public mjWarningStat_ warning0; - public mjWarningStat_ warning1; - public mjWarningStat_ warning2; - public mjWarningStat_ warning3; - public mjWarningStat_ warning4; - public mjWarningStat_ warning5; - public mjWarningStat_ warning6; - public mjWarningStat_ warning7; - public int nefc; - public int ncon; - public int nisland; - public double time; - public double* act; - public double* ctrl; - public double* xfrc_applied; - public byte* eq_active; - public double* sensordata; - public double* xpos; - public double* xquat; - public double* xmat; - public double* xipos; - public double* ximat; - public double* xanchor; - public double* xaxis; - public double* geom_xpos; - public double* geom_xmat; - public double* site_xpos; - public double* site_xmat; - public double* cam_xpos; - public double* cam_xmat; - public double* light_xpos; - public double* light_xdir; - public double* subtree_com; - public int* ten_wrapadr; - public int* ten_wrapnum; - public int* wrap_obj; - public double* ten_length; - public double* wrap_xpos; - public double* bvh_aabb_dyn; - public byte* bvh_active; - public int* island_dofadr; - public int* island_dofind; - public int* dof_island; - public int* efc_island; - public int* tendon_efcadr; - public double* flexvert_xpos; - public mjContact_* contact; - public double* efc_force; - public void* arena; -} - -[StructLayout(LayoutKind.Sequential)] -public unsafe struct mjvSceneState_ { - public int nbuffer; - public void* buffer; - public int maxgeom; - public mjvScene_ scratch; - public model model; - public data data; }public struct mjuiItem_ {}public struct mjfItemEnable {} // ----------------------------Function declarations---------------------------- @@ -6570,6 +6372,9 @@ public static unsafe extern mjData_* mj_makeData(mjModel_* m); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern mjData_* mj_copyData(mjData_* dest, mjModel_* m, mjData_* src); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern mjData_* mjv_copyData(mjData_* dest, mjModel_* m, mjData_* src); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_resetData(mjModel_* m, mjData_* d); @@ -6708,6 +6513,9 @@ public static unsafe extern void mj_transmission(mjModel_* m, mjData_* d); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_crb(mjModel_* m, mjData_* d); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_makeM(mjModel_* m, mjData_* d); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_factorM(mjModel_* m, mjData_* d); @@ -6924,15 +6732,9 @@ public static unsafe extern void mjv_alignToCamera(double* res, double* vec, dou [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_moveCamera(mjModel_* m, int action, double reldx, double reldy, mjvScene_* scn, mjvCamera_* cam); -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_moveCameraFromState(mjvSceneState_* scnstate, int action, double reldx, double reldy, mjvScene_* scn, mjvCamera_* cam); - [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_movePerturb(mjModel_* m, mjData_* d, int action, double reldx, double reldy, mjvScene_* scn, mjvPerturb_* pert); -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_movePerturbFromState(mjvSceneState_* scnstate, int action, double reldx, double reldy, mjvScene_* scn, mjvPerturb_* pert); - [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_moveModel(mjModel_* m, int action, double reldx, double reldy, double* roomup, mjvScene_* scn); @@ -6975,24 +6777,9 @@ public static unsafe extern void mjv_freeScene(mjvScene_* scn); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_updateScene(mjModel_* m, mjData_* d, mjvOption_* opt, mjvPerturb_* pert, mjvCamera_* cam, int catmask, mjvScene_* scn); -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern int mjv_updateSceneFromState(mjvSceneState_* scnstate, mjvOption_* opt, mjvPerturb_* pert, mjvCamera_* cam, int catmask, mjvScene_* scn); - [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_copyModel(mjModel_* dest, mjModel_* src); -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_defaultSceneState(mjvSceneState_* scnstate); - -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_makeSceneState(mjModel_* m, mjData_* d, mjvSceneState_* scnstate, int maxgeom); - -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_freeSceneState(mjvSceneState_* scnstate); - -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mjv_updateSceneState(mjModel_* m, mjData_* d, mjvOption_* opt, mjvSceneState_* scnstate); - [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjv_addGeoms(mjModel_* m, mjData_* d, mjvOption_* opt, mjvPerturb_* pert, int catmask, mjvScene_* scn); diff --git a/unity/package.json b/unity/package.json index 9c599fa3..c91c5d6c 100644 --- a/unity/package.json +++ b/unity/package.json @@ -1,7 +1,7 @@ { "name": "org.mujoco", "displayName": "MuJoCo", - "version": "3.3.1", + "version": "3.3.3", "description": "MuJoCo importer and runtime plug-in", "dependencies": {}, "author": {