diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..8a2aa9a2 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,20 @@ +# Run pre-commit hooks on files changed in the PR only. +name: lint + +on: + pull_request: + +jobs: + pre-commit: + name: pre-commit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - uses: pre-commit/action@v3.0.1 + with: + extra_args: --from-ref ${{ github.event.pull_request.base.sha }} --to-ref ${{ github.event.pull_request.head.sha }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..a674a832 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,6 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer diff --git a/README.md b/README.md index 9eabc498..98bee9e2 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,7 @@ Users who wish to build MuJoCo from source should consult the [build from source] section of the documentation. However, note that the commit at the tip of the `main` branch may be unstable. -### Python (>= 3.9) +### Python (>= 3.10) The native Python bindings, which come pre-packaged with a copy of MuJoCo, can be installed from [PyPI] via: diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake index 699bbb15..d2404bc1 100644 --- a/cmake/MujocoDependencies.cmake +++ b/cmake/MujocoDependencies.cmake @@ -58,11 +58,6 @@ set(MUJOCO_DEP_VERSION_benchmark CACHE STRING "Version of `benchmark` to be fetched." ) -set(MUJOCO_DEP_VERSION_TriangleMeshDistance - 2cb643de1436e1ba8e2be49b07ec5491ac604457 - CACHE STRING "Version of `TriangleMeshDistance` to be fetched." -) - mark_as_advanced(MUJOCO_DEP_VERSION_lodepng) mark_as_advanced(MUJOCO_DEP_VERSION_MarchingCubeCpp) mark_as_advanced(MUJOCO_DEP_VERSION_tinyxml2) @@ -73,7 +68,6 @@ mark_as_advanced(MUJOCO_DEP_VERSION_Eigen3) mark_as_advanced(MUJOCO_DEP_VERSION_abseil) mark_as_advanced(MUJOCO_DEP_VERSION_gtest) mark_as_advanced(MUJOCO_DEP_VERSION_benchmark) -mark_as_advanced(MUJOCO_DEP_VERSION_TriangleMeshDistance) include(FetchContent) include(FindOrFetch) @@ -202,30 +196,6 @@ if(CMAKE_POLICY_VERSION_MINIMUM_LOCALLY_DEFINED) unset(CMAKE_POLICY_VERSION_MINIMUM_LOCALLY_DEFINED) endif() -if(NOT TARGET trianglemeshdistance) - FetchContent_Declare( - trianglemeshdistance - GIT_REPOSITORY https://github.com/InteractiveComputerGraphics/TriangleMeshDistance.git - GIT_TAG ${MUJOCO_DEP_VERSION_TriangleMeshDistance} - ) - - FetchContent_GetProperties(trianglemeshdistance) - if(NOT trianglemeshdistance_POPULATED) - FetchContent_Populate(trianglemeshdistance) - # Patch the source code to silence a warning/error related to a loop variable creating a copy. - # Since this is a header only library this fix is less intrusive than disabling the warning for - # any target including the header. - set(TMD_HEADER ${trianglemeshdistance_SOURCE_DIR}/TriangleMeshDistance/include/tmd/TriangleMeshDistance.h) - file(READ ${TMD_HEADER} TMD_CONTENT) - string(REPLACE - "for (const auto edge_count : edges_count) {" - "for (const auto& edge_count : edges_count) {" - TMD_CONTENT "${TMD_CONTENT}") - file(WRITE ${TMD_HEADER} "${TMD_CONTENT}") - include_directories(${trianglemeshdistance_SOURCE_DIR}) - endif() -endif() - set(ENABLE_DOUBLE_PRECISION ON) set(CCD_HIDE_ALL_SYMBOLS ON) diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 2292549e..656c6a07 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -368,10 +368,41 @@ Numeric constants ^^^^^^^^^^^^^^^^^ Many integer constants were already documented in the primitive types above. In addition, the header files define -several other constants documented here. Unless indicated otherwise, each entry in the table below is defined in -`mjmodel.h `_. Note that some extended key -codes are defined in `mjui.h `_ which are not -shown in the table below. Their names are in the format ``mjKEY_XXX``. They correspond to GLFW key codes. +several other constants documented here. Note that some extended key codes are defined in +`mjui.h `_ which are not shown below. Their +names are in the format ``mjKEY_XXX``. They correspond to GLFW key codes. + + +.. _glNumericVersion: + +Version +~~~~~~~ + +Defined in `mujoco.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description + * - ``mjVERSION_HEADER`` + - 3005001 + - The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P" + using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001. + The API function :ref:`mj_version` returns a number with the same meaning + but for the compiled library. See + `VERSIONING.md `__ for details. + + +.. _glNumericEngine: + +Engine constants +~~~~~~~~~~~~~~~~ + +Defined in `mjmodel.h `_ unless +indicated otherwise. .. list-table:: :widths: 2 1 8 @@ -384,6 +415,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - 1E-15 - The minimal value allowed in any denominator, and in general any mathematical operation where 0 is not allowed. In almost all cases, MuJoCo silently clamps smaller values to mjMINVAL. + Defined in `mjtnum.h `_. * - ``mjPI`` - :math:`\pi` - The value of :math:`\pi`. This is used in various trigonometric functions, and also for conversion from degrees @@ -416,17 +448,38 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr is raised and ray casting may not be possible. For a balanced hierarchy, this implies 1E15 bounding volumes. * - ``mjMAXFLEXNODES`` - 27 - - Some number by Alessio that needs documentation. I guess it's related to trilinear flexes? + - The maximum number of nodes in a trilinear flex element. * - ``mjMINAWAKE`` - 10 - The minimum number of timesteps that must pass after a tree is awoken, before it is allowed to go back to sleep. + * - ``mjMAXTHREAD`` + - 128 + - Maximum number of OS threads that can be used in a thread pool. + Defined in `mjthread.h `_. + + +.. _glNumericSizes: + +Array sizes +~~~~~~~~~~~ + +Defined in `mjmodel.h `_. These constants +correspond to array sizes which we have not fully settled. There may be reasons to increase them in the future, so as to +accommodate extra parameters needed for more elaborate computations. This is why we maintain them as symbolic constants +that can be easily changed, as opposed to the array size for representing quaternions for example -- which has no reason +to change. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNEQDATA`` - 11 - The maximal number of real-valued parameters used to define each equality constraint. Determines the size of - ``mjModel.eq_data``. This and the next five constants correspond to array sizes which we have not fully settled. - There may be reasons to increase them in the future, so as to accommodate extra parameters needed for more - elaborate computations. This is why we maintain them as symbolic constants that can be easily changed, as opposed - to the array size for representing quaternions for example -- which has no reason to change. + ``mjModel.eq_data``. * - ``mjNDYN`` - 10 - The maximal number of real-valued parameters used to define the activation dynamics of each actuator. @@ -464,75 +517,104 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - The number of islands for which 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 islands for which the solver was run is given by ``mjData.nsolver_island``. + + +.. _glNumericVisualization: + +Visualization +~~~~~~~~~~~~~ + +Defined in `mjvisualize.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNGROUP`` - 6 - The number of geom, site, joint, tendon and actuator groups whose rendering can be enabled and disabled via :ref:`mjvOption`. - Defined in `mjvisualize.h `_. + * - ``mjMAXLIGHT`` + - 100 + - The maximum number of lights in a scene. * - ``mjMAXOVERLAY`` - 500 - The maximal number of characters in overlay text for rendering. - Defined in `mjvisualize.h `_. * - ``mjMAXLINE`` - 100 - The maximal number of lines per 2D figure (:ref:`mjvFigure`). - Defined in `mjvisualize.h `_. * - ``mjMAXLINEPNT`` - 1001 - The maximal number of points in each line in a 2D figure. Note that the buffer ``mjvFigure.linepnt`` has length ``2*mjMAXLINEPNT`` because each point has X and Y coordinates. - Defined in `mjvisualize.h `_. * - ``mjMAXPLANEGRID`` - 200 - The maximal number of grid lines in each dimension for rendering planes. - Defined in `mjvisualize.h `_. + + +.. _glNumericRendering: + +Rendering +~~~~~~~~~ + +Defined in `mjrender.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjNAUX`` - 10 - Number of auxiliary buffers that can be allocated in mjrContext. - Defined in `mjrender.h `_. * - ``mjMAXTEXTURE`` - 1000 - Maximum number of textures allowed. - Defined in `mjrender.h `_. - * - ``mjMAXTHREAD`` - - 128 - - Maximum number OS threads that can be used in a thread pool. - Defined in `mjthread.h `_. + * - ``mjMAXMATERIAL`` + - 1000 + - Maximum number of materials with textures. + + +.. _glNumericUI: + +UI constants +~~~~~~~~~~~~ + +Defined in `mjui.h `_. + +.. list-table:: + :widths: 2 1 8 + :header-rows: 1 + + * - symbol + - value + - description * - ``mjMAXUISECT`` - 10 - Maximum number of UI sections. - Defined in `mjui.h `_. * - ``mjMAXUIITEM`` - 200 - Maximum number of items per UI section. - Defined in `mjui.h `_. * - ``mjMAXUITEXT`` - - 500 + - 300 - Maximum number of characters in UI fields 'edittext' and 'other'. - Defined in `mjui.h `_. * - ``mjMAXUINAME`` - 40 - Maximum number of characters in any UI name. - Defined in `mjui.h `_. * - ``mjMAXUIMULTI`` - - 20 + - 35 - Maximum number of radio and select items in UI group. - Defined in `mjui.h `_. * - ``mjMAXUIEDIT`` - - 5 + - 7 - Maximum number of elements in UI edit list. - Defined in `mjui.h `_. * - ``mjMAXUIRECT`` - - 15 + - 25 - Maximum number of UI rectangles. - Defined in `mjui.h `_. - * - ``mjVERSION_HEADER`` - - 3005001 - - The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P" - using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001. - Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning - but for the compiled library. See - `VERSIONING.md `__ for details. .. _Macros: diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 5cbb95c5..bec885ae 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -1252,7 +1252,7 @@ Intersect ray ``pnt+x*vec, x >= 0`` with geoms. - If ``normal`` is not NULL, write the surface normal at the intersection point. The normal always points **out of the geometry**, regardless of the ray's direction (i.e., including rays hitting the surface from the inside). - Exclude geoms in body with id ``bodyexclude``, use -1 to include all bodies. -- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass +- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass NULL to skip geom group exclusion. - If ``flg_static`` is 0, static geoms will be excluded. @@ -2911,7 +2911,7 @@ each corresponding to one item. The last (unused) item has its type set to -1, t after the end of the last used section. There is also another version of this function (:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by -:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. .. _mjui_addToSection: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 1e0d5bbf..88e90085 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -403,7 +403,7 @@ Intersect ray ``pnt+x*vec, x >= 0`` with geoms. - If ``normal`` is not NULL, write the surface normal at the intersection point. The normal always points **out of the geometry**, regardless of the ray's direction (i.e., including rays hitting the surface from the inside). - Exclude geoms in body with id ``bodyexclude``, use -1 to include all bodies. -- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass +- ``geomgroup`` is an array of length :ref:`mjNGROUP`, where 1 means the group should be included. Pass NULL to skip geom group exclusion. - If ``flg_static`` is 0, static geoms will be excluded. @@ -444,7 +444,7 @@ each corresponding to one item. The last (unused) item has its type set to -1, t after the end of the last used section. There is also another version of this function (:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by -:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. .. _mjui_update: diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 5d08dee3..5b5392a9 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -1249,7 +1249,7 @@ The full list of processing steps applied by the compiler to each mesh is as fol transformations in ``mjModel.mesh_{pos, quat, scale}``. #. Construct the convex hull if specified; #. Find the centroid of all triangle faces, and construct the union-of-pyramids representation. Triangles whose area is - too small (below the :ref:`mjMINVAL ` value of 1E-14) result in compile error; + too small (below the :ref:`mjMINVAL ` value of 1E-14) result in compile error; #. Compute the center of mass and inertia matrix of the union-of-pyramids. Use eigenvalue decomposition to find the principal axes of inertia. Center and align the mesh, saving the translational and rotational offsets for subsequent geom-related computations. @@ -1518,7 +1518,7 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe and other geoms (except for planes and other height fields which are not supported) are computed by first selecting the sub-grid of prisms that could collide with the geom based on its bounding box, and then using the general convex collider. The number of possible contacts between a height field and a geom is limited to 50 - (:ref:`mjMAXCONPAIR `); any contacts beyond that are discarded. To avoid penetration due to discarded + (:ref:`mjMAXCONPAIR `); any contacts beyond that are discarded. To avoid penetration due to discarded contacts, the spatial features of the height field should be large compared to the geoms it collides with. .. _asset-hfield-name: @@ -2263,7 +2263,8 @@ rotations as unit quaternions. .. _body-joint-solimpfriction: :at:`solreffriction`, :at:`solimpfriction` - Constraint solver parameters for simulating dry friction. See :ref:`CSolver`. + Constraint solver parameters for simulating dry friction. + See also :ref:`Friction`. .. _body-joint-stiffness: @@ -4074,7 +4075,7 @@ friction can only be created with this element. Note that as with other :at:`solreffriction` attributes, the constraint violation is identically 0. Therefore, when using positive semantics :at:`solreffriction[1]` is ignored, while for negative semantics :at:`solreffriction[0]` is - ignored. See :ref:`CSolver` for more details. + ignored. See :ref:`Friction` for more details. .. _contact-pair-margin: @@ -4936,7 +4937,8 @@ length X, as in the clip on the right of `this example model .. _tendon-spatial-solimpfriction: :at:`solreffriction`, :at:`solimpfriction` - Constraint solver parameters for simulating dry friction in the tendon. See :ref:`CSolver`. + Constraint solver parameters for simulating dry friction in the tendon. + See also :ref:`Friction`. .. _tendon-spatial-margin: diff --git a/doc/changelog.rst b/doc/changelog.rst index 79a1929e..74178e54 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,6 +2,20 @@ Changelog ========= +Upcoming version (not yet released) +----------------------------------- + +MJX +^^^ + +- Add batch rendering support for MJX-Warp. See the :ref:`MJX-Warp batch rendering` section for details. + +Bug fixes +^^^^^^^^^ + +- Fixed a bug where :ref:`mjs_attach` silently dropped spatial tendons with wrapping geometries that had no + ``sidesite`` attribute (:issue:`3119`, reported by :github:user:`tomstewart89`). + Version 3.5.0 (February 12, 2026) --------------------------------- @@ -1449,8 +1463,8 @@ General Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. If the solver does not use islands, only row 0 is filled. - - The new constant :ref:`mjNISLAND` was set to 20. - - :ref:`mjNSOLVER` was reduced from 1000 to 200. + - The new constant :ref:`mjNISLAND` was set to 20. + - :ref:`mjNSOLVER` was reduced from 1000 to 200. - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer vectors of length ``mjNISLAND``. diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 2534e498..537d9884 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -1335,6 +1335,8 @@ It is a vector with dimensionality :math:`\nq` satisfying :math:`0` section of the +Modeling chapter. -To summarize, the user specifies the vectors of impedance coefficients :math:`0 0` -and stiffness coefficients :math:`k > 0`. The quantities :math:`R, \ar` are then computed by MuJoCo as shown above, and -the selected optimization algorithm is applied to solve problem :eq:`eq:dual`. As explained in the :ref:`solver -parameters ` section of the Modeling chapter, MuJoCo offers additional automation for setting :math:`d, b, k` -so as to achieve critical damping, or model a soft contact layer by varying :math:`d` with distance. +To summarize, the constraint behavior is determined by three per-constraint quantities: impedance :math:`0 0` and stiffness :math:`k \geq 0`. These are computed from the :at:`solimp` and :at:`solref` attributes as +described in the :ref:`solver parameters ` section of the Modeling chapter, which also offers additional +automation (e.g., achieving critical damping, or varying :math:`d` with distance to model a soft contact layer). The +quantities :math:`R, \ar` are then computed from :eq:`eq:impedance_R` and :eq:`eq:aref`, and the selected optimization +algorithm is applied to solve problem :eq:`eq:dual`. .. _soCones: @@ -1614,17 +1621,17 @@ will generate up to 1 contact or with ``multiccd`` up to 4 contacts. | **1** * - HField - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldCCD - | :ref:`mjMAXCONPAIR ` + | :ref:`mjMAXCONPAIR ` - | HFieldSDF | :ref:`sdf_initpoints ` * - Sphere @@ -1755,7 +1762,7 @@ sleeping mechanism is provided in the :ref:`Simulation chapter` but her Sleeping can occur in one of two ways: - **Automatic:** A tree whose maximum velocity in absolute value is less than the - :ref:`tolerance ` for :ref:`mjMINAWAKE ` time steps is marked as "ready to sleep". + :ref:`tolerance ` for :ref:`mjMINAWAKE ` time steps is marked as "ready to sleep". If all trees in an island are ready to sleep, they are put to sleep during state advancement. - **Initialized asleep:** By setting the :ref:`body/sleep` attribute of a tree root to "init", it is marked as "initialized-asleep" and put to sleep during :ref:`mjData` initialization. diff --git a/doc/includes/references.h b/doc/includes/references.h index 6a30e704..9bc868ab 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -289,9 +289,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_tendon int* ten_wrapadr; // start address of tendon's path (ntendon x 1) int* ten_wrapnum; // number of wrap points in path (ntendon x 1) - int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1) - int* ten_J_rowadr; // row start address in colind array (ntendon x 1) - int* ten_J_colind; // column indices in sparse Jacobian (nJten x 1) mjtNum* ten_J; // tendon Jacobian (nJten x 1) mjtNum* ten_length; // tendon lengths (ntendon x 1) int* wrap_obj; // geom id; -1: site; -2: pulley (nwrap x 2) @@ -1070,6 +1067,7 @@ struct mjModel_ { mjtSize nexclude; // number of excluded geom pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nwrap; // number of wrap objects in all tendon paths mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric custom fields @@ -1096,7 +1094,6 @@ struct mjModel_ { // sizes set after mjModel construction mjtSize nnames_map; // number of slots in the names hash map mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix - mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize ngravcomp; // number of bodies with nonzero gravcomp mjtSize nemax; // number of potential equality-constraint rows mjtSize njmax; // number of available rows in constraint Jacobian (legacy) @@ -1490,6 +1487,9 @@ struct mjModel_ { int* tendon_group; // group for visibility (ntendon x 1) int* tendon_treenum; // number of trees along tendon's path (ntendon x 1) int* tendon_treeid; // first two trees along tendon's path (ntendon x 2) + int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1) + int* ten_J_rowadr; // row start address in colind array (ntendon x 1) + int* ten_J_colind; // column indices in sparse Jacobian (nJten 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) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 4fe9dfe3..4bf4cffc 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -87,24 +87,25 @@ via Warp - if this feature is important to you, please chime in on this issue Installation ============ -The beta version of MuJoCo Warp is installed from GitHub. Please note that the beta version of MuJoCo Warp does not -support all versions of MuJoCo, Warp, CUDA, NVIDIA drivers, etc. +**From PyPI:** + +.. code-block:: shell + + pip install mujoco-warp + +**From source:** .. code-block:: shell git clone https://github.com/google-deepmind/mujoco_warp.git cd mujoco_warp - python3 -m venv env - source env/bin/activate - pip install --upgrade pip - pip install uv - uv pip install -e .[dev,cuda] + uv sync --all-extras -Test the Installation +To make sure everything is working: .. code-block:: shell - pytest + uv run pytest -n 8 .. _MJW_Usage: @@ -467,9 +468,9 @@ subset of fields. Batch Rendering =============== -MJWarp provides a high-throughput ray-tracing batch renderer built on +MJWarp provides a batch renderer for high-throughput ray tracing built on `Warp's accelerated BVHs `__ for -rendering worlds with multiple cameras in parallel on device. +rendering worlds with multiple cameras in parallel. Key features: @@ -490,8 +491,8 @@ Key features: Basic Usage ---------- -Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures -and rendering and output buffers. +Rendering or raycasting requires a :class:`mjw.RenderContext ` which contains BVH structures, +rendering specific fields, and output buffers. .. code-block:: python @@ -526,9 +527,9 @@ followed by :func:`mjw.render ` to write to output buffers. mjw.render(m, d, rc) The output buffers contain stacked pixels for all cameras with shape `(nworld, npixel)` and RGB data is -packed into one `unit32` variable. `RenderContext.rgb_adr` and `RenderContext.depth_adr` provide per-camera indexing. +packed into one `uint32` variable. `RenderContext.rgb_adr` and `RenderContext.depth_adr` provide per-camera indexing. For convenience, :func:`mjw.get_rgb ` and :func:`mjw.get_depth ` -provide per-camera batched post-processing. +return processed and reshaped RGB and depth data for a given camera batched for all worlds. .. code-block:: python @@ -555,12 +556,12 @@ For benchmark results across a variety of scenes, see the Notes ----- -- **Meshes**: Rendering computation scales with mesh complexity. A primitive is expected to have better - performance (i.e., higher throughput) compared to a similar sized :ref:`mesh` or - :ref:`heightfield `. -- **Flex**: Currently limited to 2D and 3D :ref:`flex` objects. Performance is expected to improved as +- **Meshes**: Rendering computation scales with mesh complexity, specifically the number of vertices and faces. A + primitive is expected to have better performance (i.e., higher throughput) compared to a similar-sized + :ref:`mesh` or :ref:`heightfield `. +- **Flex**: Currently limited to 2D and 3D :ref:`flex` objects. Performance is expected to improve as this feature is further developed. -- **Scaling**: Rendering scales linearly with resolution (total number of pixels) and number of cameras. +- **Scaling**: Rendering scales linearly with resolution (total pixel count) and camera count. .. _mjwFAQ: @@ -723,7 +724,7 @@ Warnings are provided when memory requirements exceed existing allocations durin setting `m.opt.contact_sensor_maxmatch`. Alternatively, refactor the contact sensor matching criteria, for example if the 2 geoms of interest are known, specify ``geom1`` and ``geom2``. - ``height field collision overflow``: The number of potential contacts generated by a height field exceeds - :ref:`mjMAXCONPAIR ` and some contacts are ignored. To resolve this warning, reduce the height field + :ref:`mjMAXCONPAIR ` and some contacts are ignored. To resolve this warning, reduce the height field resolution or reduce the size of the geom interacting with the height field. Compilation @@ -824,3 +825,80 @@ Additional MJWarp-only options are available: A new :ref:`graph capture ` may be necessary after modifying an :class:`mjw.Option ` field in order for the updated setting to take effect. + +SDF plugins +----------- + +SDF collisions support plugins. The following example for +`plugin/sdf/bowl.xml `__ illustrates how to implement the +SDF plugin implementation in `bowl.cc `__: + +.. code-block:: python + + import mujoco_warp as mjw + import warp as wp + + # distance function + @wp.func + def bowl(p: wp.vec3, attr: wp.vec3) -> float: + """Signed distance function for a bowl shape. + + attr[0] = height + attr[1] = radius + attr[2] = thickness + """ + height = attr[0] + radius = attr[1] + thick = attr[2] + width = wp.sqrt(radius * radius - height * height) + + # q = (norm_xy(p), p.z) + q0 = wp.sqrt(p[0] * p[0] + p[1] * p[1]) + q1 = p[2] + + # qdiff = q - (width, height) + qdiff0 = q0 - width + qdiff1 = q1 - height + + if height * q0 < width * q1: + dist = wp.sqrt(qdiff0 * qdiff0 + qdiff1 * qdiff1) + else: + q_norm = wp.sqrt(q0 * q0 + q1 * q1) + dist = wp.abs(q_norm - radius) + + return dist - thick + + + # gradient of distance function + @wp.func + def bowl_sdf_grad(p: wp.vec3, attr: wp.vec3) -> wp.vec3: + """Gradient of bowl SDF via finite differences.""" + eps = float(1e-6) + f0 = bowl(p, attr) + + px = wp.vec3(p[0] + eps, p[1], p[2]) + py = wp.vec3(p[0], p[1] + eps, p[2]) + pz = wp.vec3(p[0], p[1], p[2] + eps) + + grad = wp.vec3( + (bowl(px, attr) - f0) / eps, + (bowl(py, attr) - f0) / eps, + (bowl(pz, attr) - f0) / eps, + ) + return grad + + + # register the bowl SDF plugin + @wp.func + def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float: + return bowl(p, attr) + + + @wp.func + def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3: + return bowl_sdf_grad(p, attr) + + + # override the module-level hooks + mjw._src.collision_sdf.user_sdf = user_sdf + mjw._src.collision_sdf.user_sdf_grad = user_sdf_grad diff --git a/doc/mjx.rst b/doc/mjx.rst index 8825d869..a68d57b1 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -207,6 +207,8 @@ excessive graph captures in the JAX-Warp FFI layer. - 0.65M +.. _MjxWarpBatchRendering: + MJX-Warp Batch Rendering ~~~~~~~~~~~~~~~~~~~~~~~~ @@ -218,9 +220,9 @@ Note that the number of parallel worlds (``nworld``) is fixed when creating the .. code-block:: python - from mujoco.mjx import io + from mujoco.mjx import create_render_context - rc = io.create_render_context( + rc = create_render_context( mjm=m, nworld=nworld, cam_res=(width, height), @@ -247,12 +249,33 @@ volume hierarchy (BVH) and executing the raycaster: pixels, _ = mjx.render(mx, d, rc) # 3. Extract the RGB tensor for the first camera (index 0) - rgb = get_rgb(rc, pixels, 0) + rgb = get_rgb(rc, 0, pixels) # CAVEAT: Always return or use the updated `d` in your computation graph. # Otherwise, JAX's dead-code elimination will optimize away the refit_bvh call! return rgb, d +Multi-GPU rendering with ``pmap`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To render across multiple GPUs, create a render context **per device** by passing ``devices`` to +:func:`create_render_context `. + +.. code-block:: python + + ndevices = jax.local_device_count() + nworld_per_device = nworld // ndevices + + # Create one render context for all devices + rc = create_render_context( + mjm=m, + nworld=nworld_per_device, + devices=[f'cuda:{i}' for i in range(ndevices)], + cam_res=(width, height), + ) + +Then use ``jax.pmap`` to parallelize the rendering across devices. See the complete example in +`visualize_render.py `__. .. _MjxJAX: diff --git a/doc/mjx_api.rst b/doc/mjx_api.rst index e84a8d90..291674e3 100644 --- a/doc/mjx_api.rst +++ b/doc/mjx_api.rst @@ -7,6 +7,7 @@ MJX API :special-members: False :private-members: False :exclude-members: + Any, __init__, __format__, __new__, diff --git a/doc/modeling.rst b/doc/modeling.rst index fee781e9..cf9b056c 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -274,7 +274,8 @@ experiment interactively with parameter settings or implement continuation metho Here we focus on a single scalar constraint. Using slightly different notation from the Computation chapter, let :math:`\ac` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in friction dimensions), :math:`k` and :math:`b` the stiffness and damping of the virtual spring used to define the -reference acceleration :math:`\ar = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`\au` the +reference acceleration :math:`\ar = -b v - k r` (see :eq:`eq:aref`). +Let :math:`d` be the constraint impedance, and :math:`\au` the acceleration in the absence of constraint force. Our earlier analysis revealed that the dynamics in constraint space are approximately @@ -300,7 +301,7 @@ We begin by explaining the constraint impedance :math:`d`. at rest. Impedance is set using the :at:`solimp` attribute. Recall that :math:`d` must lie between 0 and 1; internally MuJoCo clamps it to the range [:ref:`mjMINIMP mjMAXIMP -`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced +`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced acceleration :math:`\au` and reference acceleration :math:`\ar`. The user can set :math:`d` to a constant, or take advantage of its interpolating property and make it position-dependent, i.e., a function of the constraint violation :math:`r`. Position-dependent impedance can be used to model soft contact layers around objects, or define @@ -344,8 +345,7 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec constraint becomes active; for contacts this margin is :ref:`margin`-:ref:`gap`. Limit and contact constraints are active when :math:`r < 0` (penetration). - For friction loss or friction dimensions of elliptic cones, the violation :math:`r` is identically zero, so - only :math:`d(0)` affects these constraints, all other :at:`solimp` values are ignored. + For frictional constraints, see :ref:`Friction`. .. _solimp0: @@ -385,25 +385,27 @@ There are two formats for this attribute, determined by the sign of the numbers. specification is considered to be in the :math:`(\text{timeconst}, \text{dampratio})` format. If negative it is in the "direct" :math:`(-\text{stiffness}, -\text{damping})` format. -Frictional constraints whose residual is identically 0 have first-order dynamics and the mass-spring-damper analysis -below does not apply. In this case the time constant is the rate of exponential decay of the constraint velocity, -and the damping ratio is ignored. Equivalently, in the direct format, the :math:`\text{stiffness}` is ignored. +For frictional constraints, the mass-spring-damper analysis below does not directly apply; +see :ref:`Friction`. **solref :** real(2), "0.02 1" We first describe the default, positive-value format where the two numbers are :math:`(\text{timeconst}, \text{dampratio})`. + .. _soRefScaling: + The idea here is to re-parameterize the model in terms of the time constant and damping ratio of a mass-spring-damper - system. By "time constant" we mean the inverse of the natural frequency times the damping ratio. In this case we use - a mass-spring-damper model to compute :math:`k, b` after suitable scaling. Note that the effective stiffness - :math:`d(r) \cdot k` and damping :math:`d(r) \cdot b` are scaled by the impedance :math:`d(r)` which is a function of - the distance :math:`r`. Thus we cannot always achieve the specified mass-spring-damper properties, unless we - completely undo the scaling by :math:`d`. But the latter is undesirable because it would ruin the interpolating - property, in particular the limit :math:`d=0` would no longer disable the constraint. Instead we scale the stiffness - and damping so that the damping ratio remains constant, while the time constant increases when :math:`d(r)` gets - smaller. The scaling formulas are + system. By "time constant" we mean the inverse of the natural frequency times the damping ratio. Now recall that the + products :math:`d \cdot k` and :math:`d \cdot b` in :eq:`eq:constraint` are the effective stiffness and damping in + constraint space. Because the impedance :math:`d(r)` varies with the + position residual :math:`r`, we cannot achieve constant mass-spring-damper properties; completely undoing the scaling + by :math:`d` is undesirable because the limit :math:`d = 0` would no longer disable the constraint. Instead, we + absorb one factor of :math:`d(r)` into :math:`k` (but not into :math:`b`), so that the damping ratio remains constant + while the time constant scales with :math:`d(r)`. The formulas are .. math:: + :label: eq:solref_standard + \begin{aligned} b &= 2 / (d_\text{width}\cdot \text{timeconst}) \\ k &= d(r) / (d_\text{width}^2 \cdot \text{timeconst}^2 \cdot \text{dampratio}^2) \\ @@ -414,7 +416,7 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: can go unstable. This is enforced internally, unless the :ref:`refsafe` attribute of :ref:`flag ` is set to false. The :math:`\text{dampratio}` parameter would normally be set to 1, corresponding to critical damping. Smaller values result in under-damped or bouncy constraints, while larger values result in - over-damped constraints. Combining the above formula with :eq:`eq:constraint`, we can derive the following result. + over-damped constraints. Combining :eq:`eq:solref_standard` with :eq:`eq:constraint`, we can derive the following If the reference acceleration is given using the positive number format and the impedance is constant :math:`d = d_0 = d_\text{width}`, then the penetration depth at rest is @@ -427,12 +429,14 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: interact. The scaling formulas are .. math:: + :label: eq:solref_direct + \begin{aligned} b &= \text{damping} / d_\text{width} \\ k &= \text{stiffness} \cdot d(r) / d_\text{width}^2 \\ \end{aligned} - Similarly to the above derivation, if the reference acceleration is given using the negative number format and the + Similarly to the derivation following :eq:`eq:solref_standard`, if the reference acceleration is given using the impedance is constant, then the penetration depth at rest is .. math:: @@ -449,6 +453,24 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: A :math:`\text{dampratio}` of 1 in the positive-value format is equivalent to :math:`\text{damping} = 2 \sqrt{ \text{stiffness} }` in the direct format. +.. _CSolverFriction: + +Friction +^^^^^^^^ + +Friction loss constraints (in joints and tendons) and friction dimensions of elliptic contact cones have zero position +violation: :math:`r \equiv 0`. This simplifies the constraint model (see also :ref:`soParameters`): + +- The **impedance** is always :math:`d_0` (:at:`solimp[0]`), since :math:`d(r)` is evaluated at :math:`r=0`. + The sigmoid shape parameters (:math:`\text{width}`, :math:`\text{midpoint}`, :math:`\text{power}`) have no effect. +- The dynamics are **first-order** (exponential decay of constraint velocity, no spring): the stiffness :math:`k` is + always 0. +- In the standard :at:`solref` format, the time constant controls exponential velocity decay. The damping ratio is + ignored (it only appears in the :math:`k` formula). +- In the direct :at:`solref` format, the damping (second value) is used but the stiffness (first value) is ignored. +- :math:`d_\text{width}` (:at:`solimp[1]`) still affects the damping :math:`b` as a scaling denominator + (:eq:`eq:solref_standard`, :eq:`eq:solref_direct`), even though it does not affect the impedance. + .. _CContact: Contact parameters @@ -1690,7 +1712,7 @@ dedicated section :ref:`therein`. 1. :ref:`Timestep`: Try to increase the simulation timestep. As explained at the end of the :ref:`Numerical Integration` section, the timestep is the single most important parameter in any model. The default value is chosen for stability rather than efficiency, and can often be increased. At some point, - increasing it further will cause diveregence, so the optimal timestep is the largest timestep at which divergence + increasing it further will cause divergence, so the optimal timestep is the largest timestep at which divergence never happens or is very rare. The actual value is model-dependent. 2. :ref:`Integrator`: Choose your integrator according to the recommendations at the end of the :ref:`Numerical Integration` section. The default recommended choice is the ``implicitfast`` diff --git a/doc/programming/index.rst b/doc/programming/index.rst index 7d91d013..5d246f78 100644 --- a/doc/programming/index.rst +++ b/doc/programming/index.rst @@ -200,7 +200,7 @@ The situation is more subtle if existing code was developed with a certain versi compiled and linked with a different version. If the definitions of the API functions used in that code have changed, either the compiler or the linker will generate errors. But even if the function definitions have not changed, it may still be a good idea to assert that the software version is the same. To this end, the main header (mujoco.h) defines -the symbol :ref:`mjVERSION_HEADER ` and the library provides the function +the symbol :ref:`mjVERSION_HEADER ` and the library provides the function :ref:`mj_version`. Thus the header and library versions can be compared with: .. code-block:: C diff --git a/doc/programming/simulation.rst b/doc/programming/simulation.rst index da6c2b16..c172b785 100644 --- a/doc/programming/simulation.rst +++ b/doc/programming/simulation.rst @@ -1098,7 +1098,7 @@ implementation details. The high level sleep state of :ref:`trees` is described by ``mjData.tree_asleep`` (though see caveat below). A negative value means a tree is awake, non-negative means asleep. Maximally awake trees are given the value - |-| (1 |-| -+ |-| :ref:`mjMINAWAKE`), and for every timestep where their velocity falls below the sleep :ref:`tolerance ++ |-| :ref:`mjMINAWAKE`), and for every timestep where their velocity falls below the sleep :ref:`tolerance `, this integer is incremented, up to -1, which means "ready to sleep". If all trees in an island are ready to sleep, they are put to sleep during state advancement and their associated values in ``tree_asleep`` are set to a (non-negative) index cycle: the "sleeping island". If any tree in the island is woken, all are woken. @@ -1213,7 +1213,7 @@ Notes **Provisional choices** Some implementation choices are provisional and subject to change. - A concrete example is the decision to hard-code the value of :ref:`mjMINAWAKE` instead of exposing it to + A concrete example is the decision to hard-code the value of :ref:`mjMINAWAKE` instead of exposing it to the user as a runtime option. This was done for two reasons. First, in our experiments, we've found that changing this value is equivalent to changing the :ref:`sleep_tolerance`, which is the more useful knob. Second, one could argue for a time-to-sleep semantic that is in units of time rather than an integer number of diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 314dfc0e..b7a1e635 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -323,9 +323,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_tendon int* ten_wrapadr; // start address of tendon's path (ntendon x 1) int* ten_wrapnum; // number of wrap points in path (ntendon x 1) - int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1) - int* ten_J_rowadr; // row start address in colind array (ntendon x 1) - int* ten_J_colind; // column indices in sparse Jacobian (nJten x 1) mjtNum* ten_J; // tendon Jacobian (nJten x 1) mjtNum* ten_length; // tendon lengths (ntendon x 1) int* wrap_obj; // geom id; -1: site; -2: pulley (nwrap x 2) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 20151f48..97de5c2f 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -728,6 +728,7 @@ struct mjModel_ { mjtSize nexclude; // number of excluded geom pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nwrap; // number of wrap objects in all tendon paths mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric custom fields @@ -754,7 +755,6 @@ struct mjModel_ { // sizes set after mjModel construction mjtSize nnames_map; // number of slots in the names hash map mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix - mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize ngravcomp; // number of bodies with nonzero gravcomp mjtSize nemax; // number of potential equality-constraint rows mjtSize njmax; // number of available rows in constraint Jacobian (legacy) @@ -1148,6 +1148,9 @@ struct mjModel_ { int* tendon_group; // group for visibility (ntendon x 1) int* tendon_treenum; // number of trees along tendon's path (ntendon x 1) int* tendon_treeid; // first two trees along tendon's path (ntendon x 2) + int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1) + int* ten_J_rowadr; // row start address in colind array (ntendon x 1) + int* ten_J_colind; // column indices in sparse Jacobian (nJten 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) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index cba09e7c..3894c898 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -215,6 +215,7 @@ X( nexclude ) \ X( neq ) \ X( ntendon ) \ + X( nJten ) \ X( nwrap ) \ X( nsensor ) \ X( nnumeric ) \ @@ -239,7 +240,6 @@ X( npaths ) \ X( nnames_map ) \ X( nJmom ) \ - X( nJten ) \ X( ngravcomp ) \ X( nemax ) \ X( njmax ) \ @@ -629,6 +629,9 @@ X ( int, tendon_group, ntendon, 1 ) \ X ( int, tendon_treenum, ntendon, 1 ) \ X ( int, tendon_treeid, ntendon, 2 ) \ + X ( int, ten_J_rownnz, ntendon, 1 ) \ + X ( int, ten_J_rowadr, ntendon, 1 ) \ + X ( int, ten_J_colind, nJten, 1 ) \ X ( mjtByte, tendon_limited, ntendon, 1 ) \ X ( mjtByte, tendon_actfrclimited, ntendon, 1 ) \ X ( mjtNum, tendon_width, ntendon, 1 ) \ @@ -848,9 +851,6 @@ X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ 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, nJten, 1 ) \ X ( mjtNum, ten_J, nJten, 1 ) \ X ( mjtNum, ten_length, ntendon, 1 ) \ X ( int, wrap_obj, nwrap, 2 ) \ diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt index deca223c..0f9ae0d5 100644 --- a/mjx/cuda_requirements.txt +++ b/mjx/cuda_requirements.txt @@ -16,8 +16,8 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \ jax-cuda12-pjrt==0.4.30; python_version == '3.9' \ --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \ --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb -warp-lang==1.11.0 \ - --hash=sha256:3a4f1c9a6e721d7de7d6dad6b242c54afaf20c6e14a767c0da03e5e963fcc13c \ - --hash=sha256:524dce20de6162ba25333552168ebf430973050e00d9f8116b8df41a60d25d6e \ - --hash=sha256:1ae6cfc226107f96e4d495b41a3dab32488e8ee8f074b0e1bcaf22e7fb8c904d \ - --hash=sha256:80d8493cbe243a3510134f3af289646d7bd7484217a30ecf565d676466ef8a5e +warp-lang==1.11.1 \ + --hash=sha256:1ad11f1fa775269e991a3d55039152c8a504baf86701c849b485cb8e66c49d15 \ + --hash=sha256:8b098f41e71d421d80ee7562e38aa8380ff6b0d3b4c6ee866cfbdef733ac5bdc \ + --hash=sha256:5d0904b0eefcc81f39ba65375427a3de99006088aa43e24a9011263f07d0cd07 \ + --hash=sha256:15dc10aa51fb0fdbe1ca16d52e5fadca35a47ffd9d0c636826506f96bb2e7c41 diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index e683633b..0ba4d41c 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -494,9 +494,12 @@ def _put_model_warp( return _strip_weak_type(model) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_model_cpp( m: mujoco.MjModel, device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" @@ -517,9 +520,11 @@ def _put_model_cpp( c_pointers_impl = types.ModelCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _model=m, ) + if keepalive_refs is not None: + keepalive_refs[addr] = m + model = types.Model( **{k: copy.copy(v) for k, v in fields.items()}, _impl=c_pointers_impl ) @@ -532,6 +537,7 @@ def put_model( device: Optional[jax.Device] = None, impl: Optional[Union[str, types.Impl]] = None, graph_mode: Optional[mjxw.types.GraphMode] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model. @@ -541,6 +547,9 @@ def put_model( impl: implementation to use graph_mode: CUDA graph capture mode (for Warp only). Use GraphMode enum from warp._src.jax_experimental.ffi. GraphMode.WARP is the default mode. + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl to keep the model alive. Returns: an mjx.Model placed on device @@ -560,7 +569,7 @@ def put_model( graph_mode = graph_mode or getattr(mjxw.types.GraphMode, 'WARP') return _put_model_warp(m, graph_mode, device) elif impl == types.Impl.CPP: - return _put_model_cpp(m, device) + return _put_model_cpp(m, device, keepalive_refs=keepalive_refs) else: raise ValueError(f'Unsupported implementation: {impl}') @@ -903,23 +912,15 @@ def _make_data_warp( data = jax.device_put(data, device=device) - with wp.ScopedDevice('cuda:0'): # pylint: disable=undefined-variable - # Warm-up the warp kernel cache. - # TODO(robotics-simulation): remove this warmup compilation once warp - # stops unloading modules during XLA graph capture for tile kernels. - # pylint: disable=undefined-variable - dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, njmax=njmax) - mw = mjwp.put_model(m) - _ = mjwp.step(mw, dw) - # pylint: enable=undefined-variable - del dw, mw - return data +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _make_data_cpp( m: Union[types.Model, mujoco.MjModel], device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data for the CPP implementation.""" if isinstance(m, mujoco.MjModel): @@ -929,7 +930,13 @@ def _make_data_cpp( m_impl = m._impl # pylint: disable=protected-access if not isinstance(m_impl, types.ModelCPP): raise ValueError(f'Expected ModelCPP impl, got {type(m_impl)}') - mj_model = m_impl._model # pylint: disable=protected-access + model_addr = int(m_impl.pointer_lo) | (int(m_impl.pointer_hi) << 32) + if keepalive_refs is None or model_addr not in keepalive_refs: + raise ValueError( + 'keepalive_refs must be provided and contain the model when calling' + ' _make_data_cpp with a types.Model.' + ) + mj_model = keepalive_refs[model_addr] # Create the raw MuJoCo data mj_data = mujoco.MjData(mj_model) @@ -944,9 +951,11 @@ def _make_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=[mj_data], ) + if keepalive_refs is not None: + keepalive_refs[addr] = [mj_data] + data = types.Data( _impl=c_pointers_impl, **fields, @@ -963,6 +972,7 @@ def make_data( nconmax: Optional[int] = None, naconmax: Optional[int] = None, njmax: Optional[int] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data. @@ -980,6 +990,9 @@ def make_data( `naconmax` argument to set the upper bound for the number of contacts across all worlds, rather than the `nconmax` argument from MuJoCo Warp. njmax: maximum number of constraints to allocate for warp across all worlds + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl when passing a types.Model. Returns: an initialized mjx.Data placed on device @@ -1008,7 +1021,7 @@ def make_data( elif impl == types.Impl.C: return _make_data_c(m, device) elif impl == types.Impl.CPP: - return _make_data_cpp(m, device) + return _make_data_cpp(m, device, keepalive_refs=keepalive_refs) elif impl == types.Impl.WARP: _check_warp_installed() naconmax = nconmax if naconmax is None else naconmax @@ -1111,17 +1124,15 @@ def _put_data_jax( impl_fields['actuator_moment'] = moment # convert ten_J to dense matrix - if mujoco.mj_isSparse(m): + if m.ntendon: ten_J = np.zeros((m.ntendon, m.nv)) mujoco.mju_sparse2dense( ten_J, d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, ) - elif m.ntendon: - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) else: ten_J = np.zeros((m.ntendon, m.nv)) impl_fields['ten_J'] = ten_J @@ -1240,6 +1251,9 @@ def _put_data_c( for f in types.DataC.fields() if hasattr(d, f.name) } + for f in types.DataC.fields(): + if not hasattr(d, f.name) and hasattr(m, f.name): + impl_fields[f.name] = getattr(m, f.name) # TODO(stunya): support islanding via C impl. impl_fields['solver_niter'] = impl_fields['solver_niter'][0] @@ -1313,11 +1327,14 @@ def _put_data_c( return _strip_weak_type(data) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_data_cpp( m: mujoco.MjModel, d: mujoco.MjData, device: Optional[jax.Device] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" @@ -1330,6 +1347,8 @@ def _put_data_cpp( mujoco.mj_copyData(new_d, m, d) data_list.append(new_d) addr = new_d._address + if keepalive_refs is not None: + keepalive_refs[addr] = new_d # To ensure that we retain the full pointer even if jax.config.enable_x64 is # set to True, we store the pointer as two 32-bit values. In the FFI call, # we combine the two values into a single pointer value. @@ -1354,7 +1373,6 @@ def _put_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=data_list, ) data = types.Data( @@ -1365,6 +1383,44 @@ def _put_data_cpp( return _strip_weak_type(data) +def _put_data_warp( + m: mujoco.MjModel, + d: mujoco.MjData, + device: Optional[jax.Device] = None, + naconmax: Optional[int] = None, + njmax: Optional[int] = None, +) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" + + with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable + dw = mjwp.put_data(m, d, nworld=1, naconmax=naconmax, njmax=njmax) # pylint: disable=undefined-variable + + fields = _put_data_public_fields(d) + for k in fields: + if not hasattr(dw, k): + continue + field = _wp_to_np_type(getattr(dw, k)) + if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access + field = field.reshape(field.shape[1:]) + fields[k] = field + + impl_fields = {} + for k in mjxw.types.DataWarp.__annotations__.keys(): + field = _get_nested_attr(dw, k, split='__') + field = _wp_to_np_type(field) + if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access + field = field.reshape(field.shape[1:]) + impl_fields[k] = field + + data = types.Data( + **fields, + _impl=mjxw.types.DataWarp(**impl_fields), + ) + + data = jax.device_put(data, device=device) + return data + + def put_data( m: mujoco.MjModel, d: mujoco.MjData, @@ -1374,6 +1430,7 @@ def put_data( naconmax: Optional[int] = None, njmax: Optional[int] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data. @@ -1390,12 +1447,13 @@ def put_data( njmax: maximum number of constraints to allocate for warp dummy_arg_for_batching: dummy argument to use for batching in cpp implementation + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Returns: an mjx.Data placed on device DeprecationWarning: if nconmax is used """ - del njmax if nconmax is not None: warnings.warn( 'nconmax will be deprecated in mujoco-mjx>=3.5. Use naconmax instead.', @@ -1410,10 +1468,16 @@ def put_data( return _put_data_c(m, d, device) elif impl == types.Impl.CPP: return _put_data_cpp( - m, d, device, dummy_arg_for_batching=dummy_arg_for_batching + m, + d, + device, + dummy_arg_for_batching=dummy_arg_for_batching, + keepalive_refs=keepalive_refs, ) - - # TODO(robotics-team): implement put_data_warp + elif impl == types.Impl.WARP: + _check_warp_installed() + naconmax = nconmax if naconmax is None else naconmax + return _put_data_warp(m, d, device, naconmax, njmax) raise NotImplementedError( f'put_data for implementation "{impl}" not implemented yet.' @@ -1601,9 +1665,6 @@ def _get_data_into( ) else: ten_j = d_i._impl.ten_J - result_i.ten_J_rownnz[:] = ten_j_rownnz - result_i.ten_J_rowadr[:] = ten_j_rowadr - result_i.ten_J_colind[:] = ten_j_colind result_i.ten_J[:] = ten_j continue @@ -1666,15 +1727,18 @@ def _get_data_into( mujoco.mj_factorM(m, result_i) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _get_data_into_cpp( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """Gets mjx.Data from CPP impl into an existing mujoco.MjData or list. For the CPP implementation, the mjx.Data wraps underlying mujoco.MjData - objects that are stored in DataCPP._data. This function simply copies the + objects that are stored in keepalive_refs. This function simply copies the data from those underlying MjData objects to the result using mj_copyData. """ @@ -1686,13 +1750,8 @@ def _get_data_into_cpp( if not isinstance(d_impl, types.DataCPP): raise ValueError(f'Expected DataCPP impl, got {type(d_impl)}') - mj_data_list = d_impl._data # pylint: disable=protected-access - - if batch_size > len(mj_data_list): - raise ValueError( - f'Batch size {batch_size} exceeds number of underlying MjData objects ' - f'({len(mj_data_list)}). Cannot copy data.' - ) + if keepalive_refs is None: + raise ValueError('keepalive_refs must be provided for CPP implementation.') # Verify that the underlying MjData state matches the mjx.Data state # Ideally we'd use mj_getState and get_state here but that requires an @@ -1702,7 +1761,19 @@ def _get_data_into_cpp( d_i: types.Data = ( jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d ) - src_data = mj_data_list[i] + result_i = result[i] if batched else result + + if batched: + addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32) + else: + addr_i = int(d_impl.pointer_lo) | (int(d_impl.pointer_hi) << 32) + + if addr_i not in keepalive_refs: + raise ValueError( + f'Address {addr_i} not found in keepalive_refs. ' + 'Ensure keepalive_refs from the original compile() is passed.' + ) + src_data = keepalive_refs[addr_i] needs_syncing = False for field in fields_to_check: @@ -1719,9 +1790,6 @@ def _get_data_into_cpp( src_data.mocap_quat[:] = d_i.mocap_quat mujoco.mj_kinematics(m, src_data) - for i in range(batch_size): - result_i = result[i] if batched else result - src_data = mj_data_list[i] mujoco.mj_copyData(result_i, m, src_data) @@ -1729,6 +1797,7 @@ def get_data_into( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" is_batched = isinstance(result, list) @@ -1744,7 +1813,7 @@ def get_data_into( return _get_data_into(result, m, d) if d.impl == types.Impl.CPP: - return _get_data_into_cpp(result, m, d) + return _get_data_into_cpp(result, m, d, keepalive_refs=keepalive_refs) if d.impl == types.Impl.WARP: return _get_data_into_warp(result, m, d) @@ -1755,7 +1824,9 @@ def get_data_into( def get_data( - m: mujoco.MjModel, d: types.Data + m: mujoco.MjModel, + d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> 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 @@ -1766,7 +1837,7 @@ def get_data( else: result = mujoco.MjData(m) - get_data_into(result, m, d) + get_data_into(result, m, d, keepalive_refs=keepalive_refs) return result @@ -1945,6 +2016,7 @@ def set_state( def create_render_context( mjm: mujoco.MjModel, nworld: int, + devices: Optional[Sequence[str]] = None, **kwargs, ): """Creates a render context. @@ -1955,6 +2027,9 @@ def create_render_context( because Warp creates arrays of size nworld that are not exposed to JAX. Thus we cannot use JAX transforms like vmap with the render context. + devices: optional list of device names (e.g. ['cuda:0', 'cuda:1']). + If provided, rendering workloads are sharded across these devices. + By default, devices is None and the default device from wp.get_device(None) is used. **kwargs: forwarded to the render context constructor. Returns: @@ -1962,4 +2037,6 @@ def create_render_context( """ _check_warp_installed() from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error - return mjxw_io.create_render_context(mjm, nworld=nworld, **kwargs) + return mjxw_io.create_render_context( + mjm, nworld=nworld, devices=devices, **kwargs + ) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 8382a6cf..9910e9b5 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -494,9 +494,15 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d._impl.contact__dist.shape[0], 9) self.assertEqual(d._impl.efc__pos.shape[0], 23) - @parameterized.parameters('jax', 'c', 'cpp') + @parameterized.parameters('jax', 'c', 'cpp', 'warp') def test_put_data(self, impl: str): """Test that put_data puts the correct data for dense and sparse.""" + if impl == 'warp': + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp is not installed.') + if not mjx_io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device.') + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) @@ -517,6 +523,15 @@ class DataIOTest(parameterized.TestCase): ) ) + # xmat, ximat, geom_xmat are all shape transformed + np.testing.assert_allclose(dx.xmat.reshape((-1, 9)), d.xmat) + np.testing.assert_allclose(dx.ximat.reshape((-1, 9)), d.ximat) + np.testing.assert_allclose(dx.geom_xmat.reshape((-1, 9)), d.geom_xmat) + np.testing.assert_allclose(dx.site_xmat.reshape((-1, 9)), d.site_xmat) + + # tendon length is correct + np.testing.assert_allclose(dx.ten_length, d.ten_length) + if impl == 'jax': # check that qM is transformed properly qm = np.zeros((m.nv, m.nv), dtype=np.float64) @@ -530,6 +545,21 @@ class DataIOTest(parameterized.TestCase): self.assertTrue(hasattr(dx._impl, 'pointer_lo')) self.assertTrue(hasattr(dx._impl, 'pointer_hi')) return # cpp does not populate other fields in _impl + elif impl == 'warp': + qm = np.zeros((m.nv, m.nv), dtype=np.float64) + mujoco.mj_fullM(m, qm, d.qM) + np.testing.assert_allclose(dx._impl.qM, qm) + # TODO(taylorhowell): test efc__J + np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3]) + + # tendon impl data is correct + 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))) + + if impl == 'warp': + return # 4 contacts, 2 for each capsule against the plane self.assertEqual(dx._impl.contact.dist.shape, (4,)) @@ -542,23 +572,6 @@ class DataIOTest(parameterized.TestCase): ) 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)) - self.assertEqual(dx.ximat.shape, (3, 3, 3)) - self.assertEqual(dx.geom_xmat.shape, (3, 3, 3)) - self.assertEqual(dx.site_xmat.shape, (1, 3, 3)) - np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat) - np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat) - np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat) - np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat) - - # tendon data is correct - np.testing.assert_allclose(dx.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._impl.efc_J.shape, (45, 8)) # nefc, nv d_efc_j = d.efc_J.reshape((-1, 8)) @@ -583,7 +596,9 @@ class DataIOTest(parameterized.TestCase): d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) dx_sparse = mjx.put_data(m, d, impl=impl) - np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8) + 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._impl.qM, d.qM, atol=1e-8) @@ -604,6 +619,31 @@ class DataIOTest(parameterized.TestCase): elif impl == 'c': np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8) + def test_put_data_warp_ndim(self): + """Tests that put_data produces expected dimensions for Warp fields.""" + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp is not installed.') + if not mjx_io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device.') + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d, impl='warp') + + def check_ndim(path, x): + k = _get_name_from_path(path) + if k not in mjxw_types._NDIM['Data']: + return + is_batched = mjxw_types._BATCH_DIM['Data'][k] + expected_ndim = mjxw_types._NDIM['Data'][k] - is_batched + if not hasattr(x, 'ndim'): + return + msg = f'Field {k} has ndim {x.ndim} but expected {expected_ndim}' + self.assertEqual(x.ndim, expected_ndim, msg) + + _ = jax.tree.map_with_path(check_ndim, dx) + @parameterized.parameters( ('jax', False), ('jax', True), ('c', False), ('c', True) ) @@ -728,9 +768,10 @@ class DataIOTest(parameterized.TestCase): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d, impl=impl) + keepalive = {} if impl == 'cpp' else None + dx = mjx.put_data(m, d, impl=impl, keepalive_refs=keepalive) d_2 = mujoco.MjData(m) - mjx.get_data_into(d_2, m, dx) + mjx.get_data_into(d_2, m, dx, keepalive_refs=keepalive) # check a few fields np.testing.assert_allclose(d_2.qpos, d.qpos) @@ -891,12 +932,11 @@ class DataIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) # Test cases for `_resolve_impl_and_device` where the device is @@ -1225,12 +1265,11 @@ class StateIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) def test_get_set_state(self): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index 58e98ef4..cfe6bf0b 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -44,7 +44,7 @@ def get_rgb( else: raise RuntimeError('Warp not installed.') - warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key] + warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)] rgb_adr = int(warp_rc.rgb_adr.numpy()[cam_id]) width = int(warp_rc.cam_res.numpy()[cam_id][0]) height = int(warp_rc.cam_res.numpy()[cam_id][1]) @@ -84,7 +84,7 @@ def get_depth( import mujoco.mjx.warp.render as mjxw_render # pylint: disable=g-import-not-at-top # pytype: disable=import-error else: raise RuntimeError('Warp not installed.') - warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key] + warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)] depth_adr = int(warp_rc.depth_adr.numpy()[cam_id]) width = int(warp_rc.cam_res.numpy()[cam_id][0]) height = int(warp_rc.cam_res.numpy()[cam_id][1]) diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py index e62658ac..ed82f5c0 100644 --- a/mjx/mujoco/mjx/_src/render_util_test.py +++ b/mjx/mujoco/mjx/_src/render_util_test.py @@ -56,7 +56,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))(rc, 0, rgb_data) @@ -70,7 +70,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): rgb = jax.jit( jax.vmap(render_util.get_rgb, in_axes=(None, None, 0)), @@ -87,7 +87,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))( rc, 0, depth_data, 5.0 @@ -103,7 +103,7 @@ class RenderUtilTest(absltest.TestCase): with mock.patch.dict( 'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS', - {0: warp_rc}, + {(0, None): warp_rc}, ): depth = jax.jit( jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None)), diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 667b3d63..7f0d8c09 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -121,8 +121,14 @@ class SmoothTest(absltest.TestCase): mujoco.mj_forward(m, d) # tendon dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d)) - # with dense jacobian mode, d.ten_J is already dense (ntendon*nv,), just reshape - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) + ten_J = np.zeros((m.ntendon, m.nv)) + mujoco.mju_sparse2dense( + ten_J, + d.ten_J, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, + ) _assert_eq(ten_J, dx._impl.ten_J, 'ten_J') _assert_attr_eq(d, dx, 'ten_length') # transmission @@ -397,19 +403,14 @@ class TendonTest(parameterized.TestCase): dx = jax.jit(mjx.forward)(mx, dx) _assert_eq(d.ten_length, dx.ten_length, 'ten_length') - # convert ten_J for comparison based on jacobian mode - if mujoco.mj_isSparse(m): - ten_J = np.zeros((m.ntendon, m.nv)) - mujoco.mju_sparse2dense( - ten_J, - d.ten_J, - d.ten_J_rownnz, - d.ten_J_rowadr, - d.ten_J_colind, - ) - else: - # dense mode: just reshape - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) + ten_J = np.zeros((m.ntendon, m.nv)) + mujoco.mju_sparse2dense( + ten_J, + d.ten_J, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, + ) _assert_eq(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') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index fcdc41d6..9f25a845 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -538,7 +538,6 @@ class ModelCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _model: mujoco.MjModel class DataCPP(PyTreeNode): @@ -548,7 +547,6 @@ class DataCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _data: list[Any] = dataclasses.field(default_factory=list, repr=False) class ModelC(PyTreeNode): diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index 9ee93d3d..e6533dd6 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -952,13 +952,12 @@ def put_data( d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float) - if mujoco.mj_isSparse(mjm): + if mjm.ntendon: ten_J = np.zeros((mjm.ntendon, mjm.nv)) - mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1)) + mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjm.ten_J_rownnz, mjm.ten_J_rowadr, mjm.ten_J_colind.reshape(-1)) d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) else: - ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv)) - d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float) + d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), 0.0), dtype=float) # TODO(taylorhowell): sparse actuator_moment actuator_moment = np.zeros((mjm.nu, mjm.nv)) @@ -1169,9 +1168,9 @@ def get_data_into( mujoco.mju_dense2sparse( result.ten_J, ten_J, - result.ten_J_rownnz, - result.ten_J_rowadr, - result.ten_J_colind, + mjm.ten_J_rownnz, + mjm.ten_J_rowadr, + mjm.ten_J_colind, ) else: result.ten_J[:] = d.ten_J.numpy()[world_id] diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py index 43f63397..f5c925dd 100644 --- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py +++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py @@ -62,11 +62,22 @@ def check_jax_version(): class GraphMode(IntEnum): - NONE = 0 # don't capture a graph - JAX = 1 # let JAX capture a graph - WARP = 2 # let Warp capture a graph - WARP_STAGED = 3 # use Warp graph with staging buffers, copy inside of the graph - WARP_STAGED_EX = 4 # use Warp graph with staging buffers, copy outside of the graph + """CUDA graph capture modes for :func:`warp.jax_experimental.jax_callable`. + + These modes control whether JAX or Warp captures a CUDA graph, and whether + staging buffers are used when capturing with Warp. + """ + + NONE = 0 + """Disable graph capture. Use when operations are not CUDA-graph compatible (for example, host synchronization).""" + JAX = 1 + """Let JAX capture the graph so the callable can be used as a subgraph within a larger JAX capture.""" + WARP = 2 + """Let Warp capture the graph and replay it for matching buffer addresses.""" + WARP_STAGED = 3 + """Capture a Warp graph using staging buffers and insert memcpy nodes inside the graph.""" + WARP_STAGED_EX = 4 + """Capture a Warp graph using staging buffers and perform memcpy outside the graph.""" class ModulePreloadMode(IntEnum): @@ -434,6 +445,7 @@ class FfiCallable: stage_out_argnames, graph_cache_max, module_preload_mode, + has_side_effect=False, ): self.func = func self.name = generate_unique_name(func) @@ -442,6 +454,7 @@ class FfiCallable: self.graph_mode = graph_mode self.output_dims = output_dims self.module_preload_mode = module_preload_mode + self.has_side_effect = has_side_effect self.first_array_arg = None self.call_id = 0 self.call_descriptors = {} @@ -613,7 +626,7 @@ class FfiCallable: out_types, vmap_method=vmap_method, input_output_aliases=self.input_output_aliases, - # has_side_effect=True, # force this function to execute even if outputs aren't used + has_side_effect=self.has_side_effect, ) # preload on the specified devices @@ -680,12 +693,13 @@ class FfiCallable: assert num_outputs == self.num_outputs cuda_stream = get_stream_from_callframe(call_frame.contents) + device_ordinal = get_device_ordinal_from_callframe(call_frame.contents) if self.graph_mode == GraphMode.WARP: # check if we already captured an identical call ip = [inputs[i].contents.data for i in self.array_input_indices] op = [outputs[i].contents.data for i in self.array_output_indices] - capture_key = hash((call_id, *ip, *op)) + capture_key = hash((device_ordinal, call_id, *ip, *op)) capture = self.captures.get(capture_key) # launch existing graph @@ -1379,6 +1393,7 @@ def jax_callable( stage_out_argnames=None, graph_cache_max: int | None = None, module_preload_mode: ModulePreloadMode = ModulePreloadMode.CURRENT_DEVICE, + has_side_effect: bool = False, ): """Create a JAX callback from an annotated Python function. @@ -1449,6 +1464,7 @@ def jax_callable( stage_out_argnames, graph_cache_max, module_preload_mode, + has_side_effect, ) _FFI_CALLABLE_REGISTRY[key] = callable else: diff --git a/mjx/mujoco/mjx/viewer.py b/mjx/mujoco/mjx/viewer.py index 075e243e..ddd88696 100644 --- a/mjx/mujoco/mjx/viewer.py +++ b/mjx/mujoco/mjx/viewer.py @@ -88,15 +88,9 @@ def _main(argv: Sequence[str]) -> None: m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value) d = mujoco.MjData(m) mx = mjx.put_model(m, impl=_IMPL.value) - if _IMPL.value == 'warp': - # TODO(btaba): use put_data. - dx = mjx.make_data( - m, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value - ) - else: - dx = mjx.put_data( - m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value - ) + dx = mjx.put_data( + m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value + ) print(f'Default backend: {jax.default_backend()}') step_fn = mjx.step diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 573c24e4..529145ee 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -88,7 +88,7 @@ def _refit_bvh_shim( _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos _d.nworld = nworld - render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + render_context = _MJX_RENDER_CONTEXT_BUFFERS[(rc_id, wp.get_device().ordinal)] dummy.zero_() mjwarp.refit_bvh(_m, _d, render_context) @@ -104,6 +104,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=True, ) out = jf( d.qpos.shape[0], @@ -122,7 +123,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): d.geom_xpos, ctx.key, ) - d = d.tree_replace({'time': d.time + out[0]}) + d = d.tree_replace({}) return d diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 35a85720..2ae01123 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _collision_shim( # Model @@ -288,6 +289,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data): ]), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index 6d0b6d7e..a485cbbb 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/mjx/mujoco/mjx/warp/ffi.py @@ -103,6 +103,7 @@ def jax_callable_variadic_tuple( in_out_argnames: Optional[Sequence[str]] = None, stage_in_argnames: Optional[Sequence[str]] = None, stage_out_argnames: Optional[Sequence[str]] = None, + has_side_effect: bool = False, ): """Wraps a JAX callable to support variadic tuples and dataclasses.""" @@ -134,6 +135,7 @@ def jax_callable_variadic_tuple( in_out_argnames=in_out_argnames, stage_in_argnames=stage_in_argnames, stage_out_argnames=stage_out_argnames, + has_side_effect=has_side_effect, ) flat_args, in_tree = jax.tree.flatten(args) diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 349e5ed2..e48219a8 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _forward_shim( # Model @@ -1278,6 +1279,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], @@ -3069,6 +3071,7 @@ def _step_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 020e924a..3eb936a7 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -149,7 +149,14 @@ class ForwardTest(parameterized.TestCase): tu.assert_attr_eq(dx, d, 'cam_xpos') tu.assert_eq(dx.cam_xmat, d.cam_xmat.reshape((-1, 3, 3)), 'cam_xmat') tu.assert_attr_eq(dx, d, 'ten_length') - ten_J = d.ten_J.reshape((m.ntendon, m.nv)) + ten_J = np.zeros((m.ntendon, m.nv)) + mujoco.mju_sparse2dense( + ten_J, + d.ten_J, + m.ten_J_rownnz, + m.ten_J_rowadr, + m.ten_J_colind, + ) tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J') tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr') tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum') diff --git a/mjx/mujoco/mjx/warp/io.py b/mjx/mujoco/mjx/warp/io.py index eb9af42c..db0f2df9 100644 --- a/mjx/mujoco/mjx/warp/io.py +++ b/mjx/mujoco/mjx/warp/io.py @@ -19,26 +19,43 @@ import threading import mujoco from mujoco.mjx.warp.types import RenderContext import mujoco.mjx.third_party.mujoco_warp as mjw +import warp as wp _MJX_RENDER_CONTEXT_COUNTER = 0 _MJX_RENDER_CONTEXT_LOCK = threading.Lock() _MJX_RENDER_CONTEXT_BUFFERS = {} +def _create_context(mjm, nworld, device, **kwargs): + with wp.ScopedDevice(device): + ctx = mjw.create_render_context(mjm=mjm, nworld=nworld, **kwargs) + ctx.rgb_data_shape = ctx.rgb_data.shape + ctx.depth_data_shape = ctx.depth_data.shape + ctx.rgb_data = None + ctx.depth_data = None + return ctx + + def create_render_context( mjm: mujoco.MjModel, nworld: int, + devices: list[str | None] | None = None, **kwargs, ): - rc = mjw.create_render_context(mjm=mjm, nworld=nworld, **kwargs) - rc.rgb_data_shape = rc.rgb_data.shape - rc.depth_data_shape = rc.depth_data.shape - rc.rgb_data = None - rc.depth_data = None - global _MJX_RENDER_CONTEXT_COUNTER + + if not devices: + devices = [None] + + contexts = [_create_context(mjm, nworld, d, **kwargs) for d in devices] + with _MJX_RENDER_CONTEXT_LOCK: _MJX_RENDER_CONTEXT_COUNTER += 1 key = _MJX_RENDER_CONTEXT_COUNTER - _MJX_RENDER_CONTEXT_BUFFERS[key] = rc + for d, ctx in zip(devices, contexts): + ordinal = wp.get_device(d).ordinal + _MJX_RENDER_CONTEXT_BUFFERS[(key, ordinal)] = ctx + if (key, None) not in _MJX_RENDER_CONTEXT_BUFFERS: + # save the first context as the default context + _MJX_RENDER_CONTEXT_BUFFERS[(key, None)] = contexts[0] return RenderContext(key, _owner=True) diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 091f72d0..3b2d3aee 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -110,14 +110,14 @@ def _render_shim( _d.light_xdir = light_xdir _d.light_xpos = light_xpos _d.nworld = nworld - render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + render_context = _MJX_RENDER_CONTEXT_BUFFERS[(rc_id, wp.get_device().ordinal)] render_context.rgb_data = rgb render_context.depth_data = depth mjwarp.render(_m, _d, render_context) def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): - render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[ctx.key] + render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[(ctx.key, None)] output_dims = { 'rgb': render_ctx.rgb_data_shape, 'depth': render_ctx.depth_data_shape, @@ -145,6 +145,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext): ]), stage_out_argnames=set([]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 217cf873..0e1ec813 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -44,6 +44,7 @@ _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) + @ffi.format_args_for_warp def _kinematics_shim( # Model @@ -213,6 +214,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): 'xquat', ]), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], @@ -407,6 +409,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): ]), stage_out_argnames=set(['ten_length']), graph_mode=m.opt._impl.graph_mode, + has_side_effect=False, ) out = jf( d.qpos.shape[0], diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index d39d4861..03d403a4 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -126,7 +126,11 @@ class RenderContext: if lock is None or buffers is None: return with lock: - buffers.pop(self.key, None) + keys_to_remove = [ + k for k in buffers.keys() if isinstance(k, tuple) and k[0] == self.key + ] + for k in keys_to_remove: + buffers.pop(k, None) class StatisticWarp(PyTreeNode): diff --git a/mjx/mujoco/mjx/warp/visualize_render.py b/mjx/mujoco/mjx/warp/visualize_render.py index 42ed02f1..6a3b5afa 100644 --- a/mjx/mujoco/mjx/warp/visualize_render.py +++ b/mjx/mujoco/mjx/warp/visualize_render.py @@ -57,6 +57,9 @@ _WP_KERNEL_CACHE_DIR = flags.DEFINE_string( '/tmp/wp_kernel_cache_dir_visualize_render', 'warp kernel cache directory', ) +_PMAP = flags.DEFINE_boolean( + 'pmap', False, 'also render with pmap across GPUs and compare' +) _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1} jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS) @@ -107,6 +110,7 @@ def _main(_: Sequence[str]): print(f' camera_id : {_CAMERA_ID.value}') print(f' use_textures: {_USE_TEXTURES.value}') print(f' use_shadows : {_USE_SHADOWS.value}') + print(f' pmap : {_PMAP.value}') print(f' output_dir : {_OUTPUT_DIR.value}\n') mx = mjx.put_model(m, impl='warp') @@ -143,7 +147,6 @@ def _main(_: Sequence[str]): enabled_geom_groups=[0, 1, 2], ) - print('rendering...') dx_batch = jax_jit(jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None)))( mx, dx_batch, rc ) @@ -187,6 +190,62 @@ def _main(_: Sequence[str]): ) _save_tiled(depth_rgb, depth_tiled_path) + if _PMAP.value: + ndevices = jax.local_device_count() + nworld = _NWORLD.value + nworld_per_device = nworld // ndevices + assert nworld >= ndevices and nworld % ndevices == 0, ( + f'--pmap requires nworld ({nworld}) divisible by device count' + f' ({ndevices})' + ) + print(f'\nrendering (pmap across {ndevices} devices)...') + + device_strs = [f'cuda:{i}' for i in range(ndevices)] + + pmap_rc = io.create_render_context( + mjm=m, + nworld=nworld_per_device, + devices=device_strs, + cam_res=(_WIDTH.value, _HEIGHT.value), + use_textures=_USE_TEXTURES.value, + use_shadows=_USE_SHADOWS.value, + render_rgb=True, + render_depth=True, + enabled_geom_groups=[0, 1, 2], + ) + + devices = jax.local_devices()[:ndevices] + mesh = jax.sharding.Mesh(np.array(devices), axis_names=('i',)) + P = jax.sharding.PartitionSpec + sharded = jax.sharding.NamedSharding(mesh, P('i')) + + def safe_shard(x, sharding): + # Go through CPU to avoid P2P DMA issues on certain machines. + x_cpu = jax.device_put(x, jax.devices('cpu')[0]) + if x_cpu.ndim > 0 and x_cpu.shape[0] == nworld: + reshaped = x_cpu.reshape(ndevices, nworld_per_device, *x_cpu.shape[1:]) + else: + reshaped = jp.stack([x_cpu] * ndevices) + return jax.device_put(reshaped, sharding) + + dx_pmap = jax.tree.map(lambda x: safe_shard(x, sharded), dx_batch) + mx_pmap = jax.tree.map(lambda x: safe_shard(x, sharded), mx) + + def inner(mx, dx): + dx = bvh.refit_bvh(mx, dx, pmap_rc) + out = render.render(mx, dx, pmap_rc) + return render_util.get_rgb(pmap_rc, _CAMERA_ID.value, out[0]) + + inner = jax.vmap(inner, in_axes=(None, 0)) + out = jax.pmap(inner)(mx_pmap, dx_pmap) + + pmap_rgb = jax.device_put(out, jax.devices('cpu')[0]).reshape(-1, *out.shape[2:]) + + pmap_tiled_path = os.path.join( + _OUTPUT_DIR.value, f'pmap_tiled_{_CAMERA_ID.value}.png' + ) + _save_tiled(pmap_rgb, pmap_tiled_path) + print('\ndone.') diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index 0348ce44..4ee26dd5 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -17,13 +17,13 @@ classifiers = [ "Intended Audience :: Science/Research", "Natural Language :: English", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ "absl-py", "etils[epath]", @@ -36,7 +36,7 @@ dependencies = [ [project.optional-dependencies] warp = [ - "warp-lang==1.11.0", + "warp-lang==1.11.1", ] [project.scripts] diff --git a/mjx/requirements.txt b/mjx/requirements.txt index f0ef43d6..4e5b56f0 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -2,15 +2,11 @@ absl-py==2.1.0 \ --hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 etils[epath]==1.10.0; python_version >= '3.10' \ --hash=sha256:0777fe60a234b4c65ca53470fc64f2dd2d0c6bca7fcc623fdaa8d7fa5a317098 -etils[epath]==1.5.2; python_version == '3.9' \ - --hash=sha256:6dc882d355e1e98a5d1a148d6323679dc47c9a5792939b9de72615aa4737eb0b jax==0.5.3; python_version >= '3.10' and (sys_platform != 'darwin' or platform_machine != 'x86_64') \ --hash=sha256:1483dc237b4f47e41755d69429e8c3c138736716147cd43bb2b99b259d4e3c41 \ --hash=sha256:f17fcb0fd61dc289394af6ce4de2dada2312f2689bb0d73642c6f026a95fbb2c jax==0.4.38; python_version >= '3.10' and sys_platform == 'darwin' and platform_machine == 'x86_64' \ --hash=sha256:78987306f7041ea8500d99df1a17c33ed92620c2268c4c3677fb24e06712be64 -jax==0.4.30; python_version == '3.9' \ - --hash=sha256:289b30ae03b52f7f4baf6ef082a9f4e3e29c1080e22d13512c5ecf02d5f1a55b jaxlib==0.5.3; python_version >= '3.10' and (sys_platform != 'darwin' or platform_machine != 'x86_64') \ --hash=sha256:48ff5c89fb8a0fe04d475e9ddc074b4879a91d7ab68a51cec5cd1e87f81e6c47 \ --hash=sha256:972400db4af6e85270d81db5e6e620d31395f0472e510c50dfcd4cb3f72b7220 \ @@ -49,12 +45,6 @@ jaxlib==0.4.38; python_version >= '3.10' and sys_platform == 'darwin' and platfo --hash=sha256:248cca3771ebf24b070f49701364ceada33e6139445b06c782cca5ac5ad92bf4 \ --hash=sha256:2ce77ba8cda9259a4bca97afc1c722e4291a6c463a63f8d372c6edc85117d625 \ --hash=sha256:4103db0b3a38a5dc132741237453c24d8547290a22079ba1b577d6c88c95300a -jaxlib==0.4.30; python_version == '3.9' \ - --hash=sha256:54987e97a22db70f3829b437b9329e4799d653634bacc8b398554d3b90c76b2a \ - --hash=sha256:f74a6b0e09df4b5e2ee399ebb9f0e01190e26e84ccb0a758fadb516415c07f18 \ - --hash=sha256:11602d5556e8baa2f16314c36518e9be4dfae0c2c256a361403fb29dc9dc79a4 \ - --hash=sha256:3d31e01191ce8052bd611aaf16ff967d8d0ec0b63f1ea4b199020cecb248d667 \ - --hash=sha256:ea3a00005faafbe3c18b178d3b534208b3b4027b2be6230227e7b87ce399fc29 pip==25.2 \ --hash=sha256:578283f006390f85bb6282dffb876454593d637f5d1be494b5202ce4877e71f2 \ --hash=sha256:6d67a2b4e7f14d8b31b8b52648866fa717f45a1eb70e83002f4331d07e953717 @@ -91,12 +81,6 @@ scipy==1.14.1; python_version >= '3.10' \ --hash=sha256:8bddf15838ba768bb5f5083c1ea012d64c9a444e16192762bd858f1e126196d0 \ --hash=sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3 \ --hash=sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389 -scipy==1.13.1; python_version == '3.9' \ - --hash=sha256:392e4ec766654852c25ebad4f64e4e584cf19820b980bc04960bca0b0cd6eaa2 \ - --hash=sha256:637e98dcf185ba7f8e663e122ebf908c4702420477ae52a04f9908707456ba4d \ - --hash=sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004 \ - --hash=sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24 \ - --hash=sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5 setuptools==78.1.1 \ --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d @@ -116,8 +100,6 @@ zipp==3.21.0 \ --hash=sha256:ac1bbe05fd2991f160ebce24ffbac5f6d11d83dc90891255885223d42b3cd931 # Transitive dependencies of jax and jaxlib -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b ml-dtypes==0.5.0 \ --hash=sha256:cb5cc7b25acabd384f75bbd78892d0c724943f3e2e1986254665a1aa10982e07 \ --hash=sha256:54415257f00eb44fbcc807454efac3356f75644f1cbfc2d4e5522a72ae1dacab \ @@ -167,15 +149,7 @@ numpy==2.1.3; python_version >= '3.10' \ --hash=sha256:6a4825252fcc430a182ac4dee5a505053d262c807f8a924603d411f6718b88fd \ --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ - --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff \ -numpy==2.0.2; python_version == '3.9' \ - --hash=sha256:a3f4ab0caa7f053f6797fcd4e1e25caee367db3112ef2b6ef82d749530768c73 \ - --hash=sha256:f26b258c385842546006213344c50655ff1555a9338e2e5e02a0756dc3e803dd \ - --hash=sha256:1e795a8be3ddbac43274f18588329c72939870a16cae810c2b73461c40718ab1 \ - --hash=sha256:97032a27bd9d8988b9a97a8c4d2c9f2c15a81f61e2f21404d7e8ef00cb5be729 \ - --hash=sha256:2b2955fa6f11907cf7a70dab0d0755159bca87755e831e47932367fc8f2f2d0b \ - --hash=sha256:423e89b23490805d2a5a96fe40ec507407b8ee786d66f7328be214f9679df6dd \ - --hash=sha256:9059e10581ce4093f735ed23f3b9d283b9d517ff46009ddd485f1747eb22653c + --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff opt-einsum==3.4.0 \ --hash=sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd diff --git a/python/build_requirements.txt b/python/build_requirements.txt index 7e2ec063..49b3750f 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -6,8 +6,6 @@ build==1.2.2.post1 \ --hash=sha256:1d61c0887fa860c01971625baae8bdd338e517b836a2f70dd1f7aa3a6b2fc5b5 etils[epath]==1.10.0; python_version >= '3.10' \ --hash=sha256:0777fe60a234b4c65ca53470fc64f2dd2d0c6bca7fcc623fdaa8d7fa5a317098 -etils[epath]==1.5.2; python_version == '3.9' \ - --hash=sha256:6dc882d355e1e98a5d1a148d6323679dc47c9a5792939b9de72615aa4737eb0b glfw==2.9.0 \ --hash=sha256:9aa3ae51601601c53838315bd2a03efb1e6bebecd072b2f64ddbd0b2556d511a \ --hash=sha256:8e4fbff88e4e953bb969b6813195d5de4641f886530cc8083897e56b00bf2c8e \ @@ -46,14 +44,6 @@ numpy==2.1.3; python_version >= '3.10' \ --hash=sha256:825656d0743699c529c5943554d223c021ff0494ff1442152ce887ef4f7561a1 \ --hash=sha256:b47fbb433d3260adcd51eb54f92a2ffbc90a4595f8970ee00e064c644ac788f5 \ --hash=sha256:c894b4305373b9c5576d7a12b473702afdf48ce5369c074ba304cc5ad8730dff -numpy==2.0.2; python_version == '3.9' \ - --hash=sha256:a3f4ab0caa7f053f6797fcd4e1e25caee367db3112ef2b6ef82d749530768c73 \ - --hash=sha256:f26b258c385842546006213344c50655ff1555a9338e2e5e02a0756dc3e803dd \ - --hash=sha256:1e795a8be3ddbac43274f18588329c72939870a16cae810c2b73461c40718ab1 \ - --hash=sha256:97032a27bd9d8988b9a97a8c4d2c9f2c15a81f61e2f21404d7e8ef00cb5be729 \ - --hash=sha256:2b2955fa6f11907cf7a70dab0d0755159bca87755e831e47932367fc8f2f2d0b \ - --hash=sha256:423e89b23490805d2a5a96fe40ec507407b8ee786d66f7328be214f9679df6dd \ - --hash=sha256:9059e10581ce4093f735ed23f3b9d283b9d517ff46009ddd485f1747eb22653c pip==24.3.1 \ --hash=sha256:3790624780082365f47549d032f3770eeb2b1e8bd1f7b2e02dace1afa361b4ed PyOpenGL==3.1.7 \ @@ -73,8 +63,6 @@ pyelftools==0.31; platform_system == 'Linux' \ # Transitive dependencies of build colorama==0.4.6; platform_system == 'Windows' \ --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6 -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b packaging==24.2 \ --hash=sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759 pyproject_hooks==1.2.0 \ diff --git a/python/make_sdist_requirements.txt b/python/make_sdist_requirements.txt index f961524c..da33bbb9 100644 --- a/python/make_sdist_requirements.txt +++ b/python/make_sdist_requirements.txt @@ -11,8 +11,6 @@ setuptools==78.1.1 \ # Transitive dependencies of build colorama==0.4.6; platform_system == 'Windows' \ --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6 -importlib-metadata==8.5.0; python_version == '3.9' \ - --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b packaging==24.2 \ --hash=sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759 pyproject_hooks==1.2.0 \ diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h index 9f26000d..524f9376 100644 --- a/python/mujoco/indexer_xmacro.h +++ b/python/mujoco/indexer_xmacro.h @@ -259,7 +259,10 @@ X( mjtNum, tendon, _length0, ntendon, 1 ) \ X( mjtNum, tendon, _invweight0, ntendon, 1 ) \ X( mjtNum, tendon, _user, ntendon, MJ_M(nuser_tendon) ) \ - X( float, tendon, _rgba, ntendon, 4 ) + X( float, tendon, _rgba, ntendon, 4 ) \ + X( int, ten_, J_rownnz, ntendon, 1 ) \ + X( int, ten_, J_rowadr, ntendon, 1 ) \ + X( int, ten_, J_colind, ntendon, MJ_M(nv) ) #define MJMODEL_TEXTURE \ X( int, tex_, type, ntex, 1 ) \ @@ -402,11 +405,7 @@ #define MJDATA_TENDON \ 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) ) \ X( mjtNum, ten_, length , ntendon, 1 ) \ - X( mjtNum, ten_, J , ntendon, MJ_M(nv) ) \ X( mjtNum, ten_, velocity, ntendon, 1 ) #define MJDATA_VIEW_GROUPS \ diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index c8f6774f..433b8167 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -1122,6 +1122,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='mjtSize'), doc='number of tendons', ), + StructFieldDecl( + name='nJten', + type=ValueType(name='mjtSize'), + doc='number of non-zeros in sparse ten_J matrix', + ), StructFieldDecl( name='nwrap', type=ValueType(name='mjtSize'), @@ -1242,11 +1247,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='mjtSize'), doc='number of non-zeros in sparse actuator_moment matrix', ), - StructFieldDecl( - name='nJten', - type=ValueType(name='mjtSize'), - doc='number of non-zeros in sparse ten_J matrix', - ), StructFieldDecl( name='ngravcomp', type=ValueType(name='mjtSize'), @@ -3927,6 +3927,30 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc="first two trees along tendon's path", array_extent=('ntendon', 2), ), + StructFieldDecl( + name='ten_J_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in Jacobian row', + array_extent=('ntendon',), + ), + StructFieldDecl( + name='ten_J_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in colind array', + array_extent=('ntendon',), + ), + StructFieldDecl( + name='ten_J_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in sparse Jacobian', + array_extent=('nJten',), + ), StructFieldDecl( name='tendon_limited', type=PointerType( @@ -5812,30 +5836,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='number of wrap points in path', array_extent=('ntendon',), ), - StructFieldDecl( - name='ten_J_rownnz', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of non-zeros in Jacobian row', - array_extent=('ntendon',), - ), - StructFieldDecl( - name='ten_J_rowadr', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='row start address in colind array', - array_extent=('ntendon',), - ), - StructFieldDecl( - name='ten_J_colind', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='column indices in sparse Jacobian', - array_extent=('nJten',), - ), StructFieldDecl( name='ten_J', type=PointerType( diff --git a/python/pyproject.toml b/python/pyproject.toml index f23b9bf5..f3f10bb7 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -9,7 +9,7 @@ authors = [ {name = "Google DeepMind", email = "mujoco@deepmind.com"}, ] description = "MuJoCo Physics Simulator" -requires-python = ">=3.9" +requires-python = ">=3.10" license = "Apache-2.0" classifiers = [ "Development Status :: 5 - Production/Stable", @@ -17,10 +17,10 @@ classifiers = [ "Intended Audience :: Science/Research", "Natural Language :: English", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", ] dependencies = [ diff --git a/sample/testspeed.cc b/sample/testspeed.cc index f14e6572..544edfc6 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -187,6 +187,12 @@ int main(int argc, char** argv) { nthread = mjMAX(1, mjMIN(maxthread, nthread)); npoolthread = mjMAX(1, mjMIN(maxthread, npoolthread)); + // load plugins from MUJOCO_PLUGIN_DIR if set + const char* plugin_dir = std::getenv("MUJOCO_PLUGIN_DIR"); + if (plugin_dir) { + mj_loadAllPluginLibraries(plugin_dir, nullptr); + } + // get filename, determine file type std::string filename(argv[1]); bool binary = (filename.find(".mjb") != std::string::npos); // NOLINT diff --git a/simulate/simulate.cc b/simulate/simulate.cc index be7318c5..022b9842 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -1584,6 +1584,7 @@ void UiEvent(mjuiState* state) { // rendering section else if (it && it->sectionid==SECT_RENDERING) { + // only update the camera when the camera itself changed if (it->pdata == &sim->camera) { if (sim->camera==0) { @@ -1603,6 +1604,7 @@ void UiEvent(mjuiState* state) { sim->cam.fixedcamid = sim->camera - 2; } } + // copy camera spec to clipboard (as MJCF element) if (it->itemid == 3) { CopyCamera(sim); diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c index c298b6d9..1339128d 100644 --- a/src/engine/engine_collision_convex.c +++ b/src/engine/engine_collision_convex.c @@ -45,44 +45,76 @@ static void ccd_free(void* data, void* buffer) { mj_freeStack((mjData*)data); } -// call libccd or nativeccd to recover penetration info -static int mjc_penetration(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, - const ccd_t* ccd, ccd_real_t* depth, ccd_vec3_t* dir, ccd_vec3_t* pos) { - // fallback to MPR - if (mjDISABLED(mjDSBL_NATIVECCD)) { - return ccdMPRPenetration(obj1, obj2, ccd, depth, dir, pos); +// ccd prism first dir +static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) { + ccdVec3Set(vec, 0, 0, 1); +} + +// wrapper around libccd; returns number of collisions found +static int libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum* dist, + mjtNum dir[3], mjtNum pos[3]) { + ccd_t ccd; + CCD_INIT(&ccd); + ccd.mpr_tolerance = m->opt.ccd_tolerance; + ccd.epa_tolerance = m->opt.ccd_tolerance; // use MPR tolerance for EPA + ccd.max_iterations = m->opt.ccd_iterations; + ccd.support1 = mjccd_support; + ccd.support2 = mjccd_support; + ccd.center1 = mjccd_center; + ccd.center2 = mjccd_center; + if (obj1->geom_type == mjGEOM_HFIELD) { + ccd.first_dir = prism_firstdir; + } + if (obj2->geom_type == mjGEOM_HFIELD) { + ccd.first_dir = prism_firstdir; } + ccd_real_t ccd_depth; + ccd_vec3_t ccd_dir, ccd_pos; + int ret = ccdMPRPenetration(obj1, obj2, &ccd, &ccd_depth, &ccd_dir, &ccd_pos); + *dist = -ccd_depth; + mji_copy3(dir, ccd_dir.v); + mji_copy3(pos, ccd_pos.v); + if (ret == 0 && dir[0] == 0 && dir[1] == 0 && dir[2] == 0) { + return 0; + } + return ret == 0; +} + + +// find penetration info between two geoms; returns number of collisions found +static int mjc_penetration(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum* dist, + mjtNum dir[3], mjtNum pos[3]) { + if (mjDISABLED(mjDSBL_NATIVECCD)) { + return libccd_wrapper(m, obj1, obj2, dist, dir, pos); + } + + // nativeccd mjCCDConfig config; mjCCDStatus status; + mjtNum d; // distance returned by mjc_ccd // set config - config.max_iterations = ccd->max_iterations; - config.tolerance = ccd->mpr_tolerance; + config.max_iterations = m->opt.ccd_iterations; + config.tolerance = m->opt.ccd_tolerance; config.max_contacts = 1; config.dist_cutoff = 0; // no geom distances needed config.context = (void*)obj1->data; config.alloc = ccd_allocate; config.free = ccd_free; - mjtNum dist = mjc_ccd(&config, &status, obj1, obj2); - if (dist < 0) { - if (depth) *depth = -dist; - if (dir) { - mju_sub3(dir->v, status.x1, status.x2); - mju_normalize3(dir->v); - } - if (pos) { - pos->v[0] = 0.5 * (status.x1[0] + status.x2[0]); - pos->v[1] = 0.5 * (status.x1[1] + status.x2[1]); - pos->v[2] = 0.5 * (status.x1[2] + status.x2[2]); - } - return 0; + if ((d = mjc_ccd(&config, &status, obj1, obj2)) < 0) { + *dist = d; + + mju_sub3(dir, status.x1, status.x2); + mju_normalize3(dir); + + pos[0] = 0.5 * (status.x1[0] + status.x2[0]); + pos[1] = 0.5 * (status.x1[1] + status.x2[1]); + pos[2] = 0.5 * (status.x1[2] + status.x2[2]); + return 1; } - if (depth) *depth = 0; - if (dir) mju_zero3(dir->v); - if (pos) mju_zero3(dir->v); - return 1; + return 0; } @@ -98,39 +130,34 @@ void mjc_center(mjtNum res[3], const mjCCDObj *obj) { int e = obj->elem; int v = obj->vert; + if (obj->geom_type == mjGEOM_HFIELD) { + mju_zero3(res); + for (int i=0; i < 6; i++) { + mji_addTo3(res, obj->prism[i]); + } + mju_scl3(res, res, 1.0/6.0); + return; + } + // return geom position if (g >= 0) { mji_copy3(res, obj->data->geom_xpos + 3*g); + return; } // return flex element position - else if (e >= 0) { + if (e >= 0) { mji_copy3(res, obj->data->flexelem_aabb + 6*(obj->model->flex_elemadr[f]+e)); + return; } // return flex vertex position - else { + if (f >= 0) { mji_copy3(res, obj->data->flexvert_xpos + 3*(obj->model->flex_vertadr[f]+v)); + return; } } - -// prism center function -static void mjc_prism_center(mjtNum res[3], const mjCCDObj* obj) { - // compute mean - mju_zero3(res); - for (int i=0; i < 6; i++) { - mji_addTo3(res, obj->prism[i]); - } - mju_scl3(res, res, 1.0/6.0); -} - - -// ccd prism center function -static void mjccd_prism_center(const void *obj, ccd_vec3_t *center) { - mjc_prism_center(center->v, (const mjCCDObj*) obj); -} - // ------------------------------------ Support functions ----------------------------------------- // transform a vector from global to local frame @@ -681,6 +708,10 @@ void mjccd_support(const void *_obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) { } break; + case mjGEOM_HFIELD: + mjc_prism_support(res, obj, dir); + return; + default: mjERROR("ccd support function is undefined for geom type %d", m->geom_type[g]); } @@ -697,12 +728,6 @@ void mjccd_support(const void *_obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) { mji_addTo3(res, d->geom_xpos+3*g); } - -// libccd prism support function -static void mjccd_prism_support(const void *obj, const ccd_vec3_t *dir, ccd_vec3_t *vec) { - mjc_prism_support(vec->v, (mjCCDObj*) obj, dir->v); -} - // ------------------------------------------------------------------------------------------------ // initialize a CCD object @@ -747,7 +772,7 @@ void mjc_initCCDObj(mjCCDObj* obj, const mjModel* m, const mjData* d, int g, mjt obj->support = mjc_boxSupport; break; case mjGEOM_HFIELD: - obj->center = mjc_prism_center; + obj->center = mjc_center; obj->support = mjc_prism_support; int hid = m->geom_dataid[g]; @@ -1140,11 +1165,6 @@ int mjc_PlaneConvex(const mjModel* m, const mjData* d, //---------------------------- heightfield collisions --------------------------------------------- -// ccd prism first dir -static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) { - ccdVec3Set(vec, 0, 0, 1); -} - // add vertex to prism static inline void addVert(mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) { @@ -1289,16 +1309,6 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d, rmin = mjMAX(0, rmin); rmax = mjMIN(nrow-1, rmax); - // CCD collision testing - - ccd_t ccd; - mjc_initCCD(&ccd, m); - ccd.first_dir = prism_firstdir; - ccd.center1 = mjccd_prism_center; - ccd.center2 = mjccd_center; - ccd.support1 = mjccd_prism_support; - ccd.support2 = mjccd_support; - // geom margin needed for actual collision test obj2.margin = margin; @@ -1325,14 +1335,12 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d, } // run penetration function, save contact - ccd_vec3_t dirccd, vecccd; - ccd_real_t depth; - if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0 - && !ccdVec3Eq(&dirccd, ccd_vec3_origin)) { + mjtNum dist; + if (mjc_penetration(m, &obj1, &obj2, &dist, dir, pos)) { // fill in contact data, transform to global coordinates - con[ncon].dist = -depth; - mji_mulMatVec3(con[ncon].frame, mat1, dirccd.v); - mji_mulMatVec3(con[ncon].pos, mat1, vecccd.v); + con[ncon].dist = dist; + mji_mulMatVec3(con[ncon].frame, mat1, dir); + mji_mulMatVec3(con[ncon].pos, mat1, pos); mji_addTo3(con[ncon].pos, pos1); mju_zero3(con[ncon].frame+3); @@ -1660,7 +1668,7 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, mjtNum xmin, xmax, ymin, ymax, zmin, zmax; int dr[2], cnt, rmin, rmax, cmin, cmax; mjCCDObj obj1; - obj1.center = mjc_prism_center; + obj1.center = mjc_center; obj1.support = mjc_prism_support; // get hfield info @@ -1682,13 +1690,9 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, mjtNum* ecenter = d->flexelem_aabb + 6*(m->flex_elemadr[f]+e); // ccd-related - ccd_vec3_t dirccd, vecccd; - ccd_real_t depth; mjCCDObj obj2; mjc_initCCDObj(&obj2, m, d, -1, margin); mjc_setCCDObjFlex(&obj2, f, e, -1); - ccd_t ccd; - //------------------------------------- AABB computation, box-box test // save elem vertices, transform to hfield frame @@ -1743,18 +1747,6 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, //------------------------------------- collision testing - // init ccd structure - CCD_INIT(&ccd); - ccd.first_dir = prism_firstdir; - ccd.center1 = mjccd_prism_center; - ccd.center2 = mjccd_center; - ccd.support1 = mjccd_prism_support; - ccd.support2 = mjccd_support; - - // set ccd parameters - ccd.max_iterations = m->opt.ccd_iterations; - ccd.mpr_tolerance = m->opt.ccd_tolerance; - // compute real-valued grid step, and triangulation direction dx = (2.0*hsize[0]) / (ncol-1); dy = (2.0*hsize[1]) / (nrow-1); @@ -1782,23 +1774,22 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, } // run ccd, save contact - if (mjc_penetration(m, &obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0) { - if (!ccdVec3Eq(&dirccd, ccd_vec3_origin)) { - // fill in contact data, transform to global coordinates - con[cnt].dist = -depth; - mji_mulMatVec3(con[cnt].frame, hmat, dirccd.v); - mji_mulMatVec3(con[cnt].pos, hmat, vecccd.v); - mji_addTo3(con[cnt].pos, hpos); - mju_zero3(con[cnt].frame+3); + mjtNum dist, dir[3], pos[3]; + if (mjc_penetration(m, &obj1, &obj2, &dist, dir, pos)) { + // fill in contact data, transform to global coordinates + con[cnt].dist = dist; + mji_mulMatVec3(con[cnt].frame, hmat, dir); + mji_mulMatVec3(con[cnt].pos, hmat, pos); + mji_addTo3(con[cnt].pos, hpos); + mju_zero3(con[cnt].frame+3); - // count, stop if max number reached - cnt++; - if (cnt >= mjMAXCONPAIR) { - r = rmax+1; - c = cmax+1; - k = 3; - break; - } + // count, stop if max number reached + cnt++; + if (cnt >= mjMAXCONPAIR) { + r = rmax+1; + c = cmax+1; + k = 3; + break; } } } diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index ac12d0e4..ceaac7ef 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -259,6 +259,78 @@ static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* mjSORT(contactSort, mjContact, contactcompare); +// filter flex contacts based on distance +static void filterFlexContacts(mjData* d, int ncon_before) { + int n = d->ncon - ncon_before; + if (n <= mjMAXCONPAIR) { + return; + } + + mjContact* contacts = d->contact + ncon_before; + + mj_markStack(d); + mjtByte* selected = mjSTACKALLOC(d, n, mjtByte); + mjtNum* min_dist = mjSTACKALLOC(d, n, mjtNum); + memset(selected, 0, n); + + for (int i = 0; i < n; i++) { + min_dist[i] = mjMAXVAL; + } + + // start with the deepest penetrating contact + int nselected = 0; + int best = 0; + mjtNum bestdist = -contacts[0].dist; + for (int i = 1; i < n; i++) { + if (-contacts[i].dist > bestdist) { + bestdist = -contacts[i].dist; + best = i; + } + } + + while (nselected < mjMAXCONPAIR && best >= 0) { + selected[best] = 1; + mjtNum* bestpos = contacts[best].pos; + + int nextbest = -1; + mjtNum nextbestdist = -1; + for (int i = 0; i < n; i++) { + if (selected[i]) continue; + + mjtNum dx = contacts[i].pos[0] - bestpos[0]; + mjtNum dy = contacts[i].pos[1] - bestpos[1]; + mjtNum dz = contacts[i].pos[2] - bestpos[2]; + mjtNum d2 = dx*dx + dy*dy + dz*dz; + if (d2 < min_dist[i]) { + min_dist[i] = d2; + } + if (min_dist[i] > nextbestdist) { + nextbestdist = min_dist[i]; + nextbest = i; + } + } + + if (nselected < mjMAXCONPAIR - 1) { + mjContact temp = contacts[nselected]; + contacts[nselected] = contacts[best]; + contacts[best] = temp; + + if (nextbest == nselected) { + nextbest = best; + } + } + + nselected++; + best = nextbest; + } + + mj_freeStack(d); + + d->ncon = ncon_before + nselected; + resetArena(d); +} + + // main collision function void mj_collision(const mjModel* m, mjData* d) { @@ -362,6 +434,12 @@ void mj_collision(const mjModel* m, mjData* d) { mj_collideTree(m, d, bf1, bf2, merged, startadr, pairadr); int ncon_after = d->ncon; + // filter flex contacts (limit per geom-flex or flex-flex pair) + if (bf1 >= nbody || bf2 >= nbody) { + filterFlexContacts(d, ncon_before); + ncon_after = d->ncon; + } + // sort contacts int n = ncon_after - ncon_before; if (n > 1) { @@ -400,15 +478,19 @@ void mj_collision(const mjModel* m, mjData* d) { // plane special processing if (m->geom_type[g] == mjGEOM_PLANE) { + int ncon_before = d->ncon; mj_collidePlaneFlex(m, d, g, f); + filterFlexContacts(d, ncon_before); continue; } // collide geom with flex elements + int ncon_before = d->ncon; int elemnum = m->flex_elemnum[f]; for (int e=0; e < elemnum; e++) { mj_collideGeomElem(m, d, g, f, e); } + filterFlexContacts(d, ncon_before); } } @@ -418,11 +500,13 @@ void mj_collision(const mjModel* m, mjData* d) { int f2 = bf2 - nbody; // collide elements of two flexes + int ncon_before = d->ncon; for (int e1=0; e1 < m->flex_elemnum[f1]; e1++) { for (int e2=0; e2 < m->flex_elemnum[f2]; e2++) { mj_collideElems(m, d, f1, e1, f2, e2); } } + filterFlexContacts(d, ncon_before); } } } @@ -439,11 +523,15 @@ void mj_collision(const mjModel* m, mjData* d) { if (!m->flex_rigid[f] && (m->flex_contype[f] & m->flex_conaffinity[f])) { // internal collisions if (m->flex_internal[f]) { + int ncon_before = d->ncon; mj_collideFlexInternal(m, d, f); + filterFlexContacts(d, ncon_before); } // active element collisions if (m->flex_selfcollide[f] != mjFLEXSELF_NONE) { + int ncon_before = d->ncon; + // element-element: midphase if (!mjDISABLED(mjDSBL_MIDPHASE) && m->flex_selfcollide[f] != mjFLEXSELF_NARROW && @@ -470,6 +558,8 @@ void mj_collision(const mjModel* m, mjData* d) { } } } + + filterFlexContacts(d, ncon_before); } } } diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index ab6abd39..982463d7 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -567,18 +567,17 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // copy Jacobian: sparse or dense if (issparse) { - // add first or second chain if (j == 0) { - NV = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); - mju_copy(jac[j], d->ten_J+d->ten_J_rowadr[id[j]], NV); + NV = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV); + mju_copy(jac[j], d->ten_J+m->ten_J_rowadr[id[j]], NV); } else { - NV2 = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); - mju_copy(jac[j], d->ten_J+d->ten_J_rowadr[id[j]], NV2); + NV2 = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain2, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV2); + mju_copy(jac[j], d->ten_J+m->ten_J_rowadr[id[j]], NV2); } } else { - mju_copy(jac[j], d->ten_J+id[j]*nv, nv); + mju_sparse2dense(jac[j], d->ten_J, 1, nv, m->ten_J_rownnz+id[j], m->ten_J_rowadr+id[j], m->ten_J_colind); } } } @@ -737,11 +736,17 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { if (m->tendon_frictionloss[i] > 0) { int efcadr = d->nefc; // add constraint - mj_addConstraint(m, d, d->ten_J + (issparse ? d->ten_J_rowadr[i] : i*nv), - 0, 0, m->tendon_frictionloss[i], - 1, mjCNSTR_FRICTION_TENDON, i, - issparse ? d->ten_J_rownnz[i] : 0, - issparse ? d->ten_J_colind+d->ten_J_rowadr[i] : NULL); + if (issparse) { + mj_addConstraint(m, d, d->ten_J + m->ten_J_rowadr[i], + 0, 0, m->tendon_frictionloss[i], + 1, mjCNSTR_FRICTION_TENDON, i, + m->ten_J_rownnz[i], + m->ten_J_colind+m->ten_J_rowadr[i]); + } else { + mju_sparse2dense(jac, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); + mj_addConstraint(m, d, jac, 0, 0, m->tendon_frictionloss[i], + 1, mjCNSTR_FRICTION_TENDON, i, 0, NULL); + } // set tendon_efcadr if (d->tendon_efcadr[i] == -1) { d->tendon_efcadr[i] = efcadr; @@ -877,19 +882,20 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { // detect tendon limit if (dist < margin) { - // prepare Jacobian: sparse or dense - if (issparse) { - mju_scl(jac, d->ten_J+d->ten_J_rowadr[i], -side, d->ten_J_rownnz[i]); - } else { - mju_scl(jac, d->ten_J+i*nv, -side, nv); - } - - // add constraint + // prepare Jacobian int efcadr = d->nefc; - mj_addConstraint(m, d, jac, &dist, &margin, 0, - 1, mjCNSTR_LIMIT_TENDON, i, - issparse ? d->ten_J_rownnz[i] : 0, - issparse ? d->ten_J_colind+d->ten_J_rowadr[i] : NULL); + if (issparse) { + mju_scl(jac, d->ten_J+m->ten_J_rowadr[i], -side, m->ten_J_rownnz[i]); + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_TENDON, i, + m->ten_J_rownnz[i], + m->ten_J_colind+m->ten_J_rowadr[i]); + } else { + mju_sparse2dense(jac, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); + mju_scl(jac, jac, -side, nv); + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_TENDON, i, 0, NULL); + } // set tendon_efcadr if (d->tendon_efcadr[i] == -1) { d->tendon_efcadr[i] = efcadr; @@ -1732,11 +1738,11 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { } } else { if (!j) { - NV = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); + NV = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV); } else { - NV2 = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); + NV2 = m->ten_J_rownnz[id[j]]; + mju_copyInt(chain2, m->ten_J_colind+m->ten_J_rowadr[id[j]], NV2); } } } @@ -1834,8 +1840,8 @@ static int mj_nf(const mjModel* m, const mjData* d, int *nnz) { for (int i=0; i < ntendon; i++) { if (m->tendon_frictionloss[i] > 0) { - nf += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); - if (nnz) *nnz += d->ten_J_rownnz[i]; + nf += mj_addConstraintCount(m, 1, m->ten_J_rownnz[i]); + if (nnz) *nnz += m->ten_J_rownnz[i]; } } @@ -1902,8 +1908,8 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { for (int i=0; i < ntendon; i++) { int count = tendonLimit(m, d->ten_length, i); for (int j = 0; j < count; j++) { - nl += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); - if (nnz) *nnz += d->ten_J_rownnz[i]; + nl += mj_addConstraintCount(m, 1, m->ten_J_rownnz[i]); + if (nnz) *nnz += m->ten_J_rownnz[i]; } } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ba5c7254..47e0fdee 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -904,8 +904,8 @@ void mj_flex(const mjModel* m, mjData* d) { // compute tendon lengths and moments void mj_tendon(const mjModel* m, mjData* d) { - int issparse = mj_isSparse(m), nv = m->nv, nten = m->ntendon; - int *rownnz = d->ten_J_rownnz, *rowadr = d->ten_J_rowadr, *colind = d->ten_J_colind; + int nv = m->nv, nten = m->ntendon; + const int *rownnz = m->ten_J_rownnz, *rowadr = m->ten_J_rowadr, *colind = m->ten_J_colind; mjtNum *L = d->ten_length, *J = d->ten_J; if (!nten) { @@ -913,28 +913,20 @@ void mj_tendon(const mjModel* m, mjData* d) { } // allocate stack arrays - int *chain = NULL, *buf_ind = NULL; - mjtNum *jac1, *jac2, *jacdif, *tmp, *sparse_buf = NULL; + int *chain; + mjtNum *jac1, *jac2, *jacdif, *tmp; mj_markStack(d); jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); tmp = mjSTACKALLOC(d, nv, mjtNum); - if (issparse) { - chain = mjSTACKALLOC(d, nv, int); - buf_ind = mjSTACKALLOC(d, nv, int); - sparse_buf = mjSTACKALLOC(d, nv, mjtNum); - } + chain = mjSTACKALLOC(d, nv, int); // clear results mju_zero(L, nten); - // clear Jacobian: sparse or dense - if (issparse) { - mju_zeroInt(rownnz, nten); - } else { - mju_zero(J, nten*nv); - } + // clear Jacobian + mju_zero(J, m->nJten); // sleep filtering int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; @@ -953,11 +945,6 @@ void mj_tendon(const mjModel* m, mjData* d) { d->ten_wrapnum[i] = 0; int tendon_num = m->tendon_num[i]; - // sparse Jacobian row init - if (issparse) { - rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); - } - // process fixed tendon if (m->wrap_type[adr] == mjWRAP_JOINT) { // process all defined joints @@ -968,18 +955,10 @@ void mj_tendon(const mjModel* m, mjData* d) { // add to length L[i] += m->wrap_prm[adr+j] * d->qpos[m->jnt_qposadr[k]]; - // add to moment - if (issparse) { - rownnz[i] = mju_combineSparse(J+rowadr[i], &m->wrap_prm[adr+j], 1, 1, - rownnz[i], 1, - colind+rowadr[i], &m->jnt_dofadr[k], - sparse_buf, buf_ind); - } - - // add to moment: dense - else { - J[i*nv + m->jnt_dofadr[k]] = m->wrap_prm[adr+j]; - } + mjtNum coef = 1; + int dofadr = m->jnt_dofadr[k]; + mju_combineSparseInc(J + rowadr[i], &coef, m->nv, 1, m->wrap_prm[adr+j], + rownnz[i], 1, colind + rowadr[i], &dofadr); } continue; @@ -1060,40 +1039,23 @@ void mj_tendon(const mjModel* m, mjData* d) { mji_sub3(dif, wpnt+3*k+3, wpnt+3*k); mju_normalize3(dif); - // sparse - if (issparse) { - // get endpoint Jacobians, subtract - int NV = mj_jacDifPair(m, d, chain, - wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, - jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); + // get endpoint Jacobians, subtract + int NV = mj_jacDifPair(m, d, chain, + wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); - // no dofs: skip - if (!NV) { - continue; - } - - // apply chain rule to compute tendon Jacobian - mju_mulMatTVec(tmp, jacdif, dif, 3, NV); - - // add to existing - rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor, - rownnz[i], NV, colind+rowadr[i], - chain, sparse_buf, buf_ind); + // no dofs: skip + if (!NV) { + continue; } - // dense - else { - // get endpoint Jacobians, subtract - mj_jac(m, d, jac1, 0, wpnt+3*k, wbody[k]); - mj_jac(m, d, jac2, 0, wpnt+3*k+3, wbody[k+1]); - mju_sub(jacdif, jac2, jac1, 3*nv); + // apply chain rule to compute tendon Jacobian + mju_mulMatTVec(tmp, jacdif, dif, 3, NV); - // apply chain rule to compute tendon Jacobian - mju_mulMatTVec(tmp, jacdif, dif, 3, nv); - - // add to existing - mju_addToScl(J + i*nv, tmp, 1/divisor, nv); - } + // add to existing + mju_combineSparseInc(J+rowadr[i], tmp, nv, 1, 1/divisor, + rownnz[i], NV, colind+rowadr[i], + chain); } } @@ -1449,28 +1411,14 @@ void mj_transmission(const mjModel* m, mjData* d) { case mjTRN_TENDON: // tendon length[i] = d->ten_length[id]*gear[0]; - // moment: sparse or dense - if (issparse) { - // sparsity - int ten_J_rownnz = d->ten_J_rownnz[id]; - int ten_J_rowadr = d->ten_J_rowadr[id]; + // moment + { + int ten_J_rownnz = m->ten_J_rownnz[id]; + int ten_J_rowadr = m->ten_J_rowadr[id]; rownnz[i] = ten_J_rownnz; - mju_copyInt(colind + adr, d->ten_J_colind + ten_J_rowadr, ten_J_rownnz); + mju_copyInt(colind + adr, m->ten_J_colind + ten_J_rowadr, ten_J_rownnz); mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz); - } else { - mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv); - - // sparsity (compress) - nnz = 0; - for (int j = 0; j < nv; j++) { - if (moment[adr+j]) { - moment[adr+nnz] = moment[adr+j]; - colind[adr+nnz] = j; - nnz++; - } - } - rownnz[i] = nnz; } break; @@ -1743,7 +1691,7 @@ void mj_transmission(const mjModel* m, mjData* d) { // add tendon armature to M void mj_tendonArmature(const mjModel* m, mjData* d) { - int nv = m->nv, ntendon = m->ntendon, issparse = mj_isSparse(m); + int nv = m->nv, ntendon = m->ntendon; const int* M_rownnz = m->M_rownnz; const int* M_rowadr = m->M_rowadr; const int* M_colind = m->M_colind; @@ -1762,47 +1710,25 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { continue; } - // dense - if (!issparse) { - // M += armature * ten_J' * ten_J - mjtNum* ten_J = d->ten_J + nv*k; - for (int i=0; i < nv; i++) { - mjtNum ten_J_i = ten_J[i]; - if (!ten_J_i) { - continue; - } + // get sparse info for tendon k + int J_rowadr = m->ten_J_rowadr[k]; + int J_rownnz = m->ten_J_rownnz[k]; + const int* J_colind = m->ten_J_colind + J_rowadr; + mjtNum* ten_J = d->ten_J + J_rowadr; - // M[i,:] += armature * ten_J[i] * ten_J - int start = M_rowadr[i]; - int end = start + M_rownnz[i]; - for (int adr = start; adr < end; adr++) { - d->M[adr] += armature * ten_J_i * ten_J[M_colind[adr]]; - } + // M += armature * ten_J' * ten_J + for (int j=0; j < J_rownnz; j++) { + mjtNum ten_J_i = ten_J[j]; + if (!ten_J_i) { + continue; } - } - // sparse - else { - // get sparse info for tendon k - int J_rowadr = d->ten_J_rowadr[k]; - int J_rownnz = d->ten_J_rownnz[k]; - const int* J_colind = d->ten_J_colind + J_rowadr; - mjtNum* ten_J = d->ten_J + J_rowadr; - - // M += armature * ten_J' * ten_J - for (int j=0; j < J_rownnz; j++) { - mjtNum ten_J_i = ten_J[j]; - if (!ten_J_i) { - continue; - } - - // M[i,:] += armature * ten_J[i] * ten_J - int i = J_colind[j]; - int M_adr = M_rowadr[i]; - mju_addToSclSparseInc(d->M + M_adr, ten_J, - M_rownnz[i], M_colind + M_adr, - J_rownnz, J_colind, armature * ten_J_i); - } + // M[i,:] += armature * ten_J[i] * ten_J + int i = J_colind[j]; + int M_adr = M_rowadr[i]; + mju_addToSclSparseInc(d->M + M_adr, ten_J, + M_rownnz[i], M_colind + M_adr, + J_rownnz, J_colind, armature * ten_J_i); } } } @@ -2686,7 +2612,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { // add bias force due to tendon armature void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; - int ntendon = m->ntendon, nv = m->nv, issparse = mj_isSparse(m); + int ntendon = m->ntendon, nv = m->nv; mjtNum* ten_Jdot = NULL; mj_markStack(d); @@ -2716,20 +2642,13 @@ void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { 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]; - } + int nnz = m->ten_J_rownnz[i]; + int adr = m->ten_J_rowadr[i]; + const int* colind = m->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]; } } } diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 6c37710f..09dd4903 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -1775,12 +1775,8 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { continue; } - // add sparse or dense - if (mj_isSparse(m)) { - addJTBJSparse(m, d, d->ten_J, &B, 1, i, d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind); - } else { - addJTBJ(m, d, d->ten_J+i*nv, &B, 1); - } + // add sparse + addJTBJSparse(m, d, d->ten_J, &B, 1, i, m->ten_J_rownnz, m->ten_J_rowadr, m->ten_J_colind); } } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 5b336890..f2a7c3c4 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -27,6 +27,7 @@ #include "engine/engine_core_constraint.h" #include "engine/engine_core_smooth.h" #include "engine/engine_derivative.h" +#include "engine/engine_core_util.h" #include "engine/engine_inverse.h" #include "engine/engine_island.h" #include "engine/engine_macro.h" @@ -229,13 +230,9 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) { mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL); - // tendon velocity: dense or sparse - if (mj_isSparse(m)) { - mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, - d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind, NULL); - } else { - mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv); - } + // tendon velocity: always sparse + mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, + m->ten_J_rownnz, m->ten_J_rowadr, m->ten_J_colind, NULL); // actuator velocity: always sparse if (!mjDISABLED(mjDSBL_ACTUATION)) { @@ -1193,14 +1190,38 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // flex: reduced dense factorization if (has_flex_interp && !sleep_filter) { + // temporary allocations for body chain + int* chain_dofs = mjSTACKALLOC(d, nv, int); + int* seen_dof = mjSTACKALLOC(d, nv, int); + mju_fillInt(seen_dof, 0, nv); + // identify flex DOFs + // For pinned nodes (body_dofnum==0): use bodyChain to include parent DOFs + // For regular flex nodes: use body_dofadr for one-way coupling for (int f=0; f < m->nflex; f++) { if (m->flex_interp[f]) { int nodenum = m->flex_nodenum[f]; int nodeadr = m->flex_nodeadr[f]; for (int n=0; n < nodenum; n++) { int b = m->flex_nodebodyid[nodeadr + n]; - nflexdofs += m->body_dofnum[b]; + int chain_nnz; + if (m->body_dofnum[b] == 0) { + // Pinned node: use bodyChain to get parent DOFs + chain_nnz = mj_bodyChain(m, b, chain_dofs); + } else { + // Regular flex node: use body's own DOFs only + chain_nnz = m->body_dofnum[b]; + for (int j = 0; j < chain_nnz; j++) { + chain_dofs[j] = m->body_dofadr[b] + j; + } + } + for (int i=0; i < chain_nnz; i++) { + int dof = chain_dofs[i]; + if (!seen_dof[dof]) { + seen_dof[dof] = 1; + nflexdofs++; + } + } } } } @@ -1211,19 +1232,34 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { int* global2local = mjSTACKALLOC(d, nv, int); mju_fillInt(global2local, -1, nv); + // collect unique DOFs in order int cnt = 0; + mju_fillInt(seen_dof, 0, nv); for (int f=0; f < m->nflex; f++) { if (m->flex_interp[f]) { int nodenum = m->flex_nodenum[f]; int nodeadr = m->flex_nodeadr[f]; for (int n=0; n < nodenum; n++) { int b = m->flex_nodebodyid[nodeadr + n]; - int dofnum = m->body_dofnum[b]; - int dofadr = m->body_dofadr[b]; - for (int j=0; j < dofnum; j++) { - flex_dof_indices[cnt] = dofadr + j; - global2local[dofadr + j] = cnt; - cnt++; + int chain_nnz; + if (m->body_dofnum[b] == 0) { + // Pinned node: use bodyChain to get parent DOFs + chain_nnz = mj_bodyChain(m, b, chain_dofs); + } else { + // Regular flex node: use body's own DOFs only + chain_nnz = m->body_dofnum[b]; + for (int j = 0; j < chain_nnz; j++) { + chain_dofs[j] = m->body_dofadr[b] + j; + } + } + for (int i=0; i < chain_nnz; i++) { + int dof = chain_dofs[i]; + if (!seen_dof[dof]) { + seen_dof[dof] = 1; + flex_dof_indices[cnt] = dof; + global2local[dof] = cnt; + cnt++; + } } } } diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 512c7661..f0cf467f 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -211,10 +211,10 @@ void mj_makeModel(mjModel** dest, mjtSize nmeshpoly, mjtSize nmeshpolyvert, mjtSize nmeshpolymap, mjtSize nskin, mjtSize nskinvert, mjtSize nskintexvert, mjtSize nskinface, mjtSize nskinbone, mjtSize nskinbonevert, mjtSize nhfield, mjtSize nhfielddata, mjtSize ntex, mjtSize ntexdata, - mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nwrap, - mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, mjtSize ntextdata, - mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, mjtSize nplugin, - mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, + mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nJten, + mjtSize nwrap, mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, + mjtSize ntextdata, mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, + mjtSize nplugin, mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, mjtSize nuser_site, mjtSize nuser_cam, mjtSize nuser_tendon, mjtSize nuser_actuator, mjtSize nuser_sensor, mjtSize nnames, mjtSize npaths) { intptr_t offset = 0; @@ -224,7 +224,7 @@ void mj_makeModel(mjModel** dest, // CHECK SIZE PARAMETERS { // dummy variables for MJMODEL_SIZES set after mjModel construction - int nnames_map = 0, nJmom = 0, nJten = 0, ngravcomp = 0, nemax = 0, njmax = 0, nconmax=0; + int nnames_map = 0, nJmom = 0, ngravcomp = 0, nemax = 0, njmax = 0, nconmax=0; int nuserdata=0, nsensordata=0, npluginstate=0, nhistory=0, narena=0, nbuffer=0; // sizes must be non-negative and fit in int, except for the byte arrays texdata and textdata @@ -243,7 +243,7 @@ void mj_makeModel(mjModel** dest, #undef X // suppress unused variable warnings - (void)nnames_map; (void)nJmom; (void)nJten; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax; + (void)nnames_map; (void)nJmom; (void)ngravcomp; (void)nemax; (void)njmax; (void)nconmax; (void)nuserdata; (void)nsensordata; (void)npluginstate; (void)nhistory; (void)narena; (void)nbuffer; } @@ -323,6 +323,7 @@ void mj_makeModel(mjModel** dest, m->nexclude = nexclude; m->neq = neq; m->ntendon = ntendon; + m->nJten = nJten; m->nwrap = nwrap; m->nsensor = nsensor; m->nnumeric = nnumeric; @@ -410,11 +411,11 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) { src->nskin, src->nskinvert, src->nskintexvert, src->nskinface, src->nskinbone, src->nskinbonevert, src->nhfield, src->nhfielddata, src->ntex, src->ntexdata, src->nmat, src->npair, src->nexclude, - src->neq, src->ntendon, src->nwrap, src->nsensor, src->nnumeric, - src->nnumericdata, src->ntext, src->ntextdata, src->ntuple, - src->ntupledata, src->nkey, src->nmocap, src->nplugin, src->npluginattr, - src->nuser_body, src->nuser_jnt, src->nuser_geom, src->nuser_site, - src->nuser_cam, src->nuser_tendon, src->nuser_actuator, + src->neq, src->ntendon, src->nJten, src->nwrap, src->nsensor, + src->nnumeric, src->nnumericdata, src->ntext, src->ntextdata, + src->ntuple, src->ntupledata, src->nkey, src->nmocap, src->nplugin, + src->npluginattr, src->nuser_body, src->nuser_jnt, src->nuser_geom, + src->nuser_site, src->nuser_cam, src->nuser_tendon, src->nuser_actuator, src->nuser_sensor, src->nnames, src->npaths); } if (!dest) { @@ -597,7 +598,8 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) { sizes[49], sizes[50], sizes[51], sizes[52], sizes[53], sizes[54], sizes[55], sizes[56], sizes[57], sizes[58], sizes[59], sizes[60], sizes[61], sizes[62], sizes[63], sizes[64], sizes[65], sizes[66], sizes[67], sizes[68], sizes[69], - sizes[70], sizes[71], sizes[72], sizes[73], sizes[74], sizes[75], sizes[76]); + sizes[70], sizes[71], sizes[72], sizes[73], sizes[74], sizes[75], sizes[76], + sizes[77]); // mj_makeModel may fail if the input buffer has invalid sizes if (!m) { diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 918ad5a3..0af1058f 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -58,10 +58,10 @@ void mj_makeModel(mjModel** dest, mjtSize nmeshpoly, mjtSize nmeshpolyvert, mjtSize nmeshpolymap, mjtSize nskin, mjtSize nskinvert, mjtSize nskintexvert, mjtSize nskinface, mjtSize nskinbone, mjtSize nskinbonevert, mjtSize nhfield, mjtSize nhfielddata, mjtSize ntex, mjtSize ntexdata, - mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nwrap, - mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, mjtSize ntextdata, - mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, mjtSize nplugin, - mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, + mjtSize nmat, mjtSize npair, mjtSize nexclude, mjtSize neq, mjtSize ntendon, mjtSize nJten, + mjtSize nwrap, mjtSize nsensor, mjtSize nnumeric, mjtSize nnumericdata, mjtSize ntext, + mjtSize ntextdata, mjtSize ntuple, mjtSize ntupledata, mjtSize nkey, mjtSize nmocap, + mjtSize nplugin, mjtSize npluginattr, mjtSize nuser_body, mjtSize nuser_jnt, mjtSize nuser_geom, mjtSize nuser_site, mjtSize nuser_cam, mjtSize nuser_tendon, mjtSize nuser_actuator, mjtSize nuser_sensor, mjtSize nnames, mjtSize npaths); diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index d52e478c..60446c13 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -119,7 +119,6 @@ static void mj_springdamper(const mjModel* m, mjData* d) { int nv = m->nv, ntendon = m->ntendon; int has_spring = !mjDISABLED(mjDSBL_SPRING); int has_damping = !mjDISABLED(mjDSBL_DAMPER); - int issparse = mj_isSparse(m); int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; int nbody = sleep_filter ? d->nbody_awake : m->nbody; @@ -472,20 +471,15 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // compute damper linear force along tendon mjtNum frc_damper = -damping * d->ten_velocity[i]; - // transform to joint torque, add to qfrc_{spring, damper}: dense or sparse - if (issparse) { - if (frc_spring || frc_damper) { - int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; - for (int j=d->ten_J_rowadr[i]; j < end; j++) { - int k = d->ten_J_colind[j]; - mjtNum J = d->ten_J[j]; - d->qfrc_spring[k] += J * frc_spring; - d->qfrc_damper[k] += J * frc_damper; - } + // transform to joint torque, add to qfrc_{spring, damper} + if (frc_spring || frc_damper) { + int end = m->ten_J_rowadr[i] + m->ten_J_rownnz[i]; + for (int j=m->ten_J_rowadr[i]; j < end; j++) { + int k = m->ten_J_colind[j]; + mjtNum J = d->ten_J[j]; + d->qfrc_spring[k] += J * frc_spring; + d->qfrc_damper[k] += J * frc_damper; } - } else { - if (frc_spring) mju_addToScl(d->qfrc_spring, d->ten_J+i*nv, frc_spring, nv); - if (frc_damper) mju_addToScl(d->qfrc_damper, d->ten_J+i*nv, frc_damper, nv); } } } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 5cee8f49..a8f407d9 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1441,16 +1441,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format); printArray2d("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format); - if (!mj_isSparse(m)) { - printArray2d("TEN_MOMENT", m->ntendon, m->nv, d->ten_J, fp, float_format); - } else { - mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, d->ten_J_rowadr, NULL, - d->ten_J_rownnz, NULL, d->ten_J_colind, fp); - printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, d->ten_J_rownnz, fp); - printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, d->ten_J_rowadr, fp); - printSparse("TEN_J", d->ten_J, m->ntendon, d->ten_J_rownnz, - d->ten_J_rowadr, d->ten_J_colind, fp, float_format); - } + mj_printSparsity("TEN_J: tendon moments", m->ntendon, m->nv, m->ten_J_rowadr, NULL, + m->ten_J_rownnz, NULL, m->ten_J_colind, fp); + printArray2dInt("TEN_J_ROWNNZ", m->ntendon, 1, m->ten_J_rownnz, fp); + printArray2dInt("TEN_J_ROWADR", m->ntendon, 1, m->ten_J_rowadr, fp); + printSparse("TEN_J", d->ten_J, m->ntendon, m->ten_J_rownnz, + m->ten_J_rowadr, m->ten_J_colind, fp, float_format); for (int i=0; i < m->ntendon; i++) { fprintf(fp, "TENDON %d: %d wrap points\n", i, d->ten_wrapnum[i]); for (int j=0; j < d->ten_wrapnum[i]; j++) { diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 7b5fe6bb..6dd01a94 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -282,6 +282,93 @@ static void setFixed(mjModel* m, mjData* d) { mj_freeStack(d); } +// compute tendon Jacobian sparsity +static void makeTendonSparse(mjModel* m) { + int ntendon = m->ntendon; + int* rownnz = m->ten_J_rownnz; + int* rowadr = m->ten_J_rowadr; + int* colind = m->ten_J_colind; + + if (!ntendon) { + return; + } + + // clear + mju_zeroInt(rownnz, ntendon); + mju_zeroInt(rowadr, ntendon); + + // compute rownnz, rowadr, and colind for each tendon + for (int i = 0; i < ntendon; i++) { + rowadr[i] = (i > 0 ? rowadr[i-1] + rownnz[i-1] : 0); + int adr = m->tendon_adr[i]; + int num = m->tendon_num[i]; + + // joint tendon: each wrap object is a joint, colind is its dofadr + if (m->wrap_type[adr] == mjWRAP_JOINT) { + for (int j = 0; j < num; j++) { + colind[rowadr[i] + j] = m->jnt_dofadr[m->wrap_objid[adr + j]]; + } + rownnz[i] = num; + } else { + // spatial tendon: collect used dofs from wrap object bodies + int nnz = 0; + for (int j = 0; j < num; j++) { + int type = m->wrap_type[adr + j]; + + // get body id from site or geom wrap object + int bodyid = -1; + if (type == mjWRAP_SITE) { + bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; + } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { + bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; + } + + // walk up the body tree, collecting used dofs + if (bodyid > 0) { + int bid = bodyid; + while (bid > 0) { + int bdofadr = m->body_dofadr[bid]; + int bdofnum = m->body_dofnum[bid]; + for (int k = 0; k < bdofnum; k++) { + int dof = bdofadr + k; + + // check if dof already in colind + int found = 0; + for (int l = 0; l < nnz; l++) { + if (colind[rowadr[i] + l] == dof) { + found = 1; + break; + } + } + + // append new dof + if (!found) { + colind[rowadr[i] + nnz] = dof; + nnz++; + } + } + bid = m->body_parentid[bid]; + } + } + } + rownnz[i] = nnz; + } + + // sort colind for this tendon + int nnz = rownnz[i]; + for (int j = 0; j < nnz - 1; j++) { + for (int k = j + 1; k < nnz; k++) { + // swap out-of-order entries + if (colind[rowadr[i] + k] < colind[rowadr[i] + j]) { + int tmp = colind[rowadr[i] + j]; + colind[rowadr[i] + j] = colind[rowadr[i] + k]; + colind[rowadr[i] + k] = tmp; + } + } + } + } +} + // compute flex sparsity: flexedge_J_{rowadr,rownnz,colind} and flexvert_J_{rowadr,rownnz} static void makeFlexSparse(mjModel* m, mjData* d) { int nv = m->nv; @@ -552,6 +639,7 @@ static void mj_alignFlex(mjModel* m, mjData* d) { // set quantities that depend on qpos0 static void set0(mjModel* m, mjData* d) { + makeTendonSparse(m); makeFlexSparse(m, d); mj_alignFlex(m, d); int nv = m->nv; @@ -756,16 +844,7 @@ static void set0(mjModel* m, mjData* d) { // compute tendon_invweight0 for (int i=0; i < m->ntendon; i++) { - // make dense vector into tmp - if (mj_isSparse(m)) { - mju_zero(tmp, nv); - int end = d->ten_J_rowadr[i] + d->ten_J_rownnz[i]; - for (int j=d->ten_J_rowadr[i]; j < end; j++) { - tmp[d->ten_J_colind[j]] = d->ten_J[j]; - } - } else { - mju_copy(tmp, d->ten_J+i*nv, nv); - } + mju_sparse2dense(tmp, d->ten_J, 1, nv, m->ten_J_rownnz+i, m->ten_J_rowadr+i, m->ten_J_colind); // solve into tmp+nv mj_solveM(m, d, tmp+nv, tmp, 1); diff --git a/src/experimental/filament/CMakeLists.txt b/src/experimental/filament/CMakeLists.txt index 99854dfc..17ebec93 100644 --- a/src/experimental/filament/CMakeLists.txt +++ b/src/experimental/filament/CMakeLists.txt @@ -21,7 +21,7 @@ target_compile_definitions(${MUJOCO_FILAMENT_TARGET_NAME} PRIVATE MJ_STATIC) target_sources(${MUJOCO_FILAMENT_TARGET_NAME} PUBLIC - stubs.cc + mjr_compat.cc render_context_filament.h render_context_filament.cc filament/buffer_util.cc diff --git a/src/experimental/filament/filament/filament_context.cc b/src/experimental/filament/filament/filament_context.cc index e02defb9..844d3cf1 100644 --- a/src/experimental/filament/filament/filament_context.cc +++ b/src/experimental/filament/filament/filament_context.cc @@ -40,6 +40,7 @@ #include "experimental/filament/filament/gui_view.h" #include "experimental/filament/filament/imgui_editor.h" #include "experimental/filament/filament/object_manager.h" +#include "experimental/filament/filament/model_util.h" #include "experimental/filament/filament/scene_view.h" #include "experimental/filament/filament/texture_util.h" #include "experimental/filament/render_context_filament.h" @@ -91,30 +92,29 @@ FilamentContext::FilamentContext(const mjrFilamentConfig* config, filament::Renderer::ClearOptions opts; opts.clear = true; opts.discard = true; - opts.clearColor = {0, 0, 0, 1}; + opts.clearColor = ReadElement(model_, "filament.clearColor", + filament::math::float4(0, 0, 0, 1)); renderer_->setClearOptions(opts); // Copy parameters from model to context. - if (model_) { - context_->shadowClip = model_->stat.extent * model_->vis.map.shadowclip; - context_->shadowScale = model_->vis.map.shadowscale; - context_->offWidth = model_->vis.global.offwidth; - context_->offHeight = model_->vis.global.offheight; - context_->offSamples = model_->vis.quality.offsamples; - context_->fogStart = - (float)(model_->stat.extent * model_->vis.map.fogstart); - context_->fogEnd = (float)(model_->stat.extent * model_->vis.map.fogend); - context_->fogRGBA[0] = model_->vis.rgba.fog[0]; - context_->fogRGBA[1] = model_->vis.rgba.fog[1]; - context_->fogRGBA[2] = model_->vis.rgba.fog[2]; - context_->fogRGBA[3] = model_->vis.rgba.fog[3]; - context_->lineWidth = model_->vis.global.linewidth; - context_->shadowSize = model_->vis.quality.shadowsize; - context_->readPixelFormat = 0x1907; // 0x1907 = GL_RGB; - context_->ntexture = model_->ntex; - for (int i = 0; i < model_->ntex; ++i) { - context_->textureType[i] = model_->tex_type[i]; - } + context_->shadowClip = model_->stat.extent * model_->vis.map.shadowclip; + context_->shadowScale = model_->vis.map.shadowscale; + context_->offWidth = model_->vis.global.offwidth; + context_->offHeight = model_->vis.global.offheight; + context_->offSamples = model_->vis.quality.offsamples; + context_->fogStart = + (float)(model_->stat.extent * model_->vis.map.fogstart); + context_->fogEnd = (float)(model_->stat.extent * model_->vis.map.fogend); + context_->fogRGBA[0] = model_->vis.rgba.fog[0]; + context_->fogRGBA[1] = model_->vis.rgba.fog[1]; + context_->fogRGBA[2] = model_->vis.rgba.fog[2]; + context_->fogRGBA[3] = model_->vis.rgba.fog[3]; + context_->lineWidth = model_->vis.global.linewidth; + context_->shadowSize = model_->vis.quality.shadowsize; + context_->readPixelFormat = 0x1907; // 0x1907 = GL_RGB; + context_->ntexture = model_->ntex; + for (int i = 0; i < model_->ntex; ++i) { + context_->textureType[i] = model_->tex_type[i]; } scene_view_ = std::make_unique(engine_, object_manager_.get()); diff --git a/src/experimental/filament/stubs.cc b/src/experimental/filament/mjr_compat.cc similarity index 65% rename from src/experimental/filament/stubs.cc rename to src/experimental/filament/mjr_compat.cc index 9d2d099c..9d06dea6 100644 --- a/src/experimental/filament/stubs.cc +++ b/src/experimental/filament/mjr_compat.cc @@ -13,9 +13,50 @@ // limitations under the License. #include +#include "experimental/filament/render_context_filament.h" + +// This library implements the entirety of mujoco's mjr API. You can link this +// library with your application (instead of the "classic" mujoco renderer) to +// use the same APIs but with Filament rendering instead. +// +// However, you should consider using filament's mjrf API directly as it will +// provide you with access to more features and optimizations. extern "C" { +// mjr functions that are supported by the filament renderer. + +void mjr_defaultContext(mjrContext* con) { + mjrf_defaultContext(con); +} +void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale) { + mjrf_makeContext(m, con, fontscale); +} +void mjr_freeContext(mjrContext* con) { + mjrf_freeContext(con); +} +void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { + mjrf_render(viewport, scn, con); +} +void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { + mjrf_uploadMesh(m, con, meshid); +} +void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { + mjrf_uploadTexture(m, con, texid); +} +void mjr_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { + mjrf_uploadHField(m, con, hfieldid); +} +void mjr_setBuffer(int framebuffer, mjrContext* con) { + mjrf_setBuffer(framebuffer, con); +} +void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, + const mjrContext* con) { + mjrf_readPixels(rgb, depth, viewport, con); +} + +// mjr functions that are NOT supported by the filament renderer. + void mjr_setAux(int index, const mjrContext* con) { mju_error("mjr_setAux not implemented."); } diff --git a/src/experimental/filament/render_context_filament.cc b/src/experimental/filament/render_context_filament.cc index 4695599f..3955a816 100644 --- a/src/experimental/filament/render_context_filament.cc +++ b/src/experimental/filament/render_context_filament.cc @@ -32,17 +32,17 @@ static mujoco::FilamentContext* g_filament_context = nullptr; static void CheckFilamentContext() { if (g_filament_context == nullptr) { - mju_error("Missing context; did you call mjr_makeFilamentContext?"); + mju_error("Missing context; did you call mjrf_makeFilamentContext?"); } } extern "C" { -void mjr_defaultFilamentConfig(mjrFilamentConfig* config) { +void mjrf_defaultFilamentConfig(mjrFilamentConfig* config) { memset(config, 0, sizeof(mjrFilamentConfig)); } -void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, +void mjrf_makeFilamentContext(const mjModel* m, mjrContext* con, const mjrFilamentConfig* config) { // TODO: Support multiple contexts and multiple threads. For now, we'll just // assume a single, global context. @@ -52,16 +52,18 @@ void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, g_filament_context = new mujoco::FilamentContext(config, m, con); } -void mjr_defaultContext(mjrContext* con) { memset(con, 0, sizeof(mjrContext)); } - -void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale) { - mjr_freeContext(con); - mjrFilamentConfig cfg; - mjr_defaultFilamentConfig(&cfg); - mjr_makeFilamentContext(m, con, &cfg); +void mjrf_defaultContext(mjrContext* con) { + memset(con, 0, sizeof(mjrContext)); } -void mjr_freeContext(mjrContext* con) { +void mjrf_makeContext(const mjModel* m, mjrContext* con, int fontscale) { + mjr_freeContext(con); + mjrFilamentConfig cfg; + mjrf_defaultFilamentConfig(&cfg); + mjrf_makeFilamentContext(m, con, &cfg); +} + +void mjrf_freeContext(mjrContext* con) { // mjr_freeContext may be called multiple times. if (g_filament_context) { delete g_filament_context; @@ -70,50 +72,50 @@ void mjr_freeContext(mjrContext* con) { mjr_defaultContext(con); } -void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { +void mjrf_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { CheckFilamentContext(); g_filament_context->Render(viewport, scn, con); } -void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { +void mjrf_uploadMesh(const mjModel* m, const mjrContext* con, int meshid) { CheckFilamentContext(); g_filament_context->UploadMesh(m, meshid); } -void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { +void mjrf_uploadTexture(const mjModel* m, const mjrContext* con, int texid) { CheckFilamentContext(); g_filament_context->UploadTexture(m, texid); } -void mjr_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { +void mjrf_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid) { CheckFilamentContext(); g_filament_context->UploadHeightField(m, hfieldid); } -void mjr_setBuffer(int framebuffer, mjrContext* con) { +void mjrf_setBuffer(int framebuffer, mjrContext* con) { CheckFilamentContext(); g_filament_context->SetFrameBuffer(framebuffer); } -void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, +void mjrf_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, const mjrContext* con) { CheckFilamentContext(); g_filament_context->ReadPixels(viewport, rgb, depth); } -uintptr_t mjr_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, +uintptr_t mjrf_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, int width, int height, int bpp, const mjrContext* con) { CheckFilamentContext(); return g_filament_context->UploadGuiImage(tex_id, pixels, width, height, bpp); } -double mjr_getFrameRate(const mjrContext* con) { +double mjrf_getFrameRate(const mjrContext* con) { CheckFilamentContext(); return g_filament_context->GetFrameRate(); } -void mjr_updateGui(const mjrContext* con) { +void mjrf_updateGui(const mjrContext* con) { if (g_filament_context != nullptr) { g_filament_context->UpdateGui(); } diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h index c0423e4a..aef8f182 100644 --- a/src/experimental/filament/render_context_filament.h +++ b/src/experimental/filament/render_context_filament.h @@ -26,6 +26,9 @@ extern "C" { #endif +// IMPORTANT: This API should still be considered experimental and is likely +// change frequently. + typedef enum mjtGraphicsApi_ { // backend graphics API to use mjGFX_DEFAULT = 0, // default based on platform mjGFX_OPENGL, // OpenGL (desktop) @@ -43,35 +46,37 @@ struct mjrFilamentConfig { bool enable_gui; }; -void mjr_defaultFilamentConfig(mjrFilamentConfig* config); +void mjrf_defaultFilamentConfig(mjrFilamentConfig* config); -void mjr_makeFilamentContext(const mjModel* m, mjrContext* con, - const mjrFilamentConfig* config); +void mjrf_makeFilamentContext(const mjModel* m, mjrContext* con, + const mjrFilamentConfig* config); -MJAPI void mjr_defaultContext(mjrContext* con); +void mjrf_defaultContext(mjrContext* con); -MJAPI void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale); +void mjrf_makeContext(const mjModel* m, mjrContext* con, int fontscale); -MJAPI void mjr_freeContext(mjrContext* con); +void mjrf_freeContext(mjrContext* con); -MJAPI void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con); +void mjrf_render(mjrRect viewport, mjvScene* scn, const mjrContext* con); -MJAPI void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid); +void mjrf_uploadMesh(const mjModel* m, const mjrContext* con, int meshid); -MJAPI void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid); +void mjrf_uploadTexture(const mjModel* m, const mjrContext* con, int texid); -MJAPI void mjr_setBuffer(int framebuffer, mjrContext* con); +void mjrf_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid); -MJAPI void mjr_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, - const mjrContext* con); +void mjrf_setBuffer(int framebuffer, mjrContext* con); -double mjr_getFrameRate(const mjrContext* con); +void mjrf_readPixels(unsigned char* rgb, float* depth, mjrRect viewport, + const mjrContext* con); -uintptr_t mjr_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, - int width, int height, int bpp, - const mjrContext* con); +double mjrf_getFrameRate(const mjrContext* con); -void mjr_updateGui(const mjrContext* con); +uintptr_t mjrf_uploadGuiImage(uintptr_t tex_id, const unsigned char* pixels, + int width, int height, int bpp, + const mjrContext* con); + +void mjrf_updateGui(const mjrContext* con); #if defined(__cplusplus) } // extern "C" diff --git a/src/experimental/platform/CMakeLists.txt b/src/experimental/platform/CMakeLists.txt index 6429db47..8ba7aa6f 100644 --- a/src/experimental/platform/CMakeLists.txt +++ b/src/experimental/platform/CMakeLists.txt @@ -42,6 +42,8 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} file_dialog.h gui.cc gui.h + gui_spec.cc + gui_spec.h helpers.cc helpers.h imgui_widgets.cc @@ -67,6 +69,13 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME} window.h ) +if(NOT WIN32) + target_sources(${MUJOCO_PLATFORM_TARGET_NAME} + PUBLIC + object_launcher_plugin.cc + ) +endif() + if(APPLE) set_source_files_properties(window_osx.mm PROPERTIES COMPILE_FLAGS "-x objective-c++") @@ -78,7 +87,7 @@ elseif(UNIX AND NOT APPLE) target_sources(${MUJOCO_PLATFORM_TARGET_NAME} PUBLIC file_dialog_zenity.cc ) -elseif(WINDOWS) +elseif(WIN32) target_sources(${MUJOCO_PLATFORM_TARGET_NAME} PUBLIC file_dialog_win.cc ) diff --git a/src/experimental/platform/gui.cc b/src/experimental/platform/gui.cc index 5cfc9f2b..2d29eeca 100644 --- a/src/experimental/platform/gui.cc +++ b/src/experimental/platform/gui.cc @@ -470,7 +470,7 @@ void StateGui(const mjModel* model, mjData* data, std::vector& state, ImVec2(0, ImGui::GetTextLineHeightWithSpacing() * 20))) { ImGui::TableSetupColumn("Index"); ImGui::TableSetupColumn("Name"); - ImGui::TableSetupColumn("Value"); + ImGui::TableSetupColumn("Value", ImGuiTableColumnFlags_WidthStretch); ImGui::TableSetupScrollFreeze(0, 1); ImGui::TableHeadersRow(); @@ -1107,111 +1107,4 @@ void StatsGui(const mjModel* model, const mjData* data, bool paused, ImGui::Columns(); } -void BodyPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsBody* body = mjs_asBody(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(body->element); - if (name.empty()) { - name = "(Body " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - - ImGui::Text("Name"); - ImGui::Text("xpos[0]"); - ImGui::Text("xpos[1]"); - ImGui::Text("xpos[2]"); - ImGui::Text("xquat[0]"); - ImGui::Text("xquat[1]"); - ImGui::Text("xquat[2]"); - ImGui::Text("xquat[3]"); - ImGui::Text("mass"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); - ImGui::Text("%f", data->xpos[3*id+0]); - ImGui::Text("%f", data->xpos[3*id+1]); - ImGui::Text("%f", data->xpos[3*id+2]); - ImGui::Text("%f", data->xquat[4*id+0]); - ImGui::Text("%f", data->xquat[4*id+1]); - ImGui::Text("%f", data->xquat[4*id+2]); - ImGui::Text("%f", data->xquat[4*id+3]); - ImGui::Text("%f", model->body_mass[id]); -} - -void JointPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsJoint* joint = mjs_asJoint(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(joint->element); - if (name.empty()) { - name = "(Joint " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - ImGui::Text("Name"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); -} - -void SitePropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id) { - const mjsSite* site = mjs_asSite(element); - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.4f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.6f); - - std::string name = *mjs_getName(site->element); - if (name.empty()) { - name = "(Joint " + std::to_string(id) + ")"; - } - - ImGui::Columns(2); - ImGui::SetColumnWidth(0, ImGui::GetWindowWidth() * 0.3f); - ImGui::SetColumnWidth(1, ImGui::GetWindowWidth() * 0.7f); - ImGui::Text("Name"); - ImGui::Text("site_xpos[0]"); - ImGui::Text("site_xpos[1]"); - ImGui::Text("site_xpos[2]"); - ImGui::Text("site_xmat[0]"); - ImGui::Text("site_xmat[1]"); - ImGui::Text("site_xmat[2]"); - ImGui::Text("site_xmat[3]"); - ImGui::Text("site_xmat[4]"); - ImGui::Text("site_xmat[5]"); - ImGui::Text("site_xmat[6]"); - ImGui::Text("site_xmat[7]"); - ImGui::Text("site_xmat[8]"); - - ImGui::NextColumn(); - ImGui::Text("%s", name.c_str()); - ImGui::Text("%f", data->site_xpos[3*id+0]); - ImGui::Text("%f", data->site_xpos[3*id+1]); - ImGui::Text("%f", data->site_xpos[3*id+2]); - ImGui::Text("%f", data->site_xmat[4*id+0]); - ImGui::Text("%f", data->site_xmat[4*id+1]); - ImGui::Text("%f", data->site_xmat[4*id+2]); - ImGui::Text("%f", data->site_xmat[4*id+3]); - ImGui::Text("%f", data->site_xmat[4*id+4]); - ImGui::Text("%f", data->site_xmat[4*id+5]); - ImGui::Text("%f", data->site_xmat[4*id+6]); - ImGui::Text("%f", data->site_xmat[4*id+7]); - ImGui::Text("%f", data->site_xmat[4*id+8]); -} - } // namespace mujoco::platform diff --git a/src/experimental/platform/gui.h b/src/experimental/platform/gui.h index 2341023e..7dbdae09 100644 --- a/src/experimental/platform/gui.h +++ b/src/experimental/platform/gui.h @@ -116,14 +116,6 @@ void CountsGui(const mjModel* model, mjData* data); // FPS needs to be tracked by the caller and passed here to be displayed. void StatsGui(const mjModel* model, const mjData* data, bool paused, float fps); -// UX for displaying properties of various mjSpec elements. -void BodyPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); -void JointPropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); -void SitePropertiesGui(const mjModel* model, const mjData* data, - mjsElement* element, int id); - } // namespace mujoco::platform #endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_H_ diff --git a/src/experimental/platform/gui_spec.cc b/src/experimental/platform/gui_spec.cc new file mode 100644 index 00000000..eeb74581 --- /dev/null +++ b/src/experimental/platform/gui_spec.cc @@ -0,0 +1,847 @@ +// Copyright 2026 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 "experimental/platform/gui_spec.h" + +#include +#include + +#include +#include +#include +#include "experimental/platform/imgui_widgets.h" + +// Define the mujoco X macros to add fields to the ImGui_DataTable. +// We limit the fields to the ones with a matching element by comparing the +// array size field (e.g. nbody) with the MATCH constexpr value. +#define X(TYPE, NAME, NELEM, SIZE) \ + if constexpr (#NELEM == MATCH) table(#NAME, ptr->NAME, SIZE); + +// Simple wrapper around MJMODEL_POINTERS that prepares values we need for the +// X macro above. +#define MJMODEL_POINTERS_X(M) \ + constexpr std::string_view MATCH(#M); \ + const auto* ptr = model; \ + MJMODEL_POINTERS + +// Simple wrapper around MJDATA_POINTERS that prepares values we need for the +// X macro above. +#define MJDATA_POINTERS_X(M) \ + constexpr std::string_view MATCH(#M); \ + const auto* ptr = data; \ + MJDATA_POINTERS + +namespace mujoco::platform { + +// Returns the index of the element in the spec. This is different from +// mjs_getId which returns the runtime ID of an element. +static int GetElementIndexInSpec(mjsElement* element) { + int n = 0; + mjSpec* spec = mjs_getSpec(element); + mjsElement* iter = mjs_firstElement(spec, element->elemtype); + while (iter) { + if (iter == element) { + return n; + } + iter = mjs_nextElement(spec, iter); + ++n; + } + return -1; +} + +// Returns a name for the element; either the element has a name, or we +// construct a unique name from the element's id (using mjs_getId) or index +// (using GetElementIndexInSpec). +static std::string ElementName(mjsElement* element) { + const mjString* name = mjs_getName(element); + std::string label = *name; + if (label.empty()) { + int id = mjs_getId(element); + if (id == -1) { + id = GetElementIndexInSpec(element); + } + const char* type_name = mju_type2Str(element->elemtype); + label = "(" + std::string(type_name) + " " + std::to_string(id) + ")"; + } + return label; +} + +static void QuatOrOrientation(ImGui_DataTable& table, const double quat[4], + const mjsOrientation& orientation, + const char* quat_name, const char* alt_name) { + auto alt = + [&](const char* label) { return std::string(alt_name) + "." + label; }; + + switch (orientation.type) { + case mjORIENTATION_QUAT: + table(quat_name, quat, 4); + break; + case mjORIENTATION_AXISANGLE: + table(alt("axisangle").c_str(), orientation.axisangle, 4); + break; + case mjORIENTATION_XYAXES: + table(alt("xyaxes").c_str(), orientation.xyaxes, 6); + break; + case mjORIENTATION_ZAXIS: + table(alt("zaxis").c_str(), orientation.zaxis, 3); + break; + case mjORIENTATION_EULER: + table(alt("euler").c_str(), orientation.euler, 3); + break; + } +} + +static void AddDeleteButton(mjsElement* element, + const SpecElementCallbackFn& on_delete) { + if (on_delete) { + // Right-align the delete button. + const float button_width = ImGui::CalcTextSize(ICON_FA_TRASH_CAN).x + + ImGui::GetStyle().FramePadding.x * 2.0f; + ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); + if (ImGui::SmallButton(ICON_FA_TRASH_CAN)) { + on_delete(element); + } + } +} + +static void SelectableElement(mjsElement* element, + mjsElement** selected_element, + const SpecElementCallbackFn& on_delete) { + constexpr ImGuiSelectableFlags flags = ImGuiSelectableFlags_AllowOverlap; + + const std::string name = ElementName(element); + const bool selected = (element == *selected_element); + if (ImGui::Selectable(name.c_str(), selected, flags)) { + *selected_element = element; + } + if (selected) { + AddDeleteButton(element, on_delete); + } +} + +static void BodyChildrenGui(const char* heading, mjtObj type, + mjsElement** element, mjsBody* body, + const SpecElementCallbackFn& on_delete) { + mjsElement* iter = mjs_firstChild(body, type, 0); + if (!iter) { + return; + } + + constexpr ImGuiTreeNodeFlags tree_flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_DrawLinesFull; + if (ImGui::TreeNodeEx(heading, tree_flags)) { + while (iter) { + SelectableElement(iter, element, on_delete); + iter = mjs_nextChild(body, iter, 0); + } + ImGui::TreePop(); + } +} + +static void ElementListGui(const char* heading, mjtObj type, + mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete) { + mjsElement* iter = mjs_firstElement(spec, type); + if (!iter) { + return; + } + + constexpr ImGuiTreeNodeFlags tree_flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + if (ImGui::TreeNodeEx(heading, tree_flags)) { + while (iter) { + SelectableElement(iter, element, on_delete); + iter = mjs_nextElement(spec, iter); + } + ImGui::TreePop(); + } +} + +static void BodyTreeGuiRecursive(mjsElement** element, mjsBody* body, + const SpecElementCallbackFn& on_delete) { + const std::string label = ElementName(body->element); + + ImGui::PushID(body); + + ImGuiTreeNodeFlags flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed | + ImGuiTreeNodeFlags_DrawLinesFull | ImGuiTreeNodeFlags_AllowOverlap; + + if (*element == body->element) { + flags |= ImGuiTreeNodeFlags_Selected; + } + + const bool tree_open = ImGui::TreeNodeEx(label.c_str(), flags); + if (ImGui::IsItemClicked()) { + *element = body->element; + } + if (*element == body->element) { + AddDeleteButton(body->element, on_delete); + } + + if (tree_open) { + mjsElement* iter = mjs_firstChild(body, mjOBJ_BODY, 0); + while (iter) { + BodyTreeGuiRecursive(element, mjs_asBody(iter), on_delete); + iter = mjs_nextChild(body, iter, 0); + } + + BodyChildrenGui("Frames", mjOBJ_FRAME, element, body, on_delete); + BodyChildrenGui("Sites", mjOBJ_SITE, element, body, on_delete); + BodyChildrenGui("Joints", mjOBJ_JOINT, element, body, on_delete); + BodyChildrenGui("Geoms", mjOBJ_GEOM, element, body, on_delete); + BodyChildrenGui("Lights", mjOBJ_LIGHT, element, body, on_delete); + BodyChildrenGui("Cameras", mjOBJ_CAMERA, element, body, on_delete); + + ImGui::TreePop(); + } + + ImGui::PopID(); +} + +void SpecExplorerGui(mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete) { + const ImGuiTreeNodeFlags flags = + ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + + if (ImGui::TreeNodeEx("Body Tree", flags)) { + mjsElement* root = mjs_firstElement(spec, mjOBJ_BODY); + if (root) { + mjsBody* body = mjs_asBody(root); + if (body) { + BodyTreeGuiRecursive(element, body, on_delete); + } + } + ImGui::TreePop(); + } + + auto list = [&](const char* heading, mjtObj type) { + ElementListGui(heading, type, element, spec, on_delete); + }; + + ImGui::PushID(spec); + + // Non-tree elements. + if (ImGui::TreeNodeEx("Elements", flags)) { + list("Actuators", mjOBJ_ACTUATOR); + list("Sensors", mjOBJ_SENSOR); + list("Flexes", mjOBJ_FLEX); + list("Tendons", mjOBJ_TENDON); + list("Pair", mjOBJ_PAIR); + list("Exclude", mjOBJ_EXCLUDE); + list("Equality", mjOBJ_EQUALITY); + list("Numeric", mjOBJ_NUMERIC); + list("Text", mjOBJ_TEXT); + list("Tuple", mjOBJ_TUPLE); + list("Key", mjOBJ_KEY); + list("Default", mjOBJ_DEFAULT); + ImGui::TreePop(); + } + + // Assets. + if (ImGui::TreeNodeEx("Assets", flags)) { + list("Meshes", mjOBJ_MESH); + list("Height Fields", mjOBJ_HFIELD); + list("Skins", mjOBJ_SKIN); + list("Textures", mjOBJ_TEXTURE); + list("Materials", mjOBJ_MATERIAL); + ImGui::TreePop(); + } + + ImGui::PopID(); +} + +void ElementSpecGui(const mjSpec* spec, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + + switch (element->elemtype) { + case mjOBJ_BODY: { + const mjsBody* body = mjs_asBody(element); + table("childclass", body->childclass, 1); // childclass name + table("pos", body->pos, 3); // frame position + QuatOrOrientation(table, body->quat, body->alt, "quat", "alt"); // frame orientation + table("ipos", body->ipos, 3); // inertial frame position + QuatOrOrientation(table, body->iquat, body->ialt, "iquat", "ialt"); // inertial frame orientation + table("mass", body->mass, 1); // mass + table("inertia", body->inertia, 3); // diagonal inertia (in i-frame) + table("fullinertia", body->fullinertia, 6); // non-axis-aligned inertia matrix + table("mocap", body->mocap, 1); // is this a mocap body + table("gravcomp", body->gravcomp, 1); // gravity compensation + table("explicitinertial", body->explicitinertial, 1); // whether to save the body with explicit inertial clause + table("sleep", body->sleep, 1); // sleep policy + table("info", body->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_JOINT: { + mjsJoint* joint = mjs_asJoint(element); + table("pos", joint->pos, 3); // anchor position + table("axis", joint->axis, 3); // joint axis + table("ref", joint->ref, 1); // value at reference configuration: qpos0 + table("align", joint->align, 1); // align free joint with body com (mjtAlignFree) + table("stiffness", joint->stiffness, 1); // stiffness coefficient + table("springref", joint->springref, 1); // spring reference value: qpos_spring + table("springdamper", joint->springdamper, 2); // timeconst, dampratio + table("limited", joint->limited, 1); // does joint have limits (mjtLimited) + table("range", joint->range, 2); // joint limits + table("margin", joint->margin, 1); // margin value for joint limit detection + table("solref_limit", joint->solref_limit, mjNREF); // solver reference: joint limits + table("solimp_limit", joint->solimp_limit, mjNIMP); // solver impedance: joint limits + table("actfrclimited", joint->actfrclimited, 1); // are actuator forces on joint limited (mjtLimited) + table("actfrcrange", joint->actfrcrange, 2); // actuator force limits + table("armature", joint->armature, 1); // armature inertia (mass for slider) + table("damping", joint->damping, 1); // damping coefficient + table("frictionloss", joint->frictionloss, 1); // friction loss + table("solref_friction", joint->solref_friction, mjNREF); // solver reference: dof friction + table("solimp_friction", joint->solimp_friction, mjNIMP); // solver impedance: dof friction + table("group", joint->group, 1); // group + table("actgravcomp", joint->actgravcomp, 1); // is gravcomp force applied via actuators + table("info", joint->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_ACTUATOR: { + mjsActuator* actuator = mjs_asActuator(element); + table("gaintype", actuator->gaintype, 1); // gain type + table("gainprm", actuator->gainprm, mjNGAIN); // gain parameters + table("biastype", actuator->biastype, 1); // bias type + table("biasprm", actuator->biasprm, mjNGAIN); // bias parameters + table("dyntype", actuator->dyntype, 1); // dynamics type + table("dynprm", actuator->dynprm, mjNDYN); // dynamics parameters + table("actdim", actuator->actdim, 1); // number of activation variables + table("actearly", actuator->actearly, 1); // apply next activations to qfrc + table("trntype", actuator->trntype, 1); // transmission type + table("gear", actuator->gear, 6); // length and transmitted force scaling + table("target", actuator->target, 1); // name of transmission target + table("refsite", actuator->refsite, 1); // reference site, for site transmission + table("slidersite", actuator->slidersite, 1); // site defining cylinder, for slider-crank + table("cranklength", actuator->cranklength, 1); // crank length, for slider-crank + table("lengthrange", actuator->lengthrange, 2); // transmission length range + table("inheritrange", actuator->inheritrange, 1); // automatic range setting for position and intvelocity + table("ctrllimited", actuator->ctrllimited, 1); // are control limits defined (mjtLimited) + table("ctrlrange", actuator->ctrlrange, 2); // control range + table("forcelimited", actuator->forcelimited, 1); // are force limits defined (mjtLimited) + table("forcerange", actuator->forcerange, 2); // force range + table("actlimited", actuator->actlimited, 1); // are activation limits defined (mjtLimited) + table("actrange", actuator->actrange, 2); // activation range + table("group", actuator->group, 1); // group + table("nsample", actuator->nsample, 1); // number of samples in history buffer + table("interp", actuator->interp, 1); // interpolation order (0=ZOH, 1=linear, 2=cubic) + table("delay", actuator->delay, 1); // delay time in seconds; 0: no delay + table("info", actuator->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SENSOR: { + mjsSensor* sensor = mjs_asSensor(element); + table("type", sensor->type, 1); // type of sensor + table("objtype", sensor->objtype, 1); // type of sensorized object + table("objname", sensor->objname, 1); // name of sensorized object + table("reftype", sensor->reftype, 1); // type of referenced object + table("refname", sensor->refname, 1); // name of referenced object + table("intprm", sensor->intprm, mjNSENS); // integer parameters + table("datatype", sensor->datatype, 1); // data type for sensor measurement + table("needstage", sensor->needstage, 1); // compute stage needed to simulate sensor + table("dim", sensor->dim, 1); // number of scalar outputs + table("cutoff", sensor->cutoff, 1); // cutoff for real and positive datatypes + table("noise", sensor->noise, 1); // noise stdev + table("nsample", sensor->nsample, 1); // number of samples in history buffer + table("interp", sensor->interp, 1); // interpolation order (0=ZOH, 1=linear, 2=cubic) + table("delay", sensor->delay, 1); // delay time in seconds + table("interval", sensor->interval, 2); // [period, time_prev] in seconds + table("info", sensor->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SITE: { + mjsSite* site = mjs_asSite(element); + table("pos", site->pos, 3); // position + QuatOrOrientation(table, site->quat, site->alt, "quat", "alt"); // orientation + table("fromto", site->fromto, 6); // alternative for capsule, cylinder, box, ellipsoid + table("size", site->size, 3); // geom size + table("type", site->type, 1); // geom type + table("material", site->material, 1); // name of material + table("group", site->group, 1); // group + table("rgba", site->rgba, 4); // rgba when material is omitted + table("info", site->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_FRAME: { + mjsFrame* frame = mjs_asFrame(element); + table("childclass", frame->childclass, 1); // childclass name + table("pos", frame->pos, 3); // position + QuatOrOrientation(table, frame->quat, frame->alt, "quat", "alt"); // orientation + table("info", frame->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_GEOM: { + mjsGeom* geom = mjs_asGeom(element); + table("type", geom->type, 1); // geom type + table("pos", geom->pos, 3); // position + QuatOrOrientation(table, geom->quat, geom->alt, "quat", "alt"); // orientation + table("fromto", geom->fromto, 6); // alternative for capsule, cylinder, box, ellipsoid + table("size", geom->size, 3); // type-specific size + table("contype", geom->contype, 1); // contact type + table("conaffinity", geom->conaffinity, 1); // contact affinity + table("condim", geom->condim, 1); // contact dimensionality + table("priority", geom->priority, 1); // contact priority + table("friction", geom->friction, 3); // one-sided friction coefficients: slide, roll, spin + table("solmix", geom->solmix, 1); // solver mixing for contact pairs + table("solref", geom->solref, mjNREF); // solver reference + table("solimp", geom->solimp, mjNIMP); // solver impedance + table("margin", geom->margin, 1); // margin for contact detection + table("gap", geom->gap, 1); // include in solver if dist < margin-gap + table("mass", geom->mass, 1); // used to compute density + table("density", geom->density, 1); // used to compute mass and inertia from volume or surface + table("typeinertia", geom->typeinertia, 1); // selects between surface and volume inertia + table("fluid_ellipsoid", geom->fluid_ellipsoid, 1); // whether ellipsoid-fluid model is active + table("fluid_coefs", geom->fluid_coefs, 5); // ellipsoid-fluid interaction coefs + table("material", geom->material, 1); // name of material + table("rgba", geom->rgba, 4); // rgba when material is omitted + table("group", geom->group, 1); // group + table("hfieldname", geom->hfieldname, 1); // heightfield attached to geom + table("meshname", geom->meshname, 1); // mesh attached to geom + table("fitscale", geom->fitscale, 1); // scale mesh uniformly + table("info", geom->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_LIGHT: { + mjsLight* light = mjs_asLight(element); + table("pos", light->pos, 3); // position + table("dir", light->dir, 3); // direction + table("mode", light->mode, 1); // tracking mode + table("targetbody", light->targetbody, 1); // target body for targeting + table("active", light->active, 1); // is light active + table("type", light->type, 1); // type of light + table("texture", light->texture, 1); // texture name for image lights + table("castshadow", light->castshadow, 1); // does light cast shadows + table("bulbradius", light->bulbradius, 1); // bulb radius, for soft shadows + table("intensity", light->intensity, 1); // intensity, in candelas + table("range", light->range, 1); // range of effectiveness + table("attenuation", light->attenuation, 3); // OpenGL attenuation (quadratic model) + table("cutoff", light->cutoff, 1); // OpenGL cutoff + table("exponent", light->exponent, 1); // OpenGL exponent + table("ambient", light->ambient, 3); // ambient color + table("diffuse", light->diffuse, 3); // diffuse color + table("specular", light->specular, 3); // specular color + table("info", light->info, 1); // message appended to compiler errorsx + break; + } + case mjOBJ_CAMERA: { + mjsCamera* camera = mjs_asCamera(element); + table("pos", camera->pos, 3); // position + QuatOrOrientation(table, camera->quat, camera->alt, "quat", "alt"); // orientation + table("mode", camera->mode, 1); // tracking mode + table("targetbody", camera->targetbody, 1); // target body for tracking/targeting + table("proj", camera->proj, 1); // camera projection type + table("resolution", camera->resolution, 2); // resolution (pixel) + table("output", camera->output, 1); // bit flags for output type + table("fovy", camera->fovy, 1); // y-field of view + table("ipd", camera->ipd, 1); // inter-pupillary distance + table("intrinsic", camera->intrinsic, 4); // camera intrinsics (length) + table("sensor_size", camera->sensor_size, 2); // sensor size (length) + table("focal_length", camera->focal_length, 2); // focal length (length) + table("focal_pixel", camera->focal_pixel, 2); // focal length (pixel) + table("principal_length", camera->principal_length, 2); // principal point (length) + table("principal_pixel", camera->principal_pixel, 2); // principal point (pixel) + table("info", camera->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_MESH: { + mjsMesh* mesh = mjs_asMesh(element); + table("content_type", mesh->content_type, 1); // content type of file + table("file", mesh->file, 1); // mesh file + table("refpos", mesh->refpos, 3); // reference position + table("refquat", mesh->refquat, 4); // reference orientation + table("scale", mesh->scale, 3); // rescale mesh + table("inertia", mesh->inertia, 1); // inertia type (convex, legacy, exact, shell) + table("smoothnormal", mesh->smoothnormal, 1); // do not exclude large-angle faces from normals + table("needsdf", mesh->needsdf, 1); // compute sdf from mesh + table("maxhullvert", mesh->maxhullvert, 1); // maximum vertex count for the convex hull + table("material", mesh->material, 1); // name of material + table("info", mesh->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_HFIELD: { + mjsHField* hfield = mjs_asHField(element); + table("content_type", hfield->content_type, 1); // content type of file + table("file", hfield->file, 1); // file: (nrow, ncol, [elevation data]) + table("size", hfield->size, 4); // hfield size (ignore referencing geom size) + table("nrow", hfield->nrow, 1); // number of rows + table("ncol", hfield->ncol, 1); // number of columns + table("info", hfield->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_SKIN: { + mjsSkin* skin = mjs_asSkin(element); + table("file", skin->file, 1); // skin file + table("material", skin->material, 1); // name of material used for rendering + table("rgba", skin->rgba, 4); // rgba when material is omitted + table("inflate", skin->inflate, 1); // inflate in normal direction + table("group", skin->group, 1); // group for visualization + table("info", skin->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_FLEX: { + mjsFlex* flex = mjs_asFlex(element); + table("contype", flex->contype, 1); // contact type + table("conaffinity", flex->conaffinity, 1); // contact affinity + table("condim", flex->condim, 1); // contact dimensionality + table("priority", flex->priority, 1); // contact priority + table("friction", flex->friction, 3); // one-sided friction coefficients: slide, roll, spin + table("solmix", flex->solmix, 1); // solver mixing for contact pairs + table("solref", flex->solref, mjNREF); // solver reference + table("solimp", flex->solimp, mjNIMP); // solver impedance + table("margin", flex->margin, 1); // margin for contact detection + table("gap", flex->gap, 1); // include in solver if distdim, 1); // element dimensionality + table("radius", flex->radius, 1); // radius around primitive element + table("size", flex->size, 3); // vertex bounding box half sizes in qpos0 + table("internal", flex->internal, 1); // enable internal collisions + table("flatskin", flex->flatskin, 1); // render flex skin with flat shading + table("selfcollide", flex->selfcollide, 1); // mode for flex self collision + table("vertcollide", flex->vertcollide, 1); // mode for vertex collision + table("passive", flex->passive, 1); // mode for passive collisions + table("activelayers", flex->activelayers, 1); // number of active element layers in 3D + table("group", flex->group, 1); // group for visualization + table("edgestiffness", flex->edgestiffness, 1); // edge stiffness + table("edgedamping", flex->edgedamping, 1); // edge damping + table("rgba", flex->rgba, 4); // rgba when material is omitted + table("material", flex->material, 1); // name of material used for rendering + table("young", flex->young, 1); // Young's modulus + table("poisson", flex->poisson, 1); // Poisson's ratio + table("damping", flex->damping, 1); // Rayleigh's damping + table("thickness", flex->thickness, 1); // thickness (2D only) + table("elastic2d", flex->elastic2d, 1); // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both + table("info", flex->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_TENDON: { + mjsTendon* tendon = mjs_asTendon(element); + table("stiffness", tendon->stiffness, 1); // stiffness coefficient + table("springlength", tendon->springlength, 2); // spring resting length; {-1, -1}: use qpos_spring + table("damping", tendon->damping, 1); // damping coefficient + table("frictionloss", tendon->frictionloss, 1); // friction loss + table("solref_friction", tendon->solref_friction, mjNREF); // solver reference: tendon friction + table("solimp_friction", tendon->solimp_friction, mjNIMP); // solver impedance: tendon friction + table("armature", tendon->armature, 1); // inertia associated with tendon velocity + table("limited", tendon->limited, 1); // does tendon have limits (mjtLimited) + table("actfrclimited", tendon->actfrclimited, 1); // does tendon have actuator force limits + table("range", tendon->range, 2); // length limits + table("actfrcrange", tendon->actfrcrange, 2); // actuator force limits + table("margin", tendon->margin, 1); // margin value for tendon limit detection + table("solref_limit", tendon->solref_limit, mjNREF); // solver reference: tendon limits + table("solimp_limit", tendon->solimp_limit, mjNIMP); // solver impedance: tendon limits + table("material", tendon->material, 1); // name of material for rendering + table("width", tendon->width, 1); // width for rendering + table("rgba", tendon->rgba, 4); // rgba when material is omitted + table("group", tendon->group, 1); // group + table("info", tendon->info, 1); // message appended to errors + break; + } + case mjOBJ_TEXTURE: { + mjsTexture* texture = mjs_asTexture(element); + table("type", texture->type, 1); // texture type + table("colorspace", texture->colorspace, 1); // colorspace + table("builtin", texture->builtin, 1); // builtin type (mjtBuiltin) + table("mark", texture->mark, 1); // mark type (mjtMark) + table("rgb1", texture->rgb1, 3); // first color for builtin + table("rgb2", texture->rgb2, 3); // second color for builtin + table("markrgb", texture->markrgb, 3); // mark color + table("random", texture->random, 1); // probability of random dots + table("height", texture->height, 1); // height in pixels (square for cube and skybox) + table("width", texture->width, 1); // width in pixels + table("nchannel", texture->nchannel, 1); // number of channels + table("content_type", texture->content_type, 1); // content type of file + table("file", texture->file, 1); // png file to load; use for all sides of cube + table("gridsize", texture->gridsize, 2); // size of grid for composite file; (1,1)-repeat + // TODO: table("gridlayout", texture->gridlayout, 12); // row-major: L,R,F,B,U,D for faces; . for unused + table("cubefiles", texture->cubefiles, 1); // different file for each side of the cube + table("hflip", texture->hflip, 1); // horizontal flip + table("vflip", texture->vflip, 1); // vertical flip + table("info", texture->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_MATERIAL: { + mjsMaterial* material = mjs_asMaterial(element); + table("textures", material->textures, 1); // names of textures (empty: none) + table("texuniform", material->texuniform, 1); // make texture cube uniform + table("texrepeat", material->texrepeat, 2); // texture repetition for 2D mapping + table("emission", material->emission, 1); // emission + table("specular", material->specular, 1); // specular + table("shininess", material->shininess, 1); // shininess + table("reflectance", material->reflectance, 1); // reflectance + table("metallic", material->metallic, 1); // metallic + table("roughness", material->roughness, 1); // roughness + table("rgba", material->rgba, 4); // rgba + table("info", material->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_PAIR: { + mjsPair* pair = mjs_asPair(element); + table("geomname1", pair->geomname1, 1); // name of geom 1 + table("geomname2", pair->geomname2, 1); // name of geom 2 + table("condim", pair->condim, 1); // contact dimensionality + table("solref", pair->solref, mjNREF); // solver reference, normal direction + table("solreffriction", pair->solreffriction, mjNREF); // solver reference, frictional directions + table("solimp", pair->solimp, mjNIMP); // solver impedance + table("margin", pair->margin, 1); // margin for contact detection + table("gap", pair->gap, 1); // include in solver if distfriction, 5); // full contact friction + table("info", pair->info, 1); // message appended to errors + break; + } + case mjOBJ_EQUALITY: { + mjsEquality* equality = mjs_asEquality(element); + table("type", equality->type, 1); // constraint type + table("data", equality->data, mjNEQDATA); // type-dependent data + table("active", equality->active, 1); // is equality initially active + table("name1", equality->name1, 1); // name of object 1 + table("name2", equality->name2, 1); // name of object 2 + table("objtype", equality->objtype, 1); // type of both objects + table("solref", equality->solref, mjNREF); // solver reference + table("solimp", equality->solimp, mjNIMP); // solver impedance + table("info", equality->info, 1); // message appended to errors + break; + } + case mjOBJ_EXCLUDE: { + mjsExclude* exclude = mjs_asExclude(element); + table("bodyname1", exclude->bodyname1, 1); // name of geom 1 + table("bodyname2", exclude->bodyname2, 1); // name of geom 2 + table("info", exclude->info, 1); // message appended to errors + break; + } + case mjOBJ_NUMERIC: { + mjsNumeric* numeric = mjs_asNumeric(element); + table("data", numeric->data, 1); // initialization data + table("size", numeric->size, 1); // array size, can be bigger than data size + table("info", numeric->info, 1); // message appended to errors + break; + } + case mjOBJ_TEXT: { + mjsText* text = mjs_asText(element); + table("data", text->data, 1); // text string + table("info", text->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_TUPLE: { + mjsTuple* tuple = mjs_asTuple(element); + table("objtype", tuple->objtype, 1); // object types + table("objname", tuple->objname, 1); // object names + table("objprm", tuple->objprm, 1); // object parameters + table("info", tuple->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_KEY: { + mjsKey* key = mjs_asKey(element); + table("time", key->time, 1); // time + table("qpos", key->qpos, 1); // qpos + table("qvel", key->qvel, 1); // qvel + table("act", key->act, 1); // act + table("mpos", key->mpos, 1); // mocap pos + table("mquat", key->mquat, 1); // mocap quat + table("ctrl", key->ctrl, 1); // ctrl + table("info", key->info, 1); // message appended to compiler errors + break; + } + case mjOBJ_PLUGIN: { + mjsPlugin* plugin = mjs_asPlugin(element); + table("name", plugin->name, 1); // instance name + table("plugin_name", plugin->plugin_name, 1); // plugin name + table("active", plugin->active, 1); // is the plugin active + table("info", plugin->info, 1); // message appended to compiler errors + break; + } + default: + // ignore other types + break; + } +} + +void ElementModelGui(const mjModel* model, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + table.SetArrayIndex(mjs_getId(element)); + + MJMODEL_POINTERS_PREAMBLE(model); + switch (element->elemtype) { + case mjOBJ_BODY: { + table.SetPrefix("body_"); + MJMODEL_POINTERS_X(nbody) + break; + } + case mjOBJ_JOINT: { + table.SetPrefix("jnt_"); + MJMODEL_POINTERS_X(njnt) + break; + } + case mjOBJ_ACTUATOR: { + table.SetPrefix("actuator_"); + MJMODEL_POINTERS_X(nu) + break; + } + case mjOBJ_SENSOR: { + table.SetPrefix("sensor_"); + MJMODEL_POINTERS_X(nsensor) + break; + } + case mjOBJ_SITE: { + table.SetPrefix("site_"); + MJMODEL_POINTERS_X(nsite) + break; + } + case mjOBJ_GEOM: { + table.SetPrefix("geom_"); + MJMODEL_POINTERS_X(ngeom) + break; + } + case mjOBJ_LIGHT: { + table.SetPrefix("light_"); + MJMODEL_POINTERS_X(nlight) + break; + } + case mjOBJ_CAMERA: { + table.SetPrefix("cam_"); + MJMODEL_POINTERS_X(ncam) + break; + } + case mjOBJ_MESH: { + table.SetPrefix("mesh_"); + MJMODEL_POINTERS_X(nmesh) + break; + } + case mjOBJ_HFIELD: { + table.SetPrefix("hfield_"); + MJMODEL_POINTERS_X(nhfield) + break; + } + case mjOBJ_SKIN: { + table.SetPrefix("skin_"); + MJMODEL_POINTERS_X(nskin) + break; + } + case mjOBJ_FLEX: { + table.SetPrefix("flex_"); + MJMODEL_POINTERS_X(nflex) + break; + } + case mjOBJ_TENDON: { + table.SetPrefix("tendon_"); + MJMODEL_POINTERS_X(ntendon) + break; + } + case mjOBJ_TEXTURE: { + table.SetPrefix("tex_"); + MJMODEL_POINTERS_X(ntex) + break; + } + case mjOBJ_MATERIAL: { + table.SetPrefix("mat_"); + MJMODEL_POINTERS_X(nmat) + break; + } + default: + // ignore other types + break; + } +} + +void ElementDataGui(const mjData* data, mjsElement* element) { + if (element == nullptr) { + return; + } + + ImGui_DataTable table; + table("Name", ElementName(element).c_str(), 1); + table.SetArrayIndex(mjs_getId(element)); + + switch (element->elemtype) { + case mjOBJ_BODY: { + MJDATA_POINTERS_X(nbody) + break; + } + case mjOBJ_JOINT: { + MJDATA_POINTERS_X(njnt) + break; + } + case mjOBJ_SITE: { + table.SetPrefix("site_"); + MJDATA_POINTERS_X(nsite) + break; + } + case mjOBJ_GEOM: { + table.SetPrefix("geom_"); + MJDATA_POINTERS_X(ngeom) + break; + } + case mjOBJ_CAMERA: { + table.SetPrefix("cam_"); + MJDATA_POINTERS_X(ncam) + break; + } + case mjOBJ_LIGHT: { + table.SetPrefix("light_"); + MJDATA_POINTERS_X(nlight) + break; + } + case mjOBJ_SENSOR: { + table.SetPrefix("sensor_"); + MJDATA_POINTERS_X(nsensor) + break; + } + case mjOBJ_MESH: { + table.SetPrefix("mesh_"); + MJDATA_POINTERS_X(nmesh) + break; + } + case mjOBJ_HFIELD: { + table.SetPrefix("hfield_"); + MJDATA_POINTERS_X(nhfield) + break; + } + case mjOBJ_SKIN: { + table.SetPrefix("skin_"); + MJDATA_POINTERS_X(nskin) + break; + } + case mjOBJ_FLEX: { + table.SetPrefix("flex_"); + MJDATA_POINTERS_X(nflex) + break; + } + case mjOBJ_TENDON: { + table.SetPrefix("ten_"); + MJDATA_POINTERS_X(ntendon) + break; + } + case mjOBJ_TEXTURE: { + table.SetPrefix("tex_"); + MJDATA_POINTERS_X(ntex) + break; + } + case mjOBJ_MATERIAL: { + table.SetPrefix("mat_"); + MJDATA_POINTERS_X(nmat) + break; + } + default: + break; + } +} +} // namespace mujoco::platform diff --git a/src/experimental/platform/gui_spec.h b/src/experimental/platform/gui_spec.h new file mode 100644 index 00000000..e1094583 --- /dev/null +++ b/src/experimental/platform/gui_spec.h @@ -0,0 +1,37 @@ +// Copyright 2026 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_PLATFORM_GUI_SPEC_H_ +#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ + +#include + +#include + +namespace mujoco::platform { + +using SpecElementCallbackFn = std::function; + +// UX for displaying the spec as a tree. +void SpecExplorerGui(mjsElement** element, mjSpec* spec, + const SpecElementCallbackFn& on_delete); + +// UX for displaying the properties of an mjSpec element. +void ElementSpecGui(const mjSpec* spec, mjsElement* element); +void ElementDataGui(const mjData* data, mjsElement* element); +void ElementModelGui(const mjModel* model, mjsElement* element); + +} // namespace mujoco::platform + +#endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ diff --git a/src/experimental/platform/imgui_widgets.cc b/src/experimental/platform/imgui_widgets.cc index 10a7b374..def513e9 100644 --- a/src/experimental/platform/imgui_widgets.cc +++ b/src/experimental/platform/imgui_widgets.cc @@ -14,9 +14,13 @@ #include "experimental/platform/imgui_widgets.h" +#include +#include #include #include +#include #include +#include #include #include @@ -57,6 +61,228 @@ KeyValues ReadIniSection(const std::string& contents, return key_values; } +ImGui_DataTable::ImGui_DataTable(float w1, float w2) { + ImGui::BeginTable("##PropertiesTable", 2); + const float width = ImGui::GetContentRegionAvail().x; + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, width * w1); + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, width * w2); +} +ImGui_DataTable::~ImGui_DataTable() { ImGui::EndTable(); } + +void ImGui_DataTable::SetArrayIndex(int index) { index_ = index; } + +void ImGui_DataTable::SetPrefix(const char* prefix) { + prefix_ = strlen(prefix); +} + +void ImGui_DataTable::operator()(const char* label, const uintptr_t* ptr, + int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("(%s)", &ptr[index_ + i] ? "[ptr]" : "null"); + } +} + +void ImGui_DataTable::operator()(const char* label, const char* ptr, int n) { + if (n == 1) { + MakeLabel(label); + ImGui::Text("%s", &ptr[index_]); + } else { + mju_error("char cannot be converted to a vector"); + } +} + +void ImGui_DataTable::operator()(const char* label, const mjtByte* ptr, int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("%s", ptr[index_ + i] ? "true" : "false"); + } +} + +void ImGui_DataTable::operator()(const char* label, const mjtByte& val, int n) { + MakeLabel(label, 0, 1); + ImGui::Text("%s", val ? "true" : "false"); +} + +void ImGui_DataTable::operator()(const char* label, const mjtSize* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const int* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const float* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const double* ptr, int n) { + Numeric(label, ptr, n); +} + +void ImGui_DataTable::operator()(const char* label, const mjtSize& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const int& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const float& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const double& val, int n) { + Scalar(label, val, n); +} + +void ImGui_DataTable::operator()(const char* label, const std::string* ptr, + int n) { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + ImGui::Text("%s", ptr[i].c_str()); + } +} + +void ImGui_DataTable::operator()(const char* label, + const std::vector* ptr, int n) { + if (n == 1) { + for (int i = 0; i < ptr->size(); ++i) { + MakeLabel(label, i, ptr->size()); + ImGui::Text("%s", ptr->at(i).c_str()); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +void ImGui_DataTable::operator()(const char* label, const std::vector* ptr, + int n) { + if (n == 1) { + const int size = ptr->size(); + if (size == 0) { + (*this)(label, "[empty]", 1); + } else { + std::string tmp = "[" + std::to_string(size) + " values]"; + (*this)(label, tmp.c_str(), 1); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +void ImGui_DataTable::operator()(const char* label, + const std::vector* ptr, int n) { + if (n == 1) { + const int size = ptr->size(); + if (size == 0) { + (*this)(label, "[empty]", 1); + } else { + std::string tmp = "[" + std::to_string(size) + " values]"; + (*this)(label, tmp.c_str(), 1); + } + } else { + mju_error("data type is vector; cannot also be an array"); + } +} + +template +void ImGui_DataTable::Scalar(const char* label, const T& value, int n) { + if (n == 1) { + Numeric(label, &value, n); + } else { + mju_error("scalar cannot be converted to a vector"); + } +} + +template +void ImGui_DataTable::Numeric(const char* label, const T* ptr, int n) { + const T* addr = ptr + index_ * n; + + using U = std::conditional_t, float, int>; + + // special treatment for NaNs. + if constexpr (std::is_same_v) { + if (*addr != *addr) { + MakeLabel(label); + ImGui::Text("nan"); + return; + } + } + + constexpr const char* fmt1 = + std::is_floating_point_v ? "%f" : "%d"; + constexpr const char* fmt2 = + std::is_floating_point_v ? "%f %f" : "%d %d"; + constexpr const char* fmt3 = + std::is_floating_point_v ? "%f %f %f" : "%d %d %d"; + constexpr const char* fmt4 = + std::is_floating_point_v ? "%f %f %f %f" : "%d %d %d %d"; + + auto text1 = [&](int offset) { + ImGui::Text(fmt1, (U)(addr[offset])); + }; + auto text2 = [&](int offset) { + ImGui::Text(fmt2, (U)(addr[offset + 0]), (U)(addr[offset + 1])); + }; + auto text3 = [&](int offset) { + ImGui::Text(fmt3, (U)(addr[offset + 0]), (U)(addr[offset + 1]), + (U)(addr[offset + 2])); + }; + auto text4 = [&](int offset) { + ImGui::Text(fmt4, (U)(addr[offset + 0]), (U)(addr[offset + 1]), + (U)(addr[offset + 2]), (U)(addr[offset + 3])); + }; + + if (n == 1) { + MakeLabel(label); + text1(0); + } else if (n == 2) { + MakeLabel(label); + text2(0); + } else if (n == 3) { + MakeLabel(label); + text3(0); + } else if (n == 4) { + MakeLabel(label); + text4(0); + } else if (n == 6) { + MakeLabel(label); + text3(0); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(3); + } else if (n == 9) { + MakeLabel(label); + text3(0); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(3); + ImGui::TableNextColumn(); + ImGui::TableNextColumn(); + text3(6); + } else { + for (int i = 0; i < n; ++i) { + MakeLabel(label, i, n); + text1(i); + } + } +} + +void ImGui_DataTable::MakeLabel(const char* label, int index, int total) { + if (total == 1) { + ImGui::TableNextColumn(); + ImGui::Text("%s", &label[prefix_]); + ImGui::TableNextColumn(); + } else { + const std::string tmp = + std::string(&label[prefix_]) + "[" + std::to_string(index) + "]"; + ImGui::TableNextColumn(); + ImGui::Text("%s", tmp.c_str()); + ImGui::TableNextColumn(); + } +} + bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max) { float f = *value; const bool res = ImGui::SliderFloat(name, &f, min, max); diff --git a/src/experimental/platform/imgui_widgets.h b/src/experimental/platform/imgui_widgets.h index e8988369..7fd64ca8 100644 --- a/src/experimental/platform/imgui_widgets.h +++ b/src/experimental/platform/imgui_widgets.h @@ -15,11 +15,13 @@ #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_IMGUI_WIDGETS_H_ #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_IMGUI_WIDGETS_H_ +#include #include #include #include #include #include +#include #include #include @@ -28,11 +30,14 @@ namespace mujoco::platform { // FontAwesome icon codes. +static constexpr const char ICON_FA_ADJUST[] = "\xEF\x81\x82"; static constexpr const char ICON_FA_ARROWS[] = "\xEF\x81\x87"; static constexpr const char ICON_FA_CAMERA[] = "\xEF\x80\xBD"; static constexpr const char ICON_FA_CARET_LEFT[] = "\xEF\x83\x99"; static constexpr const char ICON_FA_CARET_RIGHT[] = "\xEF\x83\x9A"; static constexpr const char ICON_FA_CHECK_SQUARE_O[] = "\xEF\x81\x9D"; +static constexpr const char ICON_FA_CIRCLE[] = "\xEF\x84\x91"; +static constexpr const char ICON_FA_CIRCLE_O[] = "\xEF\x84\x8C"; static constexpr const char ICON_FA_COMMENT[] = "\xEF\x83\xA5"; static constexpr const char ICON_FA_COPY[] = "\xEF\x83\x85"; static constexpr const char ICON_FA_DIAMOND[] = "\xEF\x88\x99"; @@ -109,6 +114,12 @@ struct ScopedStyle { return *this; } + ScopedStyle& Color(ImGuiCol col, ImGuiCol col2) { + ImGui::PushStyleColor(col, CurrentColor(col2)); + ++num_colors; + return *this; + } + ScopedStyle& Var(ImGuiStyleVar var, float value) { ImGui::PushStyleVar(var, value); ++num_vars; @@ -121,6 +132,10 @@ struct ScopedStyle { return *this; } + ImVec4 CurrentColor(ImGuiCol col) { + return ImGui::GetStyle().Colors[col]; + } + void Reset() { ImGui::PopStyleVar(num_vars); ImGui::PopStyleColor(num_colors); @@ -132,6 +147,74 @@ struct ScopedStyle { int num_vars = 0; }; +// Helper for displaying rows of key/value pairs in an ImGui table. +// +// Designed specifically to be used to display mjSpec, mjModel, and mjData +// values. +// +// To add a value to the table, call the operator() function with the label, +// the value, and (optionally) the dimensionality of the value (e.g. for vectors +// and matrices). To support generic code, even scalar values should be passed +// to operator() with n = 1. +class ImGui_DataTable { + public: + // Starts the table (i.e. ImGui::BeginTable()) with two columns of the + // specified widths. + ImGui_DataTable(float w1 = 0.25f, float w2 = 0.75f); + + // Ends the table (e.g. ImGui::EndTable(). + ~ImGui_DataTable(); + + ImGui_DataTable(const ImGui_DataTable& other) = delete; + ImGui_DataTable& operator=(const ImGui_DataTable& other) = delete; + + // Sets the offset into an array of values (e.g. for pointers in mjModel and + // mjData). This is only used for the display functions that take a pointer. + void SetArrayIndex(int index); + + // Sets the prefix that will be removed from all labels. Note: that we simply + // remove the first N characters of the label without actually comparing + // against this prefix. + void SetPrefix(const char* prefix); + + // Displays a labelled value in the table. + void operator()(const char* label, const uintptr_t* ptr, int n); + void operator()(const char* label, const char* ptr, int n); + void operator()(const char* label, const mjtByte* ptr, int n); + void operator()(const char* label, const mjtSize* ptr, int n); + void operator()(const char* label, const int* ptr, int n); + void operator()(const char* label, const float* ptr, int n); + void operator()(const char* label, const double* ptr, int n); + + // Displays a single scalar value in the table. Assumes n == 1. This should + // only be used for mjSpec objects and, therefore, will ignore the array index + // if set. + void operator()(const char* label, const mjtByte& val, int n); + void operator()(const char* label, const mjtSize& val, int n); + void operator()(const char* label, const int& val, int n); + void operator()(const char* label, const float& val, int n); + void operator()(const char* label, const double& val, int n); + + // Overloads for C++ container types. Assumes its only used for mjSpec objects + // and, therefore, will ignore the array index if set. + void operator()(const char* label, const std::string* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + void operator()(const char* label, const std::vector* ptr, int n); + + private: + template + void Numeric(const char* label, const T* ptr, int n); + + template + void Scalar(const char* label, const T& value, int n); + + void MakeLabel(const char* label, int index = 0, int total = 1); + + int prefix_ = 0; + int index_ = 0; +}; + // ImGui Slider that supports both float and double types. bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max); diff --git a/src/experimental/platform/object_launcher_plugin.cc b/src/experimental/platform/object_launcher_plugin.cc new file mode 100644 index 00000000..ad065bae --- /dev/null +++ b/src/experimental/platform/object_launcher_plugin.cc @@ -0,0 +1,215 @@ +// Copyright 2026 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 "experimental/platform/imgui_widgets.h" +#include "experimental/platform/plugin.h" + +namespace mujoco::studio { + +class ObjectLauncher { + public: + ObjectLauncher() : rng_(std::random_device{}()) {} + + void UpdateGui() { + using platform::ImGui_Input; + + ImGui::Checkbox("Enable Key Binding (Ctrl+Shift+Enter)", &enabled_); + ImGui_Input("Size", &size_, {0.01f, 1.0f, 0.01, 0.1}); + ImGui_Input("Speed", &speed_, {0.01f, 100.0f, 0.1, 1.0}); + ImGui_Input("Mass", &mass_, {0.01f, 100.0f, 0.01, 0.1}); + ImGui_Input("Life", &lifetime_, {0.0f, 60.0f, 0.1, 1.0}); + + int shape = type_ == mjGEOM_BOX ? 0 : 1; + const char* names[] = {"Box", "Sphere"}; + ImGui::Combo("Shape", &shape, names, 2); + type_ = shape == 0 ? mjGEOM_BOX : mjGEOM_SPHERE; + + if (ImGui::Button("Launch", ImVec2(-1.0f, 0.0f))) { + active_ = true; + } + + if (ImGui::Button("Clear")) { + for (auto& object : objects_) { + object.expiration = -1; + } + } + } + + void HandleKeyboardEvent() { if (enabled_) active_ = true; } + + bool UpdateSpecPreCompile(mjSpec* spec, const mjModel* model, + const mjData* data, const mjvCamera* camera) { + // Remove expired objects. + auto it = std::remove_if( + objects_.begin(), objects_.end(), [&](const ObjectInfo& o) { + const bool expired = + o.body_id >= 0 && o.expiration != 0 && o.expiration < data->time; + if (expired) { + mjsBody* body = mjs_findBody(spec, o.name.c_str()); + if (body) { + mjs_delete(spec, body->element); + } + } + return expired; + }); + if (it != objects_.end()) { + objects_.erase(it, objects_.end()); + return true; + } + + if (!active_) return false; + active_ = false; + + mjsBody* world = mjs_findBody(spec, "world"); + if (!world) return false; + mjsBody* body = mjs_addBody(world, nullptr); + if (!body) return false; + mjsJoint* joint = mjs_addJoint(body, nullptr); + if (!joint) return false; + mjsGeom* geom = mjs_addGeom(body, nullptr); + if (!geom) return false; + + ObjectInfo& object = objects_.emplace_back(); + object.name = "projectile" + std::to_string(counter_++);; + object.expiration = data->time + lifetime_; + + mjtNum pos[3]; + mjtNum dir[3]; + mjtNum up[3]; + mjv_cameraFrame(pos, dir, up, nullptr, data, camera); + mjs_setName(body->element, object.name.c_str()); + + joint->type = mjJNT_FREE; + body->mass = mass_; + geom->type = type_; + geom->size[0] = size_; + geom->size[1] = size_; + geom->size[2] = size_; + // Slightly in front of the camera. + body->pos[0] = pos[0] + (dir[0] * 0.1); + body->pos[1] = pos[1] + (dir[1] * 0.1); + body->pos[2] = pos[2] + (dir[2] * 0.1); + // Randomize the color. + geom->rgba[0] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[1] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[2] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); + geom->rgba[3] = 1.0; + // Launch it slightly upwards to get a nice arc. + launch_vel_[0] = (dir[0] * speed_) + up[0]; + launch_vel_[1] = (dir[1] * speed_) + up[1]; + launch_vel_[2] = (dir[2] * speed_) + up[2]; + return true; + } + + void UpdateSpecPostCompile(const mjSpec* spec, const mjModel* model, + mjData* data) { + if (objects_.empty()) { + return; + } + + ObjectInfo& object = objects_.back(); + if (object.launched) { + return; + } + object.launched = true; + + const int body_id = mj_name2id(model, mjOBJ_BODY, object.name.c_str()); + if (body_id < 0) { + return; + } + int joint_id = model->body_jntadr[body_id]; + if (joint_id < 0 || model->jnt_type[joint_id] != mjJNT_FREE) { + return; + } + int qvel_addr = model->jnt_dofadr[joint_id]; + if (qvel_addr < 0) { + return; + } + object.body_id = body_id; + data->qvel[qvel_addr + 0] = launch_vel_[0]; + data->qvel[qvel_addr + 1] = launch_vel_[1]; + data->qvel[qvel_addr + 2] = launch_vel_[2]; + } + + private: + struct ObjectInfo { + std::string name; + int body_id = -1; + mjtNum expiration = 0; + bool launched = false; + }; + + std::mt19937 rng_; + int counter_ = 0; + bool enabled_ = false; + bool active_ = false; + mjtNum size_ = 0.13365; + mjtNum speed_ = 10.0; + mjtNum mass_ = 10.0; + mjtNum lifetime_ = 5.0; + mjtGeom type_ = mjGEOM_BOX; + mjtNum launch_vel_[3] = {0, 0, 0}; + std::vector objects_; +}; + +} // namespace mujoco::studio + +mjPLUGIN_LIB_INIT { + using mujoco::studio::ObjectLauncher; + + static ObjectLauncher plugin; + + mujoco::platform::GuiPlugin gui; + gui.data = &plugin; + gui.name = "ObjectLauncher"; + gui.update = [](mujoco::platform::GuiPlugin* self) { + auto* plugin = static_cast(self->data); + plugin->UpdateGui(); + }; + mujoco::platform::RegisterPlugin(gui); + + mujoco::platform::KeyHandlerPlugin key_handler; + key_handler.data = &plugin; + key_handler.name = "ObjectLauncher"; + key_handler.key_chord = ImGuiKey_Enter | ImGuiMod_Ctrl | ImGuiMod_Shift; + key_handler.on_key_pressed = [](mujoco::platform::KeyHandlerPlugin* self) { + auto* plugin = static_cast(self->data); + plugin->HandleKeyboardEvent(); + }; + mujoco::platform::RegisterPlugin(key_handler); + + mujoco::platform::SpecEditorPlugin spec_editor; + spec_editor.data = &plugin; + spec_editor.name = "ObjectLauncher"; + spec_editor.pre_compile = [](mujoco::platform::SpecEditorPlugin* self, + mjSpec* spec, const mjModel* model, + const mjData* data, const mjvCamera* camera) { + auto* plugin = static_cast(self->data); + return plugin->UpdateSpecPreCompile(spec, model, data, camera); + }; + spec_editor.post_compile = [](mujoco::platform::SpecEditorPlugin* self, + const mjSpec* spec, const mjModel* model, + mjData* data) { + auto* plugin = static_cast(self->data); + return plugin->UpdateSpecPostCompile(spec, model, data); + }; + mujoco::platform::RegisterPlugin(spec_editor); +} diff --git a/src/experimental/platform/plugin.cc b/src/experimental/platform/plugin.cc index 13c262c1..3ddb6e1b 100644 --- a/src/experimental/platform/plugin.cc +++ b/src/experimental/platform/plugin.cc @@ -12,89 +12,66 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "experimental/platform/plugin.h" + #include #include #include -#include "experimental/platform/plugin.h" #include "engine/engine_global_table.h" +using GuiPlugin = mujoco::platform::GuiPlugin; +using ModelPlugin = mujoco::platform::ModelPlugin; +using KeyHandlerPlugin = mujoco::platform::KeyHandlerPlugin; +using SpecEditorPlugin = mujoco::platform::SpecEditorPlugin; + namespace mujoco::platform { -void RegisterGuiPlugin(const GuiPlugin* plugin) { - if (plugin->name == nullptr || plugin->name[0] == '\0') { +template +void RegisterPlugin(T plugin) { + if (plugin.name == nullptr || plugin.name[0] == '\0') { mju_error("Plugin name must not be empty or null."); } - GlobalTable::GetSingleton().AppendIfUnique(*plugin); + GlobalTable::GetSingleton().AppendIfUnique(plugin); } -void ForEachGuiPlugin(const std::function& fn) { - auto& table = GlobalTable::GetSingleton(); +template +void ForEachPlugin(const std::function& fn) { + auto& table = mujoco::GlobalTable::GetSingleton(); for (int i = 0; i < table.count(); ++i) { - const GuiPlugin* plugin = table.GetAtSlot(i); - fn(const_cast(plugin)); - } -} - -void RegisterModelPlugin(const ModelPlugin* plugin) { - if (plugin->name == nullptr || plugin->name[0] == '\0') { - mju_error("Plugin name must not be empty or null."); - } - GlobalTable::GetSingleton().AppendIfUnique(*plugin); -} - -void ForEachModelPlugin(const std::function& fn) { - auto& table = GlobalTable::GetSingleton(); - for (int i = 0; i < table.count(); ++i) { - const ModelPlugin* plugin = table.GetAtSlot(i); - fn(const_cast(plugin)); + const T* plugin = table.GetAtSlot(i); + fn(const_cast(plugin)); } } } // namespace mujoco::platform -using mujoco::GlobalTable; -using GuiPlugin = mujoco::platform::GuiPlugin; -using ModelPlugin = mujoco::platform::ModelPlugin; +#define MUJOCO_SPECIALIZE_PLUGIN(PLUGIN, NAME) \ + template <> \ + const char* mujoco::GlobalTable::HumanReadableTypeName() { \ + return NAME; \ + } \ + template <> \ + std::string_view mujoco::GlobalTable::ObjectKey(const PLUGIN& p) { \ + return std::string_view(p.name); \ + } \ + template <> \ + bool mujoco::GlobalTable::ObjectEqual(const PLUGIN& p1, \ + const PLUGIN& p2) { \ + return CaseInsensitiveEqual(p1.name, p2.name); \ + } \ + template <> \ + bool mujoco::GlobalTable::CopyObject(PLUGIN& dst, const PLUGIN& src, \ + ErrorMessage& err) { \ + dst = src; \ + return true; \ + } \ + namespace mujoco::platform { \ + template void RegisterPlugin(PLUGIN plugin); \ + template void ForEachPlugin(const std::function& fn); \ + } -template <> -const char* GlobalTable::HumanReadableTypeName() { - return "gui plugin"; -} - -template <> -std::string_view GlobalTable::ObjectKey(const GuiPlugin& plugin) { - return std::string_view(plugin.name); -} - -template <> -bool GlobalTable::ObjectEqual(const GuiPlugin& p1, const GuiPlugin& p2) { - return CaseInsensitiveEqual(p1.name, p2.name); -} - -template <> -bool GlobalTable::CopyObject(GuiPlugin& dst, const GuiPlugin& src, ErrorMessage& err) { - dst = src; - return true; -} - -template <> -const char* GlobalTable::HumanReadableTypeName() { - return "model plugin"; -} - -template <> -std::string_view GlobalTable::ObjectKey(const ModelPlugin& plugin) { - return std::string_view(plugin.name); -} - -template <> -bool GlobalTable::ObjectEqual(const ModelPlugin& p1, const ModelPlugin& p2) { - return CaseInsensitiveEqual(p1.name, p2.name); -} - -template <> -bool GlobalTable::CopyObject(ModelPlugin& dst, const ModelPlugin& src, ErrorMessage& err) { - dst = src; - return true; -} +MUJOCO_SPECIALIZE_PLUGIN(GuiPlugin, "gui plugin"); +MUJOCO_SPECIALIZE_PLUGIN(ModelPlugin, "model plugin"); +MUJOCO_SPECIALIZE_PLUGIN(KeyHandlerPlugin, "key handler plugin"); +MUJOCO_SPECIALIZE_PLUGIN(SpecEditorPlugin, "spec editor plugin"); diff --git a/src/experimental/platform/plugin.h b/src/experimental/platform/plugin.h index 824b9793..5e47fae9 100644 --- a/src/experimental/platform/plugin.h +++ b/src/experimental/platform/plugin.h @@ -20,10 +20,16 @@ namespace mujoco::platform { -// Important: Do not inherit from these plugin structs. They are copied by value -// and therefore any derived classes will be sliced. We assume plugins are -// effectively globals and so any pointers will be valid for the lifetime -// of the process. +// Registers plugins with the global registry. The plugins must have a +// case-insensitive unique name for the plugin type. Note that plugins are +// copied by value, so do not use inheritance. +template +void RegisterPlugin(T plugin); + +// Executes the given function for each registered plugin of type T. +template +void ForEachPlugin(const std::function& fn); + // Plugin for processing custom UI windows. The plugin will be listed in the // "Plugins" main menu and, when selected, an ImGui window will be opened with @@ -47,8 +53,9 @@ struct GuiPlugin final { void* data = nullptr; }; +// Plugin for loading and updating models. struct ModelPlugin final { - using GetModelToLoadFn = const char* (*)(ModelPlugin * self, int* size, + using GetModelToLoadFn = const char* (*)(ModelPlugin* self, int* size, char* content_type, int content_type_size, char* model_name, @@ -75,19 +82,44 @@ struct ModelPlugin final { void* data = nullptr; }; -// Registers a plugin with a global registry. The plugin must have a -// case-insensitive unique name. -void RegisterGuiPlugin(const GuiPlugin* plugin); +// Plugin for handling custom keyboard events. +struct KeyHandlerPlugin final { + using OnKeyPressedFn = void (*)(KeyHandlerPlugin* self); -// Executes the given function for each registered plugin. -void ForEachGuiPlugin(const std::function& fn); + // The name of the plugin; must be unique. + const char* name = ""; -// Registers a plugin with a global registry. The plugin must have a -// case-insensitive unique name. -void RegisterModelPlugin(const ModelPlugin* plugin); + // The ImGui key codes for the key combination that triggers the plugin. + int key_chord = 0; -// Executes the given function for each registered plugin. -void ForEachModelPlugin(const std::function& fn); + // The function to be called when the above key combination is pressed. + OnKeyPressedFn on_key_pressed = nullptr; + + // Optional data pointer. + void* data = nullptr; +}; + +// Plugin for editing the mjSpec. +struct SpecEditorPlugin final { + using PreCompileFn = bool (*)(SpecEditorPlugin* self, mjSpec* spec, + const mjModel* model, const mjData* data, + const mjvCamera* camera); + using PostCompileFn = void (*)(SpecEditorPlugin* self, const mjSpec* spec, + const mjModel* model, mjData* data); + + // The name of the plugin; must be unique. + const char* name = ""; + + // Callback that edits the spec. If it returns true, then the spec will be + // recompiled and `post_compile` will be called with the result. + PreCompileFn pre_compile = nullptr; + + // Callback that is called after the spec has been recompiled. + PostCompileFn post_compile = nullptr; + + // Optional data pointer. + void* data = nullptr; +}; } // namespace mujoco::platform diff --git a/src/experimental/platform/renderer.cc b/src/experimental/platform/renderer.cc index b1b39e52..2188d86f 100644 --- a/src/experimental/platform/renderer.cc +++ b/src/experimental/platform/renderer.cc @@ -56,7 +56,7 @@ void Renderer::Init(const mjModel* model) { mjr_makeContext(model, &render_context_, mjFONTSCALE_150); #else mjrFilamentConfig render_config; - mjr_defaultFilamentConfig(&render_config); + mjrf_defaultFilamentConfig(&render_config); render_config.native_window = native_window_; render_config.enable_gui = true; #if defined(MUJOCO_RENDERER_FILAMENT_OPENGL) @@ -66,7 +66,7 @@ void Renderer::Init(const mjModel* model) { #elif defined(MUJOCO_RENDERER_FILAMENT_VULKAN) render_config.graphics_api = mjGFX_VULKAN; #endif - mjr_makeFilamentContext(model, &render_context_, &render_config); + mjrf_makeFilamentContext(model, &render_context_, &render_config); #endif mjv_defaultScene(&scene_); @@ -171,9 +171,9 @@ int Renderer::UploadImage(int texture_id, const std::byte* pixels, int width, #if defined(MUJOCO_RENDERER_CLASSIC_OPENGL) return 0; #else - return mjr_uploadGuiImage(texture_id, - reinterpret_cast(pixels), - width, height, bpp, &render_context_); + return mjrf_uploadGuiImage(texture_id, + reinterpret_cast(pixels), + width, height, bpp, &render_context_); #endif } @@ -191,7 +191,7 @@ void Renderer::UpdateFps() { frames_ = 0; } #else - fps_ = mjr_getFrameRate(&render_context_); + fps_ = mjrf_getFrameRate(&render_context_); #endif } @@ -217,8 +217,8 @@ mjPLUGIN_LIB_INIT { mujoco::platform::GuiPlugin plugin; plugin.name = "Filament"; plugin.update = [](mujoco::platform::GuiPlugin* self) { - mjr_updateGui(nullptr); + mjrf_updateGui(nullptr); }; - mujoco::platform::RegisterGuiPlugin(&plugin); + mujoco::platform::RegisterPlugin(plugin); } #endif diff --git a/src/experimental/platform/step_control.cc b/src/experimental/platform/step_control.cc index 4c8b58cf..38c97fe2 100644 --- a/src/experimental/platform/step_control.cc +++ b/src/experimental/platform/step_control.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -30,6 +31,33 @@ static mjtNum Timer() { return Milliseconds(Clock::now() - start).count(); } +// Updates key viscous pause parameters restores them when done. +struct ViscousPauseState { + ViscousPauseState(mjModel* model) : model(model) { + if (model) { + mju_copy3(gravity, model->opt.gravity); + viscosity = model->opt.viscosity; + disableflags = model->opt.disableflags; + mju_zero3(model->opt.gravity); + model->opt.viscosity = 10; + model->opt.disableflags |= mjDSBL_SPRING; + } + } + + ~ViscousPauseState() { + if (model) { + mju_copy3(model->opt.gravity, gravity); + model->opt.viscosity = viscosity; + model->opt.disableflags = disableflags; + } + } + mjModel* model; + mjtNum gravity[3]; + mjtNum viscosity; + int disableflags; +}; + + StepControl::StepControl() { mjcb_time = Timer; } float StepControl::GetSpeedMeasured() const { return speed_measured_; } @@ -56,23 +84,6 @@ void StepControl::SetNoiseParameters(float ctrl_noise_scale, } void StepControl::SetPauseState(PauseState state, mjModel* m) { - if (pause_state_ == PauseState::kViscousPaused && - state != PauseState::kViscousPaused && m) { - mju_copy3(m->opt.gravity, saved_gravity_); - m->opt.viscosity = saved_viscosity_; - m->opt.disableflags = saved_disableflags_; - } - - if (state == PauseState::kViscousPaused && - pause_state_ != PauseState::kViscousPaused && m) { - mju_copy3(saved_gravity_, m->opt.gravity); - saved_viscosity_ = m->opt.viscosity; - saved_disableflags_ = m->opt.disableflags; - mju_zero3(m->opt.gravity); - m->opt.viscosity = 10; - m->opt.disableflags |= mjDSBL_SPRING; - } - pause_state_ = state; } @@ -81,6 +92,10 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d) { return Status::kOk; } + std::optional viscous_pause_state; + if (m && pause_state_ == PauseState::kViscousPaused) { + viscous_pause_state.emplace(m); + } if (pause_state_ == PauseState::kNormalPaused) { // When we eventually unpause, we need to make sure we sync to immediately diff --git a/src/experimental/platform/step_control.h b/src/experimental/platform/step_control.h index 916a23c7..7cebbd1b 100644 --- a/src/experimental/platform/step_control.h +++ b/src/experimental/platform/step_control.h @@ -111,11 +111,6 @@ class StepControl { PauseState pause_state_ = PauseState::kUnpaused; - // Viscous pause state variables - mjtNum saved_gravity_[3] = {0}; - mjtNum saved_viscosity_ = 0; - int saved_disableflags_ = 0; - // Perform only a single step on the next call to Advance() if the simulation // is paused. bool single_step_ = false; diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 4c250f36..6d33f5af 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -23,8 +23,8 @@ #include #include #include +#include #include -#include #include #include #include @@ -36,6 +36,7 @@ #include #include "experimental/platform/file_dialog.h" #include "experimental/platform/gui.h" +#include "experimental/platform/gui_spec.h" #include "experimental/platform/helpers.h" #include "experimental/platform/imgui_widgets.h" #include "experimental/platform/interaction.h" @@ -78,9 +79,9 @@ static constexpr const char* ICON_LABEL = platform::ICON_FA_COMMENT; static constexpr const char* ICON_RESET_MODEL = platform::ICON_FA_UNDO; static constexpr const char* ICON_FRAME = platform::ICON_FA_ARROWS; static constexpr const char* ICON_CAMERA = platform::ICON_FA_CAMERA; -static constexpr const char* ICON_DARKMODE = platform::ICON_FA_MOON; -static constexpr const char* ICON_LIGHTMODE = platform::ICON_FA_SUN; -static constexpr const char* ICON_CLASSICMODE = platform::ICON_FA_DIAMOND; +static constexpr const char* ICON_DARKMODE = platform::ICON_FA_CIRCLE; +static constexpr const char* ICON_LIGHTMODE = platform::ICON_FA_CIRCLE_O; +static constexpr const char* ICON_CLASSICMODE = platform::ICON_FA_ADJUST; static constexpr const char* ICON_PREV_FRAME = platform::ICON_FA_CARET_LEFT; static constexpr const char* ICON_NEXT_FRAME = platform::ICON_FA_CARET_RIGHT; static constexpr const char* ICON_CURR_FRAME = platform::ICON_FA_FAST_FORWARD; @@ -107,8 +108,7 @@ static constexpr std::array kPercentRealTime = { }; // clang-format on -App::App(Config config) - : rng_(std::random_device()()), ini_path_(std::move(config.ini_path)) { +App::App(Config config) : ini_path_(std::move(config.ini_path)) { platform::Window::Config window_config; window_config.renderer_backend = platform::Renderer::GetBackend(); window_config.offscreen_mode = config.offscreen_mode; @@ -133,6 +133,7 @@ void App::ClearModel() { tmp_ = UiTempState(); load_error_ = ""; step_error_ = ""; + edit_error_ = ""; } void App::Recompile() { @@ -222,7 +223,7 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { } } - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->post_model_loaded) { plugin->post_model_loaded(plugin, model_path_.c_str()); } @@ -258,6 +259,7 @@ void App::ResetPhysics() { mj_resetData(model(), data()); mj_forward(model(), data()); step_error_ = ""; + edit_error_ = ""; } void App::UpdatePhysics() { @@ -266,7 +268,7 @@ void App::UpdatePhysics() { } bool stepped = false; - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->do_update) { if (plugin->do_update(plugin, model(), data())) { stepped = true; @@ -391,8 +393,20 @@ void App::ProcessPendingLoads() { spec_op_ = nullptr; } + // Allow plugins to edit the spec as well. + platform::ForEachPlugin([&](auto* plugin) { + if (plugin->pre_compile) { + if (plugin->pre_compile(plugin, spec(), model(), data(), &camera_)) { + Recompile(); + if (plugin->post_compile) { + plugin->post_compile(plugin, spec(), model(), data()); + } + }; + } + }); + // Check plugins to see if we need to load a new model. - platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { if (plugin->get_model_to_load) { char model_name[1000] = ""; char content_type[1000] = ""; @@ -408,15 +422,42 @@ void App::ProcessPendingLoads() { }); } -void App::SpecDeleteSelectedElement() { - spec_op_ = [this]() { - mjs_delete(spec(), tmp_.element); - if (tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select == tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - } - tmp_.element = nullptr; +void App::SpecSelectElement(mjsElement* element) { + tmp_.element = element; + if (tmp_.element == nullptr) { tmp_.element_id = -1; + } else { + tmp_.element_id = mjs_getId(tmp_.element); + + // If we selected a body, then select the same body for perturb. + if (tmp_.element->elemtype == mjOBJ_BODY && + perturb_.select != tmp_.element_id) { + mjv_defaultPerturb(&perturb_); + perturb_.select = tmp_.element_id; + } + } +} + +void App::SpecDeleteElement(mjsElement* element) { + if (element == nullptr) { + return; + } + // Only bodies can be deleted for now... + if (element->elemtype != mjOBJ_BODY) { + edit_error_ = "WARNING: Only bodies can be deleted (for now...)"; + return; + } + spec_op_ = [this, element]() { + mjs_delete(spec(), element); + if (tmp_.element == element) { + tmp_.element = nullptr; + tmp_.element_id = -1; + } + if (element->elemtype == mjOBJ_BODY) { + if (perturb_.select == mjs_getId(element)) { + mjv_defaultPerturb(&perturb_); + } + } Recompile(); }; } @@ -624,7 +665,7 @@ void App::HandleKeyboardEvents() { } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) { ResetPhysics(); } else if (ImGui_IsChordJustPressed(ImGuiKey_Delete)) { - SpecDeleteSelectedElement(); + SpecDeleteElement(tmp_.element); } else if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) { SelectParentPerturb(model(), perturb_); } else if (ImGui_IsChordJustPressed(ImGuiKey_F1)) { @@ -715,115 +756,65 @@ void App::HandleKeyboardEvents() { ToggleFlag(vis_options_.geomgroup[4]); } else if (ImGui_IsChordJustPressed(ImGuiKey_5)) { ToggleFlag(vis_options_.geomgroup[5]); - } else if (ImGui_IsChordJustPressed(ImGuiKey_Enter | ImGuiMode_CtrlShift)) { - if (has_spec()) { - spec_op_ = [this]() { - mjsBody* world = mjs_findBody(spec(), "world"); - if (!world) return; - mjsBody* body = mjs_addBody(world, nullptr); - if (!body) return; - mjsJoint* joint = mjs_addJoint(body, nullptr); - if (!joint) return; - mjsGeom* geom = mjs_addGeom(body, nullptr); - if (!geom) return; + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_Escape)) { + ui_.camera_idx = + platform::SetCamera(model(), &camera_, platform::kTumbleCameraIdx); + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_LeftBracket)) { + ui_.camera_idx = platform::SetCamera(model(), &camera_, ui_.camera_idx - 1); + } else if (has_model() && ImGui_IsChordJustPressed(ImGuiKey_RightBracket)) { + ui_.camera_idx = platform::SetCamera(model(), &camera_, ui_.camera_idx + 1); + // WASD camera controls for free camera. + } else if (is_freecam_wasd && + (ImGui::IsKeyDown(ImGuiKey_W) || ImGui::IsKeyDown(ImGuiKey_S) || + ImGui::IsKeyDown(ImGuiKey_A) || ImGui::IsKeyDown(ImGuiKey_D) || + ImGui::IsKeyDown(ImGuiKey_Q) || ImGui::IsKeyDown(ImGuiKey_E))) { + bool moved = false; - // Set body position slightly in front of the camera. - mjtNum pos[3]; - mjtNum dir[3]; - mjtNum up[3]; - mjv_cameraFrame(pos, dir, up, nullptr, data(), &camera_); + // Move (dolly) forward/backward using W and S keys. + if (ImGui::IsKeyDown(ImGuiKey_W)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, tmp_.cam_speed); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_S)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, -tmp_.cam_speed); + moved = true; + } - static int counter = 0; - std::string name = "projectile" + std::to_string(counter++); - mjs_setName(body->element, name.c_str()); + // Strafe (truck) left/right using A and D keys. + if (ImGui::IsKeyDown(ImGuiKey_A)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, -tmp_.cam_speed, 0); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_D)) { + MoveCamera(platform::CameraMotion::TRUCK_DOLLY, tmp_.cam_speed, 0); + moved = true; + } - body->mass = 10.0; - body->pos[0] = pos[0] + dir[0] * 0.2; - body->pos[1] = pos[1] + dir[1] * 0.2; - body->pos[2] = pos[2] + dir[2] * 0.2; - geom->type = mjGEOM_BOX; - geom->size[0] = 0.13365; - geom->size[1] = 0.13365; - geom->size[2] = 0.13365; - geom->rgba[0] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[1] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[2] = std::uniform_real_distribution(0.3f, 1.0f)(rng_); - geom->rgba[3] = 1.0; + // Move (pedestal) up/down using Q and E keys. + if (ImGui::IsKeyDown(ImGuiKey_Q)) { + MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, tmp_.cam_speed); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_E)) { + MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, -tmp_.cam_speed); + moved = true; + } - joint->type = mjJNT_FREE; + if (moved) { + tmp_.cam_speed += 0.001f; - Recompile(); - - // Give the newly added body a velocity in the direction of the camera. - int bodyid = mj_name2id(model(), mjOBJ_BODY, name.c_str()); - if (bodyid >= 0) { - int jntid = model()->body_jntadr[bodyid]; - if (jntid >= 0 && model()->jnt_type[jntid] == mjJNT_FREE) { - int qveladr = model()->jnt_dofadr[jntid]; - if (qveladr >= 0) { - mjtNum speed = 10.0; // Magnitude of the initial velocity. - data()->qvel[qveladr + 0] = dir[0] * speed + up[0]; - data()->qvel[qveladr + 1] = dir[1] * speed + up[1]; - data()->qvel[qveladr + 2] = dir[2] * speed + up[2]; - } - } + const float max_speed = ImGui::GetIO().KeyShift ? 0.1 : 0.01f; + if (tmp_.cam_speed > max_speed) { + tmp_.cam_speed = max_speed; + } + } else { + tmp_.cam_speed = 0.001f; + } + } else { + platform::ForEachPlugin([&](auto* plugin) { + if (plugin->key_chord && plugin->on_key_pressed) { + if (ImGui_IsChordJustPressed(plugin->key_chord)) { + plugin->on_key_pressed(plugin); } - }; - } - } else if (has_model()) { - if (ImGui_IsChordJustPressed(ImGuiKey_Escape)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, platform::kTumbleCameraIdx); - } else if (ImGui_IsChordJustPressed(ImGuiKey_LeftBracket)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, ui_.camera_idx - 1); - } else if (ImGui_IsChordJustPressed(ImGuiKey_RightBracket)) { - ui_.camera_idx = - platform::SetCamera(model(), &camera_, ui_.camera_idx + 1); - } - - // WASD camera controls for free camera. - if (is_freecam_wasd) { - bool moved = false; - - // Move (dolly) forward/backward using W and S keys. - if (ImGui::IsKeyDown(ImGuiKey_W)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, tmp_.cam_speed); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_S)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, 0, -tmp_.cam_speed); - moved = true; } - - // Strafe (truck) left/right using A and D keys. - if (ImGui::IsKeyDown(ImGuiKey_A)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, -tmp_.cam_speed, 0); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_D)) { - MoveCamera(platform::CameraMotion::TRUCK_DOLLY, tmp_.cam_speed, 0); - moved = true; - } - - // Move (pedestal) up/down using Q and E keys. - if (ImGui::IsKeyDown(ImGuiKey_Q)) { - MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, tmp_.cam_speed); - moved = true; - } else if (ImGui::IsKeyDown(ImGuiKey_E)) { - MoveCamera(platform::CameraMotion::TRUCK_PEDESTAL, 0, -tmp_.cam_speed); - moved = true; - } - - if (moved) { - tmp_.cam_speed += 0.001f; - - const float max_speed = ImGui::GetIO().KeyShift ? 0.1 : 0.01f; - if (tmp_.cam_speed > max_speed) { - tmp_.cam_speed = max_speed; - } - } else { - tmp_.cam_speed = 0.001f; - } - } + }); } } @@ -836,7 +827,7 @@ void App::LoadSettings() { platform::KeyValues plugin_names = platform::ReadIniSection(settings, "[Studio][Plugins]"); - platform::ForEachGuiPlugin([&](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { auto it = plugin_names.find(plugin->name); if (it != plugin_names.end()) { plugin->active = std::stoi(it->second) != 0; @@ -852,7 +843,7 @@ void App::SaveSettings() { platform::AppendIniSection(settings, "[Studio][UX]", ui_.ToDict()); platform::KeyValues plugin_names; - platform::ForEachGuiPlugin([&](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([&](auto* plugin) { plugin_names[plugin->name] = std::to_string((int)plugin->active); }); platform::AppendIniSection(settings, "[Studio][Plugins]", plugin_names); @@ -891,10 +882,16 @@ void App::BuildGui() { MainMenuGui(); - if (ImGui::Begin("ToolBar")) { - ToolBarGui(); + { + ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(0, 0)); + if (ImGui::Begin("ToolBar")) { + ImGui::PopStyleVar(); + ToolBarGui(); + } else { + ImGui::PopStyleVar(); + } + ImGui::End(); } - ImGui::End(); { platform::ScopedStyle style; @@ -929,7 +926,7 @@ void App::BuildGui() { if (explorer_is_open && tmp_.element != nullptr) { if (ImGui::Begin("Properties")) { - PropertiesGui(); + SpecPropertiesGui(); } ImGui::End(); } @@ -1026,7 +1023,7 @@ void App::BuildGui() { ImGui::End(); } - platform::ForEachGuiPlugin([](platform::GuiPlugin* plugin) { + platform::ForEachPlugin([](auto* plugin) { if (!plugin->update) { return; } @@ -1162,103 +1159,62 @@ void App::DataInspectorGui() { ImGui::EndChild(); } -void DisplayElementTree(mjsElement* element) { - const mjString* name = mjs_getName(element); - if (name->empty()) { - ImGui::Text("(unnamed)"); - } else { - ImGui::Text("%s", name->c_str()); - } -} - void App::SpecExplorerGui() { if (!has_spec()) { ImGui::Text("No mjSpec loaded."); return; } - const ImGuiTreeNodeFlags flags = - ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; + auto on_delete = [this](mjsElement* element) { SpecDeleteElement(element); }; - auto display_group = [this](mjtObj type, const std::string& prefix) { - mjsElement* element = mjs_firstElement(spec(), type); - while (element) { - const int id = mjs_getId(element); - - const mjString* name = mjs_getName(element); - std::string label = *name; - if (label.empty()) { - label = "(" + prefix + " " + std::to_string(id) + ")"; - } - - const bool selected = (tmp_.element == element); - if (ImGui::Selectable(label.c_str(), selected)) { - tmp_.element = element; - tmp_.element_id = id; - } - - element = mjs_nextElement(spec(), element); - } - }; - - if (ImGui::TreeNodeEx("Bodies", flags)) { - display_group(mjOBJ_BODY, "Body"); - ImGui::TreePop(); - } - if (ImGui::TreeNodeEx("Joints", flags)) { - display_group(mjOBJ_JOINT, "Joint"); - ImGui::TreePop(); - } - if (ImGui::TreeNodeEx("Sites", flags)) { - display_group(mjOBJ_SITE, "Site"); - ImGui::TreePop(); - } - - // If we selected a body, then select the same body for the perturb object. - if (tmp_.element && tmp_.element->elemtype == mjOBJ_BODY && - perturb_.select != tmp_.element_id) { - mjv_defaultPerturb(&perturb_); - perturb_.select = tmp_.element_id; + mjsElement* element = tmp_.element; + platform::SpecExplorerGui(&element, spec(), on_delete); + if (element != tmp_.element) { + SpecSelectElement(element); } } -void App::PropertiesGui() { - if (tmp_.element == nullptr) { - ImGui::Text("No element selected."); - return; - } +void App::SpecPropertiesGui() { + platform::ScopedStyle style; - if (ImGui::BeginTable("##PropertiesHeader", 2)) { - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 20); - ImGui::TableNextColumn(); - ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); - ImGui::TableNextColumn(); - if (tmp_.element->elemtype == mjOBJ_BODY) { - if (ImGui::SmallButton(ICON_DELETE)) { - SpecDeleteSelectedElement(); - } - } - ImGui::EndTable(); + ImGui::Text("%s", mju_type2Str(tmp_.element->elemtype)); + ImGui::SameLine(); + ImGui::Text("(%d)", tmp_.element_id); + + ImGui::SameLine(120); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); } + if (ImGui::SmallButton("S")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kSpec; + } + ImGui::SetItemTooltip("Spec"); + style.Reset(); + ImGui::SameLine(); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); + } + if (ImGui::SmallButton("M")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kModel; + } + ImGui::SetItemTooltip("Model"); + style.Reset(); + ImGui::SameLine(); + if (tmp_.spec_prop_mode == SpecPropertiesMode::kData) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); + } + if (ImGui::SmallButton("D")) { + tmp_.spec_prop_mode = SpecPropertiesMode::kData; + } + ImGui::SetItemTooltip("Data"); + style.Reset(); ImGui::Separator(); - - switch (tmp_.element->elemtype) { - case mjOBJ_BODY: - platform::BodyPropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - case mjOBJ_JOINT: - platform::JointPropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - case mjOBJ_SITE: - platform::SitePropertiesGui(model(), data(), tmp_.element, - tmp_.element_id); - break; - default: - // ignore other types - break; + if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { + platform::ElementSpecGui(spec(), tmp_.element); + } else if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { + platform::ElementModelGui(model(), tmp_.element); + } else { + platform::ElementDataGui(data(), tmp_.element); } } @@ -1375,6 +1331,18 @@ void App::HelpGui() { ImGui::Columns(); } +struct SpeedStatus { + bool misaligned; + float measured; +}; + +static SpeedStatus IsSpeedMisaligned( + const platform::StepControl& step_control) { + const float desired = step_control.GetSpeed(); + const float measured = step_control.GetSpeedMeasured(); + return {std::abs(measured - desired) > 0.1f * desired, measured}; +} + void App::ToolBarGui() { if (ImGui::BeginTable("##ToolBarTable", 2)) { platform::ScopedStyle style; @@ -1384,22 +1352,30 @@ void App::ToolBarGui() { const int combo_flags = ImGuiComboFlags_NoArrowButton; const float scale = ImGui::GetWindowDpiScale(); - const float right_width = 520.f * scale; const ImVec2 button_size(48.f * scale, 32.f * scale); const ImVec2 play_button_size(80.f * scale, 32.f * scale); + const float label_width = GetExpectedLabelWidth(); + const float copy_btn_width = ImGui::CalcTextSize(ICON_COPY_CAMERA).x + + ImGui::GetStyle().FramePadding.x * 2; + const float theme_width = ImGui::CalcTextSize(ICON_LIGHTMODE).x + + ImGui::GetStyle().FramePadding.x * 2; + const float sp = ImGui::GetStyle().ItemSpacing.x; + const float right_width = label_width + sp + label_width + sp + + label_width + sp + copy_btn_width + sp + + theme_width; + const float separator_width = .2f * button_size.x; + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, right_width); ImGui::TableNextColumn(); - ImGui::Text("%s", " "); // Combined (Unload, Reload) widget { style.Var(ImGuiStyleVar_FrameRounding, 2.0f); // Unload button. - ImGui::SameLine(); { const ImColor a = red; const ImColor h(a.Value.x, a.Value.y, a.Value.z, a.Value.w * 0.6f); @@ -1421,25 +1397,19 @@ void App::ToolBarGui() { ImGui::SetItemTooltip("%s", "Reload"); } - ImGui::SameLine(0, 0); - ImGui::Text(" "); - // Reset button. - ImGui::SameLine(); + ImGui::SameLine(0, separator_width); if (ImGui::Button(ICON_RESET_MODEL, button_size)) { ResetPhysics(); } ImGui::SetItemTooltip("%s", "Reset"); - ImGui::SameLine(0, 0); - ImGui::Text(" "); - // Combined (Normal Pause, Viscous Pause, Play) widget { style.Var(ImGuiStyleVar_FrameRounding, 2.0f); // Normal pause button. - ImGui::SameLine(); + ImGui::SameLine(0, separator_width); ImColor paused_color = yellow; bool paused = step_control_.GetPauseState() == PauseState::kNormalPaused; if (platform::ImGui_ColorButton(ICON_PAUSE, paused, paused_color, @@ -1480,18 +1450,23 @@ void App::ToolBarGui() { } } - ImGui::SameLine(); - ImGui::Text("%s", " |"); - // Speed selection. ImGui::SameLine(); - ImGui::Text("%s", ICON_SPEED); - ImGui::SetItemTooltip("%s", "Playback Speed"); - - ImGui::SameLine(); - ImGui::SetNextItemWidth(50.0f * scale); - if (ImGui::BeginCombo("##Speed", kPercentRealTime[tmp_.speed_index], - combo_flags)) { + float pad_y = (button_size.y - ImGui::GetFontSize()) * 0.5f; + ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, + ImVec2(ImGui::GetStyle().FramePadding.x + 5.f, pad_y)); + const auto [misaligned, measured] = IsSpeedMisaligned(step_control_); + char speed_preview[64]; + if (misaligned) { + snprintf(speed_preview, sizeof(speed_preview), "%s%s (%-4.1f%%)", + ICON_SPEED, kPercentRealTime[tmp_.speed_index], measured); + } else { + snprintf(speed_preview, sizeof(speed_preview), "%s%s", ICON_SPEED, + kPercentRealTime[tmp_.speed_index]); + } + ImGui::SetNextItemWidth(ImGui::CalcTextSize(speed_preview).x + + ImGui::GetStyle().FramePadding.x * 2); + if (ImGui::BeginCombo("##Speed", speed_preview, combo_flags)) { for (int n = 0; n < kPercentRealTime.size(); n++) { if (ImGui::Selectable(kPercentRealTime[n], (tmp_.speed_index == n))) { SetSpeedIndex(n); @@ -1499,17 +1474,30 @@ void App::ToolBarGui() { } ImGui::EndCombo(); } - ImGui::SetItemTooltip("%s", "Playback Speed"); + ImGui::PopStyleVar(); + if (misaligned) { + ImGui::SetItemTooltip("%s", "Desired Speed (Measured Speed)"); + } else { + ImGui::SetItemTooltip("%s", "Desired Speed"); + } + + ImGui::TableNextColumn(); + ImGui::SetCursorPosY(ImGui::GetCursorPosY() + + (button_size.y - ImGui::GetFrameHeight()) * 0.5f); // Camera selection. - std::vector cameras = GetCameraNames(); - ImGui::TableNextColumn(); - ImGui::Text("%s", ICON_CAMERA); - ImGui::SetItemTooltip("%s", "Camera"); - ImGui::SameLine(); + if (ImGui::Button(ICON_COPY_CAMERA)) { + std::string camera_string = platform::CameraToString(data(), &camera_); + platform::MaybeSaveToClipboard(camera_string); + } + ImGui::SetItemTooltip("%s", "Copy Camera"); + ImGui::SameLine(0, 0); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); int camera_idx = ui_.camera_idx - platform::kTumbleCameraIdx; - if (ImGui::BeginCombo("##Camera", cameras[camera_idx], combo_flags)) { + std::vector cameras = GetCameraNames(); + std::string camera_preview = + std::string(ICON_CAMERA) + " " + cameras[camera_idx]; + if (ImGui::BeginCombo("##Camera", camera_preview.c_str(), combo_flags)) { for (int n = 0; n < cameras.size(); n++) { if (ImGui::Selectable(cameras[n], (camera_idx == n))) { ui_.camera_idx = platform::SetCamera(model(), &camera_, @@ -1519,25 +1507,13 @@ void App::ToolBarGui() { ImGui::EndCombo(); } ImGui::SetItemTooltip("%s", "Camera"); - ImGui::SameLine(); - if (ImGui::Button(ICON_COPY_CAMERA)) { - std::string camera_string = platform::CameraToString(data(), &camera_); - platform::MaybeSaveToClipboard(camera_string); - } - ImGui::SetItemTooltip("%s", "Copy Camera"); - - ImGui::SameLine(); - ImGui::Text("%s", " |"); // Label selection. - ImGui::SameLine(); - ImGui::Text("%s", ICON_LABEL); - ImGui::SetItemTooltip("%s", "Label"); - ImGui::SameLine(); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); - if (ImGui::BeginCombo("##Label", kLabelNames[vis_options_.label], - combo_flags)) { + std::string label_preview = + std::string(ICON_LABEL) + " " + kLabelNames[vis_options_.label]; + if (ImGui::BeginCombo("##Label", label_preview.c_str(), combo_flags)) { for (int n = 0; n < IM_ARRAYSIZE(kLabelNames); n++) { if (ImGui::Selectable(kLabelNames[n], (vis_options_.label == n))) { vis_options_.label = n; @@ -1547,18 +1523,12 @@ void App::ToolBarGui() { } ImGui::SetItemTooltip("%s", "Label"); - ImGui::SameLine(); - ImGui::Text("%s", " |"); - // Frame selection. - ImGui::SameLine(); - ImGui::Text("%s", ICON_FRAME); - ImGui::SetItemTooltip("%s", "Frame"); - ImGui::SameLine(); ImGui::SetNextItemWidth(GetExpectedLabelWidth()); - if (ImGui::BeginCombo("##Frame", kFrameNames[vis_options_.frame], - combo_flags)) { + std::string frame_preview = + std::string(ICON_FRAME) + " " + kFrameNames[vis_options_.frame]; + if (ImGui::BeginCombo("##Frame", frame_preview.c_str(), combo_flags)) { for (int n = 0; n < IM_ARRAYSIZE(kFrameNames); n++) { if (ImGui::Selectable(kFrameNames[n], (vis_options_.frame == n))) { vis_options_.frame = n; @@ -1568,31 +1538,31 @@ void App::ToolBarGui() { } ImGui::SetItemTooltip("%s", "Frame"); + // Theme selection. ImGui::SameLine(); - ImGui::Text("%s", " |"); - - // Style selection. - ImGui::SameLine(); - switch (ui_.theme) { - case platform::GuiTheme::kLight: - if (ImGui::Button(ICON_LIGHTMODE)) { - SetupTheme(platform::GuiTheme::kDark); + const char* theme_icons[] = {ICON_LIGHTMODE, ICON_DARKMODE, + ICON_CLASSICMODE}; + const char* theme_tooltips[] = {"Light Mode", "Dark Mode", "Classic Mode"}; + const platform::GuiTheme theme_values[] = { + platform::GuiTheme::kLight, + platform::GuiTheme::kDark, + platform::GuiTheme::kClassic, + }; + int theme_idx = static_cast(ui_.theme); + ImGui::SetNextItemWidth(ImGui::CalcTextSize(theme_icons[0]).x + + ImGui::GetStyle().FramePadding.x * 2); + if (ImGui::BeginCombo("##Theme", theme_icons[theme_idx], combo_flags)) { + for (int n = 0; n < IM_ARRAYSIZE(theme_icons); n++) { + if (ImGui::Selectable(theme_icons[n], (theme_idx == n))) { + SetupTheme(theme_values[n]); } - ImGui::SetItemTooltip("%s", "Switch to Dark Mode"); - break; - case platform::GuiTheme::kDark: - if (ImGui::Button(ICON_DARKMODE)) { - SetupTheme(platform::GuiTheme::kClassic); + if (ImGui::IsItemHovered()) { + ImGui::SetTooltip("%s", theme_tooltips[n]); } - ImGui::SetItemTooltip("%s", "Switch to Classic Mode"); - break; - case platform::GuiTheme::kClassic: - if (ImGui::Button(ICON_CLASSICMODE)) { - SetupTheme(platform::GuiTheme::kLight); - } - ImGui::SetItemTooltip("%s", "Switch to Light Mode"); - break; + } + ImGui::EndCombo(); } + ImGui::SetItemTooltip("%s", "Theme"); ImGui::EndTable(); } @@ -1613,17 +1583,7 @@ void App::StatusBarGui() { } else if (step_control_.GetPauseState() == PauseState::kNormalPaused) { ImGui::Text("Paused"); } else { - const float desired_realtime = step_control_.GetSpeed(); - const float measured_realtime = step_control_.GetSpeedMeasured(); - const float realtime_offset = - mju_abs(measured_realtime - desired_realtime); - const bool misaligned = realtime_offset > 0.1 * desired_realtime; - if (misaligned) { - ImGui::Text("Running: %g%% (%-4.1f%%)", desired_realtime, - measured_realtime); - } else { - ImGui::Text("Running: %g%%", desired_realtime); - } + ImGui::Text("Running"); } if (!step_error_.empty()) { @@ -1632,6 +1592,9 @@ void App::StatusBarGui() { } else if (!load_error_.empty()) { ImGui::SameLine(); ImGui::Text(" | Load Error: %s", load_error_.c_str()); + } else if (!edit_error_.empty()) { + ImGui::SameLine(); + ImGui::Text(" | Edit Error: %s", edit_error_.c_str()); } ImGui::TableNextColumn(); @@ -1814,6 +1777,9 @@ void App::MainMenuGui() { if (ImGui::MenuItem("ImPlot Demo")) { tmp_.implot_demo = !tmp_.implot_demo; } + ImGui::Separator(); + std::string version = "Version " + std::string(mj_versionString()); + ImGui::MenuItem(version.c_str()); ImGui::EndMenu(); } ImGui::EndMainMenuBar(); diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index 170aadd6..66332e70 100644 --- a/src/experimental/studio/app.h +++ b/src/experimental/studio/app.h @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include @@ -92,6 +91,12 @@ class App { kModelFromBuffer, }; + enum class SpecPropertiesMode { + kSpec, + kModel, + kData, + }; + // UI state that is persisted across application runs struct UiState { char watch_field[1000] = "qpos"; @@ -137,6 +142,7 @@ class App { std::vector speed_names; // Spec Properties. + SpecPropertiesMode spec_prop_mode = SpecPropertiesMode::kSpec; mjsElement* element = nullptr; int element_id = -1; @@ -210,9 +216,10 @@ class App { void ModelOptionsGui(); void DataInspectorGui(); void SpecExplorerGui(); - void PropertiesGui(); + void SpecPropertiesGui(); - void SpecDeleteSelectedElement(); + void SpecSelectElement(mjsElement* element); + void SpecDeleteElement(mjsElement* element); float GetExpectedLabelWidth(); std::vector GetCameraNames(); @@ -224,13 +231,12 @@ class App { bool has_model() const { return model_holder_ && model_holder_->model(); } bool has_data() const { return model_holder_ && model_holder_->data(); } - std::mt19937 rng_; - std::string ini_path_; std::string model_name_; // Used if model_kind_ is kModelFromBuffer. std::string model_path_; std::string load_error_; std::string step_error_; + std::string edit_error_; std::optional pending_load_; bool preserve_camera_on_load_ = false; ModelKind model_kind_ = kEmptyModel; diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 5d060117..812eb0df 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -34,7 +34,6 @@ #include #include "user/user_api.h" -#include #ifdef MUJOCO_TINYOBJLOADER_IMPL #define TINYOBJLOADER_IMPLEMENTATION @@ -780,52 +779,7 @@ void mjCMesh::TryCompile(const mjVFS* vfs) { // compute sdf coefficients if (!plugin.active) { - tmd::TriangleMeshDistance sdf(vert_.data(), nvert(), face_.data(), nface()); - - std::vector coeffs(octree_.NumVerts()); - std::vector processed(octree_.NumVerts(), false); - std::deque queue; - - if (octree_.NumNodes() > 0) { - queue.push_back(0); // start traversal from the root node - } - - while (!queue.empty()) { - int node_idx = queue.front(); - queue.pop_front(); - - for (int j = 0; j < 8; ++j) { - int vert_id = octree_.VertId(node_idx, j); - if (processed[vert_id]) { - continue; - } - if (octree_.Hang(vert_id).empty()) { - coeffs[vert_id] = sdf.signed_distance(octree_.Vert(vert_id)).distance; - } else { - double sum_coeff = 0; - for (int dep_id : octree_.Hang(vert_id)) { - sum_coeff += coeffs[dep_id]; - if (!processed[dep_id]) { - throw mjCError(this, "sdf coefficient computation failed"); - } - } - coeffs[vert_id] = sum_coeff / octree_.Hang(vert_id).size(); - } - processed[vert_id] = true; - } - - for (int child_idx : octree_.Children(node_idx)) { - if (child_idx != -1) { - queue.push_back(child_idx); - } - } - } - - for (int i = 0; i < octree_.NumNodes(); ++i) { - for (int j = 0; j < 8; j++) { - octree_.AddCoeff(i, j, coeffs[octree_.VertId(i, j)]); - } - } + octree_.ComputeSdfCoeffs(vert_.data(), nvert(), face_.data(), nface(), tree_); } } @@ -2737,17 +2691,17 @@ class MeshPolygon { public: // constructors (need starting face) MeshPolygon(const double v1[3], const double v2[3], const double v3[3], - int v1i, int v2i, int v3i); + int v1i, int v2i, int v3i, double theta, double phi); MeshPolygon() = delete; MeshPolygon(const MeshPolygon&) = delete; MeshPolygon& operator=(const MeshPolygon&) = delete; + MeshPolygon(MeshPolygon&&) = default; + MeshPolygon& operator=(MeshPolygon&&) = default; - void InsertFace(int v1, int v2, int v3); // insert a face into the polygon - std::vector> Paths() const; // return trace of the polygons - const double* Normal() const { return normal_; } // return the normal of the polygon - - // return the ith component of the normal of the polygon - double Normal(int i) const { return normal_[i]; } + void InsertFace(int v1, int v2, int v3); // insert a face into the polygon + std::vector> Paths() const; // return trace of the polygons + const double* Normal() const { return normal_; } // return the normal of the polygon + double Normal(int i) const { return normal_[i]; } // return the i-th component of the normal private: std::vector> edges_; @@ -2760,40 +2714,46 @@ class MeshPolygon { void CombineIslands(int& island1, int& island2); }; +bool MeshPolygonKey(std::pair& angles, const double v1[3], const double v2[3], + const double v3[3], double angle_tol) { + double diff12[3] = {v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2]}; + double diff13[3] = {v3[0] - v1[0], v3[1] - v1[1], v3[2] - v1[2]}; + double normal[3], norm; - -MeshPolygon::MeshPolygon(const double v1[3], const double v2[3], const double v3[3], - int v1i, int v2i, int v3i) { - mjuu_makenormal(normal_, v1, v2, v3); - edges_ = {{v1i, v2i}, {v2i, v3i}, {v3i, v1i}}; - nisland_ = 1; - islands_ = {0, 0, 0}; -} - - - -// comparison operator for std::set -bool PolygonCmp(const MeshPolygon& p1, const MeshPolygon& p2) { - const double* n1 = p1.Normal(); - const double* n2 = p2.Normal(); - double dot3 = n1[0] * n2[0] + n1[1] * n2[1] + n1[2] * n2[2]; - - // TODO(kylebayes): The tolerance should be a parameter set the user, as it should be optimized - // from mesh to mesh. - if (dot3 > 0.99999872) { + mjuu_crossvec(normal, diff12, diff13); + if ((norm = std::sqrt(mjuu_dot3(normal, normal))) < mjMINVAL) { return false; } - if (std::abs(n1[0] - n2[0]) > mjMINVAL) { - return n1[0] > n2[0]; + // atan2 is sensitive to sign of 0.0, adding 0.0 to enforcing only positive 0.0 + normal[0] = (normal[0] / norm) + 0.0; + normal[1] = (normal[1] / norm) + 0.0; + normal[2] = (normal[2] / norm) + 0.0; + double rtheta = 0.0, rphi = 0.0; + + // clamp normal to be in valid range for acos + if (std::abs(normal[2]) > 1.0 - 1e-7) { + if (normal[2] < 0) rphi = std::round(mjPI / angle_tol); + angles = std::make_pair(rtheta, rphi); + return true; } - if (std::abs(n1[1] - n2[1]) > mjMINVAL) { - return n1[1] > n2[1]; - } - if (std::abs(n1[2] - n2[2]) > mjMINVAL) { - return n1[2] > n2[2]; - } - return false; + // rounded azimuthal and polar angles + rtheta = std::round(std::atan2(normal[1], normal[0]) / angle_tol); + rphi = std::round(std::acos(normal[2]) / angle_tol); + angles = std::make_pair(rtheta, rphi); + return true; +} + + +MeshPolygon::MeshPolygon(const double v1[3], const double v2[3], const double v3[3], + int v1i, int v2i, int v3i, double theta, double phi) { + normal_[0] = std::cos(theta) * std::sin(phi); + normal_[1] = std::sin(theta) * std::sin(phi); + normal_[2] = std::cos(phi); + + edges_ = {{v1i, v2i}, {v2i, v3i}, {v3i, v1i}}; + nisland_ = 1; + islands_ = {0, 0, 0}; } @@ -2892,8 +2852,8 @@ void MeshPolygon::InsertFace(int v1, int v2, int v3) { // return the transverse vertices of the polygon, multiple paths possible if not connected -std::vector > MeshPolygon::Paths() const { - std::vector > paths; +std::vector> MeshPolygon::Paths() const { + std::vector> paths; // shortcut if polygon is just a triangular face if (edges_.size() == 3) { return {{edges_[0].first, edges_[1].first, edges_[2].first}}; @@ -2945,9 +2905,20 @@ std::vector > MeshPolygon::Paths() const { +// hash function for std::pair +struct PairHash { + template + std::size_t operator() (const std::pair& pair) const { + return std::hash()(pair.first) ^ std::hash()(pair.second); + } +}; + + + // merge coplanar mesh triangular faces into polygonal sides to represent the geometry of the mesh void mjCMesh::MakePolygons() { - std::set polygons(PolygonCmp); + constexpr double kAngleTol = 0.01; + std::unordered_map, MeshPolygon, PairHash> mesh_polygons; polygons_.clear(); polygon_normals_.clear(); polygon_map_.clear(); @@ -2969,22 +2940,30 @@ void mjCMesh::MakePolygons() { // process each face for (int i = 0; i < nfaces; i++) { - double* v1 = &vert_[3*faces[3*i + 0]]; - double* v2 = &vert_[3*faces[3*i + 1]]; - double* v3 = &vert_[3*faces[3*i + 2]]; + int vi1 = faces[3*i + 0]; + int vi2 = faces[3*i + 1]; + int vi3 = faces[3*i + 2]; + double* v1 = &vert_[3*vi1]; + double* v2 = &vert_[3*vi2]; + double* v3 = &vert_[3*vi3]; - MeshPolygon face(v1, v2, v3, faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); - auto it = polygons.find(face); - if (it == polygons.end()) { - polygons.emplace(v1, v2, v3, faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); + std::pair key; + if (!MeshPolygonKey(key, v1, v2, v3, kAngleTol)) { + continue; + } + auto it = mesh_polygons.find(key); + if (it == mesh_polygons.end()) { + double theta = kAngleTol * key.first; + double phi = kAngleTol * key.second; + mesh_polygons.emplace(key, MeshPolygon(v1, v2, v3, vi1, vi2, vi3, theta, phi)); } else { - MeshPolygon& p = const_cast(*it); - p.InsertFace(faces[3*i + 0], faces[3*i + 1], faces[3*i + 2]); + it->second.InsertFace(vi1, vi2, vi3); } } - for (const auto& polygon : polygons) { - std::vector > paths = polygon.Paths(); + for (const auto& pair : mesh_polygons) { + const MeshPolygon& polygon = pair.second; + std::vector> paths = polygon.Paths(); // separate the polygons if they were grouped together for (const auto& path : paths) { @@ -3434,15 +3413,6 @@ void mjCSkin::LoadSKN(mjResource* resource) { //-------------------------- nonlinear elasticity -------------------------------------------------- -// hash function for std::pair -struct PairHash -{ - template - std::size_t operator() (const std::pair& pair) const { - return std::hash()(pair.first) ^ std::hash()(pair.second); - } -}; - // simplex connectivity constexpr int eledge[3][6][2] = {{{ 0, 1}, {-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}}, @@ -4099,6 +4069,88 @@ void mjCFlex::ResolveReferences(const mjCModel* m) { } +std::string mjCFlex::ComputeStiffnessCacheKey() const { + std::size_t hash = 0; + auto combine = [&hash](std::size_t v) { + hash ^= v + 0x9e3779b9 + (hash << 6) + (hash >> 2); + }; + + combine(std::hash{}(young)); + combine(std::hash{}(poisson)); + combine(std::hash{}(order_)); + + // compute bounding box from vertex positions + if (!vert_.empty()) { + double minx = vert_[0], maxx = vert_[0]; + double miny = vert_[1], maxy = vert_[1]; + double minz = vert_[2], maxz = vert_[2]; + for (std::size_t i = 3; i < vert_.size(); i += 3) { + minx = std::min(minx, vert_[i]); + maxx = std::max(maxx, vert_[i]); + miny = std::min(miny, vert_[i + 1]); + maxy = std::max(maxy, vert_[i + 1]); + minz = std::min(minz, vert_[i + 2]); + maxz = std::max(maxz, vert_[i + 2]); + } + combine(std::hash{}(maxx - minx)); + combine(std::hash{}(maxy - miny)); + combine(std::hash{}(maxz - minz)); + } + + for (std::size_t i = 0; i < vert_.size(); i += std::max(1, (int)vert_.size()/100)) { + combine(std::hash{}(vert_[i])); + } + + for (std::size_t i = 0; i < shell.size(); i += std::max(1, (int)shell.size()/50)) { + combine(std::hash{}(shell[i])); + } + + return "flex_stiffness:" + std::to_string(hash); +} + + +bool mjCFlex::LoadCachedStiffness() { + mjCCache* cache = reinterpret_cast(mj_getCache()->impl_); + if (!cache) return false; + + std::string key = ComputeStiffnessCacheKey(); + + auto load_fn = [this](const void* data) { + const auto* cached = static_cast*>(data); + stiffness = *cached; + return true; + }; + + mjResource dummy_resource{}; + dummy_resource.name = const_cast(key.c_str()); + dummy_resource.timestamp[0] = '\0'; + + return cache->PopulateData(key, &dummy_resource, load_fn); +} + + +void mjCFlex::CacheStiffness() { + mjCCache* cache = reinterpret_cast(mj_getCache()->impl_); + if (!cache || stiffness.empty()) return; + + std::string key = ComputeStiffnessCacheKey(); + + auto* cached = new std::vector(stiffness); + + std::size_t size = sizeof(*cached) + sizeof(double) * stiffness.size(); + + std::shared_ptr cached_data(cached, [](const void* data) { + delete static_cast*>(data); + }); + + mjResource dummy_resource{}; + dummy_resource.name = const_cast(key.c_str()); + dummy_resource.timestamp[0] = '\0'; + + cache->Insert("", key, &dummy_resource, cached_data, size); +} + + // compiler void mjCFlex::Compile(const mjVFS* vfs) { CopyFromSpec(); @@ -4341,7 +4393,6 @@ void mjCFlex::Compile(const mjVFS* vfs) { if (min_size > nelem) { throw mjCError(this, "Trilinear dofs are require at least %d elements", "", min_size); } - ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_); } // geometrically nonlinear elasticity @@ -4391,6 +4442,16 @@ void mjCFlex::Compile(const mjVFS* vfs) { // create shell fragments and element-vertex collision pairs CreateShellPair(); + // compute linear stiffness for interpolated elements (cached) + bool stiffness_cached = false; + if (young > 0 && interpolated) { + stiffness_cached = LoadCachedStiffness(); + } + + if (!stiffness_cached && young > 0 && interpolated) { + ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_); + } + // create bounding volume hierarchy CreateBVH(); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index c2f19df6..6c642352 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -907,6 +907,7 @@ void mjCModel::ComputeSparseSizes() { // no dofs, quick return if (nv == 0) { nM = nD = nB = nC = 0; + nJten = 0; return; } @@ -1084,6 +1085,37 @@ void mjCModel::ComputeSparseSizes() { } } nC = nOD + nv; + + nJten = 0; + if (nv > 0) { + std::vector dof_bitmap(nv, false); + for (const auto* tendon : tendons_) { + if (!tendon->path.empty() && + tendon->path[0]->Type() == mjWRAP_JOINT) { + nJten += tendon->path.size(); + continue; + } + + std::fill(dof_bitmap.begin(), dof_bitmap.end(), false); + for (const auto* wrap : tendon->path) { + int bodyid = GetBodyIdFromWrap(wrap); + if (bodyid > 0) { + mjCBody* b = bodies_[bodyid]; + while (b && b->id > 0) { + for (const auto* jnt : b->joints) { + for (int k = 0; k < jnt->nv(); k++) { + dof_bitmap[jnt->dofadr_ + k] = true; + } + } + b = b->GetParent(); + } + } + } + for (int j = 0; j < nv; j++) { + nJten += dof_bitmap[j]; + } + } + } } @@ -3171,7 +3203,46 @@ void mjCModel::CopyPlugins(mjModel* m) { -// compute non-zeros in actuator_moment matrix +// compute number of dofs for a given tendon +int mjCModel::CountTendonDofs(const mjModel* m, int id) { + std::vector dof_used(m->nv, false); + int nv = m->nv; + int adr = m->tendon_adr[id]; + int num = m->tendon_num[id]; + + if (m->wrap_type[adr] == mjWRAP_JOINT) { + return num; + } + + std::fill(dof_used.begin(), dof_used.end(), false); + for (int j = 0; j < num; j++) { + int type = m->wrap_type[adr + j]; + int bodyid = -1; + if (type == mjWRAP_SITE) { + bodyid = m->site_bodyid[m->wrap_objid[adr + j]]; + } else if (type == mjWRAP_SPHERE || type == mjWRAP_CYLINDER) { + bodyid = m->geom_bodyid[m->wrap_objid[adr + j]]; + } + if (bodyid > 0) { + int bid = bodyid; + while (bid > 0) { + int bdofadr = m->body_dofadr[bid]; + int bdofnum = m->body_dofnum[bid]; + for (int k = 0; k < bdofnum; k++) { + dof_used[bdofadr + k] = true; + } + bid = m->body_parentid[bid]; + } + } + } + + int count = 0; + for (int j = 0; j < nv; j++) { + count += dof_used[j]; + } + return count; +} + int mjCModel::CountNJmom(const mjModel* m) { int nu = m->nu; int nv = m->nv; @@ -3206,7 +3277,7 @@ int mjCModel::CountNJmom(const mjModel* m) { break; case mjTRN_TENDON: - count += nv; + count += CountTendonDofs(m, id); break; case mjTRN_SITE: @@ -3228,12 +3299,13 @@ int mjCModel::CountNJmom(const mjModel* m) { // compute non-zeros in ten_J matrix int mjCModel::CountNJten(const mjModel* m) { - int nv = m->nv; int ntendon = m->ntendon; - // conservative upper bound: each tendon can have at most nv non-zeros - // TODO(taylorhowell): compute tighter bound - int count = ntendon * nv; + int count = 0; + for (int i = 0; i < ntendon; i++) { + count += CountTendonDofs(m, i); + } + return count; } @@ -5033,7 +5105,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly, nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert, nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude, - neq, ntendon, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata, + neq, ntendon, nJten, nwrap, nsensor, nnumeric, nnumericdata, ntext, ntextdata, ntuple, ntupledata, nkey, nmocap, nplugin, npluginattr, nuser_body, nuser_jnt, nuser_geom, nuser_site, nuser_cam, nuser_tendon, nuser_actuator, nuser_sensor, nnames, npaths); @@ -5065,8 +5137,6 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // compute non-zeros in actuator_moment m->nJmom = nJmom = CountNJmom(m); - // compute non-zeros in ten_J - m->nJten = nJten = CountNJten(m); // scale mass if (compiler.settotalmass > 0) { diff --git a/src/user/user_model.h b/src/user/user_model.h index 93e4d3f4..025f8097 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -73,6 +73,7 @@ class mjCModel_ : public mjsElement { mjtSize nexclude; // number of excluded body pairs mjtSize neq; // number of equality constraints mjtSize ntendon; // number of tendons + mjtSize nJten; // number of non-zeros in sparse ten_J matrix mjtSize nsensor; // number of sensors mjtSize nnumeric; // number of numeric fields mjtSize ntext; // number of text fields @@ -130,7 +131,6 @@ class mjCModel_ : public mjsElement { mjtSize nC; // number of non-zeros in reduced sparse dof-dof matrix mjtSize nD; // number of non-zeros in sparse dof-dof matrix mjtSize nJmom; // number of non-zeros in sparse actuator_moment matrix - mjtSize nJten; // number of non-zeros in sparse ten_J matrix // statistics, as computed by mj_setConst double meaninertia_auto; // mean diagonal inertia, as computed by mj_setConst @@ -368,6 +368,8 @@ class mjCModel : public mjCModel_, private mjSpec { 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 CountTendonDofs(const mjModel* m, // compute number of dofs for a given tendon + int id); int CountNJmom(const mjModel* m); // compute number of non-zeros in actuator_moment matrix int CountNJten(const mjModel* m); // compute number of non-zeros in ten_J matrix diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 9dfa3346..a5c77e9c 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -25,10 +25,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -374,6 +376,7 @@ void mjCBoundingVolumeHierarchy::RemoveInactiveVolumes(int nmax) { bvleaf_.erase(bvleaf_.begin() + nmax, bvleaf_.end()); } + const mjCBoundingVolume* mjCBoundingVolumeHierarchy::AddBoundingVolume(int id, int contype, int conaffinity, const double* pos, const double* quat, @@ -552,6 +555,7 @@ void mjCOctree::CopyLevel(int* level) const { } } + void mjCOctree::CopyChild(int* child) const { for (int i = 0; i < node_.size(); ++i) { for (int j = 0; j < 8; ++j) { @@ -560,6 +564,7 @@ void mjCOctree::CopyChild(int* child) const { } } + void mjCOctree::CopyAabb(mjtNum* aabb) const { for (int i = 0; i < node_.size(); ++i) { aabb[i * 6 + 0] = (node_[i].aamm[0] + node_[i].aamm[3]) / 2; @@ -571,6 +576,7 @@ void mjCOctree::CopyAabb(mjtNum* aabb) const { } } + void mjCOctree::CopyCoeff(mjtNum* coeff) const { for (int i = 0; i < node_.size(); ++i) { for (int j = 0; j < 8; ++j) { @@ -579,6 +585,7 @@ void mjCOctree::CopyCoeff(mjtNum* coeff) const { } } + void mjCOctree::SetFace(const std::vector& vert, const std::vector& face) { for (int i = 0; i < face.size(); i += 3) { std::array v0 = {vert[3*face[i+0]], vert[3*face[i+0]+1], vert[3*face[i+0]+2]}; @@ -621,12 +628,302 @@ void mjCOctree::CreateOctree(const double aamm[6]) { } +namespace { + +double pointBoxDistSq(const double* p, const mjtNum* aabb) { + double dist_sq = 0; + for (int i = 0; i < 3; ++i) { + double lo = aabb[i] - aabb[i + 3]; + double hi = aabb[i] + aabb[i + 3]; + if (p[i] < lo) { + dist_sq += (lo - p[i]) * (lo - p[i]); + } else if (p[i] > hi) { + dist_sq += (p[i] - hi) * (p[i] - hi); + } + } + return dist_sq; +} + + +// compute squared distance between point p and triangle (v0, v1, v2), +// and return barycentric coordinates (u,v) of the closest point +double pointTriDistSqWithUV(const double* p, const double* v0, const double* v1, + const double* v2, double& out_u, double& out_v) { + double ab[3] = {v1[0] - v0[0], v1[1] - v0[1], v1[2] - v0[2]}; + double ac[3] = {v2[0] - v0[0], v2[1] - v0[1], v2[2] - v0[2]}; + double ap[3] = {p[0] - v0[0], p[1] - v0[1], p[2] - v0[2]}; + + // the closest point on the triangle is determined by partitioning space into Voronoi regions + double d1 = ab[0]*ap[0] + ab[1]*ap[1] + ab[2]*ap[2]; + double d2 = ac[0]*ap[0] + ac[1]*ap[1] + ac[2]*ap[2]; + + // region A (vertex v0) + if (d1 <= 0 && d2 <= 0) { + out_u = 0; out_v = 0; + return ap[0]*ap[0] + ap[1]*ap[1] + ap[2]*ap[2]; + } + + double bp[3] = {p[0] - v1[0], p[1] - v1[1], p[2] - v1[2]}; + double d3 = ab[0]*bp[0] + ab[1]*bp[1] + ab[2]*bp[2]; + double d4 = ac[0]*bp[0] + ac[1]*bp[1] + ac[2]*bp[2]; + + // region B (vertex v1) + if (d3 >= 0 && d4 <= d3) { + out_u = 1; out_v = 0; + return bp[0]*bp[0] + bp[1]*bp[1] + bp[2]*bp[2]; + } + + // region AB (edge v0-v1) + double vc = d1*d4 - d3*d2; + if (vc <= 0 && d1 >= 0 && d3 <= 0) { + double u = d1 / (d1 - d3); + out_u = u; out_v = 0; + double closest[3] = {v0[0] + u*ab[0], v0[1] + u*ab[1], v0[2] + u*ab[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + double cp[3] = {p[0] - v2[0], p[1] - v2[1], p[2] - v2[2]}; + double d5 = ab[0]*cp[0] + ab[1]*cp[1] + ab[2]*cp[2]; + double d6 = ac[0]*cp[0] + ac[1]*cp[1] + ac[2]*cp[2]; + + // region C (vertex v2) + if (d6 >= 0 && d5 <= d6) { + out_u = 0; out_v = 1; + return cp[0]*cp[0] + cp[1]*cp[1] + cp[2]*cp[2]; + } + + // region AC (edge v0-v2) + double vb = d5*d2 - d1*d6; + if (vb <= 0 && d2 >= 0 && d6 <= 0) { + double v = d2 / (d2 - d6); + out_u = 0; out_v = v; + double closest[3] = {v0[0] + v*ac[0], v0[1] + v*ac[1], v0[2] + v*ac[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + // region BC (edge v1-v2) + double va = d3*d6 - d5*d4; + if (va <= 0 && (d4 - d3) >= 0 && (d5 - d6) >= 0) { + double w = (d4 - d3) / ((d4 - d3) + (d5 - d6)); + out_u = 1 - w; out_v = w; + double bc[3] = {v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2]}; + double closest[3] = {v1[0] + w*bc[0], v1[1] + w*bc[1], v1[2] + w*bc[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); + } + + // region ABC (inside triangle) + double denom = 1.0 / (va + vb + vc); + double u = vb * denom; + double v = vc * denom; + out_u = u; out_v = v; + double closest[3] = {v0[0] + u*ab[0] + v*ac[0], + v0[1] + u*ab[1] + v*ac[1], + v0[2] + u*ab[2] + v*ac[2]}; + return (p[0]-closest[0])*(p[0]-closest[0]) + + (p[1]-closest[1])*(p[1]-closest[1]) + + (p[2]-closest[2])*(p[2]-closest[2]); +} + + +// query BVH for closest face to point p, return distance, face index and barycentric coordinates +void queryClosestBVHWithFace(const mjtNum* bvh, const int* child, const int* nodeid, + const double* vert, const int* face, int node_idx, + const double* p, double& best_dist_sq, + int& best_face, double& best_u, double& best_v) { + const mjtNum* aabb = &bvh[node_idx * 6]; + if (pointBoxDistSq(p, aabb) >= best_dist_sq) return; + + int left = child[node_idx * 2]; + int right = child[node_idx * 2 + 1]; + + if (left == -1 && right == -1) { + int fi = nodeid[node_idx]; + if (fi >= 0) { + const double* v0 = vert + face[fi * 3 + 0] * 3; + const double* v1 = vert + face[fi * 3 + 1] * 3; + const double* v2 = vert + face[fi * 3 + 2] * 3; + double u, v; + double dist_sq = pointTriDistSqWithUV(p, v0, v1, v2, u, v); + if (dist_sq < best_dist_sq) { + best_dist_sq = dist_sq; + best_face = fi; + best_u = u; + best_v = v; + } + } + return; + } + + if (left >= 0) { + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, left, p, + best_dist_sq, best_face, best_u, best_v); + } + if (right >= 0) { + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, right, p, + best_dist_sq, best_face, best_u, best_v); + } +} + + +double querySignedDistance(const mjtNum* bvh, const int* child, const int* nodeid, + int nbvh, const double* point, + const double* vert, const int* face) { + if (nbvh == 0) { + return 0; + } + + double best_dist_sq = 1e20; + int best_face = -1; + double best_u = 0, best_v = 0; + queryClosestBVHWithFace(bvh, child, nodeid, vert, face, 0, point, + best_dist_sq, best_face, best_u, best_v); + double dist = std::sqrt(best_dist_sq); + + double sign = 1.0; + if (best_face >= 0) { + const double* v0 = vert + face[best_face * 3 + 0] * 3; + const double* v1 = vert + face[best_face * 3 + 1] * 3; + const double* v2 = vert + face[best_face * 3 + 2] * 3; + + double e1[3] = {v1[0]-v0[0], v1[1]-v0[1], v1[2]-v0[2]}; + double e2[3] = {v2[0]-v0[0], v2[1]-v0[1], v2[2]-v0[2]}; + double normal[3] = { + e1[1]*e2[2] - e1[2]*e2[1], + e1[2]*e2[0] - e1[0]*e2[2], + e1[0]*e2[1] - e1[1]*e2[0] + }; + + double closest[3] = { + v0[0] + best_u*(v1[0]-v0[0]) + best_v*(v2[0]-v0[0]), + v0[1] + best_u*(v1[1]-v0[1]) + best_v*(v2[1]-v0[1]), + v0[2] + best_u*(v1[2]-v0[2]) + best_v*(v2[2]-v0[2]) + }; + + double u[3] = {point[0]-closest[0], point[1]-closest[1], point[2]-closest[2]}; + double dot = u[0]*normal[0] + u[1]*normal[1] + u[2]*normal[2]; + double normal_len = mjuu_normvec(normal, 3); + double eps = 1e-12 * normal_len * dist; + sign = (dot > eps) ? 1.0 : -1.0; + } + + return sign * dist; +} + +} // namespace + + +double mjCBoundingVolumeHierarchy::QuerySignedDistance( + const double* point, const double* vert, const int* face) const { + return querySignedDistance(bvh_.data(), child_.data(), nodeid_.data(), + nbvh_, point, vert, face); +} + + +void mjCOctree::ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface, + const mjCBoundingVolumeHierarchy& tree) { + std::vector coeffs(nvert_, 0.0); + std::vector processed(nvert_, false); + std::deque queue; + + if (NumNodes() > 0) { + queue.push_back(0); + } + + while (!queue.empty()) { + int node_idx = queue.front(); + queue.pop_front(); + + // compute SDF coefficients at the 8 vertices of the octree node + for (int j = 0; j < 8; ++j) { + int vert_id = VertId(node_idx, j); + if (processed[vert_id]) { + continue; + } + if (Hang(vert_id).empty()) { + // transform from octree frame (body inertial) back to mesh frame + double p_mesh[3]; + mjuu_rotVecQuat(p_mesh, Vert(vert_id), iquat_); + p_mesh[0] += ipos_[0]; + p_mesh[1] += ipos_[1]; + p_mesh[2] += ipos_[2]; + + coeffs[vert_id] = tree.QuerySignedDistance(p_mesh, vert, face); + } else { + // hanging node: interpolate from parents + double sum_coeff = 0; + for (int dep_id : Hang(vert_id)) { + sum_coeff += coeffs[dep_id]; + } + coeffs[vert_id] = sum_coeff / Hang(vert_id).size(); + } + processed[vert_id] = true; + } + + // add children to the queue + for (int child_idx : Children(node_idx)) { + if (child_idx != -1) { + queue.push_back(child_idx); + } + } + } + + // optional Laplacian smoothing (smooths octree level transitions) + if (smoothing_iterations_ > 0) { + // build vertex neighbor graph from octree connectivity + std::vector> neighbors(nvert_); + for (int i = 0; i < NumNodes(); ++i) { + static const int edges[12][2] = { + {0, 1}, {2, 3}, {4, 5}, {6, 7}, + {0, 2}, {1, 3}, {4, 6}, {5, 7}, + {0, 4}, {1, 5}, {2, 6}, {3, 7} + }; + for (const auto& edge : edges) { + int v0 = VertId(i, edge[0]); + int v1 = VertId(i, edge[1]); + neighbors[v0].insert(v1); + neighbors[v1].insert(v0); + } + } + + // apply Laplacian smoothing + const double alpha = 0.2; + std::vector sdf_new(nvert_); + for (int iter = 0; iter < smoothing_iterations_; ++iter) { + for (int i = 0; i < nvert_; ++i) { + if (neighbors[i].empty()) { + sdf_new[i] = coeffs[i]; + } else { + double avg = 0; + for (int j : neighbors[i]) avg += coeffs[j]; + avg /= neighbors[i].size(); + sdf_new[i] = (1 - alpha) * coeffs[i] + alpha * avg; + } + } + std::swap(coeffs, sdf_new); + } + } + + // copy coefficients to the octree nodes + for (int i = 0; i < NumNodes(); ++i) { + for (int j = 0; j < 8; j++) { + AddCoeff(i, j, coeffs[VertId(i, j)]); + } + } +} + + static double dot2(const double* a, const double* b) { return a[0] * b[0] + a[1] * b[1]; } -// From M. Schwarz and H.-P. Seidel, "Fast Parallel Surface and Solid Voxelization on GPUs". +// from M. Schwarz and H.-P. Seidel, "Fast Parallel Surface and Solid Voxelization on GPUs". static bool boxTriangle(const Triangle& v, const double aamm[6]) { // bounding box tests for (int i = 0; i < 3; i++) { @@ -6267,12 +6564,16 @@ void mjCTendon::ResolveReferences(const mjCModel* m) { try { // look for wrapped element with namespace path[i]->name = prefix + pname + suffix; - path[i]->sidesite = prefix + psidesite + suffix; + if (!psidesite.empty()) { + path[i]->sidesite = prefix + psidesite + suffix; + } path[i]->ResolveReferences(m); } catch(mjCError) { // remove namespace from wrap names path[i]->name = pname; - path[i]->sidesite = psidesite; + if (!psidesite.empty()) { + path[i]->sidesite = psidesite; + } path[i]->ResolveReferences(m); nfailure++; } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 081791d7..8399a146 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -200,6 +200,10 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ { + sizeof(int) * nodeid_.size() + sizeof(int) * level_.size(); } + // query signed distance from point to mesh surface + double QuerySignedDistance(const double* point, const double* vert, + const int* face) const; + private: // internal class used during BVH construction, for partial sorting of bounding volumes struct BVElement { @@ -275,6 +279,7 @@ struct mjCOctree_ { std::vector> hang_; // hanging nodes status (nvert x 1) double ipos_[3] = {0, 0, 0}; double iquat_[4] = {1, 0, 0, 0}; + int smoothing_iterations_ = 0; // Laplacian smoothing iterations (0 = disabled) }; class mjCOctree : public mjCOctree_ { @@ -301,6 +306,15 @@ class mjCOctree : public mjCOctree_ { face_.clear(); } void AddCoeff(int n, int v, double coeff) { node_[n].coeff[v] = coeff; } + double Coeff(int n, int v) const { return node_[n].coeff[v]; } + + // Set number of Laplacian smoothing iterations (0 = disabled, default) + void SetSmoothingIterations(int iterations) { smoothing_iterations_ = iterations; } + int SmoothingIterations() const { return smoothing_iterations_; } + + // compute SDF coefficients via BVH queries, optionally with Laplacian smoothing + void ComputeSdfCoeffs(const double* vert, int nvert, const int* face, int nface, + const mjCBoundingVolumeHierarchy& tree); private: void Make(std::vector& elements); @@ -988,6 +1002,9 @@ class mjCFlex_ : public mjCBase { std::vector spec_elem_; std::vector spec_texcoord_; std::vector spec_elemtexcoord_; + + // caching + std::vector cached_stiffness_; // cached stiffness matrix }; class mjCFlex: public mjCFlex_, private mjsFlex { @@ -1036,6 +1053,11 @@ class mjCFlex: public mjCFlex_, private mjsFlex { std::vector node0_; // node Cartesian positions int order_ = 0; // interpolation order + + // stiffness caching + std::string ComputeStiffnessCacheKey() const; + bool LoadCachedStiffness(); + void CacheStiffness(); }; diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc index d06a5b02..cfcba486 100644 --- a/test/engine/engine_collision_gjk_test.cc +++ b/test/engine/engine_collision_gjk_test.cc @@ -1373,7 +1373,7 @@ TEST_F(MjGjkTest, BoxMesh) { std::vector dir, pos; mjtNum dist; int ncons = Penetration(status, dist, dir, pos, model, data, g2, g1, 0, 1000); - + EXPECT_EQ(model->nmeshpoly, 7); EXPECT_EQ(ncons, 4); mj_deleteData(data); mj_deleteModel(model); diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 53f36b38..0449cb4f 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -154,12 +154,8 @@ TEST_F(CoreSmoothTest, FixedTendonSortedIndices) { mjData* data = mj_makeData(model); mj_fwdPosition(model, data); - int rowadr = data->ten_J_rowadr[0]; - int* colind = data->ten_J_colind + rowadr; - mjtNum* J = data->ten_J + rowadr; - + mjtNum* J = data->ten_J; EXPECT_THAT(vector(J, J + 3), ElementsAre(1, 2, 3)); - EXPECT_THAT(vector(colind, colind + 3), ElementsAre(0, 1, 2)); mj_deleteData(data); mj_deleteModel(model); @@ -253,15 +249,11 @@ TEST_F(CoreSmoothTest, TendonArmature) { // 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); - } + int rowadr = m->ten_J_rowadr[j]; + int* rownnz = m->ten_J_rownnz + j; + int zero = 0; + mju_sparse2dense(ten_J.data(), d->ten_J + rowadr, 1, nv, + rownnz, &zero, m->ten_J_colind + rowadr); // get tendon inertia only, using outer product mju_mulMatMat(ten_M.data(), ten_J.data(), ten_J.data(), nv, 1, nv); diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index 58e0fe73..79df5022 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -17,6 +17,7 @@ #include "src/engine/engine_forward.h" #include "src/engine/engine_derivative.h" +#include #include #include #include @@ -1816,5 +1817,68 @@ TEST_F(ForwardTest, FlexParentCoupling) { mj_deleteModel(model); } + +TEST_F(ForwardTest, TrilinearPinnedParentWithFreejoint) { + 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); + + int parent_id = mj_name2id(m, mjOBJ_BODY, "parent"); + ASSERT_GT(parent_id, 0); + + EXPECT_EQ(m->nflexnode, 8); + EXPECT_EQ(m->body_dofnum[parent_id], 0) << "parent body should have 0 DOFs"; + + int freejoint_body = m->body_parentid[parent_id]; + EXPECT_EQ(m->body_dofnum[freejoint_body], 6) << "freejoint body has 6 DOFs"; + + mj_resetData(m, d); + mj_forward(m, d); + + for (int i = 0; i < 500; i++) { + mj_step(m, d); + + ASSERT_FALSE(mju_isBad(d->qpos[0])) + << "Simulation became unstable at step " << i; + ASSERT_FALSE(mju_isBad(d->qvel[0])) + << "Velocity became unstable at step " << i; + + for (int j = 0; j < m->nq; j++) { + ASSERT_LT(mju_abs(d->qpos[j]), 100.0) + << "Position exploded at step " << i << ", qpos[" << j + << "]=" << d->qpos[j]; + } + for (int j = 0; j < m->nv; j++) { + ASSERT_LT(mju_abs(d->qvel[j]), 1000.0) + << "Velocity exploded at step " << i << ", qvel[" << j + << "]=" << d->qvel[j]; + } + } + + mj_deleteData(d); + mj_deleteModel(m); +} + } // namespace } // namespace mujoco diff --git a/test/engine/engine_plugin_test.cc b/test/engine/engine_plugin_test.cc index 40ab13d4..180ab9c8 100644 --- a/test/engine/engine_plugin_test.cc +++ b/test/engine/engine_plugin_test.cc @@ -376,10 +376,10 @@ int RegisterNoAttributePlugin() { return mjp_registerPlugin(&plugin); } -class EnginePluginTest : public PluginTest { +class EnginePluginTest : public MujocoTest { public: // register all plugins - EnginePluginTest() : PluginTest() { + EnginePluginTest() : MujocoTest() { RegisterSensorPlugin(); for (int i = 1; i <= kNumFakePlugins; ++i) { @@ -466,7 +466,7 @@ TEST_F(MujocoTest, EmptyPluginDisallowed) { mj_deleteModel(m); } -TEST_F(PluginTest, FirstPartyPlugins) { +TEST_F(MujocoTest, FirstPartyPlugins) { EXPECT_THAT(mjp_pluginCount(), kNumTruePlugins); } diff --git a/test/engine/engine_ray_test.cc b/test/engine/engine_ray_test.cc index c55583f2..31433a33 100644 --- a/test/engine/engine_ray_test.cc +++ b/test/engine/engine_ray_test.cc @@ -82,7 +82,7 @@ using ::testing::DoubleNear; using ::testing::ElementsAre; using ::testing::NotNull; using ::testing::Pointwise; -using RayTest = PluginTest; +using RayTest = MujocoTest; TEST_F(RayTest, NoExclusions) { char error[1024]; diff --git a/test/fixture.cc b/test/fixture.cc index 4b18832d..f621330e 100644 --- a/test/fixture.cc +++ b/test/fixture.cc @@ -238,16 +238,12 @@ mjtNum CompareModel(const mjModel* m1, const mjModel* m2, MJMODEL_POINTERS_PREAMBLE(m1); // compare ints, exclude nbuffer because it hides the actual difference -// TODO(kylebayes): re-enable poly comparisons. -#define X(name) \ - if constexpr (std::string_view(#name) != "nbuffer" && \ - std::string_view(#name) != "nmeshpolymap" && \ - std::string_view(#name) != "nmeshpolyvert" && \ - std::string_view(#name) != "nmeshpoly") { \ - if (m1->name != m2->name) { \ - maxdif = std::abs((long)m1->name - (long)m2->name); \ - field = #name; \ - } \ +#define X(name) \ + if constexpr (std::string_view(#name) != "nbuffer") { \ + if (m1->name != m2->name) { \ + maxdif = std::abs((long)m1->name - (long)m2->name); \ + field = #name; \ + } \ } MJMODEL_SIZES #undef X diff --git a/test/fixture.h b/test/fixture.h index f9d0cb1f..64759b08 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -16,9 +16,12 @@ #define MUJOCO_TEST_FIXTURE_H_ #include +#include // IWYU pragma: keep +#include // IWYU pragma: keep #include #include #include +#include // IWYU pragma: keep #include #include #include @@ -54,6 +57,21 @@ class MujocoErrorTestGuard { // trigger a test failure. class MujocoTest : public ::testing::Test { public: + MujocoTest() { + static std::once_flag flag; + std::call_once(flag, []() { + const char* plugin_dir = std::getenv("MUJOCO_PLUGIN_DIR"); + if (plugin_dir) { + mj_loadAllPluginLibraries( + plugin_dir, +[](const char* filename, int first, int count) { + std::printf("Plugins registered by library '%s':\n", filename); + for (int i = first; i < first + count; ++i) { + std::printf(" %s\n", mjp_getPluginAtSlot(i)->name); + } + }); + } + }); + } ~MujocoTest() { mj_freeLastXML(); } private: @@ -183,20 +201,6 @@ class MockFilesystem { std::string dir_; // current directory }; -// Installs all plugins -class PluginTest : public MujocoTest { - public: - // load plugin library - PluginTest() : MujocoTest() { - mj_loadAllPluginLibraries( - std::string(std::getenv("MUJOCO_PLUGIN_DIR")).c_str(), +[](const char* filename, int first, int count) { - std::printf("Plugins registered by library '%s':\n", filename); - for (int i = first; i < first + count; ++i) { - std::printf(" %s\n", mjp_getPluginAtSlot(i)->name); - } - }); - } -}; } // namespace mujoco #endif // MUJOCO_TEST_FIXTURE_H_ diff --git a/test/plugin/actuator/pid_test.cc b/test/plugin/actuator/pid_test.cc index 8dfac5cf..e5fcb32c 100644 --- a/test/plugin/actuator/pid_test.cc +++ b/test/plugin/actuator/pid_test.cc @@ -29,7 +29,7 @@ namespace mujoco { namespace { -using PidTest = PluginTest; +using PidTest = MujocoTest; using ::testing::DoubleNear; using ::testing::HasSubstr; using ::testing::IsNull; diff --git a/test/plugin/elasticity/elasticity_test.cc b/test/plugin/elasticity/elasticity_test.cc index c7a06439..d2817787 100644 --- a/test/plugin/elasticity/elasticity_test.cc +++ b/test/plugin/elasticity/elasticity_test.cc @@ -26,7 +26,7 @@ namespace mujoco { namespace { -using ElasticityTest = PluginTest; +using ElasticityTest = MujocoTest; // -------------------------------- cable ----------------------------------- diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 73e36a48..82cb2a08 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -171,7 +171,7 @@ TEST_F(MujocoTest, TreeTraversal) { mj_deleteSpec(spec); } -TEST_F(PluginTest, ActivatePlugin) { +TEST_F(MujocoTest, ActivatePlugin) { mjSpec* spec = mj_makeSpec(); mjs_activatePlugin(spec, "mujoco.elasticity.cable"); @@ -196,7 +196,7 @@ TEST_F(PluginTest, ActivatePlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, DeletePlugin) { +TEST_F(MujocoTest, DeletePlugin) { mjSpec* spec = mj_makeSpec(); mjs_activatePlugin(spec, "mujoco.pid"); @@ -267,7 +267,7 @@ static constexpr char xml_plugin_2[] = R"( )"; -TEST_F(PluginTest, AttachPlugin) { +TEST_F(MujocoTest, AttachPlugin) { std::array err; mjSpec* parent = mj_parseXMLString(xml_plugin_1, 0, err.data(), err.size()); ASSERT_THAT(parent, NotNull()) << err.data(); @@ -316,7 +316,7 @@ TEST_F(PluginTest, AttachPlugin) { mj_deleteSpec(spec_3); } -TEST_F(PluginTest, DetachPlugin) { +TEST_F(MujocoTest, DetachPlugin) { std::array err; mjSpec* parent = mj_parseXMLString(xml_plugin_1, 0, err.data(), err.size()); ASSERT_THAT(parent, NotNull()) << err.data(); @@ -342,7 +342,7 @@ TEST_F(PluginTest, DetachPlugin) { mj_deleteSpec(child); } -TEST_F(PluginTest, AttachExplicitPlugin) { +TEST_F(MujocoTest, AttachExplicitPlugin) { static constexpr char xml_parent[] = R"( @@ -396,7 +396,7 @@ TEST_F(PluginTest, AttachExplicitPlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, ReplicatePlugin) { +TEST_F(MujocoTest, ReplicatePlugin) { static constexpr char xml[] = R"( @@ -429,7 +429,7 @@ TEST_F(PluginTest, ReplicatePlugin) { mj_deleteModel(model); } -TEST_F(PluginTest, ReplicateExplicitPlugin) { +TEST_F(MujocoTest, ReplicateExplicitPlugin) { static constexpr char xml[] = R"( @@ -484,7 +484,7 @@ TEST_F(MujocoTest, RecompileFails) { mj_deleteSpec(spec); } -TEST_F(PluginTest, ModifyShellInertiaFails) { +TEST_F(MujocoTest, ModifyShellInertiaFails) { static constexpr char xml[] = R"( @@ -515,7 +515,7 @@ TEST_F(PluginTest, ModifyShellInertiaFails) { } // ------------------- test recompilation multiple files ----------------------- -TEST_F(PluginTest, RecompileCompare) { +TEST_F(MujocoTest, RecompileCompare) { mjtNum tol = 0; std::string field = ""; @@ -606,7 +606,7 @@ TEST_F(PluginTest, RecompileCompare) { } } -TEST_F(PluginTest, RecompileEdit) { +TEST_F(MujocoTest, RecompileEdit) { static constexpr char xml[] = R"( @@ -640,7 +640,7 @@ TEST_F(PluginTest, RecompileEdit) { // ------------------- test cache with modified assets ------------------------- -TEST_F(PluginTest, RecompileCompareObjCache) { +TEST_F(MujocoTest, RecompileCompareObjCache) { static constexpr char xml[] = R"( @@ -716,7 +716,7 @@ static constexpr uint8_t tex2[] = { 0x82 }; -TEST_F(PluginTest, RecompileComparePngCache) { +TEST_F(MujocoTest, RecompileComparePngCache) { static constexpr char xml[] = R"( @@ -753,7 +753,7 @@ TEST_F(PluginTest, RecompileComparePngCache) { mj_deleteVFS(vfs.get()); } -TEST_F(PluginTest, DisableCache) { +TEST_F(MujocoTest, DisableCache) { static constexpr char xml[] = R"( @@ -791,7 +791,7 @@ TEST_F(PluginTest, DisableCache) { // -------------------------------- test textures ------------------------------ -TEST_F(PluginTest, TextureFromBuffer) { +TEST_F(MujocoTest, TextureFromBuffer) { mjSpec* spec = mj_makeSpec(); mjsTexture* t1 = mjs_addTexture(spec); @@ -1079,6 +1079,133 @@ TEST_F(MujocoTest, AttachSame) { mj_deleteModel(m_expected); } +TEST_F(MujocoTest, AttachSpatialTendonWithoutSidesite) { + static constexpr char xml_parent[] = R"( + + + + + + + )"; + + static constexpr char xml_child[] = R"( + + + + + + + + + + + + + + + + + + + + + + )"; + + std::array er; + mjSpec* parent = mj_parseXMLString(xml_parent, 0, er.data(), er.size()); + ASSERT_THAT(parent, NotNull()) << er.data(); + mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size()); + ASSERT_THAT(child, NotNull()) << er.data(); + + mjsBody* parent_body = mjs_findBody(parent, "parent_body"); + ASSERT_THAT(parent_body, NotNull()); + mjsSite* attach_site = mjs_addSite(parent_body, 0); + mjs_setName(attach_site->element, "attach_site"); + + mjs_attach(attach_site->element, + mjs_findBody(child, "child_body")->element, "", "_child"); + + EXPECT_THAT(mjs_findElement(parent, mjOBJ_TENDON, + "tendon_with_sidesite_child"), NotNull()); + EXPECT_THAT(mjs_findElement(parent, mjOBJ_TENDON, + "tendon_without_sidesite_child"), NotNull()); + + mjModel* model = mj_compile(parent, nullptr); + ASSERT_THAT(model, NotNull()) << mjs_getError(parent); + EXPECT_EQ(model->ntendon, 2); + + mj_deleteModel(model); + mj_deleteSpec(parent); + mj_deleteSpec(child); +} + +TEST_F(MujocoTest, AttachSpatialTendonGitHubIssue3119) { + static constexpr char parent_xml[] = R"( + + + + + + + )"; + + static constexpr char child_xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + )"; + + std::array er; + mjSpec* parent_spec = + mj_parseXMLString(parent_xml, 0, er.data(), er.size()); + ASSERT_THAT(parent_spec, NotNull()) << er.data(); + mjSpec* child_spec = + mj_parseXMLString(child_xml, 0, er.data(), er.size()); + ASSERT_THAT(child_spec, NotNull()) << er.data(); + + mjsBody* parent_body = mjs_findBody(parent_spec, "parent_body"); + ASSERT_THAT(parent_body, NotNull()); + mjsSite* attach_site = mjs_addSite(parent_body, 0); + mjs_setName(attach_site->element, "attach_site"); + + mjs_attach(attach_site->element, + mjs_findBody(child_spec, "child_body")->element, + "", "_child"); + + EXPECT_THAT(mjs_findElement(parent_spec, mjOBJ_TENDON, + "tendon_with_sidesite_child"), NotNull()); + EXPECT_THAT(mjs_findElement(parent_spec, mjOBJ_TENDON, + "tendon_without_sidesite_child"), NotNull()); + + mjModel* model = mj_compile(parent_spec, nullptr); + ASSERT_THAT(model, NotNull()) << mjs_getError(parent_spec); + EXPECT_EQ(model->ntendon, 2); + + mj_deleteModel(model); + mj_deleteSpec(parent_spec); + mj_deleteSpec(child_spec); +} + TEST_F(MujocoTest, AttachDifferent) { std::array er; mjtNum tol = 0; diff --git a/test/user/user_flex_test.cc b/test/user/user_flex_test.cc index 8db36823..841bcdac 100644 --- a/test/user/user_flex_test.cc +++ b/test/user/user_flex_test.cc @@ -456,6 +456,46 @@ TEST_F(UserFlexTest, StiffnessMatrix) { mj_deleteModel(m); } +TEST_F(UserFlexTest, StiffnessCacheDiffersByGeometry) { + std::array error; + + // Create two flexes with same material but different bounding boxes + static constexpr char xml_small[] = R"( + + + + + + + + + )"; + + static constexpr char xml_large[] = R"( + + + + + + + + + )"; + + mjModel* m_small = LoadModelFromString(xml_small, error.data(), error.size()); + ASSERT_THAT(m_small, NotNull()) << error.data(); + + mjModel* m_large = LoadModelFromString(xml_large, error.data(), error.size()); + ASSERT_THAT(m_large, NotNull()) << error.data(); + + // Same number of nodes but different stiffness due to different geometry + EXPECT_EQ(m_small->nflexnode, m_large->nflexnode); + EXPECT_NE(m_small->flex_stiffness[0], m_large->flex_stiffness[0]); + + mj_deleteModel(m_small); + mj_deleteModel(m_large); +} + TEST_F(UserFlexTest, LoadTexture) { const std::string xml_path = GetTestDataFilePath("user/testdata/textured_torus_flex.xml"); diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 8fd84caa..bfd68c0e 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -199,6 +199,48 @@ TEST_F(UserModelTest, ActuatorSparsity) { mj_deleteModel(m); } +TEST_F(UserModelTest, FixedTendonSparsity) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + )"; + mjModel* m = LoadModelFromString(xml); + ASSERT_THAT(m, NotNull()); + + EXPECT_EQ(m->nJten, 3); + EXPECT_EQ(m->ten_J_rownnz[0], 3); + EXPECT_EQ(m->ten_J_rowadr[0], 0); + EXPECT_EQ(m->wrap_type[m->tendon_adr[0]], mjWRAP_JOINT); + + int rowadr = m->ten_J_rowadr[0]; + int* colind = m->ten_J_colind + rowadr; + EXPECT_THAT(std::vector(colind, colind + 3), ElementsAre(0, 1, 2)); + + mj_deleteModel(m); +} + TEST_F(UserModelTest, NestedZeroMassBodiesOK) { static constexpr char xml[] = R"( diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc index 41528d23..a86fc6b6 100644 --- a/test/user/user_objects_test.cc +++ b/test/user/user_objects_test.cc @@ -2682,5 +2682,150 @@ TEST_F(UserObjectsTest, ZeroMass) { mj_deleteModel(model); } + +// ------------- test Octree SDF computation ----------------------------------- + +using OctreeSDFTest = MujocoTest; + +TEST_F(OctreeSDFTest, SphereSDF) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + EXPECT_GT(model->nmesh, 0); + EXPECT_EQ(model->geom_type[0], mjGEOM_SDF); + + int geom_id = 0; + int mesh_id = model->geom_dataid[geom_id]; + mjSDF sdf; + const mjpPlugin* null_plugin = nullptr; + sdf.plugin = &null_plugin; + sdf.id = &mesh_id; + sdf.type = mjSDFTYPE_SINGLE; + sdf.geomtype = (mjtGeom*)(model->geom_type + geom_id); + + // Analytic SDF for unit sphere: distance = |p| - 1 + auto analyticSdf = [](const mjtNum* p) -> double { + return mju_sqrt(p[0]*p[0] + p[1]*p[1] + p[2]*p[2]) - 1.0; + }; + + int sign_errors = 0; + int total_points = 0; + double sum_sq_error = 0.0; + + // Test grid of points + for (double x = -2.0; x <= 2.0; x += 0.5) { + for (double y = -2.0; y <= 2.0; y += 0.5) { + for (double z = -2.0; z <= 2.0; z += 0.5) { + mjtNum p[3] = {x, y, z}; + double sdf_dist = mjc_distance(model, data, &sdf, p); + double gt_dist = analyticSdf(p); + + if ((sdf_dist < 0) != (gt_dist < 0)) { + sign_errors++; + } + + double error = sdf_dist - gt_dist; + sum_sq_error += error * error; + total_points++; + } + } + } + + double rmse = mju_sqrt(sum_sq_error / total_points); + + EXPECT_LT(sign_errors, total_points / 200) + << "No more than 0.5% of points should have sign errors"; + EXPECT_LT(rmse, 0.11) << "RMSE should be less than 0.11"; + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(OctreeSDFTest, TorusSDF) { + static constexpr char xml[] = R"( + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + + EXPECT_GT(model->nmesh, 0); + EXPECT_EQ(model->geom_type[0], mjGEOM_SDF); + + int geom_id = 0; + int mesh_id = model->geom_dataid[geom_id]; + mjSDF sdf; + const mjpPlugin* null_plugin = nullptr; + sdf.plugin = &null_plugin; + sdf.id = &mesh_id; + sdf.type = mjSDFTYPE_SINGLE; + sdf.geomtype = (mjtGeom*)(model->geom_type + geom_id); + + // Analytic SDF for torus: distance = |p_proj| - r, where p_proj is + // projection of p onto circle of radius R, and r is minor radius. + // R=1, r=0.3 + auto analyticSdf = [](const mjtNum* p) -> double { + double xy = mju_sqrt(p[0]*p[0] + p[1]*p[1]); + double vec[2] = {xy - 1.0, p[2]}; + return mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]) - 0.3; + }; + + int sign_errors = 0; + int total_points = 0; + double sum_sq_error = 0.0; + + // Test grid of points + for (double x = -2.0; x <= 2.0; x += 0.5) { + for (double y = -2.0; y <= 2.0; y += 0.5) { + for (double z = -2.0; z <= 2.0; z += 0.5) { + mjtNum p[3] = {x, y, z}; + double sdf_dist = mjc_distance(model, data, &sdf, p); + double gt_dist = analyticSdf(p); + + if ((sdf_dist < 0) != (gt_dist < 0)) { + sign_errors++; + } + + double error = sdf_dist - gt_dist; + sum_sq_error += error * error; + total_points++; + } + } + } + + double rmse = mju_sqrt(sum_sq_error / total_points); + + EXPECT_LT(sign_errors, total_points / 20) + << "No more than 5% of points should have sign errors"; + EXPECT_LT(rmse, 0.52) << "RMSE should be close to 0.516"; + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index 71c250e6..20fd7a14 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -44,7 +44,7 @@ using ::testing::HasSubstr; using ::testing::Not; using ::testing::NotNull; -using XMLWriterTest = PluginTest; +using XMLWriterTest = MujocoTest; static const char* const kNonRgbTextureXMLPath = "xml/testdata/hfield_png_nonrgb.xml"; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 949dc889..0bd6e798 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4994,9 +4994,6 @@ public unsafe struct mjData_ { public double* bvh_aabb_dyn; public int* ten_wrapadr; public int* ten_wrapnum; - public int* ten_J_rownnz; - public int* ten_J_rowadr; - public int* ten_J_colind; public double* ten_J; public double* ten_length; public int* wrap_obj; @@ -5347,6 +5344,7 @@ public unsafe struct mjModel_ { public Int64 nexclude; public Int64 neq; public Int64 ntendon; + public Int64 nJten; public Int64 nwrap; public Int64 nsensor; public Int64 nnumeric; @@ -5371,7 +5369,6 @@ public unsafe struct mjModel_ { public Int64 npaths; public Int64 nnames_map; public Int64 nJmom; - public Int64 nJten; public Int64 ngravcomp; public Int64 nemax; public Int64 njmax; @@ -5712,6 +5709,9 @@ public unsafe struct mjModel_ { public int* tendon_group; public int* tendon_treenum; public int* tendon_treeid; + public int* ten_J_rownnz; + public int* ten_J_rowadr; + public int* ten_J_colind; public byte* tendon_limited; public byte* tendon_actfrclimited; public double* tendon_width; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index c49629ee..3aa71a8c 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -3934,6 +3934,12 @@ struct MjModel { void set_ntendon(int value) { ptr_->ntendon = static_cast(value); } + int nJten() const { + return static_cast(ptr_->nJten); + } + void set_nJten(int value) { + ptr_->nJten = static_cast(value); + } int nwrap() const { return static_cast(ptr_->nwrap); } @@ -4078,12 +4084,6 @@ struct MjModel { void set_nJmom(int value) { ptr_->nJmom = static_cast(value); } - int nJten() const { - return static_cast(ptr_->nJten); - } - void set_nJten(int value) { - ptr_->nJten = static_cast(value); - } int ngravcomp() const { return static_cast(ptr_->ngravcomp); } @@ -5125,6 +5125,15 @@ struct MjModel { emscripten::val tendon_treeid() const { return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon * 2, ptr_->tendon_treeid)); } + emscripten::val ten_J_rownnz() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->ten_J_rownnz)); + } + emscripten::val ten_J_rowadr() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->ten_J_rowadr)); + } + emscripten::val ten_J_colind() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nJten, ptr_->ten_J_colind)); + } emscripten::val tendon_limited() const { return emscripten::val(emscripten::typed_memory_view(ptr_->ntendon, ptr_->tendon_limited)); } @@ -6649,15 +6658,6 @@ struct MjData { emscripten::val ten_wrapnum() const { return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_wrapnum)); } - emscripten::val ten_J_rownnz() const { - return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_J_rownnz)); - } - emscripten::val ten_J_rowadr() const { - return emscripten::val(emscripten::typed_memory_view(model->ntendon, ptr_->ten_J_rowadr)); - } - emscripten::val ten_J_colind() const { - return emscripten::val(emscripten::typed_memory_view(model->nJten, ptr_->ten_J_colind)); - } emscripten::val ten_J() const { return emscripten::val(emscripten::typed_memory_view(model->nJten, ptr_->ten_J)); } @@ -8329,7 +8329,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr mj_loadXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper_1(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { @@ -8338,6 +8338,15 @@ std::unique_ptr mj_loadXML_wrapper(std::string filename) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_loadXML_wrapper_2(std::string filename, const MjVFS& vfs) { + char error[1000]; + mjModel *model = mj_loadXML(filename.c_str(), vfs.get(), error, sizeof(error)); + if (!model) { + mju_error("Loading error: %s\n", error); + } + return std::unique_ptr(new MjModel(model)); +} + void mj_saveModel_wrapper(const MjModel& m, const StringOrNull& filename, const val& buffer) { UNPACK_NULLABLE_STRING(filename); UNPACK_NULLABLE_VALUE(uint8_t, buffer); @@ -11541,9 +11550,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("subtree_com", &MjData::subtree_com) .property("subtree_linvel", &MjData::subtree_linvel) .property("ten_J", &MjData::ten_J) - .property("ten_J_colind", &MjData::ten_J_colind) - .property("ten_J_rowadr", &MjData::ten_J_rowadr) - .property("ten_J_rownnz", &MjData::ten_J_rownnz) .property("ten_length", &MjData::ten_length) .property("ten_velocity", &MjData::ten_velocity) .property("ten_wrapadr", &MjData::ten_wrapadr) @@ -11581,8 +11587,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("useexisting", &MjLROpt::useexisting, &MjLROpt::set_useexisting, reference()) .property("uselimit", &MjLROpt::uselimit, &MjLROpt::set_uselimit, reference()); emscripten::class_("MjModel") - .class_function("mj_loadXML", &mj_loadXML_wrapper, take_ownership()) - .class_function("mj_loadBinary", &mj_loadModel_wrapper, take_ownership()) + .class_function("mj_loadXML", emscripten::select_overload(std::string)>(&mj_loadXML_wrapper_1)) + .class_function("mj_loadXML", emscripten::select_overload(std::string, const MjVFS&)>(&mj_loadXML_wrapper_2)) + .class_function("mj_loadModel", &mj_loadModel_wrapper) .constructor() // Binds the functions on MjModel that return accessors. #define X_ACCESSOR(NAME, Name, OBJTYPE, field_name, nfield) \ @@ -12090,6 +12097,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("skin_vertadr", &MjModel::skin_vertadr) .property("skin_vertnum", &MjModel::skin_vertnum) .property("stat", &MjModel::stat, reference()) + .property("ten_J_colind", &MjModel::ten_J_colind) + .property("ten_J_rowadr", &MjModel::ten_J_rowadr) + .property("ten_J_rownnz", &MjModel::ten_J_rownnz) .property("tendon_actfrclimited", &MjModel::tendon_actfrclimited) .property("tendon_actfrcrange", &MjModel::tendon_actfrcrange) .property("tendon_adr", &MjModel::tendon_adr) diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py index f6282693..0517dba9 100644 --- a/wasm/codegen/generators/structs.py +++ b/wasm/codegen/generators/structs.py @@ -549,14 +549,19 @@ def _build_struct_bindings( MJDATA_ACCESSORS #undef X_ACCESSOR""".lstrip()) elif w == "MjModel": - f1 = common.wrapped_function_name( - introspect_functions.FUNCTIONS["mj_loadXML"] + builder.line( + '.class_function("mj_loadXML",' + " emscripten::select_overload(std::string)>(&mj_loadXML_wrapper_1))" + ) + builder.line( + '.class_function("mj_loadXML",' + " emscripten::select_overload(std::string," + " const MjVFS&)>(&mj_loadXML_wrapper_2))" ) - builder.line(f'.class_function("mj_loadXML", &{f1}, take_ownership())') f2 = common.wrapped_function_name( introspect_functions.FUNCTIONS["mj_loadModel"] ) - builder.line(f'.class_function("mj_loadBinary", &{f2}, take_ownership())') + builder.line(f'.class_function("mj_loadModel", &{f2})') builder.line(".constructor()") builder.line(""" // Binds the functions on MjModel that return accessors. diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index 8e29a080..8f2a1ecf 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -714,7 +714,7 @@ MjSpec::~MjSpec() { mjSpec *MjSpec::get() const { return ptr_; } void MjSpec::set(mjSpec *ptr) { ptr_ = ptr; } -std::unique_ptr mj_loadXML_wrapper(std::string filename) { +std::unique_ptr mj_loadXML_wrapper_1(std::string filename) { char error[1000]; mjModel *model = mj_loadXML(filename.c_str(), nullptr, error, sizeof(error)); if (!model) { @@ -723,6 +723,15 @@ std::unique_ptr mj_loadXML_wrapper(std::string filename) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_loadXML_wrapper_2(std::string filename, const MjVFS& vfs) { + char error[1000]; + mjModel *model = mj_loadXML(filename.c_str(), vfs.get(), error, sizeof(error)); + if (!model) { + mju_error("Loading error: %s\n", error); + } + return std::unique_ptr(new MjModel(model)); +} + void mj_saveModel_wrapper(const MjModel& m, const StringOrNull& filename, const val& buffer) { UNPACK_NULLABLE_STRING(filename); UNPACK_NULLABLE_VALUE(uint8_t, buffer); diff --git a/wasm/demo_app/vite.demo.config.ts b/wasm/demo_app/vite.demo.config.ts index b84e280f..a5a0217a 100644 --- a/wasm/demo_app/vite.demo.config.ts +++ b/wasm/demo_app/vite.demo.config.ts @@ -23,5 +23,9 @@ export default defineConfig({ }, server: { open: true, + headers: { + "Cross-Origin-Opener-Policy": "same-origin", + "Cross-Origin-Embedder-Policy": "require-corp", + }, }, }) diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index a357637c..4cff1ec9 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -2512,7 +2512,7 @@ describe('MuJoCo WASM Bindings', () => { vfs = new mujoco.MjVFS(); vfs.addBuffer(objFilename, new TextEncoder().encode(cube1)); - binaryModel = mujoco.MjModel.mj_loadBinary(mjbFilename, vfs); + binaryModel = mujoco.MjModel.mj_loadModel(mjbFilename, vfs); assertExists(binaryModel); expect(mujoco.mj_sizeModel(binaryModel)) @@ -2542,7 +2542,7 @@ describe('MuJoCo WASM Bindings', () => { const bufSize = mujoco.mj_sizeModel(model!); vfs = new mujoco.MjVFS(); - binaryModel = mujoco.MjModel.mj_loadBinary(mjbFilename, vfs); + binaryModel = mujoco.MjModel.mj_loadModel(mjbFilename, vfs); assertExists(binaryModel); expect(mujoco.mj_sizeModel(binaryModel)).toEqual(bufSize); @@ -2558,4 +2558,49 @@ describe('MuJoCo WASM Bindings', () => { } }); + it('should load XML with assets from VFS', () => { + const xml = ` + + + + + + + + `; + + const cube1 = ` + v -1 -1 1 + v 1 -1 1 + v -1 1 1 + v 1 1 1 + v -1 1 -1 + v 1 1 -1 + v -1 -1 -1 + v 1 -1 -1`; + + const xmlFilename = '/tmp/with_vfs.xml'; + writeXMLFile(xmlFilename, xml); + + let model: MjModel|null = null; + let vfs: MjVFS|null = null; + try { + vfs = new mujoco.MjVFS(); + vfs.addBuffer('cube.obj', new TextEncoder().encode(cube1)); + assertExists(vfs); + + model = mujoco.MjModel.mj_loadXML(xmlFilename, vfs); + assertExists(model); + expect(model.nmesh).toBe(1); + + const meshId = + mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH.value, 'cube'); + expect(meshId).toBeGreaterThanOrEqual(0); + } finally { + model?.delete(); + vfs?.delete(); + unlinkXMLFile(xmlFilename); + } + }); + }); diff --git a/wasm/tests/enums_test.ts b/wasm/tests/enums_test.ts index 96300469..d80fb48e 100644 --- a/wasm/tests/enums_test.ts +++ b/wasm/tests/enums_test.ts @@ -40,6 +40,10 @@ describe('Enums', () => { expect(mujoco.mjtGeom).toBeDefined(); }); + it('mjtProjection should exist', () => { + expect(mujoco.mjtProjection).toBeDefined(); + }); + it('mjtCamLight should exist', () => { expect(mujoco.mjtCamLight).toBeDefined(); }); @@ -120,6 +124,14 @@ describe('Enums', () => { expect(mujoco.mjtConDataField).toBeDefined(); }); + it('mjtRayDataField should exist', () => { + expect(mujoco.mjtRayDataField).toBeDefined(); + }); + + it('mjtCamOutBit should exist', () => { + expect(mujoco.mjtCamOutBit).toBeDefined(); + }); + it('mjtSameFrame should exist', () => { expect(mujoco.mjtSameFrame).toBeDefined(); }); diff --git a/wasm/tests/sandbox/vite.sandbox.config.ts b/wasm/tests/sandbox/vite.sandbox.config.ts index 123f412b..01312ba1 100644 --- a/wasm/tests/sandbox/vite.sandbox.config.ts +++ b/wasm/tests/sandbox/vite.sandbox.config.ts @@ -23,5 +23,9 @@ export default defineConfig({ }, server: { open: true, + headers: { + "Cross-Origin-Opener-Policy": "same-origin", + "Cross-Origin-Embedder-Policy": "require-corp", + }, }, })