diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 75f7965d..08af2d39 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -164,6 +164,7 @@ jobs: if: ${{ runner.os != 'Windows' }} working-directory: build run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && + cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin @@ -171,6 +172,7 @@ jobs: if: ${{ runner.os == 'Windows' }} working-directory: build run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && + cp bin/Release/actuator.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp bin/Release/elasticity.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin && cp bin/Release/sensor.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin - name: Configure samples diff --git a/CMakeLists.txt b/CMakeLists.txt index 7bafe59a..57b9ddd3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON) project( mujoco - VERSION 3.0.2 + VERSION 3.1.2 DESCRIPTION "MuJoCo Physics Simulator" HOMEPAGE_URL "https://mujoco.org" ) @@ -83,6 +83,7 @@ target_include_directories( ) add_subdirectory(plugin/elasticity) +add_subdirectory(plugin/actuator) add_subdirectory(plugin/sensor) add_subdirectory(plugin/sdf) add_subdirectory(src/engine) diff --git a/README.md b/README.md index 21a3cb81..cb19a639 100644 --- a/README.md +++ b/README.md @@ -164,7 +164,7 @@ These packages give users of various languages access to MuJoCo functionality: by [Manoj Velmurugan](https://github.com/vmanoj1996). - **Swift**: [swift-mujoco](https://github.com/liuliu/swift-mujoco) - **Java**: [mujoco-java](https://github.com/CommonWealthRobotics/mujoco-java) -- **Julia**: [Lyceum](https://github.com/Lyceum/MuJoCo.jl) (unmaintained) +- **Julia**: [MuJoCo.jl](https://github.com/JamieMair/MuJoCo.jl) ### Converters diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake index 87bb13a7..86c860a2 100644 --- a/cmake/MujocoDependencies.cmake +++ b/cmake/MujocoDependencies.cmake @@ -39,7 +39,7 @@ set(MUJOCO_DEP_VERSION_qhull CACHE STRING "Version of `qhull` to be fetched." ) set(MUJOCO_DEP_VERSION_Eigen3 - aa6964bf3a34fd607837dd8123bc42465185c4f8 + 454f89af9d6f3525b1df5f9ef9c86df58bf2d4d3 CACHE STRING "Version of `Eigen3` to be fetched." ) @@ -54,7 +54,7 @@ set(MUJOCO_DEP_VERSION_gtest ) set(MUJOCO_DEP_VERSION_benchmark - 344117638c8ff7e239044fd0fa7085839fc03021 # v1.8.3 + e45585a4b8e75c28479fa4107182c28172799640 # v1.8.3 CACHE STRING "Version of `benchmark` to be fetched." ) diff --git a/dist/mujoco.rc b/dist/mujoco.rc index e52e3530..2c2d58c6 100644 --- a/dist/mujoco.rc +++ b/dist/mujoco.rc @@ -1,6 +1,6 @@ 1 VERSIONINFO -FILEVERSION 3,0,2,0 -PRODUCTVERSION 3,0,2,0 +FILEVERSION 3,1,2,0 +PRODUCTVERSION 3,1,2,0 FILEOS 0x4 FILETYPE 0x1 { @@ -9,9 +9,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "3.0.2" + VALUE "ProductVersion", "3.1.2" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "3.0.2" + VALUE "FileVersion", "3.1.2" VALUE "InternalName", "mujoco.dll" VALUE "OriginalFilename", "mujoco.dll" VALUE "CompanyName", "Google DeepMind" diff --git a/dist/simulate.rc b/dist/simulate.rc index 4835c979..255cb59c 100644 --- a/dist/simulate.rc +++ b/dist/simulate.rc @@ -1,8 +1,8 @@ MUJOCO ICON "mujoco.ico" 1 VERSIONINFO -FILEVERSION 3,0,2,0 -PRODUCTVERSION 3,0,2,0 +FILEVERSION 3,1,2,0 +PRODUCTVERSION 3,1,2,0 FILEOS 0x4 FILETYPE 0x1 { @@ -11,9 +11,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "3.0.2" + VALUE "ProductVersion", "3.1.2" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "3.0.2" + VALUE "FileVersion", "3.1.2" VALUE "InternalName", "simulate.exe" VALUE "OriginalFilename", "simulate.exe" VALUE "CompanyName", "Google DeepMind" diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 33c5401c..483383c4 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - Maximum number of UI rectangles. Defined in `mjui.h `_. * - ``mjVERSION_HEADER`` - - 302 + - 312 - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning but for the compiled library. diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index 684e810b..ee8de505 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -1219,9 +1219,9 @@ a frame at the center-of-mass of the local kinematic subtree (``mjData.subtree_c This choice increases the precision of kinematic computations for mechanisms that are distant from the global origin. ``cdof``: - These 6D motion vectors describe the instantaneous axis of a degree-of-freedom and are used by all Jacobian functions. - Therefore, the minimal computation required for analytic Jacobians is :ref:`mj_kinematics` followed by - :ref:`mj_comPos`. + These 6D motion vectors (3 rotation, 3 translation) describe the instantaneous axis of a degree-of-freedom and are + used by all Jacobian functions. The minimal computation required for analytic Jacobians is :ref:`mj_kinematics` + followed by :ref:`mj_comPos`. ``cinert``: These 10-vectors describe the inertial properties of a body in the c-frame and are used by the Composite Rigid Body diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index afa607d1..04d8802a 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -583,6 +583,17 @@ mj_RungeKutta Runge-Kutta explicit order-N integrator. +.. _mj_implicit: + +mj_implicit +~~~~~~~~~~~ + +.. mujoco-include:: mj_implicit + +Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see +:ref:`Numerical Integration`), and advances simulation time. See `mjdata.h +`__ for fields computed by this function. + .. _mj_invPosition: mj_invPosition @@ -2974,7 +2985,7 @@ mju_eig3 .. mujoco-include:: mju_eig3 -Eigenvalue decomposition of symmetric 3x3 matrix. +Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'. .. _mju_boxQP: @@ -3011,10 +3022,10 @@ outputs (optional): notes: The initial value of ``res`` is used to warmstart the solver. - ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output. - ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output. + ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output. + ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output. The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures. - Only the lower triangles of H and R and are read from and written to, respectively. + Only the lower triangles of H and R are read from and written to, respectively. .. _mju_boxQPmalloc: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index b81199f5..812e6f43 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -69,6 +69,12 @@ These functions can be used to print various quantities to the screen for debugg These are components of the simulation pipeline, called internally from :ref:`mj_step`, :ref:`mj_forward` and :ref:`mj_inverse`. It is unlikely that the user will need to call them. +.. _mj_implicit: + +Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see +:ref:`Numerical Integration`), and advances simulation time. See `mjdata.h +`__ for fields computed by this function. + .. _Subcomponents: These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely @@ -486,10 +492,10 @@ outputs (optional): notes: The initial value of ``res`` is used to warmstart the solver. - ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output. - ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output. + ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output. + ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output. The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures. - Only the lower triangles of H and R and are read from and written to, respectively. + Only the lower triangles of H and R are read from and written to, respectively. .. _mju_boxQPmalloc: diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 1a168853..6100d9c1 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -103,17 +103,76 @@ In the remainder of this chapter we describe all valid MJCF elements and their a multiple contexts, in which case their meaning depends on the parent element. This is why we always show the parent as a prefix in the documentation below. +.. _meta-element: + +Meta elements +~~~~~~~~~~~~~ + +These elements are not strictly part of the low-level MJCF format definition, but rather instruct the compiler to +perform some operation on the model. A general property of meta-elements is that they disappear from the model upon +saving the XML. There are currently four meta-elements in MJCF: + +- :ref:`include` and :ref:`frame`, which are outside of the schema. +- :ref:`composite` and :ref:`flexcomp` which are part of the schema, but serve to + procedurally generate other MJCF elements. + +.. _frame: + +**frame** (R) +^^^^^^^^^^^^^ + +The frame meta-element is a pure coordinate transformation that can wrap any group of elements in the kinematic tree +(under :ref:`worldbody`). After compilation, frame elements disappear and their transformation is accumulated +in their direct children. The attributes of the frame meta-element are documented :ref:`below`. + +.. collapse:: Usage example of frame + + Loading this model and saving it: + + .. code-block:: xml + + + + + + + + + + + ... + + + + + + Results in this model: + + .. code-block:: xml + + + + + + + ... + + + + + Note that in the saved model, the frame elements have disappeared but their transformation was accumulated with those + of their child elements. .. _include: **include** (*) -~~~~~~~~~~~~~~~ +^^^^^^^^^^^^^^^ -This element does not strictly speaking belong to MJCF. Instead it is a meta-element, used to assemble multiple XML +This element does not strictly belong to MJCF. Instead it is a meta-element, used to assemble multiple XML files in a single document object model (DOM) before parsing. The included file must be a valid XML file with a unique top-level element. This top-level element is removed by the parser, and the elements below it are inserted at the location of the :el:`include` element. At least one element must be inserted as a result of this procedure. The -:el:`include` element can be used where ever an XML element is expected in the MJFC file. Nested includes are allowed, +:el:`include` element can be used where ever an XML element is expected in the MJCF file. Nested includes are allowed, however a given XML file can be included at most once in the entire model. After all the included XML files have been assembled into a single DOM, it must correspond to a valid MJCF model. Other than that, it is up to the user to decide how to use includes and how to modularize large files if desired. @@ -216,11 +275,11 @@ any effect. The settings here are global and apply to the entire model. .. _compiler-eulerseq: :at:`eulerseq`: :at-val:`string, "xyz"` - This attribute specifies the sequence of Euler rotations for all euler attributes of elements that have spatial - frames, as explained in :ref:`COrientation`. This must be a string with exactly 3 - characters from the set {'x', 'y', 'z', 'X', 'Y', 'Z'}. The character at position n determines the axis around which - the n-th rotation is performed. Lower case denotes axes that rotate with the frame, while upper case denotes axes - that remain fixed in the parent frame. The "rpy" convention used in URDF corresponds to the default "xyz" in MJCF. + This attribute specifies the sequence of Euler rotations for all :at:`euler` attributes of elements that have spatial + frames, as explained in :ref:`COrientation`. This must be a string with exactly 3 characters from the set {x, y, z, + X, Y, Z}. The character at position n determines the axis around which the n-th rotation is performed. Lower case + letters denote axes that rotate with the frame (intrinsic), while upper case letters denote axes that remain fixed in + the parent frame (extrinsic). The "rpy" convention used in URDF corresponds to "XYZ" in MJCF. .. _compiler-meshdir: @@ -247,14 +306,20 @@ any effect. The settings here are global and apply to the entire model. .. _compiler-discardvisual: :at:`discardvisual`: :at-val:`[false, true], "false" for MJCF, "true" for URDF` - This attribute instructs the parser to discard "visual geoms", defined as geoms whose contype and conaffinity - attributes are both set to 0. This functionality is useful for models that contain two sets of geoms, one for - collisions and the other for visualization. Note that URDF models are usually constructed in this way. It rarely - makes sense to have two sets of geoms in the model, especially since MuJoCo uses convex hulls for collisions, so we - recommend using this feature to discard redundant geoms. Keep in mind however that geoms considered visual per the - above definition can still participate in collisions, if they appear in the explicit list of contact - :ref:`pairs `. The parser does not check this list before discarding geoms; it relies solely on the geom - attributes to make the determination. + This attribute instructs the compiler to discard all model elements which are purely visual and have no effect on the + physics (with one exception, see below). This often enables smaller :ref:`mjModel` structs and faster simulation. + + - All materials are discarded. + - All textures are discarded. + - All geoms with :ref:`contype`=:ref:`conaffinity`=0 are discarded, if they + are not referenced in another MJCF element. If a discarded geom was used for inferring body inertia, an explicit + :ref:`inertial` element is added to the body. + - All meshes which are not referenced by any geom (in particular those discarded above) are discarded. + + The resulting compiled model will have exactly the same dynamics as the original model, with the exception of + raycasting, as used for example by :ref:`rangefinder`, since raycasting reports distances to + visual geoms. When visualizing models compiled with this flag, it is important to remember that colliding geoms are + often placed in a :ref:`group` which is invisible by default. .. _compiler-convexhull: @@ -1319,9 +1384,9 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe | For collision detection, a height field is treated as a union of triangular prisms. Collisions between height fields 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 9; 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. + 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 + contacts, the spatial features of the height field should be large compared to the geoms it collides with. .. _asset-hfield-name: @@ -1877,7 +1942,7 @@ adjust it properly through the XML. .. _option-sdf_initpoints: :at:`sdf_initpoints`: :at-val:`int, "40"` - Number of starting points used for fining contacts with Signed Distance Field collisions. + Number of starting points used for finding contacts with Signed Distance Field collisions. .. _option-actuatorgroupdisable: @@ -1924,7 +1989,7 @@ from its default. .. _option-flag-contact: :at:`contact`: :at-val:`[disable, enable], "enable"` - This flag disables all standard computations related to contact constraints. + This flag disables collision detection and all standard computations related to contact constraints. .. _option-flag-passive: @@ -3913,6 +3978,35 @@ Associate this flexcomp with an :ref:`engine plugin`. Either :at:`plug :at:`instance`: :at-val:`string, optional` Instance name, used for explicit plugin instantiation. + +.. _body-frame: + +:el-prefix:`body/` |-| **frame** (*) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Frames specify a coordinate transformation which is applied to all child elements. They disappear during compilation +and the transformation they encode is accumulated in their direct children. See :ref:`frame` for examples. + +.. _frame-pos: + +:at:`pos`: :at-val:`real(3), "0 0 0"` + The 3D position of the frame, in the parent coordinate system. + +.. _frame-quat: + +.. _frame-axisangle: + +.. _frame-xyaxes: + +.. _frame-zaxis: + +.. _frame-euler: + +:at:`quat`, :at:`axisangle`, :at:`xyaxes`, :at:`zaxis`, :at:`euler` + See :ref:`COrientation`. + + + .. _contact: **contact** (*) @@ -4532,7 +4626,7 @@ joint types (slide and hinge) can be used. :at:`polycoef`: :at-val:`real(5), "0 1 0 0 0"` Coefficients a0 ... a4 of the quartic polynomial. If the two joint values are y and x, and their reference positions (corresponding to the joint values in the initial model configuration) are y0 and x0, the constraint is: - y-y0 = a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4 + y-y0 = a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4. Omitting the second joint is equivalent to setting x = x0, in which case the constraint is y = y0 + a0. @@ -4942,21 +5036,21 @@ specify them independently. .. _actuator-general-ctrlrange: :at:`ctrlrange`: :at-val:`real(2), "0 0"` - Range for clamping the control input. The compiler expects the first value to be smaller than the second value. + Range for clamping the control input. The first value must be smaller than the second value. |br| Setting this attribute without specifying :at:`ctrllimited` is an error, unless :at:`autolimits` is set in :ref:`compiler `. .. _actuator-general-forcerange: :at:`forcerange`: :at-val:`real(2), "0 0"` - Range for clamping the force output. The compiler expects the first value to be no greater than the second value. + Range for clamping the force output. The first value must be no greater than the second value. |br| Setting this attribute without specifying :at:`forcelimited` is an error, unless :at:`autolimits` is set in :ref:`compiler `. .. _actuator-general-actrange: :at:`actrange`: :at-val:`real(2), "0 0"` - Range for clamping the activation state. The compiler expects the first value to be no greater than the second value. + Range for clamping the activation state. The first value must be no greater than the second value. See the :ref:`Activation clamping ` section for more details. |br| Setting this attribute without specifying :at:`actlimited` is an error, unless :at:`autolimits` is set in :ref:`compiler `. @@ -5229,13 +5323,13 @@ This element does not have custom attributes. It only has common attributes, whi This element creates a position servo. The underlying :el:`general` attributes are set as follows: -========= ======= ========= ======= +========= ======= ========= ========= Attribute Setting Attribute Setting -========= ======= ========= ======= +========= ======= ========= ========= dyntype none dynprm 1 0 0 gaintype fixed gainprm kp 0 0 -biastype affine biasprm 0 -kp 0 -========= ======= ========= ======= +biastype affine biasprm 0 -kp -kv +========= ======= ========= ========= This element has one custom attribute in addition to the common attributes: @@ -5289,6 +5383,11 @@ This element has one custom attribute in addition to the common attributes: :at:`kp`: :at-val:`real, "1"` Position feedback gain. +.. _actuator-position-kv: + +:at:`kv`: :at-val:`real, "0"` + Damping applied by the actuator. + When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators`. .. _actuator-velocity: @@ -5297,7 +5396,9 @@ This element has one custom attribute in addition to the common attributes: This element creates a velocity servo. Note that in order create a PD controller, one has to define two actuators: a position servo and a velocity servo. This is because MuJoCo actuators are SISO while a PD controller takes two control -inputs (reference position and reference velocity). The underlying :el:`general` attributes are set as follows: +inputs (reference position and reference velocity). +When using this actuator, it is recommended to use the implicitfast or implicit :ref:`integrators`. +The underlying :el:`general` attributes are set as follows: ========= ======= ========= ======= Attribute Setting Attribute Setting @@ -5368,14 +5469,14 @@ This element creates an integrated-velocity servo. For more information, see the :ref:`Activation clamping ` section of the Modeling chapter. The underlying :el:`general` attributes are set as follows: -========== =========== ========= ======= +========== =========== ========= ========= Attribute Setting Attribute Setting -========== =========== ========= ======= +========== =========== ========= ========= dyntype integrator dynprm 1 0 0 gaintype fixed gainprm kp 0 0 -biastype affine biasprm 0 -kp 0 +biastype affine biasprm 0 -kp -kv actlimited true -========== =========== ========= ======= +========== =========== ========= ========= This element has one custom attribute in addition to the common attributes: @@ -5430,6 +5531,11 @@ This element has one custom attribute in addition to the common attributes: :at:`kp`: :at-val:`real, "1"` Position feedback gain. +.. _actuator-intvelocity-kv: + +:at:`kv`: :at-val:`real, "0"` + Damping applied by the actuator. + When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators`. .. _actuator-damper: @@ -5437,8 +5543,9 @@ This element has one custom attribute in addition to the common attributes: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ This element is an active damper which produces a force proportional to both velocity and control: ``F = - kv * velocity -* control``, where ``kv`` must be nonnegative. :at:`ctrlrange` is required and must also be nonnegative. The underlying -:el:`general` attributes are set as follows: +* control``, where ``kv`` must be nonnegative. :at:`ctrlrange` is required and must also be nonnegative. +When using this actuator, it is recommended to use the implicitfast or implicit :ref:`integrators`. +The underlying :el:`general` attributes are set as follows: =========== ======= ========= ======= Attribute Setting Attribute Setting @@ -5784,12 +5891,29 @@ Associate this actuator with an :ref:`engine plugin`. Either :at:`plug :at:`instance`: :at-val:`string, optional` Instance name, used for explicit plugin instantiation. +.. _actuator-plugin-dyntype: + +:at:`dyntype`: :at-val:`[none, integrator, filter, filterexact, muscle, user], "none"` + Activation dynamics type for the actuator. The available dynamics types were already described in the :ref:`Actuation + model ` section. If :ref:`dyntype` is not "none", an activation variable will + be added to the actuator. This variable will be added after any activation state computed by the plugin (see + :ref:`actuator plugin activations`). + +.. _actuator-plugin-actrange: + +:at:`actrange`: :at-val:`real(2), "0 0"` + Range for clamping the activation state associated with this actuator's dyntype. The limit doesn't apply to + activations computed by the plugin. The first value must be no greater than the second value. + See the :ref:`Activation clamping ` section for more details. + .. _actuator-plugin-name: .. _actuator-plugin-class: .. _actuator-plugin-group: +.. _actuator-plugin-actlimited: + .. _actuator-plugin-ctrllimited: .. _actuator-plugin-forcelimited: @@ -5818,9 +5942,14 @@ Associate this actuator with an :ref:`engine plugin`. Either :at:`plug .. _actuator-plugin-user: -.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`ctrllimited`, +.. _actuator-plugin-dynprm: + +.. _actuator-plugin-actearly: + +.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`actlimited`, :at:`ctrllimited`, :at:`forcelimited`, :at:`ctrlrange`, :at:`forcerange`, :at:`lengthrange`, :at:`gear`, :at:`cranklength`, - :at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user` + :at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`, + :at:`dynprm`, :at:`actearly` |actuator/plugin attrib list| Same as in actuator/ :ref:`general `. @@ -7570,6 +7699,8 @@ slidersite, cranksite. .. _default-position-kp: +.. _default-position-kv: + :el-prefix:`default/` |-| **position** (?) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -7626,6 +7757,8 @@ tendon, slidersite, cranksite. .. _default-intvelocity-kp: +.. _default-intvelocity-kv: + :el-prefix:`default/` |-| **intvelocity** (?) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 84b5aff4..702adeb2 100644 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -778,7 +778,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`jointinparent` | :ref:`tendon` | :ref:`slidersite` | :ref:`cranksite` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`site` | :ref:`refsite` | :ref:`kp` | | | +| | | | :ref:`site` | :ref:`refsite` | :ref:`kp` | :ref:`kv` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| actuator |br| |_| |L| | | .. table:: | @@ -810,6 +810,8 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`cranksite` | :ref:`site` | :ref:`refsite` | :ref:`kp` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`kv` | | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| actuator |br| |_| |L| | | .. table:: | | :ref:`damper | \* | :class: mjcf-attributes | @@ -879,13 +881,15 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`name` | :ref:`class` | :ref:`plugin` | :ref:`instance` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`group` | :ref:`ctrllimited` | :ref:`forcelimited` | :ref:`ctrlrange` | | +| | | | :ref:`group` | :ref:`ctrllimited` | :ref:`forcelimited` | :ref:`actlimited` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`forcerange` | :ref:`lengthrange` | :ref:`gear` | :ref:`cranklength` | | +| | | | :ref:`ctrlrange` | :ref:`forcerange` | :ref:`actrange` | :ref:`lengthrange` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`joint` | :ref:`jointinparent` | :ref:`site` | :ref:`tendon` | | +| | | | :ref:`gear` | :ref:`cranklength` | :ref:`joint` | :ref:`jointinparent` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`cranksite` | :ref:`slidersite` | :ref:`user` | | | +| | | | :ref:`site` | :ref:`dyntype` | :ref:`dynprm` | :ref:`tendon` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`cranksite` | :ref:`slidersite` | :ref:`user` | :ref:`actearly` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| plugin |br| |_2| |L| | | .. table:: | @@ -1438,7 +1442,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`gear` | :ref:`cranklength` | :ref:`user` | :ref:`group` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`kp` | | | | | +| | | | :ref:`kp` | :ref:`kv` | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| default |br| |_| |L| | | .. table:: | @@ -1460,7 +1464,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`actrange` | :ref:`gear` | :ref:`cranklength` | :ref:`user` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`group` | :ref:`kp` | | | | +| | | | :ref:`group` | :ref:`kp` | :ref:`kv` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_| default |br| |_| |L| | | .. table:: | diff --git a/doc/changelog.rst b/doc/changelog.rst index 25d866a4..fb886273 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,6 +2,82 @@ Changelog ========= +Upcoming version (not yet released) +----------------------------------- + +General +^^^^^^^ +1. Improved the :ref:discardvisual compiler flag, which now discards all visual-only assets. See + :ref:discardvisual for details. + +MJX +^^^ +2. Added :ref:`dyntype` ``filterexact``. +3. Added :at:`site` transmission. +4. Updated MJX colab tutorial with more stable quadruped environment. +5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes. + +Bug fixes +^^^^^^^^^ +6. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes + :github:issue:`1270`. + + +Version 3.1.1 (December 18, 2023) +----------------------------------- + +Bug fixes +^^^^^^^^^ +1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the other. +2. Fixed a bug in :ref:`simulate` where the "LOADING..." message was not showing correctly. +3. Fixed a crash in the Python :ref:`passive viewer`, when used with models containing Flex objects. +4. Fixed a bug in MJX where ``site_xmat`` was ignored in ``get_data`` and ``put_data`` +5. Fixed a bug in MJX where ``efc_address`` was sometimes incorrectly calculated in ``get_data``. + + +Version 3.1.0 (December 12, 2023) +--------------------------------- + +General +^^^^^^^ +1. Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function + for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is more + robust for very small or large geom sizes. +2. Added :ref:`frame` to MJCF, a :ref:`meta-element` which defines a pure coordinate transformation + on its direct children, without requiring a :ref:`body`. +3. Added the :at:`kv` attribute to the :ref:`position` and :ref:`intvelocity` + actuators, for specifying actuator-applied damping. This can be used to implement a PD controller with 0 reference + velocity. When using this attribute, it is recommended to use the implicitfast or implicit + :ref:`integrators`. + +Plugins +^^^^^^^ + +4. Allow actuator plugins to use activation variables in ``mjData.act`` as their internal state, rather than + ``mjData.plugin_state``. Actuator plugins can now specify :ref:`callbacks` that compute activation + variables, and they can be used with built-in :ref:`dyntype` actuator dynamics. + +5. Added the `pid `__ actuator plugin, a + configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators. + +MJX +^^^ + +6. Added ``site_xpos`` and ``site_xmat`` to MJX. +7. Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be + deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as + ``efc_J``. + +Bug fixes +^^^^^^^^^ +8. Fix bug in Cartesian actuation with movable refsite, as when using body-centric Cartesian actuators on a quadruped. + Before this fix such actuators could lead to non-conservation of momentum. +9. Fix bug that prevented using flex with :ref:`simulate`. +10. Fix bug that prevented the use of elasticity plugins in combination with pinned flex vertices. +11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where SYSTEM_VERSION_COMPAT is set. The minimum + supported version is still 11.0, but we release these wheels to fix compatibility for those users. See + :github:issue:`1213`. + Version 3.0.1 (November 15, 2023) --------------------------------- diff --git a/doc/computation/fluid.rst b/doc/computation/fluid.rst index 070cb7b3..ea6f0569 100644 --- a/doc/computation/fluid.rst +++ b/doc/computation/fluid.rst @@ -520,8 +520,14 @@ in the surrounding flow a circulation of sufficient strength to hold the rear st This is the Kutta condition, a fluid dynamic phenomenon that can be observed for solid bodies with sharp corners, such as slender bodies or the trailing edges of airfoils. -.. cssclass:: caption-small .. figure:: ../images/computation/kutta_cond_plate.svg + :class: only-light + :figwidth: 95% + :align: left + +.. cssclass:: caption-small +.. figure:: ../images/computation/kutta_cond_plate_dark.svg + :class: only-dark :figwidth: 95% :align: left diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 5665fb02..840835d8 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -300,19 +300,22 @@ is attached; the possible attachment object types are :at:`joint`, :at:`tendon`, Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality constraints to the rest of the system, but that would be less efficient. -:at:`site` - :at:`site` transmission (without a :at:`refsite`, see below) and :at:`body` transmission targets have a fixed zero - length :math:`l_i(q) = 0`. They can therefore not be used to maintain a desired length, but can be used to apply - forces. Site transmissions correspond to applying a Cartsian force/torque at the site, and are useful for modeling - jets and propellors. :el:`body` transmissions correspond to applying forces at contact points belonging to a body, in +:at:`body` + :el:`body` transmission corresponds to applying forces at contact points belonging to a body, in order to model vacuum grippers and biomechanical adhesive appendages. For more information about adhesion, see the - :ref:`adhesion` actuator documentation. + :ref:`adhesion` actuator documentation. These transmission targets have a fixed zero length + :math:`l_i(q) = 0`. - If a :at:`site` transmission target is defined with the optional :at:`refsite` attribute, forces and torques are - applied in the frame of the reference site rather than the site's own frame. If a reference site is defined then - the length of the actuator is nonzero and corresponds to the pose difference of the two sites. This length can then - be controlled with a :el:`position` actuator, enabling Cartesian end-effector control. See the - :ref:`refsite` documentation for more details. +:at:`site` + Site transmissions correspond to applying a Cartsian force/torque in the frame of a site. When a :at:`refsite` is not + defined (see below), these targets have a fixed zero length :math:`l_i(q) = 0` and are useful for modeling jets and + propellors: forces and torques which are fixed to the site frame. + + If a :at:`site` transmission is defined with the optional :at:`refsite` attribute, forces and torques are applied in + the frame of the reference site rather than the site's own frame. If a reference site is defined, the length of the + actuator is nonzero and corresponds to the pose difference of the two sites, projected onto a chosen direction in the + reference frame. This length can then be controlled with a :el:`position` actuator, allowing for Cartesian + end-effector control. See the :ref:`refsite` documentation for more details. .. _geActivation: @@ -963,6 +966,12 @@ is :math:`E f`. The matrix of basis vectors is constructed as follows. .. image:: ../images/computation/contact_frame.svg :width: 700px :align: center + :class: only-light + +.. image:: ../images/computation/contact_frame_dark.svg + :width: 700px + :align: center + :class: only-dark The figure illustrates the full basis set corresponding to the case :math:`n = 6`. Otherwise we use only the first :math:`n` or :math:`2(n-1)` columns depending on the cone type. Elliptic cones are easier to understand. Since the @@ -1322,6 +1331,12 @@ representations of the constraint Jacobian and related matrices. .. image:: ../images/computation/gPGS.svg :width: 500px :align: center + :class: only-light + + .. image:: ../images/computation/gPGS_dark.svg + :width: 500px + :align: center + :class: only-dark When using pyramidal friction cones, the problem involves box constraints to which PGS has traditionally been applied. If we applied PGS directly to the conic constraints resulting from elliptic friction cones, it would get @@ -1421,12 +1436,18 @@ approximations, no matter how accurate the approximation is. The figure below il where the pyramid is not even an approximation, but represents the same constraint set as the elliptic cone. We plot the contours of the penalty/shadow for the pyramidal (red) and elliptic (dashed blue) cones, for different friction coefficients varying from left to right. Mathematically, the penalty in the pyramidal case is a quadratic spline, while -the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics - allowing +the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics -- allowing circular contours around the tip of the cone. .. image:: ../images/computation/softcontact.png :width: 600px :align: center + :class: only-light + +.. image:: ../images/computation/softcontact_dark.png + :width: 600px + :align: center + :class: only-dark In summary, elliptic and pyramidal friction cones define different soft-contact dynamics (although they are usually very close). The elliptic model is more principled and more consistent with physical intuition, and the corresponding solvers @@ -1455,48 +1476,45 @@ others can be pruned quickly without a detailed check. MuJoCo has flexible mecha checked in detail. The decision process involves two stages: generation and filtering. Generation - First we generate a list of candidate geom pairs in one of two ways: "pair" or "dynamic". The user can also specify - "all" which merges both sources (and is the default). This is done via the setting ``mjModel.opt.collision``. "Pair" - refers to an explicit list of geom pairs defined with the :ref:`pair ` element in MJCF. It gives the - user full control, however it is a static mechanism (independent of the spatial arrangement of the geoms at runtime) - and can be tedious for large models. It is normally used to supplement the output of the "dynamic" mechanism. Dynamic - generation works with bodies rather than geoms; when a body pair is included this means that all geoms attached to - one body can collide with all geoms attached to the other body. + First we generate a list of candidate geom pairs by merging from two sources: pairs of bodies that might contain + colliding geoms and the explicit list of geom pairs defined with the :ref:`pair ` element in MJCF. The body pairs are generated via broad-phase collision detection based on a modified sweep-and-prune algorithm. The modification is that the axis for sorting is chosen as the principal eigenvector of the covariance matrix of all geom - centers - which maximizes the spread. Then, for each body pair, a mid-phase collision detection using a static - bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB) is performed. Each body is - equipped with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes, + centers -- which maximizes the spread. Then, for each body pair, mid-phase collision detection is performed using a + static bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB). Each body is equipped + with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes, respectively. - Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude ` element - in MJCF. Exclusion is applied when "dynamic" or "all" are selected, but not when "pair" is selected. At the end of - this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`, but can still be - pruned further before detailed collision checking. + Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude ` element in + MJCF. At the end of this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`, + but can still be pruned further before detailed collision checking. Filtering Next we apply four filters to the list generated in the previous step. Filters 1 and 2 are applied to all geom pairs. - Filters 3 and 4 are applied only to pairs generated by the "dynamic" mechanism, thereby allowing the user to bypass + Filters 3 and 4 are applied only to pairs generated by the body-pair mechanism, thereby allowing the user to bypass those filters by specifying geom pairs explicitly. - #. The types of the two geoms must correspond to a collision function that is capable of performing the detailed + 1. The types of the two geoms must correspond to a collision function that is capable of performing the detailed check. This is usually the case but there are exceptions (for example plane-plane collisions are not supported), and furthermore the user may override the default table of collision functions with NULL pointers, effectively disabling collisions between certain geom types. - #. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a + 2. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a plane, this becomes a plane-sphere test. - #. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless + 3. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless the parent is the world body. The motivation is to avoid permanent contacts within bodies and joints. Note that if several bodies are welded together in the sense that there are no joints between them, they are treated as a single body for the purposes of this test. The parent-filter test can be disabled by the user, while the same-body test cannot be disabled. - #. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and + 4. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and ``conaffinity``. The boolean expression below must be true for the test to pass: - ``(contype1 & conaffinity2) || (contype2 & conaffinity1)`` This requires the ``contype`` of one geom and the - ``conaffinity`` of the other geom to have a common bit set to 1. This is a powerful mechanism borrowed from the - Open Dynamics Engine. The default setting for all geoms is ``contype = conaffinity = 1`` which always passes the - test, so the user can ignore this mechanism if it is confusing at first. + + ``(contype1 & conaffinity2) || (contype2 & conaffinity1)`` + + This requires the ``contype`` of one geom and the ``conaffinity`` of the other geom to have a common bit set to 1. + This is a powerful mechanism borrowed from Open Dynamics Engine. The default setting for all geoms is + ``contype = conaffinity = 1`` which always passes the test, so the user can ignore this mechanism if it is + confusing at first. .. _coChecking: @@ -1517,12 +1535,12 @@ convex hull implicitly, however pre-computing that hull can substantially improv model compiler does that by default, using the `qhull `__ library. In order to model a non-convex object other than a height field, the user must decompose it into a union of convex geoms -(which can be primitive shapes or meshes) and attach them to the same body. Tools such as the -`HACD `__ library can be used outside MuJoCo to automate this process. Finally, all -built-in collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose -"triangle soup" collision detector for example. However we do not recommend such an approach. Pre-processing the -geometry and representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both -faster and more stable simulation. +(which can be primitive shapes or meshes) and attach them to the same body. Open tools like the `CoACD library +`__ can be used outside MuJoCo to automate this process. Finally, all built-in +collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose "triangle +soup" collision detector for example. However we do not recommend such an approach. Pre-processing the geometry and +representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both faster and more +stable simulation. .. _Pipeline: @@ -1538,44 +1556,59 @@ be used to skip default steps and to enable optional steps respectively. Callbac Forward dynamics ~~~~~~~~~~~~~~~~ -The top-level function :ref:`mj_step` invokes the sequence of computations below. Alternatively one can call -:ref:`mj_forward` which invokes only steps 2-21. +The source file `engine_forward.c `__ +contains the high-level forward dynamics pipeline: -#. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence - is detected, the state is automatically reset and the corresponding warning is raised. -#. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites, - cameras and lights. It also normalizes all quaternions, just in case. -#. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding - kinematic subtrees (to improve floating-point accuracy). -#. Compute the actuator lengths and moment arms. -#. Compute the composite rigid body inertias and construct the joint-space inertia matrix. -#. Compute the sparse factorization of the joint-space inertia matrix. -#. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection. -#. Construct the constraint Jacobian and compute the constraint residuals. -#. Compute the matrices and vectors needed by the constraint solvers. -#. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial - tendons. -#. Compute sensor data that only depends on position, and the potential energy if enabled. -#. Compute the tendon and actuator velocities. -#. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at - the subtree centers of mass. -#. Compute all passive forces: spring-dampers in joints and tendons, and fluid dynamics forces. -#. Compute sensor data that depends on velocity, and the kinetic energy if enabled. - If required by sensors, call :ref:`mj_subtreeVel`. -#. Compute the reference constraint acceleration. -#. Compute the vector of Coriolis, centrifugal and gravitational forces. -#. Compute the actuator forces and activation dynamics if defined. -#. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces. -#. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the - constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics. -#. Compute sensor data that depends on force and acceleration if enabled. - If required by sensors, call :ref:`mj_rnePostConstraint`. -#. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is - automatically reset and the corresponding warning is raised. -#. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward - dynamics. This is an optional step, and is performed only when enabled. -#. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator - repeats the above sequence three more times, except for the optional computations which are performed only once. +- The top-level function :ref:`mj_step` invokes the entire sequence of computations below. +- :ref:`mj_forward` invokes only stages **2-22**, computing the continuous-time forward dynamics, ending with the + acceleration ``mjData.qacc``. +- :ref:`mj_step1` invokes stages **1-18** and :ref:`mj_step2` invokes stages **19-25**, breaking :ref:`mj_step` into two + distinct phases. This allows the user to write controllers that depend on quantities derived from the positions and + velocities (but not forces, since those have not yet been computed). Note that the :ref:`mj_step1` → :ref:`mj_step2` + pipeline does not support the Runge Kutta integrator. + +1. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence + is detected, the state is automatically reset and the corresponding warning is raised: + :ref:`mj_checkPos`, :ref:`mj_checkVel` +2. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites, + cameras and lights. It also normalizes all quaternions: :ref:`mj_kinematics`, :ref:`mj_camLight` +3. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding + kinematic subtrees: :ref:`mj_comPos` +4. Compute quantities related to :ref:`flex` objects: :ref:`mj_flex` +5. Compute the actuator lengths and moment arms: :ref:`mj_tendon` +6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb` +7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM` +8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection: + :ref:`mj_collision` +9. Construct the constraint Jacobian and compute the constraint residuals: :ref:`mj_makeConstraint` +10. Compute the matrices and vectors needed by the constraint solvers: :ref:`mj_projectConstraint` +11. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial + tendons: :ref:`mj_transmission` +12. Compute sensor data that only depends on position, and the potential energy if enabled: :ref:`mj_sensorPos`, + :ref:`mj_energyPos` +13. Compute the tendon, flex edge and actuator velocities: :ref:`mj_fwdVelocity` +14. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at + the subtree centers of mass: :ref:`mj_comVel` +15. Compute passive forces -- spring-dampers in joints and tendons, and fluid forces: :ref:`mj_passive` +16. Compute sensor data that depends on velocity, and the kinetic energy if enabled + (if required by sensors, call :ref:`mj_subtreeVel`): :ref:`mj_sensorVel` +17. Compute the reference constraint acceleration: :ref:`mj_referenceConstraint` +18. Compute the vector of Coriolis, centrifugal and gravitational forces: :ref:`mj_rne` +19. Compute the actuator forces and activation dynamics if defined: :ref:`mj_fwdActuation` +20. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces: + :ref:`mj_fwdAcceleration` +21. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the + constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics: + :ref:`mj_fwdConstraint` +22. Compute sensor data that depends on force and acceleration if enabled + (if required by sensors, call :ref:`mj_rnePostConstraint`): :ref:`mj_sensorAcc` +23. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is + automatically reset and the corresponding warning is raised: :ref:`mj_checkAcc` +24. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward + dynamics. This is an optional step, and is performed only when enabled: :ref:`mj_compareFwdInv` +25. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator + repeats the above sequence three more times, except for the optional computations which are performed only once: + one of :ref:`mj_Euler`, :ref:`mj_RungeKutta`, :ref:`mj_implicit` .. _piInverse: diff --git a/doc/images/APIreference/arrowhead.svg b/doc/images/APIreference/arrowhead.svg index 6fa78ec0..5a98995c 100644 --- a/doc/images/APIreference/arrowhead.svg +++ b/doc/images/APIreference/arrowhead.svg @@ -1 +1 @@ - + diff --git a/doc/images/computation/contact_frame_dark.svg b/doc/images/computation/contact_frame_dark.svg new file mode 100644 index 00000000..ea67cb4f --- /dev/null +++ b/doc/images/computation/contact_frame_dark.svg @@ -0,0 +1,296 @@ + + + + + + + + + + e + 1 + + + e + 2 + + + e + 3 + + + + + + + + + + + + + + + e + 4 + + + e + 5 + + + e + 6 + + x + y + z + + + + + + + + + + + + + + + e + 1 + + + e + 2 + + + e + 3 + + + e + 4 + + + + + + + + + + + + + + + + + + + + + + + + + + elliptic basis: E = I + 6 + + pyramidal basis: E = + + + + + + + + + + + e + 5 + + + + + + + + + + + + + e + 6 + + + + + + + + + + + + + e + 7 + + + + + + + + + + + + + e + 8 + + + + + + + + + + + + + e + 9 + + + + + + + + + + + + + e + 10 + + + 1 + + +m + 1 + + 0 + 0 + 0 + 0 + 1 + 0 + 0 + 0 + 0 + + -m + 1 + + 1 + + +m + 2 + + 0 + 0 + 0 + 0 + 1 + 0 + 0 + 0 + 0 + + -m + 2 + + 1 + + +m + 5 + + 0 + 0 + 0 + 0 + 1 + 0 + 0 + 0 + 0 + + -m + 5 + + ... + diff --git a/doc/images/computation/gPGS_dark.svg b/doc/images/computation/gPGS_dark.svg new file mode 100644 index 00000000..7a8cca57 --- /dev/null +++ b/doc/images/computation/gPGS_dark.svg @@ -0,0 +1,102 @@ + + + + + + + + + + cone + constraint + + + unconstrained + minimum + + + + + + + + + + continuum of + PGS local minima + + + + + + + search + ray + + + search + ellipsoid + + + + + + + + diff --git a/doc/images/computation/kutta_cond_plate_dark.svg b/doc/images/computation/kutta_cond_plate_dark.svg new file mode 100644 index 00000000..18ea52df --- /dev/null +++ b/doc/images/computation/kutta_cond_plate_dark.svg @@ -0,0 +1,124 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + v + + + + + y + x + + + + + + + + α + + + diff --git a/doc/images/computation/softcontact_dark.png b/doc/images/computation/softcontact_dark.png new file mode 100644 index 00000000..f10cf3f9 Binary files /dev/null and b/doc/images/computation/softcontact_dark.png differ diff --git a/doc/images/mjx/SPS.svg b/doc/images/mjx/SPS.svg index a35a2508..6ffa2735 100644 --- a/doc/images/mjx/SPS.svg +++ b/doc/images/mjx/SPS.svg @@ -1 +1 @@ - + diff --git a/doc/images/modeling/flexelem.png b/doc/images/modeling/flexelem.png index 39df7463..74dbe6a8 100644 Binary files a/doc/images/modeling/flexelem.png and b/doc/images/modeling/flexelem.png differ diff --git a/doc/images/modeling/impedance_dark.png b/doc/images/modeling/impedance_dark.png new file mode 100644 index 00000000..ecb40fd7 Binary files /dev/null and b/doc/images/modeling/impedance_dark.png differ diff --git a/doc/images/modeling/musclemodel_dark.png b/doc/images/modeling/musclemodel_dark.png new file mode 100644 index 00000000..c301fe66 Binary files /dev/null and b/doc/images/modeling/musclemodel_dark.png differ diff --git a/doc/images/modeling/musclerange_dark.png b/doc/images/modeling/musclerange_dark.png new file mode 100644 index 00000000..8c7bf8e6 Binary files /dev/null and b/doc/images/modeling/musclerange_dark.png differ diff --git a/doc/includes/references.h b/doc/includes/references.h index da5329cc..1c788150 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -237,7 +237,7 @@ struct mjData_ { // computed by mj_fwdPosition/mj_comPos mjtNum* subtree_com; // center of mass of each subtree (nbody x 3) - mjtNum* cdof; // com-based motion axis of each dof (nv x 6) + mjtNum* cdof; // com-based motion axis of each dof (rot:lin) (nv x 6) mjtNum* cinert; // com-based body inertia and mass (nbody x 10) // computed by mj_fwdPosition/mj_flex @@ -285,8 +285,8 @@ struct mjData_ { mjtNum* actuator_velocity; // actuator velocities (nu x 1) // computed by mj_fwdVelocity/mj_comVel - mjtNum* cvel; // com-based velocity [3D rot; 3D tran] (nbody x 6) - mjtNum* cdof_dot; // time-derivative of cdof (nv x 6) + mjtNum* cvel; // com-based velocity (rot:lin) (nbody x 6) + mjtNum* cdof_dot; // time-derivative of cdof (rot:lin) (nv x 6) // computed by mj_fwdVelocity/mj_rne (without acceleration) mjtNum* qfrc_bias; // C(qpos,qvel) (nv x 1) @@ -459,10 +459,11 @@ typedef enum mjtGeom_ { // type of geometric shape mjGEOM_ARROW1, // arrow without wedges mjGEOM_ARROW2, // arrow in both directions mjGEOM_LINE, // line + mjGEOM_LINEBOX, // box with line edges mjGEOM_FLEX, // flex mjGEOM_SKIN, // skin mjGEOM_LABEL, // text label - mjGEOM_TRIANGLE, // triangle connecting a frame + mjGEOM_TRIANGLE, // triangle mjGEOM_NONE = 1001 // missing geom type } mjtGeom; @@ -570,7 +571,9 @@ typedef enum mjtObj_ { // type of MujoCo object mjOBJ_TEXT, // text mjOBJ_TUPLE, // tuple mjOBJ_KEY, // keyframe - mjOBJ_PLUGIN // plugin instance + mjOBJ_PLUGIN, // plugin instance + + mjNOBJECT // number of object types } mjtObj; typedef enum mjtConstraint_ { // type of constraint mjCNSTR_EQUALITY = 0, // equality constraint @@ -1182,20 +1185,23 @@ struct mjModel_ { int* skin_bonebodyid; // body id of each bone (nskinbone x 1) int* skin_bonevertid; // mesh ids of vertices in each bone (nskinbonevert x 1) float* skin_bonevertweight; // weights of vertices in each bone (nskinbonevert x 1) + int* skin_pathadr; // address of asset path for skin; -1: none (nskin x 1) // height fields - mjtNum* hfield_size; // (x, y, z_top, z_bottom) (nhfield x 4) - int* hfield_nrow; // number of rows in grid (nhfield x 1) - int* hfield_ncol; // number of columns in grid (nhfield x 1) - int* hfield_adr; // address in hfield_data (nhfield x 1) - float* hfield_data; // elevation data (nhfielddata x 1) + mjtNum* hfield_size; // (x, y, z_top, z_bottom) (nhfield x 4) + int* hfield_nrow; // number of rows in grid (nhfield x 1) + int* hfield_ncol; // number of columns in grid (nhfield x 1) + int* hfield_adr; // address in hfield_data (nhfield x 1) + float* hfield_data; // elevation data (nhfielddata x 1) + int* hfield_pathadr; // address of asset path for hfield; -1: none (nhfield x 1) // textures - int* tex_type; // texture type (mjtTexture) (ntex x 1) - int* tex_height; // number of rows in texture image (ntex x 1) - int* tex_width; // number of columns in texture image (ntex x 1) - int* tex_adr; // address in rgb (ntex x 1) - mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_type; // texture type (mjtTexture) (ntex x 1) + int* tex_height; // number of rows in texture image (ntex x 1) + int* tex_width; // number of columns in texture image (ntex x 1) + int* tex_adr; // address in rgb (ntex x 1) + mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_pathadr; // address of asset path for texture; -1: none (ntex x 1) // materials int* mat_texid; // texture id; -1: none (nmat x 1) @@ -1422,6 +1428,15 @@ struct mjpPlugin_ { // called by mjv_updateScene (optional) void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance); + // methods specific to actuators (optional) + + // dimension of the actuator state for the plugin (excluding state from actuator's dyntype) + int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id); + + // updates the actuator plugin's entries in act_dot + // called after native act_dot is computed and before the compute callback + void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance); + // methods specific to signed distance fields (optional) // signed distance from the surface @@ -2155,6 +2170,7 @@ struct mjvSceneState_ { int nnames; int npaths; int nsensordata; + int narena; mjOption opt; mjVisual vis; @@ -2226,6 +2242,7 @@ struct mjvSceneState_ { int* flex_vertadr; int* flex_vertnum; int* flex_elem; + int* flex_elemlayer; int* flex_elemadr; int* flex_elemnum; int* flex_elemdataadr; @@ -2238,6 +2255,8 @@ struct mjvSceneState_ { mjtNum* flex_radius; float* flex_rgba; + int* hfield_pathadr; + int* mesh_bvhadr; int* mesh_bvhnum; int* mesh_texcoordadr; @@ -2264,6 +2283,9 @@ struct mjvSceneState_ { int* skin_bonebodyid; int* skin_bonevertid; float* skin_bonevertweight; + int* skin_pathadr; + + int* tex_pathadr; int* mat_texid; mjtByte* mat_texuniform; @@ -2362,6 +2384,7 @@ struct mjvSceneState_ { mjtNum* ten_length; mjtNum* wrap_xpos; + mjtNum* bvh_aabb_dyn; mjtByte* bvh_active; int* island_dofadr; int* island_dofind; @@ -2373,6 +2396,7 @@ struct mjvSceneState_ { mjContact* contact; mjtNum* efc_force; + void* arena; } data; }; typedef struct mjvSceneState_ mjvSceneState; @@ -2435,6 +2459,7 @@ void mj_fwdAcceleration(const mjModel* m, mjData* d); void mj_fwdConstraint(const mjModel* m, mjData* d); void mj_Euler(const mjModel* m, mjData* d); void mj_RungeKutta(const mjModel* m, mjData* d, int N); +void mj_implicit(const mjModel* m, mjData* d); void mj_invPosition(const mjModel* m, mjData* d); void mj_invVelocity(const mjModel* m, mjData* d); void mj_invConstraint(const mjModel* m, mjData* d); diff --git a/doc/mjx.rst b/doc/mjx.rst index 2bdd9a58..aafb7a86 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -181,9 +181,9 @@ The following features are **fully supported** in MJX: * - :ref:`Joint ` - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE`` * - :ref:`Transmission ` - - ``TRN_JOINT`` + - ``TRN_JOINT``, ``TRN_SITE`` * - :ref:`Actuator Dynamics ` - - ``NONE``, ``INTEGRATOR``, ``FILTER`` + - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT`` * - :ref:`Actuator Gain ` - ``FIXED``, ``AFFINE`` * - :ref:`Actuator Bias ` @@ -257,9 +257,9 @@ The following features are **unsupported**: * - Category - Feature * - :ref:`Transmission ` - - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY`` + - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY`` * - :ref:`Actuator Dynamics ` - - ``FILTEREXACT``, ``USER`` + - ``USER`` * - :ref:`Actuator Gain ` - ``USER`` * - :ref:`Actuator Bias ` @@ -288,6 +288,22 @@ Single scene simulation Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been carefully optimized for CPU. MJX works best when simulating thousands or tens of thousands of scenes in parallel. +Collisions between large meshes + MJX supports collisions between convex mesh geometries. However the convex collision algorithms + in MJX are implemented differently than in MuJoCo. MJX uses a branchless version of the + `Separating Axis Test `__ + (SAT) to determine if geometries are colliding with convex meshes, while MuJoCo uses the Minkowski Portal Refinement (MPR) + algorithm as implemented in `libccd `__. + SAT works well for smaller meshes but suffers in both runtime and memory for larger meshes. + + For + collisions between convex meshes and primitives (spheres, capsules, planes), use **3000 vertices or less** for your convex meshes. + For collisions between convex meshes and other convex meshes, use **30 vertices or less**. + With careful + tuning, MJX can simulate scenes with mesh collisions -- see the MJX + `shadow hand `__ + config for an example. Speeding up mesh collision detection is an active area of development for MJX. + Large, complex scenes with many contacts Accelerators exhibit poor performance for `branching code `__. @@ -296,7 +312,7 @@ Large, complex scenes with many contacts powerful as the one in MuJoCo. To see how this affects simulation, let us consider a physics scene with increasing numbers of humanoid bodies, - varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M1 Pro and a 64-core AMD 3995WX and time + varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M3 Max and a 64-core AMD 3995WX and time it using :ref:`testspeed`, using ``2 x numcore`` threads. We time the MJX simulation on an Nvidia A100 GPU using a batch size of 8192 and an 8-chip `v5 TPU `__ @@ -306,18 +322,10 @@ Large, complex scenes with many contacts :width: 95% :align: center - The values for a single humanoid (leftmost datapoints) for the four timed architectures are **320K**, **1.8M**, + The values for a single humanoid (leftmost datapoints) for the four timed architectures are **650K**, **1.8M**, **950K** and **2.7M** steps per second, respectively. Note that as we increase the number of humanoids (which increases the number of potential contacts in a scene), MJX throughput decreases more rapidly than MuJoCo. -Scenes with collisions between meshes with many vertices - MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of - `mesh collision algorithms `__. - These algorithms work well for smaller meshes (with hundreds of vertices) but suffer with large meshes. With careful - tuning, MJX can simulate scenes with mesh collisions well -- see the MJX - `shadow hand `__ - config for an example. - .. _MjxPerformance: Performance tuning diff --git a/doc/modeling.rst b/doc/modeling.rst index 7abdb53e..b9cc9043 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -204,16 +204,21 @@ cameras and lights. A related attribute is :ref:`compiler/angle`. It specifies whether angles in the MJCF file are expressed in degrees or radians (after compilation, angles are always expressed in radians). +Positions are specified using + +:at:`pos`: :at-val:`real(3), "0 0 0"` + Position relative to parent. + .. _COrientation: Frame orientations -~~~~~~~~~~~~~~~~~~ +^^^^^^^^^^^^^^^^^^ Several model elements have right-handed spatial frames associated with them. These are all the elements defined in the kinematic tree except for joints. A spatial frame is defined by its position and orientation. Specifying 3D positions is straightforward, but specifying 3D orientations can be challenging. This is why MJCF provides several alternative -mechanisms. No matter which mechanism the user chooses, the frame orientation is always represented as a unit quaternion -after compilation. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)` +mechanisms. No matter which mechanism the user chooses, the frame orientation is always converted internally to a unit +quaternion. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)` corresponds to the quaternion :math:`(\cos(a/2), \: \sin(a/2) \cdot (x, y, z))`. Also recall that every 3D orientation can be uniquely specified by a single 3D rotation by some angle around some axis. @@ -267,6 +272,7 @@ approximately .. math:: \ac + d \cdot (b v + k r) = (1 - d)\cdot \au + :label: eq:constraint Again, the parameters that are under the user's control are :math:`d, b, k`. The remaining quantities are functions of the system state and are computed automatically at each time step. @@ -312,7 +318,15 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec units of :math:`\text{width}`. Note that when :math:`\text{power}` is 1, the function is linear regardless of the :math:`\text{midpoint}`. - |image0| + .. image:: images/modeling/impedance.png + :width: 600px + :align: center + :class: only-light + + .. image:: images/modeling/impedance_dark.png + :width: 600px + :align: center + :class: only-dark These plots show the impedance :math:`d(r)` on the vertical axis, as a function of the constraint violation :math:`r` on the horizontal axis. @@ -338,12 +352,9 @@ Next we explain the setting of the stiffness :math:`k` and damping :math:`b` whi .. admonition:: Intuitive description of the **reference acceleration** - The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in - order to rectify violation. For example, consider a contact between a motionless free body pulled down by gravity - onto a static plane geom. Since there is no motion, the penetration will be entirely determined by the impedance - while the reference has no effect. Now imagine that the body is dropped onto the plane. Upon impact the constraint - will generate a normal force which attempts to rectify the penetration using a particular motion; this motion is - the reference acceleration. + The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in order to + rectify violation. Imagine a body dropped onto the plane. Upon impact the constraint will generate a normal force + which attempts to rectify the penetration using a particular motion; this motion is the reference acceleration. Another way of understanding the reference acceleration is to think of the unmodeled deformation variables described in the :ref:`Computation chapter`. Imagine two bodies pressed together, leading to deformation at @@ -388,7 +399,12 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: can go unstable. This is enforced internally, unless the :ref:`refsafe` 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. + over-damped constraints. Combining the above formula with :eq:`eq:constraint`, we can derive the following result. + If the reference acceleration is given using the positive number format and the impedance is constant + :math:`d = d_0 = d_\text{width}`, then the penetration depth at rest is + + .. math:: + r = \au \cdot (1 - d) \cdot \text{timeconst}^2 \cdot \text{dampratio}^2 Next we describe the direct format where the two numbers are :math:`(-\text{stiffness}, -\text{damping})`. This allows direct control over restitution in particular. We still apply some scaling so that the same numbers can be @@ -398,9 +414,15 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math: .. math:: \begin{aligned} b &= \text{damping} / d_\text{width} \\ - k &= \text{stiffness} / d_\text{width}^2 \\ + k &= \text{stiffness} \cdot d(r) / d_\text{width}^2 \\ \end{aligned} + Similarly to the above derivation, if the reference acceleration is given using the negative number format and the + impedance is constant, then the penetration depth at rest is + + .. math:: + r = \au \cdot (1 - d) \cdot \text{stiffness} + .. tip:: In the positive-value default format, the :math:`\text{timeconst}` parameter controls constraint **softness**. It is specified in units of time and means "how quickly is the constraint trying to resolve the violation". Larger @@ -851,7 +873,15 @@ The advantage of the scaled quantities is that all muscles behave similarly in t captured by the Force-Length-Velocity (:math:`\text{\small FLV}`) function measured in many experimental papers. We approximate this function as follows: -|image1| +.. image:: images/modeling/musclemodel.png + :width: 650px + :align: center + :class: only-light + +.. image:: images/modeling/musclemodel_dark.png + :width: 650px + :align: center + :class: only-dark The function is in the form: @@ -893,7 +923,15 @@ Before embarking on a mission to design more accurate :math:`\text{\small FLV}` operating range of the muscle has a bigger effect than the shape of the :math:`\text{\small FLV}` function, and in many cases this parameter is unknown. Below is a graphical illustration: -|image2| +.. image:: images/modeling/musclerange.png + :width: 500px + :align: center + :class: only-light + +.. image:: images/modeling/musclerange_dark.png + :width: 500px + :align: center + :class: only-dark This figure format is common in the biomechanics literature, showing the operating range of each muscle superimposed on the normalized :math:`\text{FL}` curve (ignore the vertical displacement). Our default range is shown in black. The blue @@ -1297,7 +1335,9 @@ A flex is a collection of MuJoCo bodies that are connected with massless stretch capsules (1D flex), triangles (2D flex), or tetrahedra (3D flex). In all cases we allow a radius, which makes the elements smooth and also volumetric in 1D and 2D. The primitive elements are illustrated below: -|flexelem| +.. image:: images/modeling/flexelem.png + :width: 600px + :align: center Thus far these look like geoms. But the key difference is that they deform: as the bodies (vertices) move independently of each other, the shape of the elements changes in real time. Collisions and contact forces are now generalized to @@ -1360,7 +1400,8 @@ In case of 3D flexes made of tetrahedra, it may be useful to examine how the fle a special visualization mode that peels off the outer layers. Below is an example with the Stanford Bunny. Note how it has smaller tetrahedra on the outside and larger ones on the inside. This mesh design makes sense, because we want the collision surface to be accurate, but on the inside we just need soft material properties - which require less spatial -resolution. +resolution. In order to convert a surface mesh to a tetrahedral mesh, we recommend open tools like the +`fTetWild library `__. |bunny1| |bunny2| @@ -1663,12 +1704,6 @@ in a visible way, and the energy fluctuates around the initial value instead of -.. |image0| image:: images/modeling/impedance.png - :width: 600px -.. |image1| image:: images/modeling/musclemodel.png - :width: 650px -.. |image2| image:: images/modeling/musclerange.png - :width: 400px .. |image3| image:: images/modeling/tendonwraps.png :width: 500px .. |image4| image:: images/modeling/particle.png @@ -1705,8 +1740,6 @@ in a visible way, and the energy fluctuates around the initial value instead of :height: 250px .. |particle| image:: images/models/particle.gif :width: 270px -.. |flexelem| image:: images/modeling/flexelem.png - :width: 400px .. |bunny1| image:: images/modeling/bunny1.png :width: 300px .. |bunny2| image:: images/modeling/bunny2.png diff --git a/doc/programming/extension.rst b/doc/programming/extension.rst index f41c9629..49ba77ff 100644 --- a/doc/programming/extension.rst +++ b/doc/programming/extension.rst @@ -199,6 +199,25 @@ When :ref:`mjData` is being copied via :ref:`mj_copyData`, MuJoCo will copy over code is responsible for setting up the plugin data for the newly copied :ref:`mjData`. To facilitate this, MuJoCo calls the ``copy`` callback from :ref:`mjpPlugin` for each plugin instance present. +.. _exActuatorAct: + +Actuator activations +"""""""""""""""""""" + +When writing stateful actuator plugins, there are two choices for where to save the actuator state. One option is using +``plugin_state`` as described above, and the other is to use ``mjData.act`` by implementing the ``actuator_actdim`` and +``actuator_act_dot`` callbacks on :ref:`mjpPlugin`. + +When using the latter option, the actuator plugin's state will be added to ``mjData.act``, and MuJoCo will +automatically integrate ``mjData.act_dot`` values between timesteps. One advantage of this approach is that +finite-differencing functions like :ref:`mjd_transitionFD` will work as they do for native actuators. The +``mjpPlugin.advance`` callback will be called after ``act_dot`` is integrated, and actuator plugins may overwrite +the ``act`` values at that point, if Euler integration isn't appropriate. + +Users may specify the :ref:`dyntype` attribute on actuator plugins, to introduce a filter or +an integrator between user inputs and actuator activations. When they do, the activation variable introduced by +``dyntype`` will be placed *after* the plugin's activation variables in the ``act`` array. + .. _exRegistration: Registration @@ -248,8 +267,11 @@ A future version of this section will include: * Things that developers need to keep in mind in order to ensure that plugins function correctly when :ref:`mjData` is copied, stepped, or reset. -Currently, there are three directories of first-party plugins: +There are several first-party plugin directories: +* **actuator:** The plugins in the `actuator/ `__ + directory implement custom actuators, so far only a PID controller. See the + `README `__ for details. * **elasticity:** The plugins in the `elasticity/ `__ directory are passive forces based on continuum mechanics for 1-dimensional and 3-dimensional bodies. The 1D model is invariant under rotations and captures @@ -269,11 +291,11 @@ Currently, there are three directories of first-party plugins: `__. The rest of this section will give more detail concerning the collision algorithm and the plugin engine interface. - Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent. - Because SDFs are non-convex, multiple starting points are required in order to converge to multiple local minima. - The number of starting points is set using :ref:`sdf_initpoints`, and are - initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes. - The number of gradient descent iterations is set using :ref:`sdf_iterations`. + Collision points are found by minimizing the function A + B + abs(max(A, B)), where A and B are the two colliding + SDFs, via gradient descent. Because SDFs are non-convex, multiple starting points are required in order to converge to + multiple local minima. The number of starting points is set using :ref:`sdf_initpoints`, and + are initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes. The number of + gradient descent iterations is set using :ref:`sdf_iterations`. While *exact* SDFs---encoding the precise signed distance to the surface---are preferred, collisions are possible with any function whose value vanishes at the surface and grows monotonically away from it, with a negative sign in the @@ -323,7 +345,6 @@ loading functions. The :ref:`mjpResourceProvider` struct stores three types of f .. _Uniform Resource Identifier: https://en.wikipedia.org/wiki/Uniform_Resource_Identifier Resource prefix - Resources are identified by prefixes in their name. The chosen prefix should have a valid `Uniform Resource Identifier`_ (URI) scheme syntax. Resource names should also have a valid URI syntax, however this isn't enforced. A resource name with the syntax ``{prefix}:{filename}`` will match a provider using the scheme ``prefix``. For diff --git a/doc/unity.rst b/doc/unity.rst index 457c2cd7..4e443f40 100644 --- a/doc/unity.rst +++ b/doc/unity.rst @@ -30,14 +30,14 @@ _____ The MuJoCo app needs to be run at least once before the native library can be used, in order to register the library as a trusted binary. Then, copy the dynamic library file from -``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.0.2.dylib`` (it can be +``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.1.2.dylib`` (it can be found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``. Linux _____ Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from -``~/.mujoco/mujoco-3.0.2/lib/libmujoco.so.3.0.2`` and rename it as ``libmujoco.so``. +``~/.mujoco/mujoco-3.1.2/lib/libmujoco.so.3.1.2`` and rename it as ``libmujoco.so``. Windows _______ diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index a3609b8b..1fc6fe19 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -265,7 +265,7 @@ struct mjData_ { // computed by mj_fwdPosition/mj_comPos mjtNum* subtree_com; // center of mass of each subtree (nbody x 3) - mjtNum* cdof; // com-based motion axis of each dof (nv x 6) + mjtNum* cdof; // com-based motion axis of each dof (rot:lin) (nv x 6) mjtNum* cinert; // com-based body inertia and mass (nbody x 10) // computed by mj_fwdPosition/mj_flex @@ -313,8 +313,8 @@ struct mjData_ { mjtNum* actuator_velocity; // actuator velocities (nu x 1) // computed by mj_fwdVelocity/mj_comVel - mjtNum* cvel; // com-based velocity [3D rot; 3D tran] (nbody x 6) - mjtNum* cdof_dot; // time-derivative of cdof (nv x 6) + mjtNum* cvel; // com-based velocity (rot:lin) (nbody x 6) + mjtNum* cdof_dot; // time-derivative of cdof (rot:lin) (nv x 6) // computed by mj_fwdVelocity/mj_rne (without acceleration) mjtNum* qfrc_bias; // C(qpos,qvel) (nv x 1) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index b9f7b8e9..d3417683 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -109,10 +109,11 @@ typedef enum mjtGeom_ { // type of geometric shape mjGEOM_ARROW1, // arrow without wedges mjGEOM_ARROW2, // arrow in both directions mjGEOM_LINE, // line + mjGEOM_LINEBOX, // box with line edges mjGEOM_FLEX, // flex mjGEOM_SKIN, // skin mjGEOM_LABEL, // text label - mjGEOM_TRIANGLE, // triangle connecting a frame + mjGEOM_TRIANGLE, // triangle mjGEOM_NONE = 1001 // missing geom type } mjtGeom; @@ -246,7 +247,9 @@ typedef enum mjtObj_ { // type of MujoCo object mjOBJ_TEXT, // text mjOBJ_TUPLE, // tuple mjOBJ_KEY, // keyframe - mjOBJ_PLUGIN // plugin instance + mjOBJ_PLUGIN, // plugin instance + + mjNOBJECT // number of object types } mjtObj; @@ -895,20 +898,23 @@ struct mjModel_ { int* skin_bonebodyid; // body id of each bone (nskinbone x 1) int* skin_bonevertid; // mesh ids of vertices in each bone (nskinbonevert x 1) float* skin_bonevertweight; // weights of vertices in each bone (nskinbonevert x 1) + int* skin_pathadr; // address of asset path for skin; -1: none (nskin x 1) // height fields - mjtNum* hfield_size; // (x, y, z_top, z_bottom) (nhfield x 4) - int* hfield_nrow; // number of rows in grid (nhfield x 1) - int* hfield_ncol; // number of columns in grid (nhfield x 1) - int* hfield_adr; // address in hfield_data (nhfield x 1) - float* hfield_data; // elevation data (nhfielddata x 1) + mjtNum* hfield_size; // (x, y, z_top, z_bottom) (nhfield x 4) + int* hfield_nrow; // number of rows in grid (nhfield x 1) + int* hfield_ncol; // number of columns in grid (nhfield x 1) + int* hfield_adr; // address in hfield_data (nhfield x 1) + float* hfield_data; // elevation data (nhfielddata x 1) + int* hfield_pathadr; // address of asset path for hfield; -1: none (nhfield x 1) // textures - int* tex_type; // texture type (mjtTexture) (ntex x 1) - int* tex_height; // number of rows in texture image (ntex x 1) - int* tex_width; // number of columns in texture image (ntex x 1) - int* tex_adr; // address in rgb (ntex x 1) - mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_type; // texture type (mjtTexture) (ntex x 1) + int* tex_height; // number of rows in texture image (ntex x 1) + int* tex_width; // number of columns in texture image (ntex x 1) + int* tex_adr; // address in rgb (ntex x 1) + mjtByte* tex_rgb; // rgb (alpha = 1) (ntexdata x 1) + int* tex_pathadr; // address of asset path for texture; -1: none (ntex x 1) // materials int* mat_texid; // texture id; -1: none (nmat x 1) diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h index 6f583c32..ca850d57 100644 --- a/include/mujoco/mjplugin.h +++ b/include/mujoco/mjplugin.h @@ -107,6 +107,15 @@ struct mjpPlugin_ { // called by mjv_updateScene (optional) void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance); + // methods specific to actuators (optional) + + // dimension of the actuator state for the plugin (excluding state from actuator's dyntype) + int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id); + + // updates the actuator plugin's entries in act_dot + // called after native act_dot is computed and before the compute callback + void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance); + // methods specific to signed distance fields (optional) // signed distance from the surface diff --git a/include/mujoco/mjvisualize.h b/include/mujoco/mjvisualize.h index 37eef03e..dc8a4bc6 100644 --- a/include/mujoco/mjvisualize.h +++ b/include/mujoco/mjvisualize.h @@ -436,6 +436,7 @@ struct mjvSceneState_ { int nnames; int npaths; int nsensordata; + int narena; mjOption opt; mjVisual vis; @@ -507,6 +508,7 @@ struct mjvSceneState_ { int* flex_vertadr; int* flex_vertnum; int* flex_elem; + int* flex_elemlayer; int* flex_elemadr; int* flex_elemnum; int* flex_elemdataadr; @@ -519,6 +521,8 @@ struct mjvSceneState_ { mjtNum* flex_radius; float* flex_rgba; + int* hfield_pathadr; + int* mesh_bvhadr; int* mesh_bvhnum; int* mesh_texcoordadr; @@ -545,6 +549,9 @@ struct mjvSceneState_ { int* skin_bonebodyid; int* skin_bonevertid; float* skin_bonevertweight; + int* skin_pathadr; + + int* tex_pathadr; int* mat_texid; mjtByte* mat_texuniform; @@ -643,6 +650,7 @@ struct mjvSceneState_ { mjtNum* ten_length; mjtNum* wrap_xpos; + mjtNum* bvh_aabb_dyn; mjtByte* bvh_active; int* island_dofadr; int* island_dofind; @@ -654,6 +662,7 @@ struct mjvSceneState_ { mjContact* contact; mjtNum* efc_force; + void* arena; } data; }; typedef struct mjvSceneState_ mjvSceneState; diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index e8879aeb..86f61e0c 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -334,7 +334,7 @@ X ( int, flex_vertbodyid, nflexvert, 1 ) \ X ( int, flex_edge, nflexedge, 2 ) \ XMJV( int, flex_elem, nflexelemdata, 1 ) \ - X ( int, flex_elemlayer, nflexelem, 1 ) \ + XMJV( int, flex_elemlayer, nflexelem, 1 ) \ XMJV( int, flex_shell, nflexshelldata,1 ) \ X ( int, flex_evpair, nflexevpair, 2 ) \ X ( mjtNum, flex_vert, nflexvert, 3 ) \ @@ -395,16 +395,19 @@ XMJV( int, skin_bonebodyid, nskinbone, 1 ) \ XMJV( int, skin_bonevertid, nskinbonevert, 1 ) \ XMJV( float, skin_bonevertweight, nskinbonevert, 1 ) \ + XMJV( int, skin_pathadr, nskin, 1 ) \ X ( mjtNum, hfield_size, nhfield, 4 ) \ X ( int, hfield_nrow, nhfield, 1 ) \ X ( int, hfield_ncol, nhfield, 1 ) \ X ( int, hfield_adr, nhfield, 1 ) \ X ( float, hfield_data, nhfielddata, 1 ) \ + XMJV( int, hfield_pathadr, nhfield, 1 ) \ X ( int, tex_type, ntex, 1 ) \ X ( int, tex_height, ntex, 1 ) \ X ( int, tex_width, ntex, 1 ) \ X ( int, tex_adr, ntex, 1 ) \ X ( mjtByte, tex_rgb, ntexdata, 1 ) \ + XMJV( int, tex_pathadr, ntex, 1 ) \ XMJV( int, mat_texid, nmat, 1 ) \ XMJV( mjtByte, mat_texuniform, nmat, 1 ) \ XMJV( float, mat_texrepeat, nmat, 2 ) \ @@ -611,7 +614,7 @@ X ( mjtNum, qLD, nM, 1 ) \ X ( mjtNum, qLDiagInv, nv, 1 ) \ X ( mjtNum, qLDiagSqrtInv, nv, 1 ) \ - X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ + XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ XMJV( mjtByte, bvh_active, nbvh, 1 ) \ X ( mjtNum, flexedge_velocity, nflexedge, 1 ) \ X ( mjtNum, ten_velocity, ntendon, 1 ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index ac1b509e..cb6962a1 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -24,7 +24,7 @@ extern "C" { #endif // header version; should match the library version as returned by mj_version() -#define mjVERSION_HEADER 302 +#define mjVERSION_HEADER 312 // needed to define size_t, fabs and log10 #include @@ -264,6 +264,9 @@ MJAPI void mj_Euler(const mjModel* m, mjData* d); // Runge-Kutta explicit order-N integrator. MJAPI void mj_RungeKutta(const mjModel* m, mjData* d, int N); +// Implicit-in-velocity integrators. +MJAPI void mj_implicit(const mjModel* m, mjData* d); + // Run position-dependent computations in inverse dynamics. MJAPI void mj_invPosition(const mjModel* m, mjData* d); @@ -1115,7 +1118,7 @@ MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, // Address of diagonal element i in band-dense matrix representation. MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense); -// Eigenvalue decomposition of symmetric 3x3 matrix. +// Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'. MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]); // minimize 0.5*x'*H*x + x'*g s.t. lower <= x <= upper, return rank or -1 if failed diff --git a/introspect/enums.py b/introspect/enums.py index bb039f65..8502dbdb 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -90,10 +90,11 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjGEOM_ARROW1', 101), ('mjGEOM_ARROW2', 102), ('mjGEOM_LINE', 103), - ('mjGEOM_FLEX', 104), - ('mjGEOM_SKIN', 105), - ('mjGEOM_LABEL', 106), - ('mjGEOM_TRIANGLE', 107), + ('mjGEOM_LINEBOX', 104), + ('mjGEOM_FLEX', 105), + ('mjGEOM_SKIN', 106), + ('mjGEOM_LABEL', 107), + ('mjGEOM_TRIANGLE', 108), ('mjGEOM_NONE', 1001), ]), )), @@ -265,6 +266,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjOBJ_TUPLE', 23), ('mjOBJ_KEY', 24), ('mjOBJ_PLUGIN', 25), + ('mjNOBJECT', 26), ]), )), ('mjtConstraint', diff --git a/introspect/enums_test.py b/introspect/enums_test.py index b4144420..3580d4c1 100644 --- a/introspect/enums_test.py +++ b/introspect/enums_test.py @@ -61,7 +61,7 @@ class EnumsTest(absltest.TestCase): self.assertEqual(enum_decl.values['mjGEOM_ARROW'], 100) self.assertEqual(enum_decl.values['mjGEOM_ARROW1'], 101) self.assertEqual(enum_decl.values['mjGEOM_ARROW2'], 102) - self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 107) + self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 108) # Skip a few... self.assertEqual(enum_decl.values['mjGEOM_NONE'], 1001) diff --git a/introspect/functions.py b/introspect/functions.py index d6b1ae56..0e9d2568 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -1156,6 +1156,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Runge-Kutta explicit order-N integrator.', )), + ('mj_implicit', + FunctionDecl( + name='mj_implicit', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + ), + doc='Implicit-in-velocity integrators.', + )), ('mj_invPosition', FunctionDecl( name='mj_invPosition', @@ -7348,7 +7368,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc='Eigenvalue decomposition of symmetric 3x3 matrix.', + doc="Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.", # pylint: disable=line-too-long )), ('mju_boxQP', FunctionDecl( diff --git a/introspect/structs.py b/introspect/structs.py index fd1f085c..43e7cdc9 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -2787,75 +2787,96 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='weights of vertices in each bone (nskinbonevert x 1)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='skin_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='address of asset path for skin; -1: none (nskin x 1)', + ), StructFieldDecl( name='hfield_size', type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='(x, y, z_top, z_bottom) (nhfield x 4)', + doc='(x, y, z_top, z_bottom) (nhfield x 4)', ), StructFieldDecl( name='hfield_nrow', type=PointerType( inner_type=ValueType(name='int'), ), - doc='number of rows in grid (nhfield x 1)', + doc='number of rows in grid (nhfield x 1)', ), StructFieldDecl( name='hfield_ncol', type=PointerType( inner_type=ValueType(name='int'), ), - doc='number of columns in grid (nhfield x 1)', + doc='number of columns in grid (nhfield x 1)', ), StructFieldDecl( name='hfield_adr', type=PointerType( inner_type=ValueType(name='int'), ), - doc='address in hfield_data (nhfield x 1)', + doc='address in hfield_data (nhfield x 1)', ), StructFieldDecl( name='hfield_data', type=PointerType( inner_type=ValueType(name='float'), ), - doc='elevation data (nhfielddata x 1)', # pylint: disable=line-too-long + doc='elevation data (nhfielddata x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='hfield_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='address of asset path for hfield; -1: none (nhfield x 1)', ), StructFieldDecl( name='tex_type', type=PointerType( inner_type=ValueType(name='int'), ), - doc='texture type (mjtTexture) (ntex x 1)', + doc='texture type (mjtTexture) (ntex x 1)', ), StructFieldDecl( name='tex_height', type=PointerType( inner_type=ValueType(name='int'), ), - doc='number of rows in texture image (ntex x 1)', + doc='number of rows in texture image (ntex x 1)', ), StructFieldDecl( name='tex_width', type=PointerType( inner_type=ValueType(name='int'), ), - doc='number of columns in texture image (ntex x 1)', + doc='number of columns in texture image (ntex x 1)', ), StructFieldDecl( name='tex_adr', type=PointerType( inner_type=ValueType(name='int'), ), - doc='address in rgb (ntex x 1)', + doc='address in rgb (ntex x 1)', ), StructFieldDecl( name='tex_rgb', type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='rgb (alpha = 1) (ntexdata x 1)', + doc='rgb (alpha = 1) (ntexdata x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='tex_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='address of asset path for texture; -1: none (ntex x 1)', ), StructFieldDecl( name='mat_texid', @@ -4472,7 +4493,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='com-based motion axis of each dof (nv x 6)', # pylint: disable=line-too-long + doc='com-based motion axis of each dof (rot:lin) (nv x 6)', # pylint: disable=line-too-long ), StructFieldDecl( name='cinert', @@ -4682,14 +4703,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='com-based velocity [3D rot; 3D tran] (nbody x 6)', # pylint: disable=line-too-long + doc='com-based velocity (rot:lin) (nbody x 6)', # pylint: disable=line-too-long ), StructFieldDecl( name='cdof_dot', type=PointerType( inner_type=ValueType(name='mjtNum'), ), - doc='time-derivative of cdof (nv x 6)', # pylint: disable=line-too-long + doc='time-derivative of cdof (rot:lin) (nv x 6)', # pylint: disable=line-too-long ), StructFieldDecl( name='qfrc_bias', @@ -6313,6 +6334,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='', ), + StructFieldDecl( + name='narena', + type=ValueType(name='int'), + doc='', + ), StructFieldDecl( name='opt', type=ValueType(name='mjOption'), @@ -6741,6 +6767,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), + StructFieldDecl( + name='flex_elemlayer', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='', + ), StructFieldDecl( name='flex_elemadr', type=PointerType( @@ -6818,6 +6851,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), + StructFieldDecl( + name='hfield_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='', + ), StructFieldDecl( name='mesh_bvhadr', type=PointerType( @@ -6993,6 +7033,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), + StructFieldDecl( + name='skin_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='', + ), + StructFieldDecl( + name='tex_pathadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='', + ), StructFieldDecl( name='mat_texid', type=PointerType( @@ -7547,6 +7601,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), + StructFieldDecl( + name='bvh_aabb_dyn', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='', + ), StructFieldDecl( name='bvh_active', type=PointerType( @@ -7610,6 +7671,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), + StructFieldDecl( + name='arena', + type=PointerType( + inner_type=ValueType(name='void'), + ), + doc='', + ), ), ), doc='', diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 3ec6312a..f4953c60 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -16,13 +16,24 @@ # pylint:disable=g-importing-member from mujoco.mjx._src.collision_driver import collision +from mujoco.mjx._src.constraint import count_constraints from mujoco.mjx._src.constraint import make_constraint from mujoco.mjx._src.device import device_get_into from mujoco.mjx._src.device import device_put +from mujoco.mjx._src.forward import euler from mujoco.mjx._src.forward import forward +from mujoco.mjx._src.forward import fwd_acceleration +from mujoco.mjx._src.forward import fwd_actuation +from mujoco.mjx._src.forward import fwd_position +from mujoco.mjx._src.forward import fwd_velocity +from mujoco.mjx._src.forward import rungekutta4 from mujoco.mjx._src.forward import step +from mujoco.mjx._src.io import get_data from mujoco.mjx._src.io import make_data +from mujoco.mjx._src.io import put_data +from mujoco.mjx._src.io import put_model from mujoco.mjx._src.passive import passive +from mujoco.mjx._src.ray import ray from mujoco.mjx._src.smooth import com_pos from mujoco.mjx._src.smooth import com_vel from mujoco.mjx._src.smooth import crb @@ -31,4 +42,5 @@ from mujoco.mjx._src.smooth import kinematics from mujoco.mjx._src.smooth import mul_m from mujoco.mjx._src.smooth import rne from mujoco.mjx._src.smooth import transmission +from mujoco.mjx._src.solver import solve from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index dde3fe92..c54278b2 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -40,7 +40,6 @@ from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import GeomType from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member -import numpy as np # pair-wise collision functions @@ -81,6 +80,12 @@ def _add_candidate( if t1 > t2: t1, t2, g1, g2 = t2, t1, g2, g1 + # MuJoCo does not collide planes with other planes or hfields + if t1 == GeomType.PLANE and t2 == GeomType.PLANE: + return + if t1 == GeomType.PLANE and t2 == GeomType.HFIELD: + return + def mesh_key(i): convex_data = [[None] * m.ngeom] * 3 if isinstance(m, Model): @@ -284,13 +289,11 @@ def _collide_geoms( solimp=params.solimp, geom1=geom1, geom2=geom2, - dim=np.array([]), - efc_address=np.array([]), ) return con -def _max_contact_points(m: Model) -> int: +def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> int: """Returns the maximum number of contact points when set as a numeric.""" for i in range(m.nnumeric): name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0] @@ -334,7 +337,7 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: return candidate_set -def ncon(m: Model) -> int: +def ncon(m: Union[Model, mujoco.MjModel]) -> int: """Returns the number of contacts computed in MJX given a model.""" if m.opt.disableflags & DisableBit.CONTACT: return 0 @@ -354,9 +357,8 @@ def ncon(m: Model) -> int: def collision(m: Model, d: Data) -> Data: """Collides geometries.""" - ncon_ = ncon(m) - if ncon_ == 0: - return d.replace(contact=Contact.zero(), ncon=0) + if ncon(m) == 0: + return d.replace(contact=Contact.zero()) candidate_set = collision_candidates(m) @@ -376,13 +378,4 @@ def collision(m: Model, d: Data) -> Data: _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact) - if ncon_ != contact.dist.shape[0]: - raise RuntimeError('Number of contacts does not match ncon.') - - # TODO(robotics-simulation): move this logic to device_put - ns = d.ne + d.nf + d.nl - contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4)) - # TODO(robotics-simulation): add support for other friction dimensions - contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32)) - - return d.replace(contact=contact, ncon=ncon_) + return d.replace(contact=contact) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index b735e014..5b58f7f8 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -52,9 +52,9 @@ def _collide( mjcf: str, assets: Optional[Dict[str, str]] = None ) -> Tuple[mujoco.MjModel, mujoco.MjData, Model, Data]: m = mujoco.MjModel.from_xml_string(mjcf, assets or {}) - mx = mjx.device_put(m) + mx = mjx.put_model(m) d = mujoco.MjData(m) - dx = mjx.device_put(d) + dx = mjx.put_data(m, d) mujoco.mj_step(m, d) collision_jit_fn = jax.jit(mjx.collision) @@ -253,7 +253,6 @@ class CapsuleCollisionTest(parameterized.TestCase): self.assertGreater(c.dist[1], 0) # extract the contact point with penetration c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) - c = c.replace(dim=c.dim[np.array([0])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) @@ -281,7 +280,6 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(-dx.contact.dist[2:], 0) # extract the contact points with penetration c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) - c = c.replace(dim=c.dim[np.array([0, 1])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2) @@ -339,7 +337,6 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(-dx.contact.dist[1:], 0) # extract the contact point with penetration c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) - c = c.replace(dim=c.dim[np.array([0])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) @@ -421,9 +418,9 @@ class BodyPairFilterTest(absltest.TestCase): def test_filter_parent_child(self): """Tests that parent-child collisions get filtered.""" m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) - mx = mjx.device_put(m) + mx = mjx.put_model(m) d = mujoco.MjData(m) - dx = mjx.device_put(d) + dx = mjx.put_data(m, d) mujoco.mj_step(m, d) collision_jit_fn = jax.jit(mjx.collision) @@ -438,9 +435,9 @@ class BodyPairFilterTest(absltest.TestCase): """Tests that filterparent flag disables parent-child filtering.""" m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_FILTERPARENT - mx = mjx.device_put(m) + mx = mjx.put_model(m) d = mujoco.MjData(m) - dx = mjx.device_put(d) + dx = mjx.put_data(m, d) mujoco.mj_step(m, d) collision_jit_fn = jax.jit(mjx.collision) @@ -457,22 +454,14 @@ class NconTest(parameterized.TestCase): """Tests ncon.""" def test_ncon(self): - m = test_util.load_test_file('ant.xml') - d = mujoco.MjData(m) - d.qpos[2] = 0.0 - - mx = mjx.device_put(m) - ncon = collision_driver.ncon(mx) - self.assertEqual(ncon, 4) + m = test_util.load_test_file('constraints.xml') + ncon = collision_driver.ncon(m) + self.assertEqual(ncon, 16) def test_disable_contact(self): - m = test_util.load_test_file('ant.xml') - d = mujoco.MjData(m) - d.qpos[2] = 0.0 - - m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT - mx = mjx.device_put(m) - ncon = collision_driver.ncon(mx) + m = test_util.load_test_file('constraints.xml') + m.opt.disableflags |= DisableBit.CONTACT + ncon = collision_driver.ncon(m) self.assertEqual(ncon, 0) @@ -503,12 +492,12 @@ class TopKContactTest(absltest.TestCase): def test_top_k_contacts(self): m = mujoco.MjModel.from_xml_string(self._CAPSULES) - mx_top_k = mjx.device_put(m) + mx_top_k = mjx.put_model(m) mx_all = mx_top_k.replace( nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([]) ) d = mujoco.MjData(m) - dx = mjx.device_put(d) + dx = mjx.put_data(m, d) collision_jit_fn = jax.jit(mjx.collision) kinematics_jit_fn = jax.jit(mjx.kinematics) @@ -517,8 +506,8 @@ class TopKContactTest(absltest.TestCase): dx_all = collision_jit_fn(mx_all, dx) dx_top_k = collision_jit_fn(mx_top_k, dx) - self.assertEqual(dx_all.ncon, 3) - self.assertEqual(dx_top_k.ncon, 2) + self.assertEqual(dx_all.contact.dist.shape, (3,)) + self.assertEqual(dx_top_k.contact.dist.shape, (2,)) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 313b43ff..50efcad3 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -14,11 +14,12 @@ # ============================================================================== """Core non-smooth constraint functions.""" -from typing import Optional, Tuple +from typing import Optional, Tuple, Union import jax from jax import numpy as jp import mujoco +from mujoco.mjx._src import collision_driver from mujoco.mjx._src import math from mujoco.mjx._src import support # pylint: disable=g-importing-member @@ -276,7 +277,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for contacts.""" - if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0: + if collision_driver.ncon(m) == 0: return None @jax.vmap @@ -313,7 +314,9 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss) -def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: +def count_constraints( + m: Union[Model, mujoco.MjModel] +) -> Tuple[int, int, int, int]: """Returns equality, friction, limit, and contact constraint counts.""" if m.opt.disableflags & DisableBit.CONSTRAINT: return 0, 0, 0, 0 @@ -333,10 +336,7 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: else: nl = int(m.jnt_limited.sum()) - if m.opt.disableflags & DisableBit.CONTACT: - nc = 0 - else: - nc = d.ncon * 4 + nc = collision_driver.ncon(m) * 4 return ne, nf, nl, nc @@ -344,10 +344,6 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: def make_constraint(m: Model, d: Data) -> Data: """Creates constraint jacobians and other supporting data.""" - ns = sum(count_constraints(m, d)[:-1]) - # TODO(robotics-simulation): make device_put set nefc/efc_address instead - d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)}) - if m.opt.disableflags & DisableBit.CONSTRAINT: efcs = () else: @@ -364,7 +360,7 @@ def make_constraint(m: Model, d: Data) -> Data: if not efcs: z = jp.empty(0) d = d.replace(efc_J=jp.empty((0, m.nv))) - d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, nefc=0) + d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z) return d efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs) @@ -378,6 +374,6 @@ def make_constraint(m: Model, d: Data) -> Data: aref, r = fn(efc) d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref) - d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0]) + d = d.replace(efc_frictionloss=efc.frictionloss) return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 8ce1bfa6..a2bd8cc9 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -15,164 +15,92 @@ """Tests for constraint functions.""" from absl.testing import absltest -from absl.testing import parameterized -import jax from jax import numpy as jp import mujoco from mujoco import mjx from mujoco.mjx._src import constraint from mujoco.mjx._src import test_util -# pylint: disable=g-importing-member -from mujoco.mjx._src.types import DisableBit -from mujoco.mjx._src.types import SolverType -# pylint: enable=g-importing-member import numpy as np -def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3): - err_msg = f'mismatch: {name} at step {step} in {fname}' - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +# tolerance for difference between MuJoCo and MJX constraint calculations, +# mostly due to float precision +_TOLERANCE = 5e-5 -class ConstraintTest(parameterized.TestCase): +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) - @parameterized.parameters(enumerate(test_util.TEST_FILES)) - def test_constraints(self, seed, fname): + +def _assert_attr_eq(a, b, attr): + _assert_eq(getattr(a, attr), getattr(b, attr), attr) + + +class ConstraintTest(absltest.TestCase): + + def test_constraints(self): """Test constraints.""" - np.random.seed(seed) - - # exclude convex.xml since convex contacts are not exactly equivalent - if fname == 'convex.xml': - return - - m = test_util.load_test_file(fname) + m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) - mx = mjx.device_put(m) - dx = mjx.make_data(mx) + mujoco.mj_step(m, d, 100) # at 100 steps mix of active/inactive constraints + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) - forward_jit_fn = jax.jit(mjx.forward) - - # give the system a little kick to ensure we have non-identity rotations - d.qvel = np.random.random(m.nv) - for i in range(100): - dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel)) - mujoco.mj_step(m, d) - dx = forward_jit_fn(mx, dx) - - nnz_filter = dx.efc_J.any(axis=1) - - mj_efc_j = d.efc_J.reshape((-1, m.nv)) - mjx_efc_j = dx.efc_J[nnz_filter] - _assert_eq(mj_efc_j, mjx_efc_j, 'efc_J', i, fname) - - mjx_efc_d = dx.efc_D[nnz_filter] - _assert_eq(d.efc_D, mjx_efc_d, 'efc_D', i, fname) - - mjx_efc_aref = dx.efc_aref[nnz_filter] - _assert_eq(d.efc_aref, mjx_efc_aref, 'efc_aref', i, fname) - - mjx_efc_frictionloss = dx.efc_frictionloss[nnz_filter] - _assert_eq( - d.efc_frictionloss, - mjx_efc_frictionloss, - 'efc_frictionloss', - i, - fname, - ) - - _JNT_RANGE = """ - - - - - - - - - - - - - """ - - def test_jnt_range(self): - """Tests that mixed joint ranges are respected.""" - # TODO(robotics-simulation): also test ball - m = mujoco.MjModel.from_xml_string(self._JNT_RANGE) - m.opt.solver = SolverType.CG.value - d = mujoco.MjData(m) - d.qpos = np.array([2.0, 15.0]) - - mx = mjx.device_put(m) - dx = mjx.device_put(d) - efc = jax.jit(constraint._instantiate_limit_slide_hinge)(mx, dx) - - # first joint is outside the joint range - np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0) - - # second joint has no range, so only one efc row - self.assertEqual(efc.J.shape[0], 1) + dx = mjx.make_constraint(mx, dx) + nnz = dx.efc_J.any(axis=1) + _assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J') + _assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D') + _assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref') + _assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss') def test_disable_refsafe(self): - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('constraints.xml') timeconst = m.opt.timestep / 4.0 # timeconst < 2 * timestep solimp = jp.array([timeconst, 1.0]) solref = jp.array([0.8, 0.99, 0.001, 0.2, 2]) pos = jp.ones(3) - m.opt.disableflags = m.opt.disableflags | DisableBit.REFSAFE + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.REFSAFE mx = mjx.device_put(m) k, *_ = constraint._kbi(mx, solimp, solref, pos) self.assertEqual(k, 1 / (0.99**2 * timeconst**2)) - m.opt.disableflags = m.opt.disableflags & ~DisableBit.REFSAFE - mx = mjx.device_put(m) - k, *_ = constraint._kbi(mx, solimp, solref, pos) - self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2)) - - def test_disableconstraint(self): - m = test_util.load_test_file('ant.xml') - d = mujoco.MjData(m) - - m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT - mx, dx = mjx.device_put(m), mjx.device_put(d) - dx = constraint.make_constraint(mx, dx) - self.assertGreater(dx.efc_J.shape[0], 1) - - m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT - mx = mjx.device_put(m) - dx = constraint.make_constraint(mx, dx) + def test_disable_constraint(self): + m = test_util.load_test_file('constraints.xml') + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT + ne, nf, nl, nc = mjx.count_constraints(m) + self.assertEqual(ne, 0) + self.assertEqual(nf, 0) + self.assertEqual(nl, 0) + self.assertEqual(nc, 0) + dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) self.assertEqual(dx.efc_J.shape[0], 0) def test_disable_equality(self): - m = test_util.load_test_file('equality.xml') - d = mujoco.MjData(m) - - m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY - mx, dx = mjx.device_put(m), mjx.device_put(d) - dx = constraint.make_constraint(mx, dx) - self.assertEqual(dx.efc_J.shape[0], 0) + m = test_util.load_test_file('constraints.xml') + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY + ne, nf, nl, nc = mjx.count_constraints(m) + self.assertEqual(ne, 0) + self.assertEqual(nf, 0) + self.assertEqual(nl, 2) + self.assertEqual(nc, 64) + dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) + self.assertEqual(dx.efc_J.shape[0], 66) # only joint range, contact def test_disable_contact(self): - m = test_util.load_test_file('ant.xml') - d = mujoco.MjData(m) - d.qpos[2] = 0.0 - mujoco.mj_forward(m, d) - - m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT - mx, dx = mjx.device_put(m), mjx.device_put(d) - dx = dx.tree_replace( - {'contact.frame': dx.contact.frame.reshape((-1, 3, 3))} - ) - efc = constraint._instantiate_contact(mx, dx) - self.assertIsNotNone(efc) - - m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT - mx, dx = mjx.device_put(m), mjx.device_put(d) - efc = constraint._instantiate_contact(mx, dx) - self.assertIsNone(efc) + m = test_util.load_test_file('constraints.xml') + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT + ne, nf, nl, nc = mjx.count_constraints(m) + self.assertEqual(ne, 10) + self.assertEqual(nf, 0) + self.assertEqual(nl, 2) + self.assertEqual(nc, 0) + dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) + self.assertEqual(dx.efc_J.shape[0], 12) # only joint range, limit if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index 0936bec6..5d95eb69 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -18,7 +18,7 @@ import copy import dataclasses import typing -from typing import Dict, Optional, Sequence, TypeVar +from typing import Any, Dict, Optional, Sequence, TypeVar import jax import numpy as np @@ -62,7 +62,7 @@ def dataclass(clz: _T) -> _T: def to_meta(field, obj): val = getattr(obj, field.name) - return to_tup(val) if isinstance(val, np.ndarray) else val + return (to_tup(val), val.dtype) if isinstance(val, np.ndarray) else val def to_data(field, obj): return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name)) @@ -75,7 +75,7 @@ def dataclass(clz: _T) -> _T: def from_meta(field, meta): if field.type is np.ndarray: - return (field.name, np.array(meta)) + return (field.name, np.array(meta[0], dtype=meta[1])) else: return (field.name, meta) @@ -113,6 +113,10 @@ class PyTreeNode: # stub for pytype raise NotImplementedError + @classmethod + def fields(cls) -> tuple[dataclasses.Field[Any], ...]: + return dataclasses.fields(cls) + def tree_replace( self, params: Dict[str, Optional[jax.typing.ArrayLike]] ) -> 'PyTreeNode': diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index c39860ef..2819d9ea 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -25,6 +25,7 @@ import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import mesh from mujoco.mjx._src import types +import numpy as np _MJ_TYPE_ATTR = { mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,), @@ -67,7 +68,7 @@ _TRANSFORMS = { (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Model, 'actuator_trnid'): lambda x: x[:, 0], + (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), (types.Contact, 'frame'): ( lambda x: x.reshape(x.shape[:-1] + (3, 3)) # pylint: disable=g-long-lambda if x is not None and x.shape[0] else jp.zeros((0, 3, 3)) @@ -78,6 +79,7 @@ _INVERSE_TRANSFORMS = { (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)), (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), (types.Contact, 'frame'): ( lambda x: x.reshape(x.shape[:-2] + (9,)) # pylint: disable=g-long-lambda if x is not None and x.shape[0] else jp.zeros((0, 9)) @@ -182,6 +184,11 @@ def device_put(value): Returns: on-device MJX struct reflecting the input value """ + warnings.warn( + 'device_put is deprecated, use put_model and put_data instead', + category=DeprecationWarning, + ) + clz = _TYPE_MAP.get(type(value)) if clz is None: raise NotImplementedError(f'{type(value)} is not supported for device_put.') @@ -240,6 +247,10 @@ def device_get_into(result, value): Raises: RuntimeError: if result length doesn't match data batch size """ + warnings.warn( + 'device_get_into is deprecated, use get_data instead', + category=DeprecationWarning, + ) value = jax.device_get(value) @@ -263,9 +274,16 @@ def device_get_into(result, value): else: if isinstance(result, mujoco.MjData): + ncon = value.contact.dist.shape[0] + nefc = value.efc_J.shape[0] mujoco._functions._realloc_con_efc( # pylint: disable=protected-access - result, ncon=value.ncon, nefc=value.nefc + result, ncon=ncon, nefc=nefc ) + result.ncon = ncon + result.nefc = nefc + efc_start = nefc - ncon * 4 + result.contact.efc_address[:] = np.arange(efc_start, nefc, 4) + result.contact.dim[:] = 3 for f in dataclasses.fields(value): # type: ignore if (type(value), f.name) in _DERIVED: diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py index e6eb7562..b4b9a0cd 100644 --- a/mjx/mujoco/mjx/_src/device_test.py +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -129,32 +129,26 @@ class ValidateInputTest(absltest.TestCase): with self.assertRaises(NotImplementedError): mjx.device_put(m) - def test_trn(self): - m = test_util.load_test_file('ant.xml') - m.actuator_trntype[0] = mujoco.mjtTrn.mjTRN_SITE - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - def test_dyn(self): - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('pendula.xml') m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE with self.assertRaises(NotImplementedError): mjx.device_put(m) def test_gain(self): - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('pendula.xml') m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE with self.assertRaises(NotImplementedError): mjx.device_put(m) def test_bias(self): - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('pendula.xml') m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE with self.assertRaises(NotImplementedError): mjx.device_put(m) def test_condim(self): - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('constraints.xml') for i in [1, 4, 6]: m.geom_condim[0] = i with self.assertRaises(NotImplementedError): diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 06bec894..1ff268e7 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -60,7 +60,7 @@ def named_scope(fn, name: str = ''): @named_scope -def _position(m: Model, d: Data) -> Data: +def fwd_position(m: Model, d: Data) -> Data: """Position-dependent computations.""" # TODO(robotics-simulation): tendon d = smooth.kinematics(m, d) @@ -74,7 +74,7 @@ def _position(m: Model, d: Data) -> Data: @named_scope -def _velocity(m: Model, d: Data) -> Data: +def fwd_velocity(m: Model, d: Data) -> Data: """Velocity-dependent computations.""" d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel) d = smooth.com_vel(m, d) @@ -84,7 +84,7 @@ def _velocity(m: Model, d: Data) -> Data: @named_scope -def _actuation(m: Model, d: Data) -> Data: +def fwd_actuation(m: Model, d: Data) -> Data: """Actuation-dependent computations.""" if not m.nu or m.opt.disableflags & DisableBit.ACTUATION: return d.replace( @@ -107,7 +107,7 @@ def _actuation(m: Model, d: Data) -> Data: act_dot = jp.array(0.0) elif dyn_typ == DynType.INTEGRATOR: act_dot = ctrl - elif dyn_typ == DynType.FILTER: + elif dyn_typ in (DynType.FILTER, DynType.FILTEREXACT): act_dot = (ctrl - act) / jp.clip(dyn_prm[0], mujoco.mjMINVAL) else: raise NotImplementedError(f'dyntype {dyn_typ.name} not implemented.') @@ -190,7 +190,7 @@ def _actuation(m: Model, d: Data) -> Data: @named_scope -def _acceleration(m: Model, d: Data) -> Data: +def fwd_acceleration(m: Model, d: Data) -> Data: """Add up all non-constraint forces, compute qacc_smooth.""" qfrc_applied = d.qfrc_applied + support.xfrc_accumulate(m, d) qfrc_smooth = d.qfrc_passive - d.qfrc_bias + d.qfrc_actuator + qfrc_applied @@ -228,6 +228,34 @@ def _integrate_pos( return jp.concatenate(qs) if qs else jp.empty((0,)) +def _next_activation(m: Model, d: Data, act_dot: jax.Array) -> jax.Array: + """Returns the next act given the current act_dot, after clamping.""" + act = d.act + + if not m.na: + return act + + actrange = jp.where( + m.actuator_actlimited[:, None], + m.actuator_actrange, + jp.array([-jp.inf, jp.inf]), + ) + + def fn(dyntype, dynprm, act, act_dot, actrange): + if dyntype == DynType.FILTEREXACT: + tau = jp.clip(dynprm[0], a_min=mujoco.mjMINVAL) + act = act + act_dot * tau * (1 - jp.exp(-m.opt.timestep / tau)) + else: + act = act + act_dot * m.opt.timestep + act = jp.clip(act, actrange[0], actrange[1]) + return act + + args = (m.actuator_dyntype, m.actuator_dynprm, act, act_dot, actrange) + act = scan.flat(m, fn, 'uuaau', 'a', *args, group_by='u') + + return act.reshape(m.na) + + @named_scope def _advance( m: Model, @@ -237,16 +265,7 @@ def _advance( qvel: Optional[jax.Array] = None, ) -> Data: """Advance state and time given activation derivatives and acceleration.""" - act = d.act - if m.na: - act = d.act + act_dot * m.opt.timestep - actrange = jp.where( - m.actuator_actlimited[:, None], - m.actuator_actrange, - jp.array([-jp.inf, jp.inf]), - ) - fn = lambda act, actrange: jp.clip(act, actrange[0], actrange[1]) - act = scan.flat(m, fn, 'au', 'a', act, actrange, group_by='u') + act = _next_activation(m, d, act_dot) # advance velocities d = d.replace(qvel=d.qvel + qacc * m.opt.timestep) @@ -263,7 +282,7 @@ def _advance( @named_scope -def _euler(m: Model, d: Data) -> Data: +def euler(m: Model, d: Data) -> Data: """Euler integrator, semi-implicit in velocity.""" # integrate damping implicitly qacc = d.qacc @@ -277,7 +296,7 @@ def _euler(m: Model, d: Data) -> Data: @named_scope -def _rungekutta4(m: Model, d: Data) -> Data: +def rungekutta4(m: Model, d: Data) -> Data: """Runge-Kutta explicit order 4 integrator.""" d_t0 = d # pylint: disable=invalid-name @@ -323,10 +342,10 @@ def _rungekutta4(m: Model, d: Data) -> Data: @named_scope def forward(m: Model, d: Data) -> Data: """Forward dynamics.""" - d = _position(m, d) - d = _velocity(m, d) - d = _actuation(m, d) - d = _acceleration(m, d) + d = fwd_position(m, d) + d = fwd_velocity(m, d) + d = fwd_actuation(m, d) + d = fwd_acceleration(m, d) if d.efc_J.size == 0: d = d.replace(qacc=d.qacc_smooth) @@ -343,9 +362,9 @@ def step(m: Model, d: Data) -> Data: d = forward(m, d) if m.opt.integrator == IntegratorType.EULER: - d = _euler(m, d) + d = euler(m, d) elif m.opt.integrator == IntegratorType.RK4: - d = _rungekutta4(m, d) + d = rungekutta4(m, d) else: raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index 40da667c..cc767ae8 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -15,77 +15,75 @@ """Tests for forward functions.""" from absl.testing import absltest -from absl.testing import parameterized import jax -from jax import numpy as jp import mujoco from mujoco import mjx -from mujoco.mjx._src import forward from mujoco.mjx._src import test_util -# pylint: disable=g-importing-member -from mujoco.mjx._src.types import DisableBit -# pylint: enable=g-importing-member import numpy as np -def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3): - err_msg = f'mismatch: {attr} at step {step} in {fname}' - a, b = getattr(a, attr), getattr(b, attr) - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +# tolerance for difference between MuJoCo and MJX forward calculations - mostly +# due to float precision +_TOLERANCE = 1e-5 -class ForwardTest(parameterized.TestCase): +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) - @parameterized.parameters( - filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES) - ) - def test_forward(self, fname): - """Test mujoco mj forward function matches mujoco_mjx forward function.""" - np.random.seed(test_util.TEST_FILES.index(fname)) - m = test_util.load_test_file(fname) +def _assert_attr_eq(a, b, attr): + _assert_eq(getattr(a, attr), getattr(b, attr), attr) + + +class ForwardTest(absltest.TestCase): + + def test_forward(self): + m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) - mx = mjx.device_put(m) - dx = mjx.make_data(mx) - forward_jit_fn = jax.jit(mjx.forward) + # apply some control and xfrc input + d.ctrl = np.array([-18, 0.59, 0.47]) + d.xfrc_applied[0, 2] = 0.1 # torque + d.xfrc_applied[1, 4] = 0.3 # linear force + mujoco.mj_step(m, d, 100) # get some dynamics going + mujoco.mj_forward(m, d) - # give the system a little kick to ensure we have non-identity rotations - d.qvel = np.random.random(m.nv) * 0.05 - for i in range(100): - qpos, qvel = d.qpos.copy(), d.qvel.copy() - mujoco.mj_step(m, d) - dx = forward_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + mx = mjx.put_model(m) - _assert_attr_eq(d, dx, 'qfrc_smooth', i, fname) - _assert_attr_eq(d, dx, 'qacc_smooth', i, fname) + # fwd_actuation + dx = jax.jit(mjx.fwd_actuation)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'act_dot') + _assert_attr_eq(d, dx, 'qfrc_actuator') - @parameterized.parameters( - filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES) - ) - def test_step(self, fname): - """Test mujoco mj step matches mujoco_mjx step.""" - np.random.seed(test_util.TEST_FILES.index(fname)) - m = test_util.load_test_file(fname) - step_jit_fn = jax.jit(forward.step) + # fwd_accleration (fwd_position and fwd_velocity already tested elsewhere) + dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_smooth') + _assert_attr_eq(d, dx, 'qacc_smooth') - mx = mjx.device_put(m) + # euler + dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d)) + mujoco.mj_Euler(m, d) + _assert_attr_eq(d, dx, 'act') + _assert_attr_eq(d, dx, 'qpos') + _assert_attr_eq(d, dx, 'time') + + def test_step(self): + m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) - # give the system a little kick to ensure we have non-identity rotations - d.qvel = np.random.normal(m.nv) * 0.05 - for i in range(100): - # in order to avoid re-jitting, reuse the same mj_data shape - qpos, qvel = d.qpos, d.qvel - d = mujoco.MjData(m) - d.qpos, d.qvel = qpos, qvel - dx = mjx.device_put(d) + # apply some control and xfrc input + d.ctrl = np.array([-18, 0.59, 0.47]) + d.xfrc_applied[0, 2] = 0.1 # torque + d.xfrc_applied[1, 4] = 0.3 # linear force + mujoco.mj_step(m, d, 100) # get some dynamics going - mujoco.mj_step(m, d) - dx = step_jit_fn(mx, dx) - - _assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2) - _assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2) - _assert_attr_eq(d, dx, 'act', i, fname) - _assert_attr_eq(d, dx, 'time', i, fname) + mx = mjx.put_model(m) + dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d)) + mujoco.mj_step(m, d) + _assert_attr_eq(d, dx, 'act') + _assert_attr_eq(d, dx, 'time') + _assert_attr_eq(d, dx, 'qvel') + _assert_attr_eq(d, dx, 'qpos') def test_rk4(self): m = mujoco.MjModel.from_xml_string(""" @@ -94,7 +92,6 @@ class ForwardTest(parameterized.TestCase): - @@ -107,41 +104,75 @@ class ForwardTest(parameterized.TestCase): """) - step_jit_fn = jax.jit(forward.step) - mx = mjx.device_put(m) d = mujoco.MjData(m) # give the system a little kick to ensure we have non-identity rotations - d.qvel = np.random.normal(m.nv) * 0.05 - for i in range(100): - # in order to avoid re-jitting, reuse the same mj_data shape - qpos, qvel = d.qpos, d.qvel - d = mujoco.MjData(m) - d.qpos, d.qvel = qpos, qvel - dx = mjx.device_put(d) + d.qvel = np.array([0.2, -0.1]) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) - mujoco.mj_step(m, d) - dx = step_jit_fn(mx, dx) + mx = mjx.put_model(m) + dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d)) + mujoco.mj_RungeKutta(m, d, 4) - _assert_attr_eq(d, dx, 'qvel', i, 'test_rk4', atol=1e-2) - _assert_attr_eq(d, dx, 'qpos', i, 'test_rk4', atol=1e-2) - _assert_attr_eq(d, dx, 'act', i, 'test_rk4') - _assert_attr_eq(d, dx, 'time', i, 'test_rk4') + _assert_attr_eq(d, dx, 'qvel') + _assert_attr_eq(d, dx, 'qpos') + _assert_attr_eq(d, dx, 'act') + _assert_attr_eq(d, dx, 'time') def test_disable_eulerdamp(self): - m = test_util.load_test_file('ant.xml') - m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP + m = test_util.load_test_file('pendula.xml') + self.assertTrue((m.dof_damping > 0).any()) + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EULERDAMP d = mujoco.MjData(m) - mx = mjx.device_put(m) - self.assertTrue((mx.dof_damping > 0).any()) - dx = mjx.device_put(d) - dx = jax.jit(forward.forward)(mx, dx) + d.qvel[:] = 1.0 + d.qacc[:] = 1.0 + mx = mjx.put_model(m) + dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d)) - dx = dx.replace(qvel=jp.ones_like(dx.qvel), qacc=jp.ones_like(dx.qacc)) - dx = jax.jit(forward._euler)(mx, dx) np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep) +class ActuatorTest(absltest.TestCase): + _DYN_XML = """ + + + + + + + + + + + + + + + + + + + """ + + def test_dyntype(self): + m = mujoco.MjModel.from_xml_string(self._DYN_XML) + d = mujoco.MjData(m) + d.ctrl = np.array([1.5, 1.5, 1.5, 1.5]) + d.act = np.array([0.5, 0.5, 0.5]) + + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + + mujoco.mj_fwdActuation(m, d) + dx = jax.jit(mjx.fwd_actuation)(mx, dx) + _assert_attr_eq(d, dx, 'act_dot') + + mujoco.mj_Euler(m, d) + dx = jax.jit(mjx.euler)(mx, dx) + _assert_attr_eq(d, dx, 'act') + + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 67e50a7e..5b406185 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -14,92 +14,339 @@ # ============================================================================== """Functions to initialize, load, or save data.""" +import copy +from typing import List, Union + +import jax from jax import numpy as jp +import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint -# pylint: disable=g-importing-member -from mujoco.mjx._src.types import Contact -from mujoco.mjx._src.types import Data -from mujoco.mjx._src.types import Model -# pylint: enable=g-importing-member +from mujoco.mjx._src import mesh +from mujoco.mjx._src import types import numpy as np -def make_data(m: Model) -> Data: - """Allocate and initialize Data.""" +def _put_option(o: mujoco.MjOption, device=None) -> types.Option: + """Puts mujoco.MjOption onto a device, resulting in mjx.Option.""" + if o.integrator not in set(types.IntegratorType): + raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') - # create first d to get num contacts and nc - d = Data( - solver_niter=jp.array(0, dtype=jp.int32), - ne=0, - nf=0, - nl=0, - nefc=0, - ncon=0, - time=jp.zeros((), dtype=jp.float32), - qpos=m.qpos0, - qvel=jp.zeros(m.nv, dtype=jp.float32), - act=jp.zeros(m.na, dtype=jp.float32), - qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32), - ctrl=jp.zeros(m.nu, dtype=jp.float32), - qfrc_applied=jp.zeros(m.nv, dtype=jp.float32), - xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32), - eq_active=jp.zeros(m.neq, dtype=jp.int32), - qacc=jp.zeros(m.nv, dtype=jp.float32), - act_dot=jp.zeros(m.na, dtype=jp.float32), - xpos=jp.zeros((m.nbody, 3), dtype=jp.float32), - xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), - xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), - xipos=jp.zeros((m.nbody, 3), dtype=jp.float32), - ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), - xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32), - xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32), - geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32), - geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32), - subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32), - cdof=jp.zeros((m.nv, 6), dtype=jp.float32), - cinert=jp.zeros((m.nbody, 10), dtype=jp.float32), - actuator_length=jp.zeros(m.nu, dtype=jp.float32), - actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), - crb=jp.zeros((m.nbody, 10), dtype=jp.float32), - qM=jp.zeros(m.nM, dtype=jp.float32), - qLD=jp.zeros(m.nM, dtype=jp.float32), - qLDiagInv=jp.zeros(m.nv, dtype=jp.float32), - qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32), - contact=Contact.zero(), - efc_J=jp.zeros((), dtype=jp.float32), - efc_frictionloss=jp.zeros((), dtype=jp.float32), - efc_D=jp.zeros((), dtype=jp.float32), - actuator_velocity=jp.zeros(m.nu, dtype=jp.float32), - cvel=jp.zeros((m.nbody, 6), dtype=jp.float32), - cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32), - qfrc_bias=jp.zeros(m.nv, dtype=jp.float32), - qfrc_passive=jp.zeros(m.nv, dtype=jp.float32), - efc_aref=jp.zeros((), dtype=jp.float32), - actuator_force=jp.zeros(m.nu, dtype=jp.float32), - qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32), - qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32), - qacc_smooth=jp.zeros(m.nv, dtype=jp.float32), - qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32), - qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32), - efc_force=jp.zeros((), dtype=jp.float32), + if o.cone not in set(types.ConeType): + raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') + + if o.solver not in set(types.SolverType): + raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') + + for i in range(mujoco.mjtEnableBit.mjNENABLE): + if o.enableflags & 2**i: + raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') + + static_fields = { + f.name: copy.copy(getattr(o, f.name)) + for f in types.Option.fields() + if f.type in (int, bytes, np.ndarray) + } + static_fields['integrator'] = types.IntegratorType(o.integrator) + static_fields['cone'] = types.ConeType(o.cone) + static_fields['solver'] = types.SolverType(o.solver) + static_fields['disableflags'] = types.DisableBit(o.disableflags) + + device_fields = { + f.name: copy.copy(getattr(o, f.name)) + for f in types.Option.fields() + if f.type is jax.Array + } + device_fields = jax.device_put(device_fields, device=device) + + has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() + + return types.Option( + has_fluid_params=has_fluid_params, + **static_fields, + **device_fields, ) - # get contact data with correct shapes - ncon = collision_driver.ncon(m) - d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon) - d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)}) - ne, nf, nl, nc = constraint.count_constraints(m, d) - d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc) - ns = ne + nf + nl - d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)}) - d = d.replace( - efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32), - efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32), - efc_D=jp.zeros(d.nefc, dtype=jp.float32), - efc_aref=jp.zeros(d.nefc, dtype=jp.float32), - efc_force=jp.zeros(d.nefc, dtype=jp.float32), +def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic: + """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" + return types.Statistic( + meaninertia=jax.device_put(s.meaninertia, device=device) + ) + + +def put_model(m: mujoco.MjModel, device=None) -> types.Model: + """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" + + if m.ntendon: + raise NotImplementedError('tendons are not supported') + + if (m.geom_condim != 3).any() or (m.pair_dim != 3).any(): + raise NotImplementedError('only condim=3 is supported') + + # check collision geom types + for g1, g2, *_ in collision_driver.collision_candidates(m): + if collision_driver.get_collision_fn((g1, g2)) is None: + g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) + raise NotImplementedError(f'({g1}, {g2}) has no collision function') + + for enum_field, enum_type, mj_type in ( + (m.actuator_biastype, types.BiasType, mujoco.mjtBias), + (m.actuator_dyntype, types.DynType, mujoco.mjtDyn), + (m.actuator_gaintype, types.GainType, mujoco.mjtGain), + (m.actuator_trntype, types.TrnType, mujoco.mjtTrn), + (m.eq_type, types.EqType, mujoco.mjtEq), + ): + missing = set(enum_field) - set(enum_type) + if missing: + raise NotImplementedError( + f'{[mj_type(m) for m in missing]} not supported' + ) + + opt = _put_option(m.opt, device=device) + stat = _put_statistic(m.stat, device=device) + + static_fields = { + f.name: getattr(m, f.name) + for f in types.Model.fields() + if f.type in (int, bytes, np.ndarray) + } + static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4)) + static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4)) + + device_fields = { + f.name: copy.copy(getattr(m, f.name)) # copy because device_put is async + for f in types.Model.fields() + if f.type is jax.Array + } + device_fields.update(mesh.get(m)) + device_fields = jax.device_put(device_fields, device=device) + + return types.Model( + opt=opt, + stat=stat, + **static_fields, + **device_fields, + ) + + +def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: + """Allocate and initialize Data.""" + + ncon = collision_driver.ncon(m) + ne, nf, nl, nc = constraint.count_constraints(m) + nefc = ne + nf + nl + nc + + zero_nv = jp.zeros(m.nv, dtype=jp.float32) + zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32) + zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32) + zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32) + zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32) + zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32) + zero_nefc = jp.zeros(nefc, dtype=jp.float32) + zero_na = jp.zeros(m.na, dtype=jp.float32) + zero_nu = jp.zeros(m.nu, dtype=jp.float32) + zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32) + zero_nm = jp.zeros(m.nM, dtype=jp.float32) + + # create first d to get num contacts and nc + d = types.Data( + solver_niter=jp.array(0, dtype=jp.int32), + time=jp.array(0.0), + qpos=jp.array(m.qpos0), + qvel=zero_nv, + act=zero_na, + qacc_warmstart=zero_nv, + ctrl=zero_nu, + qfrc_applied=zero_nv, + xfrc_applied=zero_nbody_6, + eq_active=jp.zeros(m.neq, dtype=jp.int32), + qacc=zero_nv, + act_dot=zero_na, + xpos=zero_nbody_3, + xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), + xmat=zero_nbody_3_3, + xipos=zero_nbody_3, + ximat=zero_nbody_3_3, + xanchor=zero_njnt_3, + xaxis=zero_njnt_3, + geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32), + geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32), + site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32), + site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32), + subtree_com=zero_nbody_3, + cdof=zero_nv_6, + cinert=zero_nbody_10, + actuator_length=zero_nu, + actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), + crb=zero_nbody_10, + qM=zero_nm, + qLD=zero_nm, + qLDiagInv=zero_nv, + qLDiagSqrtInv=zero_nv, + contact=types.Contact.zero(ncon), + efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32), + efc_frictionloss=zero_nefc, + efc_D=zero_nefc, + actuator_velocity=zero_nu, + cvel=zero_nbody_6, + cdof_dot=zero_nv_6, + qfrc_bias=zero_nv, + qfrc_passive=zero_nv, + efc_aref=zero_nefc, + qfrc_actuator=zero_nv, + qfrc_smooth=zero_nv, + qacc_smooth=zero_nv, + qfrc_constraint=zero_nv, + qfrc_inverse=zero_nv, + efc_force=zero_nefc, ) return d + + +def _get_contact( + c: mujoco._structs._MjContactList, + cx: types.Contact, + efc_start: int, +): + """Converts mjx.Contact to mujoco._structs._MjContactList.""" + con_id = np.nonzero(cx.dist <= 0)[0] + for field in types.Contact.fields(): + value = getattr(cx, field.name)[con_id] + if field.name == 'frame': + value = value.reshape((-1, 9)) + getattr(c, field.name)[:] = value + + ncon = cx.dist.shape[0] + c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id] + + +def get_data( + m: mujoco.MjModel, d: types.Data +) -> Union[mujoco.MjData, List[mujoco.MjData]]: + """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" + dx = jax.device_get(d) + batched = len(d.qpos.shape) > 1 + batch_size = d.qpos.shape[0] if batched else 1 + ne, nf, nl, nc = constraint.count_constraints(m) + efc_type = np.array([ + mujoco.mjtConstraint.mjCNSTR_EQUALITY, + mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF, + mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT, + mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL, + ]).repeat([ne, nf, nl, nc]) + + ds = [] + for i in range(batch_size): + dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d + ncon = (dx_i.contact.dist <= 0).sum() + efc_active = (dx_i.efc_J != 0).any(axis=1) + efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL + nefc, nc = efc_active.sum(), (efc_active & efc_con).sum() + d_i = mujoco.MjData(m) + d_i.nnzJ = nefc * m.nv + mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access + d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc) + d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv) + d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) + + for field in types.Data.fields(): + if field.name == 'contact': + _get_contact(d_i.contact, dx_i.contact, nefc - nc) + continue + + value = getattr(dx_i, field.name) + + if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'): + value = value.reshape((-1, 9)) + + if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'): + value = value[efc_active] + + if field.name == 'efc_J': + value = value[efc_active].reshape(-1) + + if value.shape: + getattr(d_i, field.name)[:] = value + else: + setattr(d_i, field.name, value) + + d_i.efc_type[:] = efc_type[efc_active] + ds.append(d_i) + + return ds if batched else ds[0] + + +def _put_contact( + c: mujoco._structs._MjContactList, ncon: int, device=None +) -> types.Contact: + """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact.""" + fields = { + f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields() + } + fields['frame'] = fields['frame'].reshape((-1, 3, 3)) + pad_size = ncon - c.dist.shape[0] + pad_fn = lambda x: np.concatenate( + (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype)) + ) + fields = jax.tree_map(pad_fn, fields) + fields['dist'][-pad_size:] = np.inf + fields = jax.device_put(fields, device=device) + + return types.Contact(**fields) + + +def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" + ncon = collision_driver.ncon(m) + ne, nf, nl, nc = constraint.count_constraints(m) + nefc = ne + nf + nl + nc + + for d_val, val, name in ( + (d.ncon, ncon, 'ncon'), + (d.ne, ne, 'ne'), + (d.nf, nf, 'nf'), + (d.nl, nl, 'nl'), + (d.nefc, nefc, 'nefc'), + ): + if d_val > val: + raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') + + fields = { + f.name: copy.copy(getattr(d, f.name)) # copy because device_put is async + for f in types.Data.fields() + if f.type is jax.Array + } + + for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'): + fields[fname] = fields[fname].reshape((-1, 3, 3)) + + # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints. + # efc_J is also optionally column-sparse (typically for large nv). MJX is + # neither: it contains zeros for inactive constraints, and efc_J is always + # (nefc, nv). this may change in the future. + if mujoco.mj_isSparse(m): + nr = d.efc_J_rownnz.shape[0] + efc_j = np.zeros((nr, m.nv)) + for i in range(nr): + rowadr = d.efc_J_rowadr[i] + for j in range(d.efc_J_rownnz[i]): + efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j] + fields['efc_J'] = efc_j + else: + fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv)) + + for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'): + value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc) + for i in range(4): + value_beg = sum([ne, nf, nl][:i]) + d_beg = sum([d.ne, d.nf, d.nl][:i]) + size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i] + value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size] + fields[fname] = value + + fields = jax.device_put(fields, device=device) + fields['contact'] = _put_contact(d.contact, ncon, device=device) + + return types.Data(**fields) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 269509a8..8b25f359 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -17,25 +17,384 @@ from absl.testing import absltest from absl.testing import parameterized import jax +from jax import numpy as jp +import mujoco from mujoco import mjx -from mujoco.mjx._src import test_util +import numpy as np -class IoTest(parameterized.TestCase): +_MULTIPLE_CONVEX_OBJECTS = """ + + +""" - @parameterized.parameters(test_util.TEST_FILES) - def test_make_data(self, fname): - """Test that data created by make_data matches data returned by step.""" +_MULTIPLE_CONSTRAINTS = """ + + + + + + + + + + + + + + + + + + +""" - m = test_util.load_test_file(fname) - mx = mjx.device_put(m) - dx = mjx.make_data(mx) - dx_step = mjx.step(mx, dx) - _, dx_treedef = jax.tree_util.tree_flatten(dx) - _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step) +class ModelIOTest(parameterized.TestCase): + """IO tests for mjx.Model.""" - self.assertEqual(dx_treedef, dx_step_treedef) + def test_put_model(self): + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) + mx = mjx.put_model(m) + self.assertEqual(mx.nq, m.nq) + self.assertEqual(mx.nv, m.nv) + self.assertEqual(mx.nu, m.nu) + self.assertEqual(mx.na, m.na) + self.assertEqual(mx.nbody, m.nbody) + self.assertEqual(mx.njnt, m.njnt) + self.assertEqual(mx.ngeom, m.ngeom) + self.assertEqual(mx.nmesh, m.nmesh) + self.assertEqual(mx.npair, m.npair) + self.assertEqual(mx.nexclude, m.nexclude) + self.assertEqual(mx.neq, m.neq) + self.assertEqual(mx.nnumeric, m.nnumeric) + self.assertEqual(mx.nM, m.nM) + self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep) + + np.testing.assert_allclose(mx.body_parentid, m.body_parentid) + np.testing.assert_allclose(mx.geom_type, m.geom_type) + np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid) + np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref) + np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos) + self.assertLen(mx.geom_convex_face, 6) + self.assertLen(mx.geom_convex_vert, 6) + self.assertLen(mx.geom_convex_edge, 6) + self.assertLen(mx.geom_convex_facenormal, 6) + + np.testing.assert_allclose(mx.jnt_type, m.jnt_type) + np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr) + np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid) + np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited) + np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis) + + np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype) + np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype) + np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype) + np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype) + np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid) + + def test_fluid_params(self): + """Test that has_fluid_params is set when fluid params are present.""" + m = mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + self.assertTrue(m.opt.has_fluid_params) + + def test_implicit_not_implemented(self): + """Test that MJX guards against models with unimplemented features.""" + + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_cone_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_pgs_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_tendon_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + """)) + + def test_condim_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + """)) + + def test_cylinder_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + """)) + + +class DataIOTest(parameterized.TestCase): + """IO tests for mjx.Data.""" + + def test_make_data(self): + """Test that make_data returns the correct shapes.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) + d = mjx.make_data(m) + + nq = 22 + nbody = 5 + ncon = 46 + nv = 19 + nefc = 185 + nm = 64 + + self.assertEqual(d.qpos.shape, (nq,)) + self.assertEqual(d.qvel.shape, (nv,)) + self.assertEqual(d.act.shape, (0,)) + self.assertEqual(d.qacc_warmstart.shape, (nv,)) + self.assertEqual(d.ctrl.shape, (1,)) + self.assertEqual(d.qfrc_applied.shape, (nv,)) + self.assertEqual(d.xfrc_applied.shape, (nbody, 6)) + self.assertEqual(d.eq_active.shape, (0,)) + self.assertEqual(d.qacc.shape, (nv,)) + self.assertEqual(d.act_dot.shape, (0,)) + self.assertEqual(d.xpos.shape, (nbody, 3)) + self.assertEqual(d.xquat.shape, (nbody, 4)) + self.assertEqual(d.xmat.shape, (nbody, 3, 3)) + self.assertEqual(d.xipos.shape, (nbody, 3)) + self.assertEqual(d.ximat.shape, (nbody, 3, 3)) + self.assertEqual(d.xanchor.shape, (4, 3)) + self.assertEqual(d.xaxis.shape, (4, 3)) + self.assertEqual(d.geom_xpos.shape, (6, 3)) + self.assertEqual(d.geom_xmat.shape, (6, 3, 3)) + self.assertEqual(d.subtree_com.shape, (nbody, 3)) + self.assertEqual(d.cdof.shape, (nv, 6)) + self.assertEqual(d.cinert.shape, (nbody, 10)) + self.assertEqual(d.crb.shape, (nbody, 10)) + self.assertEqual(d.actuator_length.shape, (1,)) + self.assertEqual(d.actuator_moment.shape, (1, nv)) + self.assertEqual(d.qM.shape, (nm,)) + self.assertEqual(d.qLD.shape, (nm,)) + self.assertEqual(d.qLDiagInv.shape, (nv,)) + self.assertEqual(d.qLDiagSqrtInv.shape, (nv,)) + self.assertEqual(d.contact.dist.shape, (ncon,)) + self.assertEqual(d.contact.pos.shape, (ncon, 3)) + self.assertEqual(d.contact.frame.shape, (ncon, 3, 3)) + self.assertEqual(d.contact.solref.shape, (ncon, 2)) + self.assertEqual(d.contact.solimp.shape, (ncon, 5)) + self.assertEqual(d.contact.geom1.shape, (ncon,)) + self.assertEqual(d.contact.geom2.shape, (ncon,)) + self.assertEqual(d.efc_J.shape, (nefc, nv)) + self.assertEqual(d.efc_frictionloss.shape, (nefc,)) + self.assertEqual(d.efc_D.shape, (nefc,)) + self.assertEqual(d.actuator_velocity.shape, (1,)) + self.assertEqual(d.cvel.shape, (nbody, 6)) + self.assertEqual(d.cdof_dot.shape, (nv, 6)) + self.assertEqual(d.qfrc_bias.shape, (nv,)) + self.assertEqual(d.qfrc_passive.shape, (nv,)) + self.assertEqual(d.efc_aref.shape, (nefc,)) + self.assertEqual(d.qfrc_actuator.shape, (nv,)) + self.assertEqual(d.qfrc_smooth.shape, (nv,)) + self.assertEqual(d.qacc_smooth.shape, (nv,)) + self.assertEqual(d.qfrc_constraint.shape, (nv,)) + self.assertEqual(d.qfrc_inverse.shape, (nv,)) + self.assertEqual(d.efc_force.shape, (nefc,)) + + def test_put_data(self): + """Test that put_data puts the correct data for dense and sparse.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + + # check a few fields + np.testing.assert_allclose(dx.qpos, d.qpos) + np.testing.assert_allclose(dx.xpos, d.xpos) + np.testing.assert_allclose(dx.cvel, d.cvel) + np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot) + np.testing.assert_allclose(dx.qM, d.qM) + + # 4 contacts, 2 for each capsule against the plane + self.assertEqual(dx.contact.dist.shape, (4,)) + self.assertEqual(d.ncon, 1) # however only 1 contact in this step + np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0]) + self.assertTrue(np.isinf(dx.contact.dist[1:]).all()) + self.assertEqual(dx.contact.frame.shape, (4, 3, 3)) + np.testing.assert_allclose( + dx.contact.frame[0].reshape(9), d.contact.frame[0] + ) + np.testing.assert_allclose(dx.contact.frame[1:], 0) + + # xmat, ximat, geom_xmat are all shape transformed + self.assertEqual(dx.xmat.shape, (3, 3, 3)) + self.assertEqual(dx.ximat.shape, (3, 3, 3)) + self.assertEqual(dx.geom_xmat.shape, (3, 3, 3)) + self.assertEqual(dx.site_xmat.shape, (1, 3, 3)) + np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat) + np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat) + np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat) + np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat) + + # efc_ are also shape transformed and padded + self.assertEqual(dx.efc_J.shape, (21, 8)) # nefc, nv + d_efc_j = d.efc_J.reshape((-1, 8)) + np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq + np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit + np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit + np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8]) # contact + np.testing.assert_allclose(dx.efc_J[9:], 0) # no contact + + # check another efc_ too + self.assertEqual(dx.efc_aref.shape, (21,)) # nefc + np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3]) + np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3]) + np.testing.assert_allclose(dx.efc_aref[4], 0) + np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8]) + np.testing.assert_allclose(dx.efc_aref[9:], 0) + + # check sparse transform is correct + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx_from_sparse = mjx.put_data(m, d) + np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8) + + def test_get_data(self): + """Test that get_data makes correct MjData.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + d_2: mujoco.MjData = mjx.get_data(m, dx) + + # check a few fields + np.testing.assert_allclose(d_2.qpos, d.qpos) + np.testing.assert_allclose(d_2.xpos, d.xpos) + np.testing.assert_allclose(d_2.cvel, d.cvel) + np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot) + np.testing.assert_allclose(d_2.qM, d.qM) + + # only 1 contact active + self.assertEqual(d_2.contact.dist.shape, (1,)) + self.assertEqual(d_2.ncon, 1) + np.testing.assert_allclose(d_2.contact.dist, d.contact.dist) + self.assertEqual(d_2.contact.frame.shape, (1, 9)) + np.testing.assert_allclose(d_2.contact.frame, d.contact.frame) + + # xmat, ximat, geom_xmat, site_xmat are all shape transformed + self.assertEqual(d_2.xmat.shape, (3, 9)) + self.assertEqual(d_2.ximat.shape, (3, 9)) + self.assertEqual(d_2.geom_xmat.shape, (3, 9)) + self.assertEqual(d_2.site_xmat.shape, (1, 9)) + np.testing.assert_allclose(d_2.xmat, d.xmat) + np.testing.assert_allclose(d_2.ximat, d.ximat) + np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat) + np.testing.assert_allclose(d_2.site_xmat, d.site_xmat) + + # efc_* are also shape transformed and filtered + self.assertEqual(d_2.efc_J.shape, (64,)) # nefc * nv + np.testing.assert_allclose(d_2.efc_J, d.efc_J) + self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc + np.testing.assert_allclose(d_2.efc_aref, d.efc_aref) + + # efc_address is created on demand + np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address) + + def test_get_data_batched(self): + """Test that get_data makes correct List[MjData] for batched Data.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + # second data in batch has contact dist > 0, disables contact + dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) + ds = mjx.get_data(m, dx_b) + self.assertLen(ds, 2) + np.testing.assert_allclose(ds[0].qpos, d.qpos) + np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8) + self.assertEqual(ds[0].ncon, 1) + self.assertEqual(ds[1].ncon, 0) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py index 9b0b41f0..268de126 100644 --- a/mjx/mujoco/mjx/_src/passive.py +++ b/mjx/mujoco/mjx/_src/passive.py @@ -76,7 +76,7 @@ def _inertia_box_fluid_model( def passive(m: Model, d: Data) -> Data: """Adds all passive forces.""" if m.opt.disableflags & DisableBit.PASSIVE: - return d + return d.replace(qfrc_passive=jp.zeros(m.nv)) # joint-level springs def fn(jnt_typs, stiffness, qpos_spring, qpos): diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py index 4a5026c8..6ff5a596 100644 --- a/mjx/mujoco/mjx/_src/passive_test.py +++ b/mjx/mujoco/mjx/_src/passive_test.py @@ -14,100 +14,65 @@ # ============================================================================== """Tests passive forces.""" -import itertools - from absl.testing import absltest -from absl.testing import parameterized -from etils import epath import jax -import jax.numpy as jp import mujoco from mujoco import mjx +from mujoco.mjx._src import test_util import numpy as np - -def _assert_attr_eq(a, b, attr, step, fname, atol=1e-4, rtol=1e-4): - err_msg = f'mismatch: {attr} at step {step} in {fname}' - a, b = getattr(a, attr), getattr(b, attr) - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +# tolerance for difference between MuJoCo and MJX passive calculations - mostly +# due to float precision +_TOLERANCE = 1e-7 -class PassiveTest(parameterized.TestCase): +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) - @parameterized.parameters(enumerate(('ant.xml', 'pendula.xml'))) - def test_stiffness_damping(self, seed, fname): - """Tests stiffness and damping on Ant.""" - np.random.seed(seed) - path = epath.resource_path('mujoco.mjx') / 'test_data' - path /= fname - m = mujoco.MjModel.from_xml_string(path.read_text()) - # set stiffness/damping - m.jnt_stiffness = np.random.uniform(size=m.njnt) - m.dof_damping = np.random.uniform(size=m.nv) +def _assert_attr_eq(a, b, attr): + _assert_eq(getattr(a, attr), getattr(b, attr), attr) + + +class PassiveTest(absltest.TestCase): + + def test_passive(self): + m = test_util.load_test_file('pendula.xml') d = mujoco.MjData(m) - d.qvel = np.random.random(m.nv) # random kick + # give the system a little kick to ensure we have non-identity rotations + d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4]) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) - mx = mjx.device_put(m) - dx = mjx.make_data(mx) + dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_passive') - passive_jit_fn = jax.jit(mjx.passive) + # test with fluid forces + m.opt.density = 0.01 + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_passive') - for i in range(100): - qpos, qvel = d.qpos.copy(), d.qvel.copy() - mujoco.mj_step(m, d) - dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) - _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) + m.opt.viscosity = 0.02 + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_passive') - @parameterized.parameters( - itertools.product(range(3), ('pendula.xml',)) - ) - def test_fluid(self, seed, fname): - np.random.seed(seed) - path = epath.resource_path('mujoco.mjx') / 'test_data' - path /= fname - m = mujoco.MjModel.from_xml_string(path.read_text()) + m.opt.wind = np.array([0.03, 0.04, 0.05]) + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_passive') - # set density/viscosity/wind - m.opt.density = np.random.uniform() - m.opt.viscosity = np.random.uniform() - m.opt.wind = np.random.uniform() - - passive_jit_fn = jax.jit(mjx.passive) - - mx = mjx.device_put(m) - d = mujoco.MjData(m) - d.qvel = np.random.random(m.nv) # random kick - - for i in range(100): - mujoco.mj_step(m, d) - dx = mjx.device_put(d) - mujoco.mj_passive(m, d) - dx = passive_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) - - def test_disable_passive(self): - m = mujoco.MjModel.from_xml_string(""" - - - - - - - - - - """) - mx = mjx.device_put(m) - d = mujoco.MjData(m) - dx = mjx.device_put(d) - dx = dx.replace(qvel=jp.ones(mx.nv)) - - passive_jit_fn = jax.jit(mjx.passive) - dx = passive_jit_fn(mx, dx) - np.testing.assert_equal(dx.qfrc_passive, np.zeros(mx.nv)) + # test disable passive + mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE}) + dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d)) + np.testing.assert_allclose(dx.qfrc_passive, 0) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py new file mode 100644 index 00000000..d093e926 --- /dev/null +++ b/mjx/mujoco/mjx/_src/ray.py @@ -0,0 +1,208 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Functions for ray interesection testing.""" + +from typing import Sequence, Tuple + +import jax +from jax import numpy as jp +import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +def _ray_quad( + a: jax.Array, b: jax.Array, c: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0.""" + det = b * b - a * c + det_2 = jp.sqrt(det) + + x0, x1 = (-b - det_2) / a, (-b + det_2) / a + x0 = jp.where((det < mujoco.mjMINVAL) | (x0 < 0), jp.inf, x0) + x1 = jp.where((det < mujoco.mjMINVAL) | (x1 < 0), jp.inf, x1) + + return x0, x1 + + +def _ray_plane( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with a plane.""" + x = -pnt[2] / vec[2] + + valid = vec[2] <= -mujoco.mjMINVAL # z-vec pointing towards front face + valid &= x >= 0 + # only within rendered rectangle + p = pnt[0:2] + x * vec[0:2] + valid &= jp.all((size[0:2] <= 0) | (jp.abs(p) <= size[0:2])) + + return jp.where(valid, x, jp.inf) + + +def _ray_sphere( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with a sphere.""" + x0, x1 = _ray_quad(vec @ vec, vec @ pnt, pnt @ pnt - size[0] * size[0]) + x = jp.where(jp.isinf(x0), x1, x0) + + return x + + +def _ray_capsule( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with a capsule.""" + + # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0] + a = vec[0:2] @ vec[0:2] + b = vec[0:2] @ pnt[0:2] + c = pnt[0:2] @ pnt[0:2] - size[0] * size[0] + + # solve a*x^2 + 2*b*x + c = 0 + x0, x1 = _ray_quad(a, b, c) + x = jp.where(jp.isinf(x0), x1, x0) + + # make sure round solution is between flat sides + x = jp.where(jp.abs(pnt[2] + x * vec[2]) <= size[1], x, jp.inf) + + # top cap + dif = pnt - jp.array([0, 0, size[1]]) + x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0]) + # accept only top half of sphere + x = jp.where((pnt[2] + x0 * vec[2] >= size[1]) & (x0 < x), x0, x) + x = jp.where((pnt[2] + x1 * vec[2] >= size[1]) & (x1 < x), x1, x) + + # bottom cap + dif = pnt + jp.array([0, 0, size[1]]) + x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0]) + + # accept only bottom half of sphere + x = jp.where((pnt[2] + x0 * vec[2] <= -size[1]) & (x0 < x), x0, x) + x = jp.where((pnt[2] + x1 * vec[2] <= -size[1]) & (x1 < x), x1, x) + + return x + + +def _ray_box( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with a box.""" + + iface = jp.array([(1, 2), (0, 2), (0, 1), (1, 2), (0, 2), (0, 1)]) + + # side +1, -1 + # solution of pnt[i] + x * vec[i] = side * size[i] + x = jp.concatenate([(size - pnt) / vec, (-size - pnt) / vec]) + + # intersection with face + p0 = pnt[iface[:, 0]] + x * vec[iface[:, 0]] + p1 = pnt[iface[:, 1]] + x * vec[iface[:, 1]] + valid = jp.abs(p0) <= size[iface[:, 0]] + valid &= jp.abs(p1) <= size[iface[:, 1]] + + return jp.min(jp.where(valid, x, jp.inf)) + + +def _ray_mesh( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with a mesh.""" + del size, pnt, vec + raise NotImplementedError("ray <> mesh not implemented yet") + + +_RAY_FUNC = { + GeomType.PLANE: _ray_plane, + GeomType.SPHERE: _ray_sphere, + GeomType.CAPSULE: _ray_capsule, + GeomType.BOX: _ray_box, + # GeomType.MESH: _ray_mesh, +} + + +def ray( + m: Model, + d: Data, + pnt: jax.Array, + vec: jax.Array, + geomgroup: Sequence[int] = (), + flg_static: bool = True, + bodyexclude: int = -1, +) -> Tuple[jax.Array, jax.Array]: + """Returns the geom id and distance at which a ray intersects with a geom. + + Args: + m: MJX model + d: MJX data + pnt: ray origin point (3,) + vec: ray direction (3,) + geomgroup: group inclusion/exclusion mask, or empty to ignore + flg_static: if True, allows rays to intersect with static geoms + bodyexclude: ignore geoms on specified body id + + Returns: + dist: distance from ray origin to geom surface (or -1.0 for no intersection) + id: id of intersected geom (or -1 for no intersection) + """ + + dists, ids = [], [] + geom_filter = m.geom_bodyid != bodyexclude + geom_filter &= (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0) + geom_filter &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0) + geom_filter &= flg_static | (m.body_weldid[m.geom_bodyid] != 0) + if geomgroup: + geomgroup = np.array(geomgroup, dtype=bool) + geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)] + + # map ray to local geom frames + geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos) + geom_vecs = jax.vmap(lambda x: x.T @ vec)(d.geom_xmat) + + for geom_type, fn in _RAY_FUNC.items(): + id_, = np.nonzero(geom_filter & (m.geom_type == geom_type)) + + if id_.size == 0: + continue + + size, pnt, vec = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_] + dist = jax.vmap(fn)(size, pnt, vec) + dists, ids = dists + [dist], ids + [id_] + + if not ids: + return jp.array(-1), jp.array(-1.0) + + dists = jp.concatenate(dists) + ids = jp.concatenate(ids) + min_id = jp.argmin(dists) + dist = jp.where(jp.isinf(dists[min_id]), -1, dists[min_id]) + id_ = jp.where(jp.isinf(dists[min_id]), -1, ids[min_id]) + + return dist, id_ diff --git a/mjx/mujoco/mjx/_src/ray_test.py b/mjx/mujoco/mjx/_src/ray_test.py new file mode 100644 index 00000000..0556c9a3 --- /dev/null +++ b/mjx/mujoco/mjx/_src/ray_test.py @@ -0,0 +1,220 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for ray functions.""" + +from absl.testing import absltest +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import test_util +import numpy as np + +# tolerance for difference between MuJoCo and MJX ray calculations - mostly +# due to float precision +_TOLERANCE = 5e-5 + + +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) + + +class RayTest(absltest.TestCase): + + def test_ray_nothing(self): + """Tests that MJX ray returns -1 when nothing is hit.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + pnt, vec = jp.array([12.146, 1.865, 3.895]), jp.array([0, 0, -1.0]) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + def test_ray_plane(self): + """Tests MJX ray<>plane matches MuJoCo.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + # looking down at a slight angle + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 0, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + # looking on wrong side of plane + pnt = jp.array([0, 0, -0.5]) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + def test_ray_sphere(self): + """Tests MJX ray<>sphere matches MuJoCo.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + # looking down at sphere at a slight angle + pnt, vec = jp.array([0, 0, 1.6]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 1, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + def test_ray_capsule(self): + """Tests MJX ray<>capsule matches MuJoCo.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + # looking down at capsule at a slight angle + pnt, vec = jp.array([0.5, 1, 1.6]), jp.array([0, 0.05, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 2, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + # looking up at capsule from below + pnt, vec = jp.array([-0.5, 1, 0.05]), jp.array([0, 0.05, 1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 2, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + # looking at cylinder of capsule from the side + pnt, vec = jp.array([0, 1, 0.75]), jp.array([1, 0, 0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 2, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + def test_ray_box(self): + """Tests MJX ray<>box matches MuJoCo.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + # looking down at box at a slight angle + pnt, vec = jp.array([1, 0, 1.6]), jp.array([0, 0.05, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 3, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + # looking up at box from below + pnt, vec = jp.array([1, 0, 0.05]), jp.array([0, 0.05, 1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, 3, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + def test_ray_geomgroup(self): + """Tests ray geomgroup filter.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + ray_fn = jax.jit(mjx.ray, static_argnums=(4,)) + + # hits plane with geom_group[0] = 1 + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + geomgroup = (1, 0, 0, 0, 0, 0) + dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup) + _assert_eq(geomid, 0, 'geom_id') + pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32) + mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused) + _assert_eq(dist, mj_dist, 'dist') + + # nothing hit with geom_group[0] = 0 + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + geomgroup = (0, 0, 0, 0, 0, 0) + dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + def test_ray_flg_static(self): + """Tests ray flg_static filter.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + ray_fn = jax.jit(mjx.ray, static_argnames=('flg_static',)) + + # nothing hit with flg_static = False + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = ray_fn(mx, dx, pnt, vec, flg_static=False) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + def test_ray_bodyexclude(self): + """Tests ray bodyexclude filter.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + ray_fn = jax.jit(mjx.ray, static_argnames=('bodyexclude',)) + + # nothing hit with bodyexclude = 0 (world body) + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = ray_fn(mx, dx, pnt, vec, bodyexclude=0) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + def test_ray_invisible(self): + """Tests ray doesn't hit transparent geoms.""" + m = test_util.load_test_file('ray.xml') + # nothing hit with transparent geoms: + m.geom_rgba = 0 + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec) + _assert_eq(geomid, -1, 'geom_id') + _assert_eq(dist, -1, 'dist') + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index de904e14..89f360cb 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -49,7 +49,9 @@ def _take(obj: Y, idx: np.ndarray) -> Y: def take(x): # TODO(erikfrey): if this helps perf, add support for striding too - if ( + if not x.shape[0]: + return x + elif ( len(idx.shape) == 1 and idx.size > 0 and (idx == np.arange(idx[0], idx[0] + idx.size)).all() @@ -113,8 +115,12 @@ def _nvmap(f: Callable[..., Y], *args) -> Y: if isinstance(arg, np.ndarray) and not np.all(arg == arg[0]): raise RuntimeError(f'numpy arg elements do not match: {arg}') + # split out numpy and jax args np_args = [a[0] if isinstance(a, np.ndarray) else None for a in args] args = [a if n is None else None for n, a in zip(np_args, args)] + + # remove empty args that we should not vmap over + args = jax.tree_map(lambda a: a if a.shape[0] else None, args) in_axes = [None if a is None else 0 for a in args] def outer_f(*args, np_args=np_args): @@ -126,7 +132,15 @@ def _nvmap(f: Callable[..., Y], *args) -> Y: def _check_input(m: Model, args: Any, in_types: str) -> None: """Checks that scan input has the right shape.""" - size = {'b': m.nbody, 'j': m.njnt, 'q': m.nq, 'v': m.nv, 'u': m.nu, 'a': m.na} + size = { + 'b': m.nbody, + 'j': m.njnt, + 'q': m.nq, + 'v': m.nv, + 'u': m.nu, + 'a': m.na, + 's': m.nsite, + } for idx, (arg, typ) in enumerate(zip(args, in_types)): if len(arg) != size[typ]: raise IndexError( @@ -162,7 +176,7 @@ def flat( ) -> Y: r"""Scan a function across bodies or actuators. - Scan group data according to type and batch shape then calls vmap(f) on it. + Scan group data according to type and batch shape then calls vmap(f) on it.\ Args: m: an mjx model @@ -206,6 +220,7 @@ def flat( m.actuator_dyntype[ids_u], m.actuator_trntype[ids_u], m.jnt_type[ids_j], + m.actuator_trnid[ids_u, 1] == -1, # key by refsite being present ) def type_ids_j(m, i): @@ -221,16 +236,24 @@ def flat( 'u': i, 'a': m.actuator_actadr[i], 'j': ( - m.actuator_trnid[i] + m.actuator_trnid[i, 0] if m.actuator_trntype[i] == TrnType.JOINT - else np.array(-1) + else -1 + ), + 's': ( + m.actuator_trnid[i] + if m.actuator_trntype[i] == TrnType.SITE + else np.array([-1, -1]) ), } - # v/q associated with joint transmissions - typ_ids.update({ - 'v': np.nonzero(m.dof_jntid == typ_ids['j'])[0], - 'q': np.nonzero(_q_jointid(m) == typ_ids['j'])[0], - }) + v, q = np.array([-1]), np.array([-1]) + if m.actuator_trntype[i] == TrnType.JOINT: + # v/q are associated with the joint transmissions only + v = np.nonzero(m.dof_jntid == typ_ids['j'])[0] + q = np.nonzero(_q_jointid(m) == typ_ids['j'])[0] + + typ_ids.update({'v': v, 'q': q}) + return typ_ids # build up a grouping of type take-ids in body/actuator order diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py index 9b845de1..456067d6 100644 --- a/mjx/mujoco/mjx/_src/scan_test.py +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -193,7 +193,7 @@ class ScanTest(absltest.TestCase): """ - def testscan_actuators(self): + def test_scan_actuators(self): """Tests scanning over actuators.""" m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML) m = mjx.device_put(m) @@ -210,15 +210,16 @@ class ScanTest(absltest.TestCase): m, fn, 'ujqva', 'ujqva', *args, group_by='u' ) + actuator_trnid = m.actuator_trnid[:, 0] np.testing.assert_array_equal(gear, m.actuator_gear) - np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid]) + np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid]) np.testing.assert_array_equal(act, jp.array([1.4, 1.1])) expected_vadr = np.concatenate( - [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid] + [np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid] ) np.testing.assert_array_equal(vadr, expected_vadr) expected_qadr = np.concatenate( - [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid] + [np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid] ) np.testing.assert_array_equal(qadr, expected_qadr) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 2b421f2f..c1531192 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -19,12 +19,15 @@ from jax import numpy as jp import mujoco from mujoco.mjx._src import math from mujoco.mjx._src import scan +from mujoco.mjx._src import support # pylint: disable=g-importing-member from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import JointType from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import TrnType # pylint: enable=g-importing-member +import numpy as np def kinematics(m: Model, d: Data) -> Data: @@ -101,13 +104,20 @@ def kinematics(m: Model, d: Data) -> Data: # TODO(erikfrey): confirm that quats are more performant for mjx than mats xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat) - geom_xpos, geom_xmat = local_to_global( - xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat - ) - d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos) d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat) - d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat) + + if m.ngeom: + geom_xpos, geom_xmat = local_to_global( + xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat + ) + d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat) + + if m.nsite: + site_xpos, site_xmat = local_to_global( + xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat + ) + d = d.replace(site_xpos=site_xpos, site_xmat=site_xmat) return d @@ -423,45 +433,118 @@ def rne(m: Model, d: Data) -> Data: return d +def _site_dof_mask(m: Model) -> np.ndarray: + """Creates a dof mask for site transmissions.""" + mask = np.ones((m.nu, m.nv)) + for i in np.nonzero(m.actuator_trnid[:, 1] != -1)[0]: + id_, refid = m.actuator_trnid[i] + # intialize last dof address for each body + b0 = m.body_weldid[m.site_bodyid[id_]] + b1 = m.body_weldid[m.site_bodyid[refid]] + dofadr0 = m.body_dofadr[b0] + m.body_dofnum[b0] - 1 + dofadr1 = m.body_dofadr[b1] + m.body_dofnum[b1] - 1 + + # find common ancestral dof, if any + while dofadr0 != dofadr1: + if dofadr0 < dofadr1: + dofadr1 = m.dof_parentid[dofadr1] + else: + dofadr0 = m.dof_parentid[dofadr0] + if dofadr0 == -1 or dofadr1 == -1: + break + + # if common ancestral dof was found, clear the columns of its parental chain + da = dofadr0 if dofadr0 == dofadr1 else -1 + while da >= 0: + mask[i, da] = 0.0 + da = m.dof_parentid[da] + + return mask + + def transmission(m: Model, d: Data) -> Data: """Computes actuator/transmission lengths and moments.""" + # TODO: consider combining transmission calculation into fwd_actuation. if not m.nu: return d - def fn(gear, jnt_typ, m_i, m_j, qpos): - # handles joint transmissions only - if jnt_typ == JointType.FREE: - length = jp.zeros(1) - moment = gear - m_i = jp.repeat(m_i, 6) - m_j = m_j + jp.arange(6) - elif jnt_typ == JointType.BALL: - axis, _ = math.quat_to_axis_angle(qpos) - length = jp.dot(axis, gear[:3])[None] - moment = gear[:3] - m_i = jp.repeat(m_i, 3) - m_j = m_j + jp.arange(3) - elif jnt_typ in (JointType.SLIDE, JointType.HINGE): - length = qpos * gear[0] - moment = gear[:1] - m_i, m_j = m_i[None], m_j[None] - else: - raise RuntimeError(f'unrecognized joint type: {jnt_typ}') - return length, moment, m_i, m_j + def fn( + trntype, + trnid, + gear, + jnt_typ, + m_j, + qpos, + has_refsite, + site_dof_mask, + site_xpos, + site_xmat, + site_quat, + ): + if trntype == TrnType.JOINT: + if jnt_typ == JointType.FREE: + length = jp.zeros(1) + moment = gear + m_j = m_j + jp.arange(6) + elif jnt_typ == JointType.BALL: + axis, angle = math.quat_to_axis_angle(qpos) + length = jp.dot(axis * angle, gear[:3])[None] + moment = gear[:3] + m_j = m_j + jp.arange(3) + elif jnt_typ in (JointType.SLIDE, JointType.HINGE): + length = qpos * gear[0] + moment = gear[:1] + m_j = m_j[None] + else: + raise RuntimeError(f'unrecognized joint type: {JointType(jnt_typ)}') - length, m_val, m_i, m_j = scan.flat( + moment = jp.zeros((m.nv,)).at[m_j].set(moment) + elif trntype == TrnType.SITE: + length = jp.zeros(1) + id_, refid = jp.array(m.site_bodyid)[trnid] + jacp, jacr = support.jac(m, d, site_xpos[0], id_) + frame_xmat = site_xmat[0] + if has_refsite: + vecp = site_xmat[1].T @ (site_xpos[0] - site_xpos[1]) + vecr = math.quat_sub(site_quat[0], site_quat[1]) + length += jp.dot(jp.concatenate([vecp, vecr]), gear) + jacrefp, jacrefr = support.jac(m, d, site_xpos[1], refid) + jacp, jacr = jacp - jacrefp, jacr - jacrefr + frame_xmat = site_xmat[1] + + jac = jp.concatenate((jacp, jacr), axis=1) * site_dof_mask[:, None] + wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:])) + moment = jac @ wrench + else: + raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}') + + return length, moment + + # pre-compute values for site transmissions + has_refsite = m.actuator_trnid[:, 1] != -1 + site_dof_mask = _site_dof_mask(m) + site_quat = jax.vmap(math.quat_mul)(m.site_quat, d.xquat[m.site_bodyid]) + + length, moment = scan.flat( m, fn, - 'ujujq', - 'uvvv', + 'uuujjquusss', + 'uu', + m.actuator_trntype, + jp.array(m.actuator_trnid), m.actuator_gear, m.jnt_type, - jp.arange(m.nu), jp.array(m.jnt_dofadr), d.qpos, + has_refsite, + jp.array(site_dof_mask), + d.site_xpos, + d.site_xmat, + site_quat, group_by='u', ) - moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val) length = length.reshape((m.nu,)) + moment = moment.reshape((m.nu, m.nv)) + d = d.replace(actuator_length=length, actuator_moment=moment) return d diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 9e518672..b7d62eb9 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -15,122 +15,107 @@ """Tests for smooth dynamics functions.""" from absl.testing import absltest -from absl.testing import parameterized import jax from jax import numpy as jp import mujoco from mujoco import mjx from mujoco.mjx._src import test_util -# pylint: disable=g-importing-member -from mujoco.mjx._src.types import DisableBit -# pylint: enable=g-importing-member import numpy as np - -def _assert_eq(a, b, name, step, fname, atol=5e-4, rtol=5e-4): - err_msg = f'mismatch: {name} at step {step} in {fname}' - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +# tolerance for difference between MuJoCo and MJX smooth calculations - mostly +# due to float precision +_TOLERANCE = 5e-5 -def _assert_attr_eq(a, b, attr, step, fname, atol=5e-4, rtol=5e-4): - err_msg = f'mismatch: {attr} at step {step} in {fname}' - a, b = getattr(a, attr), getattr(b, attr) - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +def _assert_eq(a, b, name): + tol = _TOLERANCE * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) -class SmoothTest(parameterized.TestCase): +def _assert_attr_eq(a, b, attr): + _assert_eq(getattr(a, attr), getattr(b, attr), attr) - @parameterized.parameters(enumerate(test_util.TEST_FILES)) - def test_smooth(self, seed, fname): - """Tests mujoco mj smooth functions match mujoco_mjx smooth functions.""" - if fname in ('convex.xml', 'equality.xml'): - return - np.random.seed(seed) +class SmoothTest(absltest.TestCase): - m = test_util.load_test_file(fname) + def setUp(self): + super().setUp() + # although we already have generous padding of thresholds, it doesn't hurt + # to also fix the seed to reduce test flakiness + np.random.seed(0) + + def test_smooth(self): + """Tests MJX smooth functions match MuJoCo smooth functions.""" + + m = test_util.load_test_file('pendula.xml') d = mujoco.MjData(m) - - kinematics_jit_fn = jax.jit(mjx.kinematics) - com_pos_jit_fn = jax.jit(mjx.com_pos) - crb_jit_fn = jax.jit(mjx.crb) - factor_m_fn = jax.jit(mjx.factor_m) - com_vel_jit_fn = jax.jit(mjx.com_vel) - rne_jit_fn = jax.jit(mjx.rne) - mul_m_jit_fn = jax.jit(mjx.mul_m) - transmission_jit_fn = jax.jit(mjx.transmission) - - mx = mjx.device_put(m) - dx = mjx.make_data(mx) - # give the system a little kick to ensure we have non-identity rotations d.qvel = np.random.random(m.nv) - for i in range(100): - qpos, qvel = d.qpos.copy(), d.qvel.copy() - mujoco.mj_step(m, d) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) - # kinematics - dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) - _assert_attr_eq(d, dx, 'xanchor', i, fname) - _assert_attr_eq(d, dx, 'xaxis', i, fname) - _assert_attr_eq(d, dx, 'xpos', i, fname) - _assert_attr_eq(d, dx, 'xquat', i, fname) - _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname) - _assert_attr_eq(d, dx, 'xipos', i, fname) - _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname) - _assert_attr_eq(d, dx, 'geom_xpos', i, fname) - _assert_eq( - d.geom_xmat.reshape((-1, 3, 3)), - dx.geom_xmat, - 'geom_xmat', - i, - fname, - ) + # kinematics + dx = jax.jit(mjx.kinematics)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'xanchor') + _assert_attr_eq(d, dx, 'xaxis') + _assert_attr_eq(d, dx, 'xpos') + _assert_attr_eq(d, dx, 'xquat') + _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat') + _assert_attr_eq(d, dx, 'xipos') + _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat') + _assert_attr_eq(d, dx, 'geom_xpos') + _assert_eq(d.geom_xmat.reshape((-1, 3, 3)), dx.geom_xmat, 'geom_xmat') + _assert_attr_eq(d, dx, 'site_xpos') + _assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat') + # com_pos + dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'subtree_com') + _assert_attr_eq(d, dx, 'cinert') + _assert_attr_eq(d, dx, 'cdof') + # crb + dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'crb') + _assert_attr_eq(d, dx, 'qM') + # factor_m + dx = mjx.put_data(m, d) + dx = jax.jit(mjx.factor_m)(mx, dx, dx.qM) + _assert_attr_eq(d, dx, 'qLD') + _assert_attr_eq(d, dx, 'qLDiagInv') + # com_vel + dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'cvel') + _assert_attr_eq(d, dx, 'cdof_dot') + # rne + dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qfrc_bias') + # transmission + dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'actuator_length') + _assert_attr_eq(d, dx, 'actuator_moment') - # com_pos - dx = com_pos_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'subtree_com', i, fname) - _assert_attr_eq(d, dx, 'cinert', i, fname) - _assert_attr_eq(d, dx, 'cdof', i, fname) + def test_mul_m(self): + m = test_util.load_test_file('pendula.xml') + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + vec = np.random.random(m.nv) + mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec)) + mj_vec = np.zeros(m.nv) + mujoco.mj_mulM(m, d, mj_vec, vec) + _assert_eq(mj_vec, mjx_vec, 'mul_m') - # crb - dx = crb_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'crb', i, fname) - _assert_attr_eq(d, dx, 'qM', i, fname) - - # factor_m - dx = factor_m_fn(mx, dx, dx.qM) - _assert_attr_eq(d, dx, 'qLD', i, fname, atol=1e-3) - _assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3) - - # com_vel - dx = com_vel_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'cvel', i, fname) - _assert_attr_eq(d, dx, 'cdof_dot', i, fname) - - # rne - dx = rne_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'qfrc_bias', i, fname) - - # mul_m (auxilliary function, not part of smooth step) - vec = np.random.random(m.nv) - mjx_vec = mul_m_jit_fn(mx, dx, jp.array(vec)) - mj_vec = np.zeros(m.nv) - mujoco.mj_mulM(m, d, mj_vec, vec) - _assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname) - - # transmission - dx = transmission_jit_fn(mx, dx) - _assert_attr_eq(d, dx, 'actuator_length', i, fname) - _assert_attr_eq(d, dx, 'actuator_moment', i, fname) - - -class DisableGravityTest(absltest.TestCase): - - def test_disabled(self): + def test_disable_gravity(self): m = mujoco.MjModel.from_xml_string(""" - @@ -139,27 +124,52 @@ class DisableGravityTest(absltest.TestCase): """) - mx = mjx.device_put(m) d = mujoco.MjData(m) - dx = mjx.device_put(d) + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) - # test with gravity - step_jit_fn = jax.jit(mjx.step) - dx = step_jit_fn(mx, dx) - np.testing.assert_array_almost_equal( - dx.qpos, np.array([0.0, 0.0, -9.81e-4, 1.0, 0.0, 0.0, 0.0]), decimal=7 - ) + dx = jax.jit(mjx.rne)(mx, dx) + np.testing.assert_allclose(dx.qfrc_bias, 0) - # test with gravity disabled - mx = mx.tree_replace( - {'opt.disableflags': mx.opt.disableflags | DisableBit.GRAVITY} - ) - dx = mjx.device_put(d) - step_jit_fn = jax.jit(mjx.step) - dx = step_jit_fn(mx, dx) - np.testing.assert_equal( - dx.qpos, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) - ) + def test_site_transmission(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + + + + + + + + + + + + + + + """) + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + + mujoco.mj_transmission(m, d) + dx = jax.jit(mjx.transmission)(mx, dx) + _assert_attr_eq(d, dx, 'actuator_length') + _assert_attr_eq(d, dx, 'actuator_moment') if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index 567cb709..f8ef9bed 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -19,6 +19,7 @@ from typing import Optional import jax from jax import numpy as jp import mujoco +from mujoco.mjx._src import constraint from mujoco.mjx._src import math from mujoco.mjx._src import smooth # pylint: disable=g-importing-member @@ -69,12 +70,12 @@ class _Context(PyTreeNode): # TODO(robotics-team): determine nv at which sparse mul is faster M = smooth.dense_m(m, d) if m.nv < 100 else None # pylint: disable=invalid-name ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc - nv_0 = jp.zeros((m.nv,)) + nv_0 = jp.zeros(m.nv) ctx = _Context( qacc=d.qacc, qfrc_constraint=d.qfrc_constraint, Jaref=jaref, - efc_force=jp.zeros(d.nefc), + efc_force=d.efc_force, M=M, Ma=ma, grad=nv_0, @@ -111,7 +112,7 @@ class _LSPoint(PyTreeNode): @classmethod def create( cls, - d: Data, + m: Model, ctx: _Context, alpha: jax.Array, jv: jax.Array, @@ -122,13 +123,14 @@ class _LSPoint(PyTreeNode): # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c # TODO(robotics-team): change this to support friction constraints - active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True) quad = jax.vmap(jp.multiply)(quad, active) # only active quad_total = quad_gauss + jp.sum(quad, axis=0) cost = alpha * alpha * quad_total[2] + alpha * quad_total[1] + quad_total[0] deriv_0 = 2 * alpha * quad_total[2] + quad_total[1] - deriv_1 = 2 * quad_total[2] + deriv_1 = 2 * quad_total[2] + (quad_total[2] == 0) * mujoco.mjMINVAL return _LSPoint(alpha=alpha, cost=cost, deriv_0=deriv_0, deriv_1=deriv_1) @@ -177,12 +179,11 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: Returns: context with new constraint force and costs """ - del m - # TODO(robotics-team): add friction constraints # only count active constraints - active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = (ctx.Jaref < 0).at[:ne + nf].set(True) efc_force = d.efc_D * -ctx.Jaref * active qfrc_constraint = d.efc_J.T @ efc_force @@ -221,7 +222,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: if m.opt.solver == SolverType.CG: mgrad = smooth.solve_m(m, d, grad) elif m.opt.solver == SolverType.NEWTON: - active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = (ctx.Jaref < 0).at[:ne + nf].set(True) h = (d.efc_J.T * d.efc_D * active) @ d.efc_J h = smooth.dense_m(m, d) + h h_ = jax.scipy.linalg.cho_factor(h) @@ -265,7 +267,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv)) quad = (quad * d.efc_D).T - point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss) + point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss) def cond(ctx: _LSContext) -> jax.Array: done = ctx.ls_iter >= m.opt.ls_iterations diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py index 6a4f7792..d1677ff9 100644 --- a/mjx/mujoco/mjx/_src/solver_test.py +++ b/mjx/mujoco/mjx/_src/solver_test.py @@ -12,119 +12,64 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for forward functions.""" +"""Tests for constraint functions.""" from absl.testing import absltest -from absl.testing import parameterized -from etils import epath import jax import mujoco from mujoco import mjx +from mujoco.mjx._src import test_util import numpy as np -def _assert_attr_eq(a, b, attr, step, fname, atol=1e-2, rtol=1e-2): - err_msg = f'mismatch: {attr} at step {step} in {fname}' - a, b = getattr(a, attr), getattr(b, attr) - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) +# tolerance for difference between MuJoCo and MJX constraint calculations, +# mostly due to float precision +_TOLERANCE = 5e-5 -class Solver64Test(parameterized.TestCase): - """Tests solvers at 64 bit precision.""" +def _assert_eq(a, b, name, tol=_TOLERANCE): + tol = tol * 10 # avoid test noise + err_msg = f'mismatch: {name}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) - def setUp(self): - super().setUp() - jax.config.update('jax_enable_x64', True) - def tearDown(self): - super().tearDown() - jax.config.update('jax_enable_x64', False) +def _assert_attr_eq(a, b, attr): + _assert_eq(getattr(a, attr), getattr(b, attr), attr) - @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) - def test_cg(self, seed, fname): - """Test mjx cg solver matches mujoco cg solver at 64 bit precision.""" - f = epath.resource_path('mujoco.mjx') / 'test_data' / fname - m = mujoco.MjModel.from_xml_string(f.read_text()) + +class SolverTest(absltest.TestCase): + + def test_solver(self): + """Test solver.""" + m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) - mx = mjx.device_put(m) + mujoco.mj_step(m, d, 100) # at 100 steps mix of active/inactive constraints + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) - jax.config.update('jax_enable_x64', True) - forward_jit_fn = jax.jit(mjx.forward) + dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qacc_warmstart') + _assert_attr_eq(d, dx, 'qacc') + _assert_attr_eq(d, dx, 'qfrc_constraint') + nnz = dx.efc_J.any(axis=1) + _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') - # give the system a little kick to ensure we have non-identity rotations - np.random.seed(seed) - d.qvel = 0.01 * np.random.random(m.nv) - - for i in range(100): - # in order to avoid re-jitting, reuse the same mj_data shape - save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth - d = mujoco.MjData(m) - d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save - dx = mjx.device_put(d) - - mujoco.mj_step(m, d) - dx = forward_jit_fn(mx, dx) - - # at 64 bits the solutions returned by the two solvers are quite close - self.assertLessEqual(dx.solver_niter[0], d.solver_niter[0]) - _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname) - _assert_attr_eq(d, dx, 'qacc', i, fname) - - -class SolverTest(parameterized.TestCase): - - @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) - def test_cg(self, seed, fname): - """Test mjx cg solver is close to mj at 32 bit precision. - - Args: - seed: int - fname: file to test - - At lower float resolution there's wiggle room in valid forces that satisfy - constraints. So instead let's mainly validate that mjx is finding solutions - with as good cost as mujoco, even if the resulting forces/accelerations - are not quite the same. - """ - f = epath.resource_path('mujoco.mjx') / 'test_data' / fname - m = mujoco.MjModel.from_xml_string(f.read_text()) - d = mujoco.MjData(m) - mx = mjx.device_put(m) - - forward_jit_fn = jax.jit(mjx.forward) - - # give the system a little kick to ensure we have non-identity rotations - np.random.seed(seed) - d.qvel = 0.01 * np.random.random(m.nv) - - for i in range(100): - # in order to avoid re-jitting, reuse the same mj_data shape - save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth - d = mujoco.MjData(m) - d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save - dx = mjx.device_put(d) - - mujoco.mj_step(m, d) - dx = forward_jit_fn(mx, dx) - - def cost(qacc): - jaref = np.zeros(d.nefc) - mujoco.mj_mulJacVec(m, d, jaref, qacc) - jaref -= d.efc_aref - cost = np.array([0.0]) - mujoco.mj_constraintUpdate(m, d, jaref, cost, 0) - return cost[0] - - cost_mj, cost_mjx = cost(d.qacc), cost(dx.qacc) - - self.assertLessEqual( - cost_mjx, - cost_mj * 1.01, - msg=f'mismatch: {fname} at step {i}, cost too high', - ) - _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname, atol=1e-1, rtol=1e-1) - _assert_attr_eq(d, dx, 'qacc', i, fname, atol=1e-1, rtol=1e-1) + # also test normal CG + m.opt.solver = mujoco.mjtSolver.mjSOL_CG + mujoco.mj_forward(m, d) + dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'qacc_warmstart') + _assert_attr_eq(d, dx, 'qacc') + _assert_attr_eq(d, dx, 'qfrc_constraint') + _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + # without warmstart, the solution is not as close + m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON + m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d)) + _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2) if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index fe88fc84..fb3a5389 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -34,8 +34,8 @@ class SupportTest(parameterized.TestCase): m = test_util.load_test_file(fname) d = mujoco.MjData(m) mujoco.mj_step(m, d) - mx = mjx.device_put(m) - dx = mjx.device_put(d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) point = np.random.randn(3) body = np.random.choice(m.nbody) jacp, jacr = jax.jit(support.jac)(mx, dx, point, body) @@ -49,11 +49,11 @@ class SupportTest(parameterized.TestCase): """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate.""" np.random.seed(0) - m = test_util.load_test_file('ant.xml') + m = test_util.load_test_file('pendula.xml') d = mujoco.MjData(m) mujoco.mj_step(m, d) - mx = mjx.device_put(m) - dx = mjx.device_put(d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) self.assertFalse((dx.xipos == 0.0).all()) xfrc = np.random.rand(*dx.xfrc_applied.shape) diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index 8d765644..d350c11a 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -23,20 +23,21 @@ import mujoco import numpy as np TEST_FILES: List[str] = [ - 'ant.xml', + 'constraints.xml', 'convex.xml', - 'equality.xml', - 'humanoid.xml', 'pendula.xml', + 'ray.xml', ] _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] +_DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact'] +_DYN_PRMS = ['0.189', '2.1'] _JOINT_TYPES = ['free', 'hinge', 'slide', 'ball'] _JOINT_AXES = ['1 0 0', '0 1 0', '0 0 1'] _FRICTIONS = ['1.2 0.003 0.0002', '0.2 0.0001 0.0005'] _KP_POS = ['1', '2'] _KP_INTVEL = ['10000', '2000'] -_KV_VEL = ['123', '1'] +_KV_VEL = ['12', '1', '0', '0.1'] _PAIR_FRICTIONS = ['1.2 0.9 0.003 0.0002 0.0001'] _SOLREFS = ['0.04 1.01', '0.05 1.02', '0.03 1.1', '0.015 1.0'] _SOLIMPS = [ @@ -47,7 +48,7 @@ _SOLIMPS = [ _DIMS = ['3'] _MARGINS = ['0.0', '0.01', '0.02'] _GAPS = ['0.0', '0.005'] -_GEARS = ['20', '50', '100'] +_GEARS = ['2.1 0.0 3.3 0 2.3 0', '5.0 3.1 0 2.3 0.0 1.1'] def p(pct: int) -> bool: @@ -123,13 +124,29 @@ def _make_geom( return attr -def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]: +def _make_actuator( + actuator_type: str, + joint: str | None = None, + site: str | None = None, + refsite: str | None = None, +) -> Dict[str, str]: """Returns attributes for an actuator.""" - attr = {'joint': joint} - if actuator_type == 'motor': - attr['gear'] = np.random.choice(_GEARS) - elif actuator_type == 'position': + if joint: + attr = {'joint': joint} + elif site: + attr = {'site': site} + else: + raise ValueError('must provide a joint or site name') + + if refsite: + attr['refsite'] = refsite + + attr['gear'] = np.random.choice(_GEARS) + + # set actuator type + if actuator_type == 'position': attr['kp'] = np.random.choice(_KP_POS) + attr['kv'] = np.random.choice(_KV_VEL) elif actuator_type == 'general': attr['biastype'] = 'affine' attr['gainprm'] = '35 0 0' @@ -141,10 +158,18 @@ def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]: elif actuator_type == 'velocity': attr['kv'] = np.random.choice(_KV_VEL) + # set dyntype + if actuator_type == 'general': + attr['dyntype'] = np.random.choice(_DYN_TYPES) + if attr['dyntype'] != 'none': + attr['dynprm'] = np.random.choice(_DYN_PRMS) + + # ctrlrange if p(50) and actuator_type != 'intvelocity': lb, ub = -np.random.uniform(), np.random.uniform() attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}' + # forcerange if p(50): lb, ub = -np.random.uniform(), np.random.uniform() attr['forcerange'] = f'{lb*10:.2f} {ub*10:.2f}' @@ -235,6 +260,7 @@ def create_mjcf( pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}' n_bodies = len(list(mjcf.iter('body'))) child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'}) + ET.SubElement(child, 'site', {'name': f'site{n_bodies}'}) n_joints = len(list(mjcf.iter('joint'))) for nj in range(np.random.randint(1, max_stacked_joints + 1)): @@ -272,17 +298,35 @@ def create_mjcf( for _ in range(num_trees): make_tree(world, 0) + bodies = list(mjcf.iter('body')) + n_bodies = len(bodies) + # actuators if add_actuators: actuator = ET.SubElement(mjcf, 'actuator') n_joints = len(list(mjcf.iter('joint'))) nu = np.random.randint(1, n_joints + 1) actuators = [] + + # joint transmission for i in range(nu): actuator_type = np.random.choice(_ACTUATOR_TYPES) attr = _make_actuator(actuator_type, joint=f'joint{i}') actuators.append((actuator_type, attr)) + # site transmission + for i in range(np.random.randint(0, n_bodies)): + actuator_type = np.random.choice(_ACTUATOR_TYPES) + attr = _make_actuator(actuator_type, site=f'site{i}') + actuators.append((actuator_type, attr)) + + # site transmission with refsite + for i in range(np.random.randint(0, n_bodies)): + j = np.random.randint(0, n_bodies) + actuator_type = np.random.choice(_ACTUATOR_TYPES) + attr = _make_actuator(actuator_type, site=f'site{i}', refsite=f'site{j}') + actuators.append((actuator_type, attr)) + np.random.shuffle(actuators) for typ, attr in actuators: ET.SubElement(actuator, typ, attr) @@ -310,9 +354,7 @@ def create_mjcf( ET.SubElement(contact, 'pair', attr) # exclude contacts - bodies = list(mjcf.iter('body')) body_names = [b.get('name') for b in bodies] - n_bodies = len(bodies) for _ in range(min(max_contact_excludes, (n_bodies * (n_bodies - 1) // 2))): if p(50): continue diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 5875ad7b..79cecaf5 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -20,9 +20,7 @@ from typing import Sequence import jax import jax.numpy as jp import mujoco -# pylint: disable=g-importing-member -from mujoco.mjx._src.dataclasses import PyTreeNode -# pylint: enable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member import numpy as np @@ -155,9 +153,11 @@ class TrnType(enum.IntEnum): Attributes: JOINT: force on joint + SITE: force on site """ JOINT = mujoco.mjtTrn.mjTRN_JOINT - # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, SITE, BODY + SITE = mujoco.mjtTrn.mjTRN_SITE + # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, BODY class DynType(enum.IntEnum): @@ -166,11 +166,14 @@ class DynType(enum.IntEnum): Attributes: NONE: no internal dynamics; ctrl specifies force INTEGRATOR: integrator: da/dt = u + FILTER: linear filter: da/dt = (u-a) / tau + FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration """ NONE = mujoco.mjtDyn.mjDYN_NONE INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR FILTER = mujoco.mjtDyn.mjDYN_FILTER - # unsupported: FILTEREXACT, MUSCLE, USER + FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT + # unsupported: MUSCLE, USER class GainType(enum.IntEnum): @@ -213,7 +216,6 @@ class Option(PyTreeNode): integrator: integration mode cone: type of friction cone solver: solver algorithm - integrator: integration mode iterations: number of main solver iterations ls_iterations: maximum number of CG/Newton linesearch iterations disableflags: bit flags for disabling standard features @@ -260,7 +262,9 @@ class Model(PyTreeNode): nbody: number of bodies njnt: number of joints ngeom: number of geoms + nsite: number of sites nmesh: number of meshes + nmat: number of materials npair: number of predefined geom pairs nexclude: number of excluded geom pairs neq: number of equality constraints @@ -317,6 +321,8 @@ class Model(PyTreeNode): geom_conaffinity: geom contact affinity (ngeom,) geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) geom_bodyid: id of geom's body (ngeom,) + geom_group: group for visibility (ngeom,) + geom_matid: material id for rendering (ngeom,) geom_priority: geom contact priority (ngeom,) geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) geom_solref: constraint solver reference: contact (ngeom, mjNREF) @@ -327,6 +333,11 @@ class Model(PyTreeNode): geom_friction: friction for (slide, spin, roll) (ngeom, 3) geom_margin: include in solver if dist 'Contact': + def zero(cls, ncon: int = 0) -> 'Contact': """Returns a contact filled with zeros.""" return Contact( - dist=jp.zeros(shape), - pos=jp.zeros(shape + (3,)), - frame=jp.zeros(shape + (3, 3)), - includemargin=jp.zeros(shape), - friction=jp.zeros(shape + (5,)), - solref=jp.zeros(shape + (mujoco.mjNREF,)), - solreffriction=jp.zeros(shape + (mujoco.mjNREF,)), - solimp=jp.zeros(shape + (mujoco.mjNIMP,)), - dim=np.zeros(shape, dtype=np.int32), - geom1=jp.zeros(shape, dtype=jp.int32), - geom2=jp.zeros(shape, dtype=jp.int32), - efc_address=np.zeros(shape, dtype=np.int32), + dist=jp.zeros(ncon), + pos=jp.zeros((ncon, 3,)), + frame=jp.zeros((ncon, 3, 3)), + includemargin=jp.zeros(ncon), + friction=jp.zeros((ncon, 5)), + solref=jp.zeros((ncon, mujoco.mjNREF)), + solreffriction=jp.zeros((ncon, mujoco.mjNREF)), + solimp=jp.zeros((ncon, mujoco.mjNIMP,)), + geom1=jp.zeros(ncon, dtype=jp.int32), + geom2=jp.zeros(ncon, dtype=jp.int32), ) @@ -544,11 +558,6 @@ class Data(PyTreeNode): Attributes: solver_niter: number of solver iterations, per island (mjNISLAND,) - ne: number of equality constraints - nf: number of friction constraints - nl: number of limit constraints - nefc: number of constraints - ncon: nubmer of contacts time: simulation time qpos: position (nq,) qvel: velocity (nv,) @@ -569,6 +578,8 @@ class Data(PyTreeNode): xaxis: Cartesian joint axis (njnt, 3) geom_xpos: Cartesian geom position (ngeom, 3) geom_xmat: Cartesian geom orientation (ngeom, 3, 3) + site_xpos: Cartesian site position (nsite, 3) + site_xmat: Cartesian site orientation (nsite, 9) subtree_com: center of mass of each subtree (nbody, 3) cdof: com-based motion axis of each dof (nv, 6) cinert: com-based body inertia and mass (nbody, 10) @@ -589,7 +600,6 @@ class Data(PyTreeNode): qfrc_bias: C(qpos,qvel) (nv,) qfrc_passive: passive force (nv,) efc_aref: reference pseudo-acceleration (nefc,) - actuator_force: actuator force in actuation space (nu,) qfrc_actuator: actuator force (nv,) qfrc_smooth: net unconstrained force (nv,) qacc_smooth: unconstrained acceleration (nv,) @@ -600,12 +610,6 @@ class Data(PyTreeNode): """ # solver statistics: solver_niter: jax.Array - # sizes (variable in MJ, constant in MJX) - ne: int - nf: int - nl: int - nefc: int - ncon: int # global properties: time: jax.Array # state: @@ -631,6 +635,8 @@ class Data(PyTreeNode): xaxis: jax.Array geom_xpos: jax.Array geom_xmat: jax.Array + site_xpos: jax.Array + site_xmat: jax.Array subtree_com: jax.Array cdof: jax.Array cinert: jax.Array @@ -653,7 +659,6 @@ class Data(PyTreeNode): qfrc_passive: jax.Array efc_aref: jax.Array # position, velcoity, control & acceleration dependent: - actuator_force: jax.Array qfrc_actuator: jax.Array qfrc_smooth: jax.Array qacc_smooth: jax.Array diff --git a/mjx/mujoco/mjx/integration_test/collision_driver_test.py b/mjx/mujoco/mjx/integration_test/collision_driver_test.py index d49ef10a..a9891656 100644 --- a/mjx/mujoco/mjx/integration_test/collision_driver_test.py +++ b/mjx/mujoco/mjx/integration_test/collision_driver_test.py @@ -58,9 +58,9 @@ class CollisionDriverIntegrationTest(parameterized.TestCase): ) m = mujoco.MjModel.from_xml_string(mjcf) - mx = mjx.device_put(m) + mx = mjx.put_model(m) d = mujoco.MjData(m) - dx = mjx.device_put(d) + dx = mjx.put_data(m, d) mujoco.mj_step(m, d) collision_jit_fn = jax.jit(mjx.collision) @@ -82,7 +82,6 @@ class CollisionDriverIntegrationTest(parameterized.TestCase): mjx_contact = jax.tree_map( lambda x: x.take(np.array(idx), axis=0), dx.contact ) - mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx]) for field in dataclasses.fields(Contact): _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7) diff --git a/mjx/mujoco/mjx/integration_test/forward_test.py b/mjx/mujoco/mjx/integration_test/forward_test.py index 6fa913c4..67f20371 100644 --- a/mjx/mujoco/mjx/integration_test/forward_test.py +++ b/mjx/mujoco/mjx/integration_test/forward_test.py @@ -19,7 +19,6 @@ from absl.testing import parameterized import jax import mujoco from mujoco import mjx -from mujoco.mjx._src import forward from mujoco.mjx._src import test_util import numpy as np @@ -46,7 +45,7 @@ class ActuationIntegrationTest(parameterized.TestCase): enable_contact=False, ) m = mujoco.MjModel.from_xml_string(mjcf) - actuation_jit_fn = jax.jit(forward._actuation) + actuation_jit_fn = jax.jit(mjx.fwd_actuation) # init d = mujoco.MjData(m) @@ -57,8 +56,8 @@ class ActuationIntegrationTest(parameterized.TestCase): mujoco.mj_fwdVelocity(m, d) # put on device - mx = mjx.device_put(m) - dx = mjx.device_put(d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) mujoco.mj_fwdActuation(m, d) dx = actuation_jit_fn(mx, dx) diff --git a/mjx/mujoco/mjx/integration_test/smooth_test.py b/mjx/mujoco/mjx/integration_test/smooth_test.py index c924e005..c4d3bfda 100644 --- a/mjx/mujoco/mjx/integration_test/smooth_test.py +++ b/mjx/mujoco/mjx/integration_test/smooth_test.py @@ -57,17 +57,21 @@ class TransmissionIntegrationTest(parameterized.TestCase): d = mujoco.MjData(m) d.ctrl = np.random.normal(scale=10, size=m.nu) d.act = np.random.normal(scale=10, size=m.na) + d.qpos = np.random.normal(m.nq) d.qvel = np.random.random(m.nv) + mujoco.mj_forward(m, d) # put on device - mx = mjx.device_put(m) - dx = mjx.device_put(d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) mujoco.mj_transmission(m, d) dx = transmission_jit_fn(mx, dx) _assert_attr_eq(d, dx, 'actuator_length', seed, f'transmission{seed}') - _assert_attr_eq(d, dx, 'actuator_moment', seed, f'transmission{seed}') + _assert_attr_eq( + d, dx, 'actuator_moment', seed, f'transmission{seed}', atol=1e-4 + ) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/test_data/ant.xml b/mjx/mujoco/mjx/test_data/ant.xml deleted file mode 100644 index 7417c3ae..00000000 --- a/mjx/mujoco/mjx/test_data/ant.xml +++ /dev/null @@ -1,82 +0,0 @@ - - - diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml new file mode 100644 index 00000000..52eb95d8 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -0,0 +1,53 @@ + + + diff --git a/mjx/mujoco/mjx/test_data/equality.xml b/mjx/mujoco/mjx/test_data/equality.xml deleted file mode 100644 index e5c9d184..00000000 --- a/mjx/mujoco/mjx/test_data/equality.xml +++ /dev/null @@ -1,71 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/mjx/mujoco/mjx/test_data/humanoid.xml b/mjx/mujoco/mjx/test_data/humanoid.xml deleted file mode 100644 index 2d7158ee..00000000 --- a/mjx/mujoco/mjx/test_data/humanoid.xml +++ /dev/null @@ -1,109 +0,0 @@ - - - - - - - - diff --git a/mjx/mujoco/mjx/test_data/pendula.xml b/mjx/mujoco/mjx/test_data/pendula.xml index 2363a3f8..0dc476ab 100644 --- a/mjx/mujoco/mjx/test_data/pendula.xml +++ b/mjx/mujoco/mjx/test_data/pendula.xml @@ -18,6 +18,8 @@ + + @@ -26,45 +28,49 @@ - + + - + + - + + - - + + - + - + - + - + + @@ -72,14 +78,14 @@ - + - + - + @@ -89,14 +95,21 @@ - + - + - + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/ray.xml b/mjx/mujoco/mjx/test_data/ray.xml new file mode 100644 index 00000000..a6424ec4 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/ray.xml @@ -0,0 +1,16 @@ + + + + + + + + + + + + + + + + diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index a385e9af..b0672ead 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name="mujoco-mjx" -version = "3.0.2" +version = "3.1.2" authors = [ {name = "Google DeepMind", email = "mujoco@deepmind.com"}, ] @@ -31,13 +31,13 @@ dependencies = [ "etils[epath]", "jax", "jaxlib", - "mujoco>=3.0.2.dev0", + "mujoco>=3.1.2.dev0", "scipy", "trimesh", ] [project.urls] Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx" -Documentation = "https://mujoco.readthedocs.io/en/3.0.2" +Documentation = "https://mujoco.readthedocs.io/en/3.1.2" Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx" -Changelog = "https://mujoco.readthedocs.io/en/3.0.2/changelog.html" +Changelog = "https://mujoco.readthedocs.io/en/3.1.2/changelog.html" diff --git a/mjx/requirements.txt b/mjx/requirements.txt index d5cab687..9ac5f83d 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -65,7 +65,8 @@ scipy==1.11.3; python_version >= '3.9' \ --hash=sha256:c77da50c9a91e23beb63c2a711ef9e9ca9a2060442757dffee34ea41847d8156 \ --hash=sha256:9ea7f579182d83d00fed0e5c11a4aa5ffe01460444219dedc448a36adf0c3917 \ --hash=sha256:5305792c7110e32ff155aed0df46aa60a60fc6e52cd4ee02cdeb67eaccd5356e \ - --hash=sha256:a63d1ec9cadecce838467ce0631c17c15c7197ae61e49429434ba01d618caa83 + --hash=sha256:a63d1ec9cadecce838467ce0631c17c15c7197ae61e49429434ba01d618caa83 \ + --hash=sha256:715c9966eb8906bc67e450e962bd07a5254420077178f98258904da4004a172f setuptools==68.2.2 \ --hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a trimesh==4.0.0 \ diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb index 16aa9c32..8644de41 100644 --- a/mjx/tutorial.ipynb +++ b/mjx/tutorial.ipynb @@ -157,20 +157,25 @@ "source": [ "#@title Import MuJoCo, MJX, and Brax\n", "\n", + "\n", "from datetime import datetime\n", "import functools\n", + "from IPython.display import HTML\n", "import jax\n", "from jax import numpy as jp\n", "import numpy as np\n", - "from typing import Any, Dict, Tuple, Union\n", + "from typing import Any, Dict, Sequence, Tuple, Union\n", "\n", + "from brax import base\n", "from brax import envs\n", "from brax import math\n", "from brax.base import Base, Motion, Transform\n", - "from brax.envs.base import Env, State\n", + "from brax.envs.base import Env, MjxEnv, State\n", + "from brax.mjx.base import State as MjxState\n", "from brax.training.agents.ppo import train as ppo\n", "from brax.training.agents.ppo import networks as ppo_networks\n", - "from brax.io import model\n", + "from brax.io import html, mjcf, model\n", + "\n", "from etils import epath\n", "from flax import struct\n", "from matplotlib import pyplot as plt\n", @@ -180,6 +185,211 @@ "from mujoco import mjx\n" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "Nj4-Xmx4DFaq" + }, + "source": [ + "# Introduction to MJX\n", + "\n", + "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we will demonstrate how to train RL policies with MJX.\n", + "\n", + "Before we get into hefty RL workloads, let's get started with a simpler example! The entrypoint into MJX is through MuJoCo, so first we load a MuJoCo model:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bNus3mbbDz6a" + }, + "outputs": [], + "source": [ + "xml = \"\"\"\n", + "\u003cmujoco\u003e\n", + " \u003cworldbody\u003e\n", + " \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n", + " \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n", + " \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n", + " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n", + " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n", + " \u003c/body\u003e\n", + " \u003c/worldbody\u003e\n", + "\u003c/mujoco\u003e\n", + "\"\"\"\n", + "\n", + "# Make model, data, and renderer\n", + "mj_model = mujoco.MjModel.from_xml_string(xml)\n", + "mj_data = mujoco.MjData(mj_model)\n", + "renderer = mujoco.Renderer(mj_model)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Po5oykJbFQbj" + }, + "source": [ + "Next we take the MuJoCo model and data, and place them on the GPU device using MJX." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "TSpoOWqeEC3P" + }, + "outputs": [], + "source": [ + "mjx_model = mjx.put_model(mj_model)\n", + "mjx_data = mjx.put_data(mj_model, mj_data)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6rxMMSs4OJJf" + }, + "source": [ + "Below, we print the `qpos` from MuJoCo and MJX. Notice that the `qpos` for the mjData is a numpy array living on the CPU, while the `qpos` for `mjx.Data` is a JAX Array living on the GPU device." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZOD582pfOLP-" + }, + "outputs": [], + "source": [ + "print(mj_data.qpos, type(mj_data.qpos))\n", + "print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ZShF9-o_JLm3" + }, + "source": [ + "Let's run the simulation in MuJoCo and render the trajectory. This example is taken from the [MuJoCo tutorial](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "HDlPlX05I3m-" + }, + "outputs": [], + "source": [ + "# enable joint visualization option:\n", + "scene_option = mujoco.MjvOption()\n", + "scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True\n", + "\n", + "duration = 3.8 # (seconds)\n", + "framerate = 60 # (Hz)\n", + "\n", + "frames = []\n", + "mujoco.mj_resetData(mj_model, mj_data)\n", + "while mj_data.time \u003c duration:\n", + " mujoco.mj_step(mj_model, mj_data)\n", + " if len(frames) \u003c mj_data.time * framerate:\n", + " renderer.update_scene(mj_data, scene_option=scene_option)\n", + " pixels = renderer.render()\n", + " frames.append(pixels)\n", + "\n", + "# Simulate and display video.\n", + "media.show_video(frames, fps=framerate)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "m70b_RxBJOyd" + }, + "source": [ + "Now let's run the same exact simulation on the GPU device using MJX!\n", + "\n", + "In the example below, we use `mjx.step` instead of `mujoco.mj_step`, and we also [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) the `mjx.step` so that it runs efficiently on the GPU. After each step, we convert the `mjx.Data` back to `mjData` so that we can use the MuJoCo renderer.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Pr29xq0-JRQv" + }, + "outputs": [], + "source": [ + "\n", + "jit_step = jax.jit(mjx.step)\n", + "\n", + "frames = []\n", + "mujoco.mj_resetData(mj_model, mj_data)\n", + "mjx_data = mjx.put_data(mj_model, mj_data)\n", + "while mjx_data.time \u003c duration:\n", + " mjx_data = jit_step(mjx_model, mjx_data)\n", + " if len(frames) \u003c mjx_data.time * framerate:\n", + " mj_data = mjx.get_data(mj_model, mjx_data)\n", + " renderer.update_scene(mj_data, scene_option=scene_option)\n", + " pixels = renderer.render()\n", + " frames.append(pixels)\n", + "\n", + "media.show_video(frames, fps=framerate)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "wXsQ4qO2KO3Q" + }, + "source": [ + "Running single threaded physics simulation on the GPU is not very [efficient](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits). The advantage with MJX is that we can run environments in parallel on a hardware accelerated device. Let's try it out!\n", + "\n", + "In the example below, we create 4096 copies of the `mjx.Data` and we run the `mjx.step` over the batched data. Since MJX is implemented in JAX, we take advantage of [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html) to run the `mjx.step` in parallel over all `mjx.Data`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "rrdrcKRVK6w9" + }, + "outputs": [], + "source": [ + "rng = jax.random.PRNGKey(0)\n", + "rng = jax.random.split(rng, 4096)\n", + "batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)\n", + "\n", + "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n", + "batch = jit_step(mjx_model, batch)\n", + "\n", + "print(batch.qpos)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "x4lL220cOj0q" + }, + "source": [ + "We can copy the batched `mjx.Data` back to MuJoCo like we did before:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Jtz7j1PDOnw5" + }, + "outputs": [], + "source": [ + "batched_mj_data = mjx.get_data(mj_model, batch)\n", + "print([d.qpos for d in batched_mj_data])" + ] + }, { "cell_type": "markdown", "metadata": { @@ -187,145 +397,10 @@ }, "source": [ "# Training a Policy with MJX\n", - "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate how to train RL policies with MJX.\n", "\n", - "First, we implement an environment `State` so that we can plug into the [Brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "7DQ_rW4CkIB_" - }, - "outputs": [], - "source": [ - "#@title State\n", + "Running large batch physics simulation is useful for training RL policies. Here we demonstrate training RL policies with MJX using the RL library from [Brax](https://github.com/google/brax).\n", "\n", - "@struct.dataclass\n", - "class State(Base):\n", - " \"\"\"Environment state for training and inference with brax.\n", - "\n", - " Args:\n", - " pipeline_state: the physics state, mjx.Data\n", - " obs: environment observations\n", - " reward: environment reward\n", - " done: boolean, True if the current episode has terminated\n", - " metrics: metrics that get tracked per environment step\n", - " info: environment variables defined and updated by the environment reset\n", - " and step functions\n", - " \"\"\"\n", - "\n", - " pipeline_state: mjx.Data\n", - " obs: jax.Array\n", - " reward: jax.Array\n", - " done: jax.Array\n", - " metrics: Dict[str, jax.Array] = struct.field(default_factory=dict)\n", - " info: Dict[str, Any] = struct.field(default_factory=dict)\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "acpXtDLNXLV9" - }, - "source": [ - "\n", - "Next, we implement `MjxEnv`, an environment class we'll use through the notebook. `MjxEnv` initializes a `mjx.Model` and `mjx.Data` object. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`.\n", - "\n", - "`MjxEnv` also inherits from `brax.envs.base.Env` which allows us to use the training agents implemented in brax." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "ccujYeJ5XOhx" - }, - "outputs": [], - "source": [ - "#@title MjxEnv\n", - "\n", - "class MjxEnv(Env):\n", - " \"\"\"API for driving an MJX system for training and inference in brax.\"\"\"\n", - "\n", - " def __init__(\n", - " self,\n", - " mj_model: mujoco.MjModel,\n", - " physics_steps_per_control_step: int = 1,\n", - " ):\n", - " \"\"\"Initializes MjxEnv.\n", - "\n", - " Args:\n", - " mj_model: mujoco.MjModel\n", - " physics_steps_per_control_step: the number of times to step the physics\n", - " pipeline for each environment step\n", - " \"\"\"\n", - " self.model = mj_model\n", - " self.data = mujoco.MjData(mj_model)\n", - " self.sys = mjx.device_put(mj_model)\n", - " self._physics_steps_per_control_step = physics_steps_per_control_step\n", - "\n", - " def pipeline_init(\n", - " self, qpos: jax.Array, qvel: jax.Array\n", - " ) -\u003e mjx.Data:\n", - " \"\"\"Initializes the physics state.\"\"\"\n", - " data = mjx.device_put(self.data)\n", - " data = data.replace(qpos=qpos, qvel=qvel, ctrl=jp.zeros(self.sys.nu))\n", - " data = mjx.forward(self.sys, data)\n", - " return data\n", - "\n", - " def pipeline_step(\n", - " self, data: mjx.Data, ctrl: jax.Array\n", - " ) -\u003e mjx.Data:\n", - " \"\"\"Takes a physics step using the physics pipeline.\"\"\"\n", - " def f(data, _):\n", - " data = data.replace(ctrl=ctrl)\n", - " return (\n", - " mjx.step(self.sys, data),\n", - " None,\n", - " )\n", - " data, _ = jax.lax.scan(f, data, (), self._physics_steps_per_control_step)\n", - " return data\n", - "\n", - " @property\n", - " def dt(self) -\u003e jax.Array:\n", - " \"\"\"The timestep used for each env step.\"\"\"\n", - " return self.sys.opt.timestep * self._physics_steps_per_control_step\n", - "\n", - " @property\n", - " def observation_size(self) -\u003e int:\n", - " rng = jax.random.PRNGKey(0)\n", - " reset_state = self.unwrapped.reset(rng)\n", - " return reset_state.obs.shape[-1]\n", - "\n", - " @property\n", - " def action_size(self) -\u003e int:\n", - " return self.sys.nu\n", - "\n", - " @property\n", - " def backend(self) -\u003e str:\n", - " return 'mjx'\n", - "\n", - " def _pos_vel(\n", - " self, data: mjx.Data\n", - " ) -\u003e Tuple[Transform, Motion]:\n", - " \"\"\"Returns 6d spatial transform and 6d velocity for all bodies.\"\"\"\n", - " x = Transform(pos=data.xpos[1:, :], rot=data.xquat[1:, :])\n", - " cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])\n", - " offset = data.xpos[1:, :] - data.subtree_com[\n", - " self.model.body_rootid[np.arange(1, self.model.nbody)]]\n", - " xd = Transform.create(pos=offset).vmap().do(cvel)\n", - " return x, xd\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "iPlFu4CiIgBN" - }, - "source": [ - "Finally we can implement a real environment. We choose to first implement the Humanoid environment. Notice that `reset` initializes a `State`, and `step` steps through the physics step and reward logic. The reward and stepping logic train the Humanoid to run forwards." + "Below, we implement the classic Humanoid environment using MJX and Brax. We inherit from the `MjxEnv` implementation in Brax so that we can step the physics with MJX while training with Brax RL implementations.\n" ] }, { @@ -361,10 +436,10 @@ " mj_model.opt.ls_iterations = 6\n", "\n", " physics_steps_per_control_step = 5\n", - " kwargs['physics_steps_per_control_step'] = kwargs.get(\n", - " 'physics_steps_per_control_step', physics_steps_per_control_step)\n", + " kwargs['n_frames'] = kwargs.get(\n", + " 'n_frames', physics_steps_per_control_step)\n", "\n", - " super().__init__(mj_model=mj_model, **kwargs)\n", + " super().__init__(model=mj_model, **kwargs)\n", "\n", " self._forward_reward_weight = forward_reward_weight\n", " self._ctrl_cost_weight = ctrl_cost_weight\n", @@ -390,7 +465,7 @@ "\n", " data = self.pipeline_init(qpos, qvel)\n", "\n", - " obs = self._get_obs(data, jp.zeros(self.sys.nu))\n", + " obs = self._get_obs(data.data, jp.zeros(self.sys.nu))\n", " reward, done, zero = jp.zeros(3)\n", " metrics = {\n", " 'forward_reward': zero,\n", @@ -410,16 +485,14 @@ " data0 = state.pipeline_state\n", " data = self.pipeline_step(data0, action)\n", "\n", - " com_before = data0.subtree_com[1]\n", - " com_after = data.subtree_com[1]\n", + " com_before = data0.data.subtree_com[1]\n", + " com_after = data.data.subtree_com[1]\n", " velocity = (com_after - com_before) / self.dt\n", " forward_reward = self._forward_reward_weight * velocity[0]\n", "\n", " min_z, max_z = self._healthy_z_range\n", - " is_healthy = jp.where(data.qpos[2] \u003c min_z, x=0.0, y=1.0)\n", - " is_healthy = jp.where(\n", - " data.qpos[2] \u003e max_z, x=0.0, y=is_healthy\n", - " )\n", + " is_healthy = jp.where(data.q[2] \u003c min_z, 0.0, 1.0)\n", + " is_healthy = jp.where(data.q[2] \u003e max_z, 0.0, is_healthy)\n", " if self._terminate_when_unhealthy:\n", " healthy_reward = self._healthy_reward\n", " else:\n", @@ -427,7 +500,7 @@ "\n", " ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n", "\n", - " obs = self._get_obs(data, action)\n", + " obs = self._get_obs(data.data, action)\n", " reward = forward_reward + healthy_reward - ctrl_cost\n", " done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n", " state.metrics.update(\n", @@ -494,31 +567,7 @@ "\n", "# define the jit reset/step functions\n", "jit_reset = jax.jit(env.reset)\n", - "jit_step = jax.jit(env.step)\n", - "\n", - "# instantiate the renderer\n", - "renderer = mujoco.Renderer(env.model)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "9f2ME2WbA5Ip" - }, - "outputs": [], - "source": [ - "#@title Define a render utility function\n", - "\n", - "def get_image(state: State, camera: str) -\u003e np.ndarray:\n", - " \"\"\"Renders the environment state.\"\"\"\n", - " d = mujoco.MjData(env.model)\n", - " # write the mjx.Data into an mjData object\n", - " mjx.device_get_into(d, state.pipeline_state)\n", - " mujoco.mj_forward(env.model, d)\n", - " # use the mjData object to update the renderer\n", - " renderer.update_scene(d, camera=camera)\n", - " return renderer.render()\n" + "jit_step = jax.jit(env.step)\n" ] }, { @@ -531,17 +580,15 @@ "source": [ "# initialize the state\n", "state = jit_reset(jax.random.PRNGKey(0))\n", - "rollout = [state]\n", - "images = [get_image(state, camera='side')]\n", + "rollout = [state.pipeline_state]\n", "\n", "# grab a trajectory\n", "for i in range(10):\n", " ctrl = -0.1 * jp.ones(env.sys.nu)\n", " state = jit_step(state, ctrl)\n", - " rollout.append(state)\n", - " images.append(get_image(state, camera='side'))\n", + " rollout.append(state.pipeline_state)\n", "\n", - "media.show_video(images, fps=1.0 / env.dt)" + "media.show_video(env.render(rollout, camera='side'), fps=1.0 / env.dt)" ] }, { @@ -552,7 +599,7 @@ "source": [ "## Train Humanoid Policy\n", "\n", - "Let's finally train a policy with PPO to make the Humanoid run forwards. Training takes about 13-14 minutes on a Tesla V100 GPU." + "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 9-10 minutes on a Tesla A100 GPU." ] }, { @@ -606,7 +653,7 @@ "id": "YYIch0HEApBx" }, "source": [ - "## Save and Load Policy\n", + "\u003c!-- ## Save and Load Policy --\u003e\n", "\n", "We can save and load the policy using the brax model API." ] @@ -675,8 +722,7 @@ "# initialize the state\n", "rng = jax.random.PRNGKey(0)\n", "state = jit_reset(rng)\n", - "rollout = [state]\n", - "images = [get_image(state, camera='side')]\n", + "rollout = [state.pipeline_state]\n", "\n", "# grab a trajectory\n", "n_steps = 500\n", @@ -686,14 +732,12 @@ " act_rng, rng = jax.random.split(rng)\n", " ctrl, _ = jit_inference_fn(state.obs, act_rng)\n", " state = jit_step(state, ctrl)\n", - " rollout.append(state)\n", - " if i % render_every == 0:\n", - " images.append(get_image(state, camera='side'))\n", + " rollout.append(state.pipeline_state)\n", "\n", " if state.done:\n", " break\n", "\n", - "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" + "media.show_video(env.render(rollout[::render_every], camera='side'), fps=1.0 / env.dt / render_every)" ] }, { @@ -704,7 +748,7 @@ "source": [ "# MJX Policy in MuJoCo\n", "\n", - "Note that we can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo." + "We can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo." ] }, { @@ -715,7 +759,7 @@ }, "outputs": [], "source": [ - "mj_model = eval_env.model\n", + "mj_model = eval_env._model\n", "mj_data = mujoco.MjData(mj_model)\n", "\n", "renderer = mujoco.Renderer(mj_model)\n", @@ -725,11 +769,11 @@ "for i in range(n_steps):\n", " act_rng, rng = jax.random.split(rng)\n", "\n", - " obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n", + " obs = eval_env._get_obs(mjx.put_data(mj_model, mj_data), ctrl)\n", " ctrl, _ = jit_inference_fn(obs, act_rng)\n", "\n", " mj_data.ctrl = ctrl\n", - " for _ in range(eval_env._physics_steps_per_control_step):\n", + " for _ in range(eval_env._n_frames):\n", " mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n", "\n", " if i % render_every == 0:\n", @@ -745,7 +789,7 @@ "id": "65mIPj6DQNNa" }, "source": [ - "# Domain Randomization\n", + "# Training a Policy with Domain Randomization\n", "\n", "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias." ] @@ -768,7 +812,7 @@ " friction = sys.geom_friction.at[:, 0].set(friction)\n", " # actuator\n", " _, key = jax.random.split(key, 2)\n", - " gain_range = (-10, -5)\n", + " gain_range = (-5, 5)\n", " param = jax.random.uniform(\n", " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n", " ) + sys.actuator_gainprm[:, 0]\n", @@ -800,7 +844,7 @@ "id": "gnsZo-GWSYYj" }, "source": [ - "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjModel` along with a dictionary specifying the axes that are batched." + "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjx.Model` along with a dictionary specifying the axes that are batched." ] }, { @@ -830,7 +874,18 @@ "source": [ "## Quadruped Env\n", "\n", - "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0) and an environment that trains a joystick policy." + "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour vb Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_vb) from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie). We implement an environment that trains a joystick policy with Brax." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VfyK73gtRXid" + }, + "outputs": [], + "source": [ + "!git clone https://github.com/google-deepmind/mujoco_menagerie" ] }, { @@ -841,7 +896,7 @@ }, "outputs": [], "source": [ - "#@title Barkour v0 Quadruped Env\n", + "#@title Barkour vb Quadruped Env\n", "\n", "def get_config():\n", " \"\"\"Returns reward config for barkour quadruped environment.\"\"\"\n", @@ -871,11 +926,10 @@ " # Penalize non-zero roll and pitch angles. L2 penalty.\n", " orientation=-5.0,\n", " # L2 regularization of joint torques, |tau|^2.\n", - " # torques=-0.0002,\n", - " torques=-0.002,\n", + " torques=-0.0002,\n", " # Penalize the change in the action and encourage smooth\n", " # actions. L2 regularization |action - last_action|^2\n", - " action_rate=-0.1,\n", + " action_rate=-0.01,\n", " # Encourage long swing steps. However, it does not\n", " # encourage high clearances.\n", " feet_air_time=0.2,\n", @@ -895,7 +949,10 @@ " return default_config\n", "\n", " default_config = config_dict.ConfigDict(\n", - " dict(rewards=get_default_rewards_config(),))\n", + " dict(\n", + " rewards=get_default_rewards_config(),\n", + " )\n", + " )\n", "\n", " return default_config\n", "\n", @@ -906,43 +963,69 @@ " def __init__(\n", " self,\n", " obs_noise: float = 0.05,\n", - " action_scale: float=0.3,\n", + " action_scale: float = 0.3,\n", + " kick_vel: float = 0.05,\n", " **kwargs,\n", " ):\n", - " path = epath.Path(epath.resource_path('mujoco')) / (\n", - " 'mjx/benchmark/model/barkour_v0/assets'\n", + " path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n", + " self._dt = 0.02 # this environment is 50 fps\n", + " self.brax_sys = mjcf.load(path).replace(dt=self._dt)\n", + " model = self.brax_sys.get_model()\n", + " model.opt.timestep = 0.004\n", + "\n", + " # override menagerie params for smoother policy\n", + " model.dof_damping[6:] = 0.5239\n", + " model.actuator_gainprm[:, 0] = 35.0\n", + " model.actuator_biasprm[:, 1] = -35.0\n", + "\n", + " n_frames = kwargs.pop('n_frames', int(self._dt / model.opt.timestep))\n", + " super().__init__(model=model, n_frames=n_frames)\n", + "\n", + " self.reward_config = get_config()\n", + " # set custom from kwargs\n", + " for k, v in kwargs.items():\n", + " if k.endswith('_scale'):\n", + " self.reward_config.rewards.scales[k[:-6]] = v\n", + "\n", + " self._torso_idx = mujoco.mj_name2id(\n", + " model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n", " )\n", - " mj_model = mujoco.MjModel.from_xml_path(\n", - " (path / 'barkour_v0_mjx.xml').as_posix())\n", - " mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n", - " mj_model.opt.iterations = 4\n", - " mj_model.opt.ls_iterations = 6\n", - "\n", - " physics_steps_per_control_step = 10\n", - " kwargs['physics_steps_per_control_step'] = kwargs.get(\n", - " 'physics_steps_per_control_step', physics_steps_per_control_step)\n", - " super().__init__(mj_model=mj_model, **kwargs)\n", - "\n", - " self.torso_idx = 1\n", " self._action_scale = action_scale\n", " self._obs_noise = obs_noise\n", - " self._reset_horizon = 500\n", - " self._feet_index = jp.array([3, 6, 9, 12])\n", - " # local positions for each foot\n", - " self._feet_pos = jp.array([\n", - " [-0.191284, -0.0191638, 0.013],\n", - " [-0.191284, -0.0191638, -0.013],\n", - " [-0.191284, -0.0191638, 0.013],\n", - " [-0.191284, -0.0191638, -0.013],\n", - " ])\n", - " self._init_q = mj_model.keyframe('standing').qpos\n", - " self._default_ap_pose = mj_model.keyframe('standing').qpos[7:]\n", - " self.reward_config = get_config()\n", - " self.lowers = self._default_ap_pose - jp.array([0.2, 0.8, 0.8] * 4)\n", - " self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n", + " self._kick_vel = kick_vel\n", + " self._init_q = jp.array(model.keyframe('home').qpos)\n", + " self._default_pose = model.keyframe('home').qpos[7:]\n", + " self.lowers = jp.array([-0.7, -1.0, 0.05] * 4)\n", + " self.uppers = jp.array([0.52, 2.1, 2.1] * 4)\n", + " feet_site = [\n", + " 'foot_front_left',\n", + " 'foot_hind_left',\n", + " 'foot_front_right',\n", + " 'foot_hind_right',\n", + " ]\n", + " feet_site_id = [\n", + " mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n", + " for f in feet_site\n", + " ]\n", + " assert not any(id_ == -1 for id_ in feet_site_id), 'Site not found.'\n", + " self._feet_site_id = np.array(feet_site_id)\n", + " lower_leg_body = [\n", + " 'lower_leg_front_left',\n", + " 'lower_leg_hind_left',\n", + " 'lower_leg_front_right',\n", + " 'lower_leg_hind_right',\n", + " ]\n", + " lower_leg_body_id = [\n", + " mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n", + " for l in lower_leg_body\n", + " ]\n", + " assert not any(id_ == -1 for id_ in lower_leg_body_id), 'Body not found.'\n", + " self._lower_leg_body_id = np.array(lower_leg_body_id)\n", + " self._foot_radius = 0.0175\n", + " self._nv = model.nv\n", "\n", " def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n", - " lin_vel_x = [-0.6, 1.0] # min max [m/s]\n", + " lin_vel_x = [-0.6, 1.5] # min max [m/s]\n", " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n", " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n", "\n", @@ -959,216 +1042,156 @@ " new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n", " return new_cmd\n", "\n", - " def reset(self, rng: jax.Array) -\u003e State:\n", + " def reset(self, rng: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n", " rng, key = jax.random.split(rng)\n", "\n", - " qpos = jp.array(self._init_q)\n", - " qvel = jp.zeros(self.model.nv)\n", - " new_cmd = self.sample_command(key)\n", - " data = self.pipeline_init(qpos, qvel)\n", + " pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n", "\n", " state_info = {\n", " 'rng': rng,\n", " 'last_act': jp.zeros(12),\n", " 'last_vel': jp.zeros(12),\n", - " 'last_contact_buffer': jp.zeros((20, 4), dtype=bool),\n", - " 'command': new_cmd,\n", + " 'command': self.sample_command(key),\n", " 'last_contact': jp.zeros(4, dtype=bool),\n", " 'feet_air_time': jp.zeros(4),\n", - " 'obs_history': jp.zeros(15 * 31),\n", - " 'reward_tuple': {\n", - " 'tracking_lin_vel': 0.0,\n", - " 'tracking_ang_vel': 0.0,\n", - " 'lin_vel_z': 0.0,\n", - " 'ang_vel_xy': 0.0,\n", - " 'orientation': 0.0,\n", - " 'torque': 0.0,\n", - " 'action_rate': 0.0,\n", - " 'stand_still': 0.0,\n", - " 'feet_air_time': 0.0,\n", - " 'foot_slip': 0.0,\n", - " },\n", + " 'rewards': {k: 0.0 for k in self.reward_config.rewards.scales.keys()},\n", + " 'kick': jp.array([0.0, 0.0]),\n", " 'step': 0,\n", " }\n", "\n", - " x, xd = self._pos_vel(data)\n", - " obs = self._get_obs(data.qpos, x, xd, state_info)\n", + " obs_history = jp.zeros(15 * 31) # store 15 steps of history\n", + " obs = self._get_obs(pipeline_state, state_info, obs_history)\n", " reward, done = jp.zeros(2)\n", " metrics = {'total_dist': 0.0}\n", - " for k in state_info['reward_tuple']:\n", - " metrics[k] = state_info['reward_tuple'][k]\n", - " state = State(data, obs, reward, done, metrics, state_info)\n", + " for k in state_info['rewards']:\n", + " metrics[k] = state_info['rewards'][k]\n", + " state = State(pipeline_state, obs, reward, done, metrics, state_info) # pytype: disable=wrong-arg-types\n", " return state\n", "\n", - " def step(self, state: State, action: jax.Array) -\u003e State:\n", - " rng, rng_noise, cmd_rng = jax.random.split(\n", - " state.info['rng'], 3\n", - " )\n", + " def step(self, state: State, action: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n", + " rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n", + "\n", + " # kick\n", + " push_interval = 10\n", + " kick_theta = jax.random.uniform(kick_noise_2, maxval=2 * jp.pi)\n", + " kick = jp.array([jp.cos(kick_theta), jp.sin(kick_theta)])\n", + " kick *= jp.mod(state.info['step'], push_interval) == 0\n", + " qvel = state.pipeline_state.data.qvel # pytype: disable=attribute-error\n", + " qvel = qvel.at[:2].set(kick * self._kick_vel + qvel[:2])\n", + " state = state.tree_replace({'pipeline_state.data.qvel': qvel})\n", "\n", " # physics step\n", - " cur_action = jp.array(action)\n", - " action = action[:12] * self._action_scale\n", - " motor_targets = jp.clip(\n", - " action + self._default_ap_pose, self.lowers, self.uppers\n", - " )\n", - " data = self.pipeline_step(state.pipeline_state, motor_targets)\n", + " motor_targets = self._default_pose + action * self._action_scale\n", + " motor_targets = jp.clip(motor_targets, self.lowers, self.uppers)\n", + " pipeline_state = self.pipeline_step(state.pipeline_state, motor_targets)\n", + " x, xd = pipeline_state.x, pipeline_state.xd\n", "\n", " # observation data\n", - " x, xd = self._pos_vel(data)\n", - " obs = self._get_obs(data.qpos, x, xd, state.info)\n", - " obs_noise = self._obs_noise * jax.random.uniform(\n", - " rng_noise, obs.shape, minval=-1, maxval=1)\n", - " qpos, qvel = data.qpos, data.qvel\n", - " joint_angles = qpos[7:]\n", - " joint_vel = qvel[6:]\n", + " obs = self._get_obs(pipeline_state, state.info, state.obs)\n", + " joint_angles = pipeline_state.q[7:]\n", + " joint_vel = pipeline_state.qd[6:]\n", "\n", " # foot contact data based on z-position\n", - " foot_contact = 0.017 - self._get_feet_pos_vel(x, xd)[0][:, 2]\n", - " contact = foot_contact \u003e -1e-3 # a mm or less off the floor\n", - " contact_filt_mm = jp.logical_or(contact, state.info['last_contact'])\n", - " contact_filt_cm = jp.logical_or(\n", - " foot_contact \u003e -1e-2, state.info['last_contact']\n", - " )\n", - " first_contact = (state.info['feet_air_time'] \u003e 0) * (contact_filt_mm)\n", + " foot_pos = pipeline_state.data.site_xpos[self._feet_site_id] # pytype: disable=attribute-error\n", + " foot_contact_z = foot_pos[:, 2] - self._foot_radius\n", + " contact = foot_contact_z \u003c 1e-3 # a mm or less off the floor\n", + " contact_filt_mm = contact | state.info['last_contact']\n", + " contact_filt_cm = (foot_contact_z \u003c 3e-2) | state.info['last_contact']\n", + " first_contact = (state.info['feet_air_time'] \u003e 0) * contact_filt_mm\n", " state.info['feet_air_time'] += self.dt\n", "\n", + " # done if joint limits are reached or robot is falling\n", + " up = jp.array([0.0, 0.0, 1.0])\n", + " done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) \u003c 0\n", + " done |= jp.any(joint_angles \u003c self.lowers)\n", + " done |= jp.any(joint_angles \u003e self.uppers)\n", + " done |= pipeline_state.x.pos[self._torso_idx - 1, 2] \u003c 0.18\n", + "\n", " # reward\n", - " reward_tuple = {\n", + " rewards = {\n", " 'tracking_lin_vel': (\n", " self._reward_tracking_lin_vel(state.info['command'], x, xd)\n", - " * self.reward_config.rewards.scales.tracking_lin_vel\n", " ),\n", " 'tracking_ang_vel': (\n", " self._reward_tracking_ang_vel(state.info['command'], x, xd)\n", - " * self.reward_config.rewards.scales.tracking_ang_vel\n", " ),\n", - " 'lin_vel_z': (\n", - " self._reward_lin_vel_z(xd)\n", - " * self.reward_config.rewards.scales.lin_vel_z\n", + " 'lin_vel_z': self._reward_lin_vel_z(xd),\n", + " 'ang_vel_xy': self._reward_ang_vel_xy(xd),\n", + " 'orientation': self._reward_orientation(x),\n", + " 'torques': self._reward_torques(pipeline_state.data.qfrc_actuator), # pytype: disable=attribute-error\n", + " 'action_rate': self._reward_action_rate(action, state.info['last_act']),\n", + " 'stand_still': self._reward_stand_still(\n", + " state.info['command'], joint_angles,\n", " ),\n", - " 'ang_vel_xy': (\n", - " self._reward_ang_vel_xy(xd)\n", - " * self.reward_config.rewards.scales.ang_vel_xy\n", - " ),\n", - " 'orientation': (\n", - " self._reward_orientation(x)\n", - " * self.reward_config.rewards.scales.orientation\n", - " ),\n", - " 'torque': (\n", - " self._reward_torques(data.qfrc_actuator)\n", - " * self.reward_config.rewards.scales.torques\n", - " ),\n", - " 'action_rate': (\n", - " self._reward_action_rate(cur_action, state.info['last_act'])\n", - " * self.reward_config.rewards.scales.action_rate\n", - " ),\n", - " 'stand_still': (\n", - " self._reward_stand_still(\n", - " state.info['command'], joint_angles, self._default_ap_pose\n", - " )\n", - " * self.reward_config.rewards.scales.stand_still\n", - " ),\n", - " 'feet_air_time': (\n", - " self._reward_feet_air_time(\n", - " state.info['feet_air_time'],\n", - " first_contact,\n", - " state.info['command'],\n", - " )\n", - " * self.reward_config.rewards.scales.feet_air_time\n", - " ),\n", - " 'foot_slip': (\n", - " self._reward_foot_slip(x, xd, contact_filt_cm)\n", - " * self.reward_config.rewards.scales.foot_slip\n", + " 'feet_air_time': self._reward_feet_air_time(\n", + " state.info['feet_air_time'],\n", + " first_contact,\n", + " state.info['command'],\n", " ),\n", + " 'foot_slip': self._reward_foot_slip(pipeline_state, contact_filt_cm),\n", + " 'termination': self._reward_termination(done, state.info['step']),\n", " }\n", - " reward = sum(reward_tuple.values())\n", - " reward = jp.clip(reward * self.dt, 0.0, 10000.0)\n", + " rewards = {\n", + " k: v * self.reward_config.rewards.scales[k] for k, v in rewards.items()\n", + " }\n", + " reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)\n", "\n", " # state management\n", - " state.info['last_act'] = cur_action\n", + " state.info['kick'] = kick\n", + " state.info['last_act'] = action\n", " state.info['last_vel'] = joint_vel\n", " state.info['feet_air_time'] *= ~contact_filt_mm\n", " state.info['last_contact'] = contact\n", - " state.info['last_contact_buffer'] = jp.roll(\n", - " state.info['last_contact_buffer'], 1, axis=0\n", - " )\n", - " state.info['last_contact_buffer'] = (\n", - " state.info['last_contact_buffer'].at[0].set(contact)\n", - " )\n", - " state.info['reward_tuple'] = reward_tuple\n", + " state.info['rewards'] = rewards\n", " state.info['step'] += 1\n", - " state.info.update(rng=rng)\n", + " state.info['rng'] = rng\n", "\n", - " # resetting logic if joint limits are reached or robot is falling\n", - " done = 0.0\n", - " up = jp.array([0.0, 0.0, 1.0])\n", - " done = jp.where(jp.dot(math.rotate(up, x.rot[0]), up) \u003c 0, 1.0, done)\n", - " done = jp.where(jp.logical_or(\n", - " jp.any(joint_angles \u003c .98 * self.lowers),\n", - " jp.any(joint_angles \u003e .98 * self.uppers)), 1.0, done)\n", - " done = jp.where(x.pos[self.torso_idx, 2] \u003c 0.18, 1.0, done)\n", - "\n", - " # termination reward\n", - " reward += jp.where(\n", - " (done == 1.0) \u0026 (state.info['step'] \u003c self._reset_horizon),\n", - " self.reward_config.rewards.scales.termination,\n", - " 0.0,\n", - " )\n", - "\n", - " # when done, sample new command if more than _reset_horizon timesteps\n", - " # achieved\n", + " # sample new command if more than 500 timesteps achieved\n", " state.info['command'] = jp.where(\n", - " (done == 1.0) \u0026 (state.info['step'] \u003e self._reset_horizon),\n", - " self.sample_command(cmd_rng), state.info['command'])\n", + " state.info['step'] \u003e 500,\n", + " self.sample_command(cmd_rng),\n", + " state.info['command'],\n", + " )\n", " # reset the step counter when done\n", " state.info['step'] = jp.where(\n", - " (done == 1.0) | (state.info['step'] \u003e self._reset_horizon), 0,\n", - " state.info['step']\n", + " done | (state.info['step'] \u003e 500), 0, state.info['step']\n", " )\n", "\n", " # log total displacement as a proxy metric\n", - " state.metrics['total_dist'] = math.normalize(x.pos[self.torso_idx])[1]\n", - " for k in state.info['reward_tuple'].keys():\n", - " state.metrics[k] = state.info['reward_tuple'][k]\n", + " state.metrics['total_dist'] = math.normalize(x.pos[self._torso_idx - 1])[1]\n", + " state.metrics.update(state.info['rewards'])\n", "\n", + " done = jp.float32(done)\n", " state = state.replace(\n", - " pipeline_state=data, obs=obs + obs_noise, reward=reward,\n", - " done=done)\n", + " pipeline_state=pipeline_state, obs=obs, reward=reward, done=done\n", + " )\n", " return state\n", "\n", - " def _get_obs(self, qpos: jax.Array, x: Transform, xd: Motion,\n", - " state_info: Dict[str, Any]) -\u003e jax.Array:\n", - " # Get observations:\n", - " # yaw_rate, projected_gravity, command, motor_angles, last_action\n", + " def _get_obs(\n", + " self,\n", + " pipeline_state: base.State,\n", + " state_info: dict[str, Any],\n", + " obs_history: jax.Array,\n", + " ) -\u003e jax.Array:\n", + " inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n", + " local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n", "\n", - " inv_base_orientation = math.quat_inv(x.rot[0])\n", - " local_rpyrate = math.rotate(xd.ang[0], inv_base_orientation)\n", - " cmd = state_info['command']\n", + " obs = jp.concatenate([\n", + " jp.array([local_rpyrate[2]]) * 0.25, # yaw rate\n", + " math.rotate(jp.array([0, 0, -1]), inv_torso_rot), # projected gravity\n", + " state_info['command'] * jp.array([2.0, 2.0, 0.25]), # command\n", + " pipeline_state.q[7:] - self._default_pose, # motor angles\n", + " state_info['last_act'], # last action\n", + " ])\n", "\n", - " obs_list = []\n", - " # yaw rate\n", - " obs_list.append(jp.array([local_rpyrate[2]]) * 0.25)\n", - " # projected gravity\n", - " obs_list.append(\n", - " math.rotate(jp.array([0.0, 0.0, -1.0]), inv_base_orientation))\n", - " # command\n", - " obs_list.append(cmd * jp.array([2.0, 2.0, 0.25]))\n", - " # motor angles\n", - " angles = qpos[7:19]\n", - " obs_list.append(angles - self._default_ap_pose)\n", - " # last action\n", - " obs_list.append(state_info['last_act'])\n", - "\n", - " obs = jp.clip(jp.concatenate(obs_list), -100.0, 100.0)\n", - "\n", - " # stack observations through time\n", - " single_obs_size = len(obs)\n", - " state_info['obs_history'] = jp.roll(\n", - " state_info['obs_history'], single_obs_size\n", + " # clip, noise\n", + " obs = jp.clip(obs, -100.0, 100.0) + self._obs_noise * jax.random.uniform(\n", + " state_info['rng'], obs.shape, minval=-1, maxval=1\n", " )\n", - " state_info['obs_history'] = jp.array(\n", - " state_info['obs_history']).at[:single_obs_size].set(obs)\n", - " return state_info['obs_history']\n", + " # stack observations through time\n", + " obs = jp.roll(obs_history, obs.size).at[:obs.size].set(obs)\n", + "\n", + " return obs\n", "\n", " # ------------ reward functions----------------\n", " def _reward_lin_vel_z(self, xd: Motion) -\u003e jax.Array:\n", @@ -1190,12 +1213,14 @@ " return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n", "\n", " def _reward_action_rate(\n", - " self, act: jax.Array, last_act: jax.Array) -\u003e jax.Array:\n", + " self, act: jax.Array, last_act: jax.Array\n", + " ) -\u003e jax.Array:\n", " # Penalize changes in actions\n", " return jp.sum(jp.square(act - last_act))\n", "\n", " def _reward_tracking_lin_vel(\n", - " self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n", + " self, commands: jax.Array, x: Transform, xd: Motion\n", + " ) -\u003e jax.Array:\n", " # Tracking of linear velocity commands (xy axes)\n", " local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\n", " lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n", @@ -1205,15 +1230,16 @@ " return lin_vel_reward\n", "\n", " def _reward_tracking_ang_vel(\n", - " self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n", + " self, commands: jax.Array, x: Transform, xd: Motion\n", + " ) -\u003e jax.Array:\n", " # Tracking of angular velocity commands (yaw)\n", " base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\n", " ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n", - " return jp.exp(-ang_vel_error/self.reward_config.rewards.tracking_sigma)\n", + " return jp.exp(-ang_vel_error / self.reward_config.rewards.tracking_sigma)\n", "\n", " def _reward_feet_air_time(\n", - " self, air_time: jax.Array, first_contact: jax.Array,\n", - " commands: jax.Array) -\u003e jax.Array:\n", + " self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n", + " ) -\u003e jax.Array:\n", " # Reward air time.\n", " rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n", " rew_air_time *= (\n", @@ -1222,29 +1248,38 @@ " return rew_air_time\n", "\n", " def _reward_stand_still(\n", - " self, commands: jax.Array, joint_angles: jax.Array,\n", - " default_angles: jax.Array) -\u003e jax.Array:\n", + " self,\n", + " commands: jax.Array,\n", + " joint_angles: jax.Array,\n", + " ) -\u003e jax.Array:\n", " # Penalize motion at zero commands\n", - " return jp.sum(jp.abs(joint_angles - default_angles)) * (\n", + " return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n", " math.normalize(commands[:2])[1] \u003c 0.1\n", " )\n", "\n", - " def _get_feet_pos_vel(\n", - " self, x: Transform, xd: Motion) -\u003e Tuple[jax.Array, jax.Array]:\n", - " offset = Transform.create(pos=self._feet_pos)\n", - " pos = x.take(self._feet_index).vmap().do(offset).pos\n", - " vel = offset.vmap().do(xd.take(self._feet_index)).vel\n", - " return pos, vel\n", - "\n", " def _reward_foot_slip(\n", - " self, x: Transform, xd: Motion, contact_filt: jax.Array) -\u003e jax.Array:\n", - " # Get feet velocities\n", - " _, foot_world_vel = self._get_feet_pos_vel(x, xd)\n", - " # Penalize large feet velocity for feet that are in contact with the ground.\n", - " return jp.sum(\n", - " jp.square(foot_world_vel[:, :2]) * contact_filt.reshape((-1, 1))\n", - " )\n", + " self, pipeline_state: base.State, contact_filt: jax.Array\n", + " ) -\u003e jax.Array:\n", + " # get velocities at feet which are offset from lower legs\n", + " # pytype: disable=attribute-error\n", + " pos = pipeline_state.data.site_xpos[self._feet_site_id] # feet position\n", + " feet_offset = pos - pipeline_state.data.xpos[self._lower_leg_body_id]\n", + " # pytype: enable=attribute-error\n", + " offset = base.Transform.create(pos=feet_offset)\n", + " foot_indices = self._lower_leg_body_id - 1 # we got rid of the world body\n", + " foot_vel = offset.vmap().do(pipeline_state.xd.take(foot_indices)).vel\n", "\n", + " # Penalize large feet velocity for feet that are in contact with the ground.\n", + " return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n", + "\n", + " def _reward_termination(self, done: jax.Array, step: jax.Array) -\u003e jax.Array:\n", + " return done \u0026 (step \u003c 500)\n", + "\n", + " def render(\n", + " self, trajectory: List[base.State], camera: str | None = None\n", + " ) -\u003e Sequence[np.ndarray]:\n", + " camera = camera or 'track'\n", + " return super().render(trajectory, camera)\n", "\n", "envs.register_environment('barkour', BarkourEnv)" ] @@ -1258,10 +1293,7 @@ "outputs": [], "source": [ "env_name = 'barkour'\n", - "env = envs.get_environment(env_name)\n", - "\n", - "# re-instantiate the renderer\n", - "renderer = mujoco.Renderer(env.model)" + "env = envs.get_environment(env_name)" ] }, { @@ -1272,7 +1304,7 @@ "source": [ "## Train Policy\n", "\n", - "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes about 14 minutes on a Tesla V100 GPU." + "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes 8-9 minutes on a Tesla A100 GPU." ] }, { @@ -1287,22 +1319,19 @@ " ppo_networks.make_ppo_networks,\n", " policy_hidden_layer_sizes=(128, 128, 128, 128))\n", "train_fn = functools.partial(\n", - " ppo.train,\n", - " num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n", - " episode_length=1000, normalize_observations=True,\n", - " action_repeat=1, unroll_length=20, num_minibatches=8, gae_lambda=0.95,\n", - " num_updates_per_batch=4, discounting=0.99, learning_rate=3e-4,\n", - " entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n", + " ppo.train, num_timesteps=100_000_000, num_evals=10,\n", + " reward_scaling=1, episode_length=1000, normalize_observations=True,\n", + " action_repeat=1, unroll_length=20, num_minibatches=32,\n", + " num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n", + " entropy_cost=1e-2, num_envs=8192, batch_size=256,\n", " network_factory=make_networks_factory,\n", - " num_resets_per_eval=10,\n", " randomization_fn=domain_randomize, seed=0)\n", "\n", - "\n", "x_data = []\n", "y_data = []\n", "ydataerr = []\n", "times = [datetime.now()]\n", - "max_y, min_y = 30, 0\n", + "max_y, min_y = 40, 0\n", "\n", "# Reset environments since internals may be overwritten by tracers from the\n", "# domain randomization function.\n", @@ -1366,7 +1395,6 @@ }, "outputs": [], "source": [ - "\n", "# @markdown Commands **only used for Barkour Env**:\n", "x_vel = 1.0 #@param {type: \"number\"}\n", "y_vel = 0.0 #@param {type: \"number\"}\n", @@ -1378,8 +1406,7 @@ "rng = jax.random.PRNGKey(0)\n", "state = jit_reset(rng)\n", "state.info['command'] = the_command\n", - "rollout = [state]\n", - "images = [get_image(state, camera='track')]\n", + "rollout = [state.pipeline_state]\n", "\n", "# grab a trajectory\n", "n_steps = 500\n", @@ -1389,11 +1416,31 @@ " act_rng, rng = jax.random.split(rng)\n", " ctrl, _ = jit_inference_fn(state.obs, act_rng)\n", " state = jit_step(state, ctrl)\n", - " rollout.append(state)\n", - " if i % render_every == 0:\n", - " images.append(get_image(state, camera='track'))\n", + " rollout.append(state.pipeline_state)\n", "\n", - "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" + "media.show_video(\n", + " eval_env.render(rollout[::render_every], camera='track'),\n", + " fps=1.0 / eval_env.dt / render_every)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "aD6H6WD0915X" + }, + "source": [ + "We can also render the rollout using the Brax renderer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "V7jqv08X95u4" + }, + "outputs": [], + "source": [ + "HTML(html.render(eval_env.brax_sys, rollout))" ] } ], @@ -1401,12 +1448,13 @@ "accelerator": "GPU", "colab": { "gpuClass": "premium", - "gpuType": "V100", + "gpuType": "A100", + "machine_shape": "hm", "private_outputs": true, "provenance": [ { - "file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2", - "timestamp": 1697763780236 + "file_id": "11cFRVCJ8Kn71tlQFbFcw4JzQZ00F8BRG", + "timestamp": 1704355889284 } ], "toc_visible": true diff --git a/model/humanoid/README.md b/model/humanoid/README.md index 28ab7e7a..de4c7069 100644 --- a/model/humanoid/README.md +++ b/model/humanoid/README.md @@ -1,17 +1,34 @@ -Humanoid -======== - -Degrees of Freedom: 27 -Actuators: 21 +# Humanoid This simplified humanoid model, introduced in [1], is designed for bipedal locomotion behaviours. While several variants of it exist in the wild, this version is based on the model in the DeepMind Control Suite [2], which has fairly realistic actuator gains. -[1] [Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization] - (https://doi.org/10.1109/IROS.2012.6386025). +* Degrees of Freedom: 27 +* Actuators: 21 + +

+ +

+ +## Changelog + +* 02-01-2024: Add more keyframes. +* 27-11-2023: Move humanoid geoms to group 1. +* 05-04-2023: Fix typo in texture size. +* 20-09-2022: Use default class for left_upper_arm geom. +* 17-09-2022: Increase offscreen render buffer resolution of the humanoid to 2560x1440. +* 12-09-2022: + * Increased maximum hip flexion. + * Symmetrised shoulder and ankle joints. + * Added hamstring tendons which couple the hip and knee at large flexion values. + * Moved duplicated values into defaults. + * Added two keyframes. + * Improved lighting. + * Changed naming convention. + +## References + +[1] [Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization](https://doi.org/10.1109/IROS.2012.6386025). [2] [DeepMind Control Suite](https://arxiv.org/abs/1801.00690). - - -![humanoid](humanoid.png) diff --git a/model/humanoid/humanoid.xml b/model/humanoid/humanoid.xml index 3f2f7436..013ebe23 100644 --- a/model/humanoid/humanoid.xml +++ b/model/humanoid/humanoid.xml @@ -37,7 +37,7 @@ - + @@ -233,17 +233,35 @@ left leg arms --> - - + + + + diff --git a/model/plugin/actuator/pid.xml b/model/plugin/actuator/pid.xml new file mode 100644 index 00000000..357ae90a --- /dev/null +++ b/model/plugin/actuator/pid.xml @@ -0,0 +1,68 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/model/plugin/elasticity/trampoline_flex.xml b/model/plugin/elasticity/trampoline_flex.xml index 874e12fa..106e37cc 100644 --- a/model/plugin/elasticity/trampoline_flex.xml +++ b/model/plugin/elasticity/trampoline_flex.xml @@ -44,6 +44,7 @@ radius=".001" mass="10" name="plate" dim="2"> + @@ -52,11 +53,4 @@ - - - - - - - diff --git a/model/plugin/sdf/nutbolt.xml b/model/plugin/sdf/nutbolt.xml index d3d80ff7..51e2b554 100644 --- a/model/plugin/sdf/nutbolt.xml +++ b/model/plugin/sdf/nutbolt.xml @@ -8,7 +8,7 @@ - + @@ -30,7 +30,7 @@ -