From 8c7f6ce5a096c60973f7fe9c934cf217c8708572 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 15 Mar 2023 13:16:45 -0700 Subject: [PATCH] New `implicitfast` integrator and sparse RNE derivatives for `implicit`. PiperOrigin-RevId: 516910733 Change-Id: I29a0465c0f0b1749a73e3d7e01925200d025ddd0 --- doc/changelog.rst | 28 +- doc/computation.rst | 191 ++++++----- doc/includes/references.h | 13 +- include/mujoco/mjdata.h | 7 +- include/mujoco/mjmodel.h | 6 +- include/mujoco/mjxmacro.h | 4 + introspect/enums.py | 1 + simulate/simulate.cc | 2 +- src/engine/engine_derivative.c | 472 ++++++++++++++++++++------ src/engine/engine_derivative.h | 18 +- src/engine/engine_forward.c | 71 ++-- src/engine/engine_io.c | 192 ++++++++++- src/engine/engine_print.c | 21 ++ src/engine/engine_support.c | 103 ++---- src/engine/engine_support.h | 12 +- src/user/user_model.cc | 36 +- src/user/user_model.h | 3 +- src/user/user_objects.cc | 1 + src/user/user_objects.h | 4 +- src/xml/xml_native_reader.cc | 9 +- test/engine/engine_derivative_test.cc | 92 +++-- test/engine/engine_io_test.cc | 2 + unity/Runtime/Bindings/MjBindings.cs | 5 + 23 files changed, 961 insertions(+), 332 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 4051732c..e911fb2c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -8,17 +8,39 @@ Upcoming version (not yet released) General ^^^^^^^ +- Improvements to implicit integration: + + - The derivatives of the RNE algorithm are now computed using sparse math, leading to significant speed + improvements for large models when using the :ref:`implicit integrator`. + - A new integrator called ``implicitfast`` was added. It is similar to the existing implicit integrator, but skips the + derivatives of Coriolis and centripetal forces. See the :ref:`numerical integration` section for a + detailed motivation and discussion. The implicitfast integrator is recommended for all new models and will + become the default integrator in a future version. + + The table below shows the compute cost of the 627-DoF `humanoid100 + `_ model using different integrators. + "implicit (old)" uses dense RNE derivatives, "implicit (new)" is after the sparsification mentioned above. + Timings were measured on a single core of an AMD 3995WX CPU. + + .. csv-table:: + :header: "timing", "Euler", "implicitfast", "implicit (new)", "implicit (old)" + :widths: auto + :align: left + + one step (ms), 0.5, 0.53, 0.77, 5.0 + steps/second, 2000, 1900, 1300, 200 + - The ``mjd_transitionFD`` function no longer triggers sensor calculation unless explicitly requested. - Corrected the spelling of the ``inteval`` attribute to ``interval`` in the :ref:`mjLROpt` struct. - Mesh texture and normal mappings are now 3-per-triangle rather than 1-per-vertex. Mesh vertices are no longer duplicated in order to circumvent this limitation as they previously were. - The non-zeros for the sparse constraint Jacobian matrix are now precounted and used for matrix memory allocation. - For instance, the constraint Jacobian matrix from the `humanoid100.xml + For instance, the constraint Jacobian matrix from the `humanoid100 `_ model, which previously required ~500,000 ``mjtNum``'s, now only requires ~6000. Very large models can now load and run with the CG solver. - Modified :ref:`mju_error` and :ref:`mju_warning` to be variadic functions (support for printf-like arguments). The functions :ref:`mju_error_i`, :ref:`mju_error_s`, :ref:`mju_warning_i`, and :ref:`mju_warning_s` are now deprecated. -- Implemented a performant :ref:`mju_sqrMatTDSparse` function that doesn't require dense memory allocation. +- Implemented a performant ``mju_sqrMatTDSparse`` function that doesn't require dense memory allocation. Python bindings @@ -28,7 +50,7 @@ Python bindings continuation of an IPython interactive shell session, and is no longer considered experimental feature. - Remove ``efc_`` fields from joint indexers. Since the introduction of arena memory, these fields now have dynamic sizes that change between time steps depending on the number of active constraints, breaking strict correspondence - between joints and `efc_` rows. + between joints and ``efc_`` rows. Simulate ^^^^^^^^ diff --git a/doc/computation.rst b/doc/computation.rst index 5ecc733c..aecd293d 100644 --- a/doc/computation.rst +++ b/doc/computation.rst @@ -421,118 +421,131 @@ Numerical integration MuJoCo computes forward and inverse dynamics in continuous time. The end result of forward dynamics is the joint acceleration :math:`a=\dot{v}` as well as the actuator activations :math:`\dot{w}` when present in the model. These are used to advance the simulation time from :math:`t` to :math:`t+h`, and to update the state variables :math:`q, v, w`. -Three numerical integrators are available: -Semi-implicit Euler method (Euler) - This method updates the activation and velocity with the usual Euler method, however the position is updated using - the *new* velocity; this is known as a "semi-implicit" update: +Four numerical integrators are available, three single-step integrators and the multi-step 4th order Runge-Kutta +integrator. Before describing the integrators, we begin with a general description of single-step Euler integrators: +*explicit*, *semi-implicit* and *implicit-in-velocity*. The *explicit* Euler method is not supported by MuJoCo but has +pedagogical value. It can be written as: - .. math:: - \begin{aligned} - \textrm{activation: }w_{t+h} &= w_t + h \dot{w}_t \\ - \textrm{velocity: }v_{t+h} &= v_t + h a_t \\ - \textrm{position: }q_{t+h} &= q_t + h v_{t+h} - \end{aligned} +.. math:: + :label: eq_explicit - Using the new velocity in the position update improves stability, and is standard in physics simulation. The - summation in the position update generally involves vectors with different dimensionality, and is done by taking into - account the properties of quaternions. + \begin{aligned} + \textrm{activation:}\quad w_{t+h} &= w_t + h \dot{w}_t \\ + \textrm{velocity:}\quad v_{t+h} &= v_t + h a_t \\ + \textrm{position:}\quad q_{t+h} &= q_t + h v_t + \end{aligned} - When joint damping is defined in the model, the Euler method applies a correction to the inertia matrix that - corresponds to implicit integration of damping forces. Let :math:`B` be the diagonal matrix of negative joint damping - coefficients. Letting :math:`\widehat{M} = M-h B`, we then update the velocity using a corrected acceleration as - follows: +Note that in the presence of quaternions, the operation :math:`q_t + h v_t` is more involved than a simple summation, as +the dimensionalities of :math:`q` and :math:`v` are different. The reason explicit Euler is not implemented is that the +following formulation, known as *semi-implicit* Euler is `strictly better `_, and standard in physics simulation: - .. math:: - \begin{aligned} - v_{t+h} &= v_t + h \widehat{M}^{-1} M a_t - \end{aligned} +.. math:: + :label: eq_semimplicit - This correction is temporary, and the original acceleration computed by forward dynamics and saved in ``mjData.qacc`` - is not modified. We explain below how this correction is derived, and why it is a special case of implicit - integration. + \begin{aligned} + v_{t+h} &= v_t + h a_t \\ + q_{t+h} &= q_t + h v_{\color{red}t+h} + \end{aligned} -Implicit-in-velocity Euler method (implicit) - This method approximates the following discrete-time update: +Comparing :eq:`eq_explicit` and :eq:`eq_semimplicit`, we see that in semi-implicit Euler, the position is updated using +the *new* velocity. *Implicit* Euler means: - .. math:: - \begin{aligned} - w_{t+h} &= w_t + h \dot{w}_t \\ - v_{t+h} &= v_t + h a_{t+h} \\ - q_{t+h} &= q_t + h v_{t+h} - \end{aligned} +.. math:: + :label: eq_implicit - Note the acceleration :math:`a_{t+h}=\dot{v}_{t+h}` on the right hand side of the velocity update is evaluated at - :math:`t+h`, making the integrator fully implicit in velocity. When applied to Hamiltonian systems, such integrators - are called symplectic, and there is a large mathematical literature on them. Hamiltonian systems have the property - that certain quantities (symplectic forms) remain constant over time in the true continuous-time system. Symplectic - integrators have the unique property that they preserve the same quantities exactly, despite using a discrete-time - approximation. Few MuJoCo models used in practice correspond to Hamiltonian systems; these are systems without - activation dynamics, contacts, friction, driving forces etc. Nevertheless this type of integrator has appealing - properties in terms of accuracy and stability, achieved at the expense of added computation. It is particularly - effective in systems where instabilities (of the regular Euler integrator) are caused by velocity-dependent forces: - multi-joint pendulums, bodies tumbling through space, systems with substantial lift and drag forces, systems with - substantial damping in tendons and actuators (as well as joints, but joint damping is already handled implicitly in - the regular Euler integrator). + \begin{aligned} + v_{t+h} &= v_t + h a_{\color{red}t+h} \\ + q_{t+h} &= q_t + h v_{t+h} + \end{aligned} - Writing the acceleration (i.e. the forward dynamics) more explicitly as a function of velocity: :math:`a_t=\dot - {v}_t = a(v_t)`, the velocity update we aim to approximate is +Comparing :eq:`eq_semimplicit` and :eq:`eq_implicit`, we see that the acceleration :math:`a_{t+h}=\dot{v}_{t+h}` on the +right hand side of the velocity update is evaluated at the *next time step*. While evaluating the next acceleration +is not possible without stepping, we can use a first-order Taylor expansion to approximate this quantity, and +take a single step of Newton's method. When the expansion is only with respect to velocity (and not position), the +integrator is known as *implicit-in-velocity* Euler. This approach is particularly effective in systems where +instabilities are caused by velocity-dependent forces: multi-joint pendulums, bodies tumbling through space, systems +with lift and drag forces, and systems with substantial damping in tendons and actuators. Writing the +acceleration as a function of velocity: :math:`a_t = a(v_t)`, the velocity update we aim to approximate is - .. math:: v_{t+h} = v_t + h a(v_{t+h}) +.. math:: v_{t+h} = v_t + h a(v_{t+h}) - This is a non-linear equation in the unknown vector :math:`v_{t+h}`. It must be solved numerically at each time step - in order to implement an implicit integrator. We approximate the solution using a single step of Newton's method, - based on the first-order Taylor series expansion of :math:`a(v_{t+h})` around :math:`v_t`. Recall that the forward - dynamics are +This is a non-linear equation in the unknown vector :math:`v_{t+h}` and can be solved numerically at each time step +using a first-order expansion of :math:`a(v_{t+h})` around :math:`v_t`. Recall that the forward dynamics are - .. math:: a(v) = M^{-1} \big(\tau(v) - c(v) + J^T f(v)\big) +.. math:: a(v) = M^{-1} \big(\tau(v) - c(v) + J^T f(v)\big) - Here we will ignore the velocity dependence of the constraint forces :math:`J^T f(v)`, because differentiating them - is complicated, and furthermore MuJoCo uses soft constraints whose integration is very stable even without implicit - integration. Thus we define the approximate derivative +Thus we define the derivative - .. math:: - \begin{aligned} - {\partial a(v) \over \partial v} &\approx M^{-1} D(v) \\ - D(v) &= {\partial\big(\tau(v) - c (v)\big) \over \partial v} - \end{aligned} +.. math:: + \begin{aligned} + {\partial a(v) \over \partial v} &= M^{-1} D \\ + D &\equiv {\partial \over \partial v} \Big(\tau(v) - c (v) + J^T f(v)\Big) + \end{aligned} - MuJoCo computes :math:`D(v)` analytically, by differentating the Recursive Newton-Euler algorithm as well as the code - that computes applied and bias forces. A further approximation is that we restrict :math:`D` to have the same - sparsity pattern as :math:`M`, for computational efficiency (see below). This restriction will exclude damping in - tendons which connect bodies that are on different branches of the kinematic tree. The velocity update corresponding - to Newton's method is as follows. First, we expand the right hand side to first order +The velocity update corresponding to Newton's method is as follows. First, we expand the right hand side to first order - .. math:: - \begin{aligned} - v_{t+h} &= v_t + h a(v_{t+h}) \\ - &= v_t + h a(v_t + v_{t+h}-v_t) \\ - &\approx v_t + h a(v_t) + h M^{-1} D \cdot (v_{t+h}-v_t) - \end{aligned} +.. math:: + \begin{aligned} + v_{t+h} &= v_t + h a(v_{t+h}) \\ + &\approx v_t + h \big( a(v_t) + {\partial a(v) \over \partial v} \cdot (v_{t+h}-v_t) \big) \\ + &= v_t + h a(v_t) + h M^{-1} D \cdot (v_{t+h}-v_t) + \end{aligned} - Premultiplying by :math:`M` and rearranging yields +Premultiplying by :math:`M` and rearranging yields - .. math:: (M-h D) v_{t+h} = (M-h D) v_t + h M a(v_t) +.. math:: (M-h D) v_{t+h} = (M-h D) v_t + h M a(v_t) - Now letting :math:`\widehat{M} = M-h D`, we obtain the implicit update +Solving for :math:`v_{t+h}`, we obtain the implicit-in-velocity update - .. math:: v_{t+h} = v_t + h \widehat{M}^{-1} M a(v_t) +.. math:: + :label: eq_implicit_update - Comparing to the regular Euler method with damping correction, we see that the matrix :math:`B` in the regular Euler - method corresponds to :math:`D` used here, but restricted to joint damping. Since :math:`D` and :math:`M` have the - same sparsity pattern corresponding to the topology of the kinematic tree, reverse-order LU factorization of - :math:`\widehat{M}` is guaranteed to have no fill-in, which is the computational speed-up mentioned above. This - factorization is stored ``mjData.qLU``. It is now also clear why the Euler method uses :math:`B` rather than - :math:`D`: since :math:`B` is diagonal, :math:`\widehat{M}` remains symmetric and can be factorized with Cholesky - rather than LU. + v_{t+h} = v_t + h (M-h D)^{-1} M a(v_t) -4th-order Runge-Kutta method (RK4) - One advantage of our continuous-time formulation is that we can use more advanced integrators such as Runge-Kutta or - multistep methods. The only such integrator currently implemented is the fixed-step 4th-order Runge-Kutta method. We - have observed that for energy-conserving systems it is qualitatively better than the Euler method, both in terms of - stability and accuracy, even when the timestep of the Euler method is decreased by a factor of 4 (so the - computational effort is identical). In the presence of contacts we have not observed significant benefits, although a - more systematic investigation remains to be performed. +All three single-step integrators in MuJoCo use the update :eq:`eq_implicit_update`, with different definitions of the +:math:`D` matrix, which is computed analytically. + +Semi-implicit Euler with implicit joint damping (``Euler``) + For this method :math:`D` only includes derivatives of joint damping. Note that in this case :math:`D` is diagonal + and :math:`M-h D` is symmetric, so Cholesky decomposition can be used. + +Implicit-in-velocity Euler (``implicit``) + For this method :math:`D` includes derivatives of all forces except the constraint forces :math:`J^T f(v)`. These are + currently ignored since even though computing them is possible, it is complicated, and numerical tests show that + including them does not confer much benefit. That said, analytical derivatives of constraint forces are planned for a + future version. Additionally, we restrict :math:`D` to have the same sparsity pattern as :math:`M`, for computational + efficiency. This restriction will exclude damping in tendons which connect bodies that are on different branches of + the kinematic tree. Since :math:`D` is not symmetric, we cannot use Cholesky factorization, but because :math:`D` and + :math:`M` have the same sparsity pattern corresponding to the topology of the kinematic tree, reverse-order LU + factorization of :math:`M-h D` is `guaranteed to have no fill-in + `_. This factorization is stored ``mjData.qLU``. + +Fast implicit-in-velocity (``implicitfast``) + For this method :math:`D` includes derivatives of all forces used in the implicit method, with the exception of the + centripetal and Coriolis forces :math:`c (v)` computed by the RNE algorithm. Additionally, it is symmetrized :math:`D + \leftarrow (D + D^T)/2`. One reason for dropping the RNE derivatives is that they are the most expensive to compute. + Second, these forces change rapidly only at high rotational velocities of complex pendula and spinning bodies, + scenarios which are not common and already well-handled by the Runge-Kutta integrator (see below). Because the RNE + derivatives are also the main source of asymmetry of :math:`D`, by dropping them and symmetrizing, we can use the + faster Cholesky rather than LU decomposition. + + .. tip:: + The implicitfast integrator has similar computational cost to Euler, yet provides increased stability, and is + therefore a strict improvement. It is the recommended integrator and will become the default in a future version. + +4th-order Runge-Kutta (``RK4``) + One advantage of our continuous-time formulation is that we can use higher order integrators such as Runge-Kutta or + multistep methods. The only such integrator currently implemented is the fixed-step `4th-order Runge-Kutta method + `_, though users + can easily implement other integrators by calling :ref:`mj_forward` and integrating accelerations themselves. We have + observed that for energy-conserving systems (`example + `_) RK4 + is qualitatively better than the single-step methods, both in terms of stability and accuracy, even when the timestep + is decreased by a factor of 4 (so the computational effort is identical). In the presence of large velocity- + dependent forces, if the chosen single-step method integrates those forces implicitly, single-step methods can be + significantly more stable than RK4. .. note:: The accuracy and stability of all integrators can be improved by reducing the time step :math:`h` which is stored in diff --git a/doc/includes/references.h b/doc/includes/references.h index 0b39e000..e258436b 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -245,14 +245,17 @@ struct mjData_ { mjtNum* subtree_linvel; // linear velocity of subtree com (nbody x 3) mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) - // computed by mj_Euler + // computed by mj_Euler or mj_implicit mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) - // computed by mj_implicit + // computed by mj_resetData int* D_rownnz; // non-zeros in each row (nv x 1) int* D_rowadr; // address of each row in D_colind (nv x 1) int* D_colind; // column indices of non-zeros (nD x 1) + int* B_rownnz; // non-zeros in each row (nbody x 1) + int* B_rowadr; // address of each row in B_colind (nbody x 1) + int* B_colind; // column indices of non-zeros (nB x 1) // computed by mj_implicit/mj_derivative mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1) @@ -391,7 +394,8 @@ typedef enum mjtTexture_ { // type of texture typedef enum mjtIntegrator_ { // integrator mode mjINT_EULER = 0, // semi-implicit Euler mjINT_RK4, // 4th-order Runge Kutta - mjINT_IMPLICIT // implicit in velocity + mjINT_IMPLICIT, // implicit in velocity + mjINT_IMPLICITFAST // implicit in velocity, no rne derivative } mjtIntegrator; typedef enum mjtCollision_ { // collision mode for selecting geom pairs mjCOL_ALL = 0, // test precomputed and dynamic pairs @@ -791,7 +795,8 @@ struct mjModel_ { // sizes set after mjModel construction (only affect mjData) int nM; // number of non-zeros in sparse inertia matrix - int nD; // number of non-zeros in sparse derivative matrix + int nD; // number of non-zeros in sparse dof-dof matrix + int nB; // number of non-zeros in sparse body-dof matrix int nemax; // number of potential equality-constraint rows int njmax; // number of available rows in constraint Jacobian int nconmax; // number of potential contacts in contact list diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 3f0c22c2..d005a603 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -270,14 +270,17 @@ struct mjData_ { mjtNum* subtree_linvel; // linear velocity of subtree com (nbody x 3) mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) - // computed by mj_Euler + // computed by mj_Euler or mj_implicit mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) - // computed by mj_implicit + // computed by mj_resetData int* D_rownnz; // non-zeros in each row (nv x 1) int* D_rowadr; // address of each row in D_colind (nv x 1) int* D_colind; // column indices of non-zeros (nD x 1) + int* B_rownnz; // non-zeros in each row (nbody x 1) + int* B_rowadr; // address of each row in B_colind (nbody x 1) + int* B_colind; // column indices of non-zeros (nB x 1) // computed by mj_implicit/mj_derivative mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 45258575..90c60fc0 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -125,7 +125,8 @@ typedef enum mjtTexture_ { // type of texture typedef enum mjtIntegrator_ { // integrator mode mjINT_EULER = 0, // semi-implicit Euler mjINT_RK4, // 4th-order Runge Kutta - mjINT_IMPLICIT // implicit in velocity + mjINT_IMPLICIT, // implicit in velocity + mjINT_IMPLICITFAST // implicit in velocity, no rne derivative } mjtIntegrator; @@ -583,7 +584,8 @@ struct mjModel_ { // sizes set after mjModel construction (only affect mjData) int nM; // number of non-zeros in sparse inertia matrix - int nD; // number of non-zeros in sparse derivative matrix + int nD; // number of non-zeros in sparse dof-dof matrix + int nB; // number of non-zeros in sparse body-dof matrix int nemax; // number of potential equality-constraint rows int njmax; // number of available rows in constraint Jacobian int nconmax; // number of potential contacts in contact list diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 04ec723f..82dfe873 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -117,6 +117,7 @@ X( nnames_map ) \ X( nM ) \ X( nD ) \ + X( nB ) \ X( nemax ) \ X( njmax ) \ X( nconmax ) \ @@ -514,6 +515,9 @@ X( int, D_rownnz, nv, 1 ) \ X( int, D_rowadr, nv, 1 ) \ X( int, D_colind, nD, 1 ) \ + X( int, B_rownnz, nbody, 1 ) \ + X( int, B_rowadr, nbody, 1 ) \ + X( int, B_colind, nB, 1 ) \ X( mjtNum, qDeriv, nD, 1 ) \ X( mjtNum, qLU, nD, 1 ) \ X( mjtNum, actuator_force, nu, 1 ) \ diff --git a/introspect/enums.py b/introspect/enums.py index e1752d4e..eaaf8d65 100755 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -120,6 +120,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjINT_EULER', 0), ('mjINT_RK4', 1), ('mjINT_IMPLICIT', 2), + ('mjINT_IMPLICITFAST', 3), ]), )), ('mjtCollision', diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 1a40196f..cfb20801 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -547,7 +547,7 @@ void makephysics(mj::Simulate* sim, int oldstate) { mjuiDef defPhysics[] = { {mjITEM_SECTION, "Physics", oldstate, nullptr, "AP"}, - {mjITEM_SELECT, "Integrator", 2, &(opt.integrator), "Euler\nRK4\nimplicit"}, + {mjITEM_SELECT, "Integrator", 2, &(opt.integrator), "Euler\nRK4\nimplicit\nimplicitfast"}, {mjITEM_SELECT, "Collision", 2, &(opt.collision), "All\nPair\nDynamic"}, {mjITEM_SELECT, "Cone", 2, &(opt.cone), "Pyramidal\nElliptic"}, {mjITEM_SELECT, "Jacobian", 2, &(opt.jacobian), "Dense\nSparse\nAuto"}, diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 9c822a6b..eaf8b1e1 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -358,6 +358,7 @@ void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) { break; case mjINT_IMPLICIT: + case mjINT_IMPLICITFAST: mj_implicitSkip(m, d, skipstage >= mjSTAGE_VEL); break; @@ -370,10 +371,11 @@ void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) { -//------------------------- derivatives of component functions ------------------------------------- +//------------------------- dense derivatives of component functions ------------------------------- +// no longer used, for comparison only -// derivative of cvel, cdof_dot w.r.t qvel -static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* Dcdofdot) +// derivative of cvel, cdof_dot w.r.t qvel (dense version) +static void mjd_comVel_vel_dense(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* Dcdofdot) { int nv = m->nv, nbody = m->nbody; mjtNum mat[36]; @@ -388,8 +390,7 @@ static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* D // Dcvel += D(cdof * qvel), Dcdofdot = D(cvel x cdof) for (int j=m->body_dofadr[i]; jbody_dofadr[i]+m->body_dofnum[i]; j++) { - switch (m->jnt_type[m->dof_jntid[j]]) - { + switch (m->jnt_type[m->dof_jntid[j]]) { case mjJNT_FREE: // Dcdofdot = 0 mju_zero(Dcdofdot+j*6*nv, 18*nv); @@ -439,8 +440,8 @@ static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* D -// subtract (d qfrc_bias / d qvel) from DfDv -static void mjd_rne_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { +// subtract (d qfrc_bias / d qvel) from qDeriv (dense version) +void mjd_rne_vel_dense(const mjModel* m, mjData* d) { int nv = m->nv, nbody = m->nbody; mjtNum mat[36], mat1[36], mat2[36], dmul[36], tmp[6]; @@ -449,9 +450,10 @@ static void mjd_rne_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { mjtNum* Dcdofdot = mj_stackAlloc(d, nv*6*nv); mjtNum* Dcacc = mj_stackAlloc(d, nbody*6*nv); mjtNum* Dcfrcbody = mj_stackAlloc(d, nbody*6*nv); + mjtNum* row = mj_stackAlloc(d, nv); - // compute Dcdofdot and Dcvel - mjd_comVel_vel(m, d, Dcvel, Dcdofdot); + // compute Dcvel and Dcdofdot + mjd_comVel_vel_dense(m, d, Dcvel, Dcdofdot); // clear Dcacc mju_zero(Dcacc, nbody*6*nv); @@ -500,10 +502,17 @@ static void mjd_rne_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { } } - // DfDv -= D(cdof * cfrc_body) + // qDeriv -= D(cdof * cfrc_body) for (int i=0; idof_bodyid[i]*6+k)*nv, -d->cdof[i*6+k], nv); + // compute D(cdof * cfrc_body), store in row + mju_scl(row, Dcfrcbody + (m->dof_bodyid[i]*6+k)*nv, d->cdof[i*6+k], nv); + + // dense to sparse: qDeriv -= row + int end = d->D_rowadr[i] + d->D_rownnz[i]; + for (int adr=d->D_rowadr[i]; adrqDeriv[adr] -= row[d->D_colind[adr]]; + } } } @@ -512,6 +521,223 @@ static void mjd_rne_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { +//------------------------- sparse derivatives of component functions ------------------------------ +// internal sparse format: dense body/dof x sparse dof x 6 (inner size is 6) + +// copy sparse B-row from parent, shared ancestors only +static void copyFromParent(const mjModel* m, mjData* d, mjtNum* mat, int n) { + // return if this is world or parent is world + if (n == 0 || m->body_weldid[m->body_parentid[n]] == 0) { + return; + } + + // count dofs in ancestors + int ndof = 0; + int np = m->body_weldid[m->body_parentid[n]]; + while (np>0) { + // add self dofs + ndof += m->body_dofnum[np]; + + // advance to parent + np = m->body_weldid[m->body_parentid[np]]; + } + + // copy: guaranteed to be at beginning of sparse array, due to sorting + mju_copy(mat + 6*d->B_rowadr[n], mat + 6*d->B_rowadr[m->body_parentid[n]], 6*ndof); +} + + + +// add sparse B-row to parent, all overlapping nonzeros +static void addToParent(const mjModel* m, mjData* d, mjtNum* mat, int n) { + // return if this is world or parent is world + if (n == 0 || m->body_weldid[m->body_parentid[n]] == 0) { + return; + } + + // find matching nonzeros + int np = m->body_parentid[n]; + int i = 0, ip = 0; + while (iB_rownnz[n] && ipB_rownnz[np]) { + // columns match + if (d->B_colind[d->B_rowadr[n] + i] == d->B_colind[d->B_rowadr[np] + ip]) { + mju_addTo(mat + 6*(d->B_rowadr[np] + ip), mat + 6*(d->B_rowadr[n] + i), 6); + + // advance both + i++; + ip++; + } + + // mismatch columns: advance parent + else if (d->B_colind[d->B_rowadr[n] + i] > d->B_colind[d->B_rowadr[np] + ip]) { + ip++; + } + + // child nonzeroes must be subset of parent; SHOULD NOT OCCUR + else { + mju_error("Error in addToParent: child nonzeroes must be subset of parent"); + } + } +} + + + +// derivative of cvel, cdof_dot w.r.t qvel +static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* Dcdofdot) { + int nv = m->nv, nbody = m->nbody; + int* Badr = d->B_rowadr, * Dadr = d->D_rowadr; + mjtNum mat[36], matT[36]; // 6x6 matrices + + // forward pass over bodies: accumulate Dcvel, set Dcdofdot + for (int i = 1; ibody_dofadr[i] + m->body_dofnum[i]; + for (int j = m->body_dofadr[i]; jdof_Madr[j + 1] : m->nM) - (m->dof_Madr[j] + 1); + + // Dcvel += D(cdof * qvel), Dcdofdot = D(cvel x cdof) + switch (m->jnt_type[m->dof_jntid[j]]) { + case mjJNT_FREE: + // Dcdofdot = 0 (already cleared) + + // Dcvel += cdof * D(qvel) + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 0), d->cdof + 6*(j + 0), 6); + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 1), d->cdof + 6*(j + 1), 6); + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 2), d->cdof + 6*(j + 2), 6); + + // continue with rotations + j += 3; + Jadr += 3; + mjFALLTHROUGH; + + case mjJNT_BALL: + // Dcdofdot = Dcvel * D crossMotion(cvel, cdof) + for (int dj=0; dj<3; dj++) { + mjd_crossMotion_vel(mat, d->cdof + 6 * (j + dj)); + mju_transpose(matT, mat, 6, 6); + mju_mulMatMat(Dcdofdot + 6*Dadr[j + dj], Dcvel + 6*Badr[i], matT, Jadr + dj, 6, 6); + } + + // Dcvel += cdof * (D qvel) + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 0), d->cdof + 6*(j + 0), 6); + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 1), d->cdof + 6*(j + 1), 6); + mju_addTo(Dcvel + 6*(Badr[i] + Jadr + 2), d->cdof + 6*(j + 2), 6); + + // adjust for 3-dof joint + j += 2; + break; + + case mjJNT_HINGE: + case mjJNT_SLIDE: + // Dcdofdot = D crossMotion(cvel, cdof) * Dcvel + mjd_crossMotion_vel(mat, d->cdof + 6 * j); + mju_transpose(matT, mat, 6, 6); + mju_mulMatMat(Dcdofdot + 6*Dadr[j], Dcvel + 6*Badr[i], matT, Jadr, 6, 6); + + // Dcvel += cdof * (D qvel) + mju_addTo(Dcvel + 6*(Badr[i] + Jadr), d->cdof + 6*j, 6); + break; + + default: + mju_error("mjd_comVel_vel: Unknown joint type"); + } + } + } +} + + + +// subtract d qfrc_bias / d qvel from qDeriv +static void mjd_rne_vel(const mjModel* m, mjData* d) { + int nv = m->nv, nbody = m->nbody; + const int* Badr = d->B_rowadr; + const int* Dadr = d->D_rowadr; + const int* Bnnz = d->B_rownnz; + + mjtNum mat[36], mat1[36], mat2[36], dmul[36], tmp[6]; + + mjMARKSTACK; + mjtNum* Dcdofdot = mj_stackAlloc(d, 6*m->nD); + mjtNum* Dcvel = mj_stackAlloc(d, 6*m->nB); + mjtNum* Dcacc = mj_stackAlloc(d, 6*m->nB); + mjtNum* Dcfrcbody = mj_stackAlloc(d, 6*m->nB); + mjtNum* row = mj_stackAlloc(d, nv); + + // clear + mju_zero(Dcdofdot, 6*m->nD); + mju_zero(Dcvel, 6*m->nB); + mju_zero(Dcacc, 6*m->nB); + mju_zero(Dcfrcbody, 6*m->nB); + + // compute Dcvel and Dcdofdot + mjd_comVel_vel(m, d, Dcvel, Dcdofdot); + + // forward pass over bodies: accumulate Dcacc, set Dcfrcbody + for (int i=1; ibody_dofadr[i] + m->body_dofnum[i]; + for (int j=m->body_dofadr[i]; jdof_Madr[j + 1] : m->nM) - (m->dof_Madr[j] + 1); + + // Dcacc += cdofdot * (D qvel) + mju_addTo(Dcacc + 6*(Badr[i] + Jadr), d->cdof_dot + 6*j, 6); + + // Dcacc += (D cdofdot) * qvel + // Dcacc[row i] and Dcdofdot[row j] have identical sparsity + mju_addToScl(Dcacc + 6*Badr[i], Dcdofdot + 6*Dadr[j], d->qvel[j], 6*Bnnz[i]); + } + + //---------- Dcfrcbody = D(cinert * cacc + cvel x (cinert * cvel)) + + // Dcfrcbody = (D mul / D cacc) * Dcacc + mjd_mulInertVec_vel(dmul, d->cinert + 10*i); + mju_transpose(mat1, dmul, 6, 6); + mju_mulMatMat(Dcfrcbody + 6*Badr[i], Dcacc + 6*Badr[i], mat1, Bnnz[i], 6, 6); + + // mat = (D cross / D cvel) + (D cross / D mul) * (D mul / D cvel) + mju_mulInertVec(tmp, d->cinert + 10*i, d->cvel + i*6); + mjd_crossForce_vel(mat, tmp); + mjd_crossForce_frc(mat1, d->cvel + i*6); + mju_mulMatMat(mat2, mat1, dmul, 6, 6, 6); + mju_addTo(mat, mat2, 36); + + // Dcfrcbody += mat * Dcvel (use worldbody as temp) + mju_transpose(mat1, mat, 6, 6); + mju_mulMatMat(Dcfrcbody, Dcvel + 6*Badr[i], mat1, Bnnz[i], 6, 6); + mju_addTo(Dcfrcbody + 6*Badr[i], Dcfrcbody, 6*Bnnz[i]); + } + + // clear worldbody Dcfrcbody + mju_zero(Dcfrcbody, 6*Bnnz[0]); + + // backward pass over bodies: accumulate Dcfrcbody + for (int i=m->nbody-1; i>0; i--) { + addToParent(m, d, Dcfrcbody, i); + } + + // process all dofs, update qDeriv + for (int j=0; jdof_bodyid[j]; + + // qDeriv -= D(cdof * cfrc_body) + mju_mulMatVec(row, Dcfrcbody + 6*Badr[i], d->cdof + 6*j, Bnnz[i], 6); + mju_subFrom(d->qDeriv + Dadr[j], row, Bnnz[i]); + } + + mjFREESTACK; +} + + + //--------------------- utility functions for (d force / d vel) Jacobians -------------------------- // construct sparse Jacobian structure of body; return nnz @@ -548,51 +774,84 @@ static int bodyJacSparse(const mjModel* m, int body, int* ind) { -// add J'*B*J to DfDv -static void addJTBJ(mjtNum* DfDv, const mjtNum* J, const mjtNum* B, int n, int nv) { + +// add J'*B*J to qDeriv +static void addJTBJ(const mjModel* m, mjData* d, const mjtNum* J, const mjtNum* B, int n) { + int nv = m->nv; + + // allocate dense row + mjMARKSTACK; + mjtNum* row = mj_stackAlloc(d, nv); + // process non-zero elements of B for (int i=0; iD_rownnz[k]; + for (int s=0; sD_rowadr[k] + s; + d->qDeriv[adr] += row[d->D_colind[adr]]; } } } } } + + mjFREESTACK; } -// add J'*B*J to DfDv, sparse version -static void addJTBJSparse(mjtNum* DfDv, const mjtNum* J, const mjtNum* B, - int n, int nv, int offset, - const int* rownnz, const int* rowadr, const int* colind) { +// add J'*B*J to qDeriv, sparse version +static void addJTBJSparse(const mjModel* m, mjData* d, const mjtNum* J, + const mjtNum* B, int n, int offset, + const int* rownnz, const int* rowadr, const int* colind) { + int nv = m->nv; + + // allocate row + mjMARKSTACK; + mjtNum* row = mj_stackAlloc(d, nv); + // process non-zero elements of B for (int i=0; iD_rownnz[k]; s++) { + int adr = d->D_rowadr[k] + s; + d->qDeriv[adr] += row[d->D_colind[adr]]; } } } } + + // free space + mjFREESTACK; } @@ -667,8 +926,8 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang -// add (d qfrc_actuator / d qvel) to DfDv -static void mjd_actuator_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { +// add (d qfrc_actuator / d qvel) to qDeriv +void mjd_actuator_vel(const mjModel* m, mjData* d) { int nv = m->nv; // disabled: nothing to add @@ -712,7 +971,7 @@ static void mjd_actuator_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { // add if (bias_vel!=0) { - addJTBJ(DfDv, d->actuator_moment+i*nv, &bias_vel, 1, nv); + addJTBJ(m, d, d->actuator_moment+i*nv, &bias_vel, 1); } } } @@ -1007,7 +1266,7 @@ static inline void mjd_magnus_force( //----------------- fluid force derivatives, ellipsoid and inertia-box models ---------------------- // fluid forces based on ellipsoid approximation -void mjd_ellipsoidFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int bodyid) { +void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) { mjMARKSTACK; int nv = m->nv; @@ -1100,10 +1359,22 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int bodyid) { mjd_addedMassForces(B, lvel, m->opt.density, virtual_mass, virtual_inertia); + // make B symmetric if integrator is IMPLICITFAST + if (m->opt.integrator == mjINT_IMPLICITFAST) { + for (int i=0; i<5; i++) { + for (j=i+1; j<6; j++) { + mjtNum tmp = 0.5 * (B[6 * i + j] + B[6 * j + i]); + B[6 * i + j] = tmp; + B[6 * j + i] = tmp; + } + } + } + if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, B, 6, nv, 0, rownnz, rowadr, colind_compressed); - } else { - addJTBJ(DfDv, J, B, 6, nv); + addJTBJSparse(m, d, J, B, 6, 0, rownnz, rowadr, colind_compressed); + } + else { + addJTBJ(m, d, J, B, 6); } } @@ -1112,7 +1383,7 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int bodyid) { // fluid forces based on inertia-box approximation -void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) +void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, int i) { mjMARKSTACK; @@ -1127,11 +1398,11 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) // equivalent inertia box box[0] = mju_sqrt(mju_max(mjMINVAL, - (inertia[1] + inertia[2] - inertia[0])) / m->body_mass[i] * 6.0); + (inertia[1] + inertia[2] - inertia[0])) / m->body_mass[i] * 6.0); box[1] = mju_sqrt(mju_max(mjMINVAL, - (inertia[0] + inertia[2] - inertia[1])) / m->body_mass[i] * 6.0); + (inertia[0] + inertia[2] - inertia[1])) / m->body_mass[i] * 6.0); box[2] = mju_sqrt(mju_max(mjMINVAL, - (inertia[0] + inertia[1] - inertia[2])) / m->body_mass[i] * 6.0); + (inertia[0] + inertia[1] - inertia[2])) / m->body_mass[i] * 6.0); // map from CoM-centered to local body-centered 6D velocity mj_objectVelocity(m, d, mjOBJ_BODY, i, lvel, 1); @@ -1190,9 +1461,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) B = -mjPI*diam*diam*diam*m->opt.viscosity; for (int j=0; j<3; j++) { if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, j, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, j, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+j*nv, &B, 1, nv); + addJTBJ(m, d, J+j*nv, &B, 1); } } @@ -1200,9 +1471,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) B = -3.0*mjPI*diam*m->opt.viscosity; for (int j=0; j<3; j++) { if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 3+j, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 3+j, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+3*nv+j*nv, &B, 1, nv); + addJTBJ(m, d, J+3*nv+j*nv, &B, 1); } } } @@ -1214,9 +1485,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) B = -m->opt.density*box[0]*(box[1]*box[1]*box[1]*box[1]+box[2]*box[2]*box[2]*box[2])* 2*mju_abs(lvel[0])/64.0; if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 0, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 0, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J, &B, 1, nv); + addJTBJ(m, d, J, &B, 1); } // lfrc[1] -= m->opt.density*box[1]*(box[0]*box[0]*box[0]*box[0]+box[2]*box[2]*box[2]*box[2])* @@ -1224,9 +1495,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) B = -m->opt.density*box[1]*(box[0]*box[0]*box[0]*box[0]+box[2]*box[2]*box[2]*box[2])* 2*mju_abs(lvel[1])/64.0; if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 1, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 1, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+nv, &B, 1, nv); + addJTBJ(m, d, J+nv, &B, 1); } // lfrc[2] -= m->opt.density*box[2]*(box[0]*box[0]*box[0]*box[0]+box[1]*box[1]*box[1]*box[1])* @@ -1234,33 +1505,33 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) B = -m->opt.density*box[2]*(box[0]*box[0]*box[0]*box[0]+box[1]*box[1]*box[1]*box[1])* 2*mju_abs(lvel[2])/64.0; if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 2, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 2, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+2*nv, &B, 1, nv); + addJTBJ(m, d, J+2*nv, &B, 1); } // lfrc[3] -= 0.5*m->opt.density*box[1]*box[2]*mju_abs(lvel[3])*lvel[3]; B = -0.5*m->opt.density*box[1]*box[2]*2*mju_abs(lvel[3]); if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 3, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 3, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+3*nv, &B, 1, nv); + addJTBJ(m, d, J+3*nv, &B, 1); } // lfrc[4] -= 0.5*m->opt.density*box[0]*box[2]*mju_abs(lvel[4])*lvel[4]; B = -0.5*m->opt.density*box[0]*box[2]*2*mju_abs(lvel[4]); if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 4, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 4, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+4*nv, &B, 1, nv); + addJTBJ(m, d, J+4*nv, &B, 1); } // lfrc[5] -= 0.5*m->opt.density*box[0]*box[1]*mju_abs(lvel[5])*lvel[5]; B = -0.5*m->opt.density*box[0]*box[1]*2*mju_abs(lvel[5]); if (mj_isSparse(m)) { - addJTBJSparse(DfDv, J, &B, 1, nv, 5, rownnz, rowadr, colind); + addJTBJSparse(m, d, J, &B, 1, 5, rownnz, rowadr, colind); } else { - addJTBJ(DfDv, J+5*nv, &B, 1, nv); + addJTBJ(m, d, J+5*nv, &B, 1); } } @@ -1271,9 +1542,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, mjtNum* DfDv, int i) //------------------------- derivatives of passive forces ------------------------------------------ -// add (d qfrc_passive / d qvel) to DfDv -void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { - int nv = m->nv; +// add (d qfrc_passive / d qvel) to qDeriv +void mjd_passive_vel(const mjModel* m, mjData* d) { + int nv = m->nv, nbody = m->nbody; // disabled: nothing to add if (mjDISABLED(mjDSBL_PASSIVE)) { @@ -1282,7 +1553,16 @@ void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { // dof damping for (int i=0; idof_damping[i]; + int nnz_i = d->D_rownnz[i]; + for (int j=0; jD_rowadr[i] + j; + + // identify diagonal element + if (d->D_colind[ij] == i) { + d->qDeriv[ij] -= m->dof_damping[i]; + break; + } + } } // tendon damping @@ -1292,17 +1572,17 @@ void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { // add sparse or dense if (mj_isSparse(m)) { - addJTBJSparse(DfDv, d->ten_J, &B, 1, nv, i, + addJTBJSparse(m, d, d->ten_J, &B, 1, i, d->ten_J_rownnz, d->ten_J_rowadr, d->ten_J_colind); } else { - addJTBJ(DfDv, d->ten_J+i*nv, &B, 1, nv); + addJTBJ(m, d, d->ten_J+i*nv, &B, 1); } } } // fluid drag model, either body-level (inertia box) or geom-level (ellipsoid) if (m->opt.viscosity>0 || m->opt.density>0) { - for (int i=1; inbody; i++) { + for (int i=1; ibody_mass[i]geom_fluid[mjNFLUID*geomid] > 0); } if (use_ellipsoid_model) { - mjd_ellipsoidFluid(m, d, DfDv, i); + mjd_ellipsoidFluid(m, d, i); } else { - mjd_inertiaBoxFluid(m, d, DfDv, i); + mjd_inertiaBoxFluid(m, d, i); } } } @@ -1324,13 +1604,17 @@ void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { -// add forward fin-diff approximation of (d qfrc_passive / d qvel) to DfDv -void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* DfDv) { +// add forward fin-diff approximation of (d qfrc_passive / d qvel) to qDeriv +void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps) { int nv = m->nv; mjMARKSTACK; mjtNum* qfrc_passive = mj_stackAlloc(d, nv); mjtNum* fd = mj_stackAlloc(d, nv); + int* cnt = (int*)mj_stackAlloc(d, nv); + + // clear row counters + memset(cnt, 0, nv*sizeof(int)); // save qfrc_passive, assume mj_fwdVelocity was called mju_copy(qfrc_passive, d->qfrc_passive, nv); @@ -1351,9 +1635,13 @@ void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* DfDv) { mju_sub(fd, d->qfrc_passive, qfrc_passive, nv); mju_scl(fd, fd, 1/eps, nv); - // copy to i-th column of DfDv + // copy to i-th column of qDeriv for (int j=0; jD_rowadr[j] + cnt[j]; + if (cnt[j]D_rownnz[j] && d->D_colind[adr] == i) { + d->qDeriv[adr] = fd[j]; + cnt[j]++; + } } } @@ -1365,10 +1653,8 @@ void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* DfDv) { - //-------------------- derivatives of all smooth (unconstrained) forces ---------------------------- - // centered finite difference approximation to mjd_smooth_vel void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) { int nv = m->nv; @@ -1436,31 +1722,21 @@ void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) { //------------------------- main entry points ------------------------------------------------------ // analytical derivative of smooth forces w.r.t velocities: -// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel -void mjd_smooth_vel(const mjModel *m, mjData *d) { - int nv = m->nv; +// d->qDeriv = d (qfrc_actuator + qfrc_passive - [qfrc_bias]) / d qvel +void mjd_smooth_vel(const mjModel* m, mjData* d, int flg_bias) { + // clear qDeriv + mju_zero(d->qDeriv, m->nD); - // allocate space - mjMARKSTACK; - mjtNum *DfDv = mj_stackAlloc(d, nv*nv); + // qDeriv += d qfrc_actuator / d qvel + mjd_actuator_vel(m, d); - // clear DfDv - mju_zero(DfDv, nv*nv); + // qDeriv += d qfrc_passive / d qvel + mjd_passive_vel(m, d); - // DfDv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel - mjd_actuator_vel(m, d, DfDv); - mjd_passive_vel(m, d, DfDv); - mjd_rne_vel(m, d, DfDv); - - // copy dense DfDv to sparse qDeriv - for (int i=0; iD_rownnz[i]; j++) { - int adr = d->D_rowadr[i] + j; - d->qDeriv[adr] = DfDv[i*nv + d->D_colind[adr]]; - } + // qDeriv -= d qfrc_bias / d qvel; optional + if (flg_bias) { + mjd_rne_vel(m, d); } - - mjFREESTACK; } diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index 0cc92bfe..ab093415 100644 --- a/src/engine/engine_derivative.h +++ b/src/engine/engine_derivative.h @@ -24,17 +24,23 @@ extern "C" { #endif // analytical derivative of smooth forces w.r.t velocities: -// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel -MJAPI void mjd_smooth_vel(const mjModel* m, mjData* d); +// d->qDeriv = d (qfrc_actuator + qfrc_passive - [qfrc_bias]) / d qvel +MJAPI void mjd_smooth_vel(const mjModel* m, mjData* d, int flg_bias); // centered finite difference approximation to mjd_smooth_vel MJAPI void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps); -// add (d qfrc_passive / d qvel) to DfDv -MJAPI void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv); +// add (d qfrc_actuator / d qvel) to qDeriv +MJAPI void mjd_actuator_vel(const mjModel* m, mjData* d); -// add forward finite difference approximation of (d qfrc_passive / d qvel) to DfDv -MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* DfDv); +// add (d qfrc_passive / d qvel) to qDeriv +MJAPI void mjd_passive_vel(const mjModel* m, mjData* d); + +// subtract (d qfrc_bias / d qvel) from qDeriv (dense version) +MJAPI void mjd_rne_vel_dense(const mjModel* m, mjData* d); + +// add forward finite difference approximation of (d qfrc_passive / d qvel) to qDeriv +MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps); // advance simulation using control callback, skipstage is mjtStage MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor); diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index d7747001..9209d91a 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -704,33 +704,59 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { // fully implicit in velocity, possibly skipping factorization -void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor) { +void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { int nv = m->nv; mjMARKSTACK; - mjtNum *qfrc = mj_stackAlloc(d, nv); - mjtNum *qacc = mj_stackAlloc(d, nv); - - if (!skipfactor) { - // construct sparse structure in d->D_xxx - mj_makeMSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_colind); - - // compute analytical derivative qDeriv - mjd_smooth_vel(m, d); - - // set qLU = qM - dt*qDeriv - mj_setMSparse(m, d, d->qLU, d->D_rownnz, d->D_rowadr, d->D_colind); - mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); - - // factorize qLU, use qacc as scratch space - mju_factorLUSparse(d->qLU, nv, (int*)qacc, d->D_rownnz, d->D_rowadr, d->D_colind); - } + mjtNum* qfrc = mj_stackAlloc(d, nv); + mjtNum* qacc = mj_stackAlloc(d, nv); // set qfrc = qfrc_smooth + qfrc_constraint mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); - // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc - mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_colind); + // IMPLICIT + if (m->opt.integrator == mjINT_IMPLICIT) { + if (!skipfactor) { + // compute analytical derivative qDeriv + mjd_smooth_vel(m, d, /* flg_bias = */ 1); + + // set qLU = qM + mj_copyM2DSparse(m, d, d->qLU, d->qM); + + // set qLU = qM - dt*qDeriv + mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); + + // factorize qLU, use qacc as scratch space + mju_factorLUSparse(d->qLU, nv, (int*)qacc, d->D_rownnz, d->D_rowadr, d->D_colind); + } + + // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc + mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_colind); + } + + // IMPLICITFAST + else if (m->opt.integrator == mjINT_IMPLICITFAST) { + if (!skipfactor) { + // compute analytical derivative qDeriv; skip rne derivative + mjd_smooth_vel(m, d, /* flg_bias = */ 0); + + // modified mass matrix MhB = qDeriv[Lower] + mjtNum* MhB = mj_stackAlloc(d, m->nM); + mj_copyD2MSparse(m, d, MhB, d->qDeriv); + + // set MhB = M - dt*qDeriv + mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, m->nM); + + // factorize + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL); + } + + // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc + mju_copy(qacc, qfrc, m->nv); + mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); + } else { + mju_error("mj_implicitSkip: integrator must be implicit or implicitfast"); + } // advance state and time mj_advance(m, d, d->act_dot, qacc, NULL); @@ -741,7 +767,7 @@ void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor) { // fully implicit in velocity -void mj_implicit(const mjModel *m, mjData *d) { +void mj_implicit(const mjModel* m, mjData* d) { mj_implicitSkip(m, d, 0); } @@ -825,6 +851,7 @@ void mj_step(const mjModel* m, mjData* d) { break; case mjINT_IMPLICIT: + case mjINT_IMPLICITFAST: mj_implicit(m, d); break; @@ -873,7 +900,7 @@ void mj_step2(const mjModel* m, mjData* d) { } // integrate with Euler or implicit; RK4 defaults to Euler - if (m->opt.integrator==mjINT_IMPLICIT) { + if (m->opt.integrator == mjINT_IMPLICIT || m->opt.integrator == mjINT_IMPLICITFAST) { mj_implicit(m, d); } else { mj_Euler(m, d); diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index b6293861..6c348b25 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -29,6 +29,7 @@ #include "engine/engine_plugin.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" +#include "engine/engine_util_misc.h" #include "engine/engine_vfs.h" #ifdef _MSC_VER @@ -796,6 +797,183 @@ int mj_sizeModel(const mjModel* m) { + +//-------------------------- sparse system matrix construction ------------------------------------- + +// construct sparse representation of dof-dof matrix +static void makeDSparse(const mjModel* m, mjData* d) { + int nv = m->nv; + int* rownnz = d->D_rownnz; + int* rowadr = d->D_rowadr; + int* colind = d->D_colind; + + mjMARKSTACK; + int* remaining = (int*)mj_stackAlloc(d, nv); + + // compute rownnz + memset(rownnz, 0, nv * sizeof(int)); + for (int i = nv - 1; i >= 0; i--) { + // init at diagonal + int j = i; + rownnz[i]++; + + // process below diagonal + while ((j = m->dof_parentid[j]) >= 0) { + rownnz[i]++; + rownnz[j]++; + } + } + + // accumulate rowadr + rowadr[0] = 0; + for (int i = 1; i < nv; i++) { + rowadr[i] = rowadr[i - 1] + rownnz[i - 1]; + } + + // populate colind + memcpy(remaining, rownnz, nv * sizeof(int)); + for (int i = nv - 1; i >= 0; i--) { + // init at diagonal + remaining[i]--; + colind[rowadr[i] + remaining[i]] = i; + + // process below diagonal + int j = i; + while ((j = m->dof_parentid[j]) >= 0) { + remaining[i]--; + colind[rowadr[i] + remaining[i]] = j; + + remaining[j]--; + colind[rowadr[j] + remaining[j]] = i; + } + } + + // sanity check; SHOULD NOT OCCUR + for (int i = 0; i < nv; i++) { + if (remaining[i] != 0) { + mju_error("Error in mj_makeDSparse: unexpected remaining"); + } + } + + mjFREESTACK; +} + + + +// construct sparse representation of body-dof matrix +static void makeBSparse(const mjModel* m, mjData* d) { + int nv = m->nv, nbody = m->nbody; + int* rownnz = d->B_rownnz; + int* rowadr = d->B_rowadr; + int* colind = d->B_colind; + + // set rownnz to subtree dofs counts, including self + memset(rownnz, 0, sizeof(int) * nbody); + for (int i = nbody - 1; i > 0; i--) { + rownnz[i] += m->body_dofnum[i]; + rownnz[m->body_parentid[i]] += rownnz[i]; + } + + // sanity check; SHOULD NOT OCCUR + if (rownnz[0] != nv) { + mju_error("Error in mj_makeBSparse: rownnz[0] different from nv"); + } + + // add dofs in ancestors bodies + for (int i = 0; i < nbody; i++) { + int j = m->body_parentid[i]; + while (j > 0) { + rownnz[i] += m->body_dofnum[j]; + j = m->body_parentid[j]; + } + } + + // compute rowadr + rowadr[0] = 0; + for (int i = 1; i < nbody; i++) { + rowadr[i] = rowadr[i - 1] + rownnz[i - 1]; + } + + // sanity check; SHOULD NOT OCCUR + if (m->nB != rowadr[nbody - 1] + rownnz[nbody - 1]) { + mju_error("Error in mj_makeBSparse: sum of rownnz different from nB"); + } + + // allocate and clear incremental row counts + mjMARKSTACK; + int* cnt = (int*)mj_stackAlloc(d, nbody); + memset(cnt, 0, sizeof(int) * nbody); + + // add subtree dofs to colind + for (int i = nbody - 1; i > 0; i--) { + // add this body's dofs to subtree + for (int n = 0; n < m->body_dofnum[i]; n++) { + colind[rowadr[i] + cnt[i]] = m->body_dofadr[i] + n; + cnt[i]++; + } + + // add body subtree to parent + int par = m->body_parentid[i]; + for (int n = 0; n < cnt[i]; n++) { + colind[rowadr[par] + cnt[par]] = colind[rowadr[i] + n]; + cnt[par]++; + } + } + + // add all ancestor dofs + for (int i = 0; i < nbody; i++) { + int par = m->body_parentid[i]; + while (par > 0) { + // add ancestor body dofs + for (int n = 0; n < m->body_dofnum[par]; n++) { + colind[rowadr[i] + cnt[i]] = m->body_dofadr[par] + n; + cnt[i]++; + } + + // advance to parent + par = m->body_parentid[par]; + } + } + + // process all bodies + for (int i = 0; i < nbody; i++) { + // make sure cnt = rownnz; SHOULD NOT OCCUR + if (rownnz[i] != cnt[i]) { + mju_error("Error in mj_makeBSparse: cnt different from rownnz"); + } + + // sort colind in each row + if (cnt[i] > 1) { + mju_insertionSortInt(colind + rowadr[i], cnt[i]); + } + } + + mjFREESTACK; +} + + + +// check D and B sparsity for consistency +static void checkDBSparse(const mjModel* m, mjData* d) { + // process all dofs + for (int j = 0; j < m->nv; j++) { + // get body for this dof + int i = m->dof_bodyid[j]; + + // D[row j] and B[row i] should be identical + if (d->D_rownnz[j] != d->B_rownnz[i]) { + mju_error("Error in checkDBSparse: rows have different nnz"); + } + for (int k = 0; k < d->D_rownnz[j]; k++) { + if (d->D_colind[d->D_rowadr[j] + k] != d->B_colind[d->B_rowadr[i] + k]) { + mju_error("Error in checkDBSparse: rows have different colind"); + } + } + } +} + + + //----------------------------------- mjData construction ------------------------------------------ // set pointers into mjData buffer @@ -1142,6 +1320,13 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { } } + // construct sparse matrix representations + if (m->body_dofadr) { + makeDSparse(m, d); + makeBSparse(m, d); + checkDBSparse(m, d); + } + // restore pluginstate and plugindata memcpy(d->plugin_state, plugin_state, sizeof(mjtNum) * m->npluginstate); mju_free(plugin_state); @@ -1507,6 +1692,7 @@ const char* mj_validateReferences(const mjModel* m) { return "Invalid model: eq_obj2id out of bounds."; } break; + case mjEQ_TENDON: if (obj1id >= m->ntendon || obj1id < 0) { return "Invalid model: eq_obj1id out of bounds."; @@ -1516,8 +1702,7 @@ const char* mj_validateReferences(const mjModel* m) { return "Invalid model: eq_obj2id out of bounds."; } break; - case mjEQ_DISTANCE: - return "distance equality constraints are no longer supported"; + case mjEQ_WELD: case mjEQ_CONNECT: if (obj1id >= m->nbody || obj1id < 0) { @@ -1527,6 +1712,9 @@ const char* mj_validateReferences(const mjModel* m) { return "Invalid model: eq_obj2id out of bounds."; } break; + + default: + mju_error("mj_validateReferences: unknown equality constraint type."); } } for (int i=0; inwrap; i++) { diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 78724273..563fdf72 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -933,6 +933,27 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, } fprintf(fp, "\n\n"); + // B_rownnz + fprintf(fp, NAME_FORMAT, "B_rownnz"); + for (int i = 0; i < m->nbody; i++) { + fprintf(fp, " %d", d->B_rownnz[i]); + } + fprintf(fp, "\n\n"); + + // B_rowadr + fprintf(fp, NAME_FORMAT, "B_rowadr"); + for (int i = 0; i < m->nbody; i++) { + fprintf(fp, " %d", d->B_rowadr[i]); + } + fprintf(fp, "\n\n"); + + // B_colind + fprintf(fp, NAME_FORMAT, "B_colind"); + for (int i = 0; i < m->nB; i++) { + fprintf(fp, " %d", d->B_colind[i]); + } + fprintf(fp, "\n\n"); + // print qDeriv mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind); printArray("QDERIV", m->nv, m->nv, M, fp, float_format); diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 77315a30..774d6835 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -975,94 +975,63 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, -// construct sparse matrix representations matching qM -void mj_makeMSparse(const mjModel* m, mjData* d, int* rownnz, int* rowadr, int* colind) { +//-------------------------- sparse system matrix conversion --------------------------------------- + +// dst[D] = src[M], handle different sparsity representations +void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) { int nv = m->nv; - mjMARKSTACK; - int *remaining = (int*) mj_stackAlloc(d, nv); - // compute rownnz - memset(rownnz, 0, nv*sizeof(int)); - for (int i=nv-1; i>=0; i--) { - // init at diagonal - int j = i; - rownnz[i]++; - - // process below diagonal - while ((j=m->dof_parentid[j]) >= 0) { - rownnz[i]++; - rownnz[j]++; - } - } - - // accumulate rowadr - rowadr[0] = 0; - for (int i=1; i=0; i--) { - // init at diagonal - remaining[i]--; - colind[rowadr[i] + remaining[i]] = i; - - // process below diagonal - int j = i; - while ((j = m->dof_parentid[j]) >= 0) { - remaining[i]--; - colind[rowadr[i] + remaining[i]] = j; - - remaining[j]--; - colind[rowadr[j] + remaining[j]] = i; - } - } - - // sanity check; SHOULD NOT OCCUR - for (int i=0; inv; - - mjMARKSTACK; - int *remaining = (int*) mj_stackAlloc(d, nv); + // init remaining + int* remaining = (int*)mj_stackAlloc(d, nv); + memcpy(remaining, d->D_rownnz, nv * sizeof(int)); // copy data - memcpy(remaining, rownnz, nv*sizeof(int)); - for (int i=nv-1; i>=0; i--) { + for (int i = nv - 1; i >= 0; i--) { // init at diagonal int adr = m->dof_Madr[i]; remaining[i]--; - dst[rowadr[i] + remaining[i]] = d->qM[adr]; + dst[d->D_rowadr[i] + remaining[i]] = src[adr]; adr++; // process below diagonal int j = i; while ((j = m->dof_parentid[j]) >= 0) { remaining[i]--; - dst[rowadr[i] + remaining[i]] = d->qM[adr]; + dst[d->D_rowadr[i] + remaining[i]] = src[adr]; remaining[j]--; - dst[rowadr[j] + remaining[j]] = d->qM[adr]; + dst[d->D_rowadr[j] + remaining[j]] = src[adr]; adr++; } } - mjFREESTACK; + mjFREESTACK +} + + + +// dst[M] = src[D lower], handle different sparsity representations +void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) { + int nv = m->nv; + + // copy data + for (int i = nv - 1; i >= 0; i--) { + // find diagonal in qDeriv + int j = 0; + while (d->D_colind[d->D_rowadr[i] + j] < i) { + j++; + } + + // copy + int adr = m->dof_Madr[i]; + while (j >= 0) { + dst[adr] = src[d->D_rowadr[i] + j]; + adr++; + j--; + } + } } diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index 203a8478..35e0b64c 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -104,12 +104,14 @@ MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind); -// construct sparse matrix representations matching qM -MJAPI void mj_makeMSparse(const mjModel* m, mjData* d, int *rownnz, int *rowadr, int *colind); -// set dst = qM, handle different sparsity representations -MJAPI void mj_setMSparse(const mjModel* m, mjData* d, mjtNum* dst, - const int *rownnz, const int *rowadr, const int *colind); +//-------------------------- sparse system matrix conversion --------------------------------------- + +// dst[D] = src[M], handle different sparsity representations +MJAPI void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src); + +// dst[M] = src[D lower], handle different sparsity representations +MJAPI void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src); //-------------------------- perturbations --------------------------------------------------------- diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 304f0d50..37abc90a 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -281,6 +281,7 @@ void mjCModel::Clear(void) { nemax = 0; nM = 0; nD = 0; + nB = 0; njmax = -1; nconmax = -1; @@ -1622,10 +1623,39 @@ void mjCModel::CopyTree(mjModel* m) { } m->nM = nM; - // set nD - nD = 2*nM - nv; + // compute nD + nD = 2 * nM - nv; m->nD = nD; + // compute subtreedofs in backward pass over bodies + for (i = nbody - 1; i > 0; i--) { + // add body dofs to self count + bodies[i]->subtreedofs += bodies[i]->dofnum; + + // add to parent count + bodies[bodies[i]->parentid]->subtreedofs += bodies[i]->subtreedofs; + } + + // make sure all dofs are in world "subtree", SHOULD NOT OCCUR + if (bodies[0]->subtreedofs != nv) { + throw mjCError(0, "all DOFs should be in world subtree"); + } + + // compute nB + nB = 0; + for (i = 0; i < nbody; i++) { + // add subtree dofs (including self) + nB += bodies[i]->subtreedofs; + + // add dofs in ancestor bodies + j = bodies[i]->parentid; + while (j > 0) { + nB += bodies[j]->dofnum; + j = bodies[j]->parentid; + } + } + m->nB = nB; + // set dof_simplenum int scnt = 0; for (i=nv-1; i>=0; i--) { @@ -2846,7 +2876,7 @@ bool mjCModel::CopyBack(const mjModel* m) { neq!=m->neq || ntendon!=m->ntendon || nwrap!=m->nwrap || nsensor!=m->nsensor || nnumeric!=m->nnumeric || nnumericdata!=m->nnumericdata || ntext!=m->ntext || ntextdata!=m->ntextdata || nnames!=m->nnames || nM!=m->nM || nD!=m->nD || - nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax) { + nB!=m->nB || nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax) { errInfo = mjCError(0, "incompatible models in CopyBack"); return false; } diff --git a/src/user/user_model.h b/src/user/user_model.h index 7d2b44fc..5bc1a86f 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -235,7 +235,8 @@ class mjCModel { int npluginattr; // number of chars in all plugin config attributes int nnames; // number of chars in all names int nM; // number of non-zeros in sparse inertia matrix - int nD; // number of non-zeros in sparse derivative matrix + int nD; // number of non-zeros in sparse dof-dof matrix + int nB; // number of non-zeros in sparse body-dof matrix //------------------------ object lists // objects created here diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 3396f562..99049f30 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -327,6 +327,7 @@ mjCBody::mjCBody(mjCModel* _model) { weldid = -1; dofnum = 0; lastdof = -1; + subtreedofs = 0; gravcomp = 0; userdata.clear(); diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 80bf7189..573a3324 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -209,7 +209,9 @@ class mjCBody : public mjCBase { int mocapid; // mocap id, -1: not mocap bool explicitinertial; // whether to save the body with an explicit inertial clause - int lastdof; // id of last dof (used by compiler) + // used internally by compiler + int lastdof; // id of last dof + int subtreedofs; // number of dofs in subtree, including self // objects allocated by Add functions std::vector bodies; // child bodies diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 0244f556..6f476ecb 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -511,11 +511,12 @@ const mjMap camlight_map[camlight_sz] = { // integrator type -const int integrator_sz = 3; +const int integrator_sz = 4; const mjMap integrator_map[integrator_sz] = { - {"Euler", mjINT_EULER}, - {"RK4", mjINT_RK4}, - {"implicit", mjINT_IMPLICIT} + {"Euler", mjINT_EULER}, + {"RK4", mjINT_RK4}, + {"implicit", mjINT_IMPLICIT}, + {"implicitfast", mjINT_IMPLICITFAST} }; diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 476a344d..826866d9 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -24,7 +24,6 @@ #include "src/engine/engine_core_smooth.h" #include "src/engine/engine_derivative.h" #include "src/engine/engine_io.h" -#include "src/engine/engine_support.h" #include "src/engine/engine_util_blas.h" #include "src/engine/engine_util_errmem.h" #include "test/fixture.h" @@ -101,6 +100,7 @@ static const char* const kDampedPendulumPath = static const char* const kLinearPath = "engine/testdata/derivative/linear.xml"; static const char* const kModelPath = "testdata/model.xml"; + // compare analytic and finite-difference d_smooth/d_qvel TEST_F(DerivativeTest, SmoothDvel) { // run test on all models @@ -110,6 +110,7 @@ TEST_F(DerivativeTest, SmoothDvel) { kDamperActuatorsPath}) { const std::string xml_path = GetTestDataFilePath(local_path); mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int nD = model->nD; mjData* data = mj_makeData(model); for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { @@ -127,24 +128,24 @@ TEST_F(DerivativeTest, SmoothDvel) { mj_forward(model, data); // construct sparse structure in d->D_xxx, compute analytical qDeriv - mj_makeMSparse(model, data, - data->D_rownnz, data->D_rowadr, data->D_colind); - mjd_smooth_vel(model, data); + mju_zero(data->qDeriv, nD); + mjd_smooth_vel(model, data, /*flg_bias=*/true); // expect derivatives to be non-zero, make copy of qDeriv as a vector - EXPECT_GT(mju_norm(data->qDeriv, model->nD), 0); - std::vector qDerivAnalytic = AsVector(data->qDeriv, model->nD); + EXPECT_GT(mju_norm(data->qDeriv, nD), 0); + std::vector qDerivAnalytic = AsVector(data->qDeriv, nD); // compute finite-difference derivatives mjtNum eps = 1e-7; + mju_zero(data->qDeriv, nD); mjd_smooth_velFD(model, data, eps); // expect FD and analytic derivatives to be numerically different - EXPECT_NE(mju_norm(data->qDeriv, model->nD), - mju_norm(qDerivAnalytic.data(), model->nD)); + EXPECT_NE(mju_norm(data->qDeriv, nD), + mju_norm(qDerivAnalytic.data(), nD)); // expect FD and analytic derivatives to be similar to eps precision - EXPECT_THAT(AsVector(data->qDeriv, model->nD), + EXPECT_THAT(AsVector(data->qDeriv, nD), Pointwise(DoubleNear(eps), qDerivAnalytic)); } mj_deleteData(data); @@ -159,11 +160,11 @@ TEST_F(DerivativeTest, PassiveDvel) { // load model const std::string xml_path = GetTestDataFilePath(local_path); mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - int nv = model->nv; + int nD = model->nD; mjData* data = mj_makeData(model); // allocate Jacobians - mjtNum* DfDv_analytic = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv*nv); - mjtNum* DfDv_FD = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv*nv); + mjtNum* qDerivAnalytic = (mjtNum*) mju_malloc(sizeof(mjtNum)*nD); + mjtNum* qDerivFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*nD); for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { // set sparsity @@ -176,22 +177,23 @@ TEST_F(DerivativeTest, PassiveDvel) { } mj_forward(model, data); - // clear DfDv, get analytic derivatives - mju_zero(DfDv_analytic, nv*nv); - mjd_passive_vel(model, data, DfDv_analytic); + // get analytic derivatives + mju_copy(qDerivAnalytic, data->qDeriv, nD); - // clear DfDv, get finite-difference derivatives - mju_zero(DfDv_FD, nv*nv); + // clear qDeriv, get finite-difference derivatives + mju_zero(data->qDeriv, nD); + mju_zero(qDerivFD, nD); mjtNum eps = 1e-6; - mjd_passive_velFD(model, data, eps, DfDv_FD); + mjd_passive_velFD(model, data, eps); // expect FD and analytic derivatives to be similar to tol precision mjtNum tol = 1e-4; - CompareMatrices(DfDv_analytic, DfDv_FD, nv, nv, tol); + EXPECT_THAT(AsVector(data->qDeriv, nD), + Pointwise(DoubleNear(tol), AsVector(qDerivAnalytic, nD))); } - mju_free(DfDv_FD); - mju_free(DfDv_analytic); + mju_free(qDerivFD); + mju_free(qDerivAnalytic); mj_deleteData(data); mj_deleteModel(model); } @@ -210,7 +212,9 @@ TEST_F(DerivativeTest, StepSkip) { // disable warmstarts so we don't need to save qacc_warmstart model->opt.disableflags |= mjDSBL_WARMSTART; - for (const mjtIntegrator integrator : {mjINT_EULER, mjINT_IMPLICIT}) { + for (const mjtIntegrator integrator : {mjINT_EULER, + mjINT_IMPLICIT, + mjINT_IMPLICITFAST}) { model->opt.integrator = integrator; // reset, take 20 steps, save initial state @@ -598,5 +602,49 @@ TEST_F(DerivativeTest, NoStateMutation) { mj_deleteModel(model); } +// compare dense and sparse derivatives of qfrc_bias (RNE) +TEST_F(DerivativeTest, DenseSparseRneEquivalent) { + // run test on all models + for (const char* local_path : {kEnergyConservingPendulumPath, + kTumblingThinObjectPath, + kDampedActuatorsPath, + kDamperActuatorsPath}) { + const std::string xml_path = GetTestDataFilePath(local_path); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int nD = model->nD; + mjtNum* qDeriv = (mjtNum*) mju_malloc(sizeof(mjtNum)*nD); + mjData* data = mj_makeData(model); + + // take 100 steps so we have some velocities, then call forward + mj_resetData(model, data); + if (model->nu) { + data->ctrl[0] = 0.1; + } + for (int i=0; i < 100; i++) { + mj_step(model, data); + } + mj_forward(model, data); + + // compute qDeriv with sparse function, make local copy + mjd_smooth_vel(model, data, /*flg_bias=*/1); + mju_copy(qDeriv, data->qDeriv, nD); + + // re-compute with dense function + mju_zero(data->qDeriv, model->nD); + mjd_actuator_vel(model, data); + mjd_passive_vel(model, data); + mjd_rne_vel_dense(model, data); + + // expect dense and sparse derivatives to be similar to eps precision + mjtNum eps = 1e-12; + EXPECT_THAT(AsVector(data->qDeriv, nD), + Pointwise(DoubleNear(eps), AsVector(qDeriv, nD))); + + mj_deleteData(data); + mju_free(qDeriv); + mj_deleteModel(model); + } +} + } // namespace } // namespace mujoco diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index e2a6c8da..9517097d 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -77,6 +77,8 @@ TEST_F(EngineIoTest, MakeDataFromPartialModel) { { MJDATA_POINTERS_PREAMBLE((&partial_model)) #define X(type, name, nr, nc) \ + if (strcmp(#name, "D_rownnz") && strcmp(#name, "D_rowadr") && \ + strcmp(#name, "B_rownnz") && strcmp(#name, "B_rowadr")) \ EXPECT_EQ(std::memcmp(data_from_partial->name, data_from_model->name, \ sizeof(type)*(partial_model.nr)*(nc)), \ 0) << "mjData::" #name " differs"; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index d02042c6..021bd781 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -196,6 +196,7 @@ public enum mjtIntegrator : int{ mjINT_EULER = 0, mjINT_RK4 = 1, mjINT_IMPLICIT = 2, + mjINT_IMPLICITFAST = 3, } public enum mjtCollision : int{ mjCOL_ALL = 0, @@ -1658,6 +1659,9 @@ public unsafe struct mjData_ { public int* D_rownnz; public int* D_rowadr; public int* D_colind; + public int* B_rownnz; + public int* B_rowadr; + public int* B_colind; public double* qDeriv; public double* qLU; public double* actuator_force; @@ -1918,6 +1922,7 @@ public unsafe struct mjModel_ { public int nnames_map; public int nM; public int nD; + public int nB; public int nemax; public int njmax; public int nconmax;