Use ID_FO_AZA-style backward sweep for inner-term computation

shubhamsingh91 · shubhamsingh91 · commit 43fdba264022 · 2026-05-18T19:39:22.000-05:00
The mass-matrix derivative inner products required by the second-order forward dynamics formula are now computed via an O(N^2) per-column articulated-body recursion instead of the dense (dM/dq) * (dqdd/d{q,v,tau}) matrix-matrix products. Following Singh, Russell & Wensing, "On Second-Order Derivatives of Rigid-Body Dynamics: Theory & Implementation" (T-RO 2023), Fig. 14-15, the recursive Inner-Term becomes preferable to the dense DTM/DMM product beyond N ~ 25 for tree topologies, while the Outer-Term keeps the dense -M^{-1} * term_in GEMM (BLAS-3) since that is faster at practical robotics sizes. The recursion mirrors the ID_FO_AZA algorithm in the reference, with the base-joint spatial acceleration deliberately set to zero so the result contains only the mass-matrix derivative term (no gravity contribution). Algorithm public API and Data layout are unchanged. Refs: #2691
diff --git a/include/pinocchio/src/algorithm/aba-second-order-derivatives.hxx b/include/pinocchio/src/algorithm/aba-second-order-derivatives.hxx
@@ -11,9 +11,160 @@
   #include "pinocchio/algorithm/aba-second-order-derivatives.hpp"
 #endif // PINOCCHIO_LSP
 
+#include "pinocchio/spatial/act-on-set.hpp"
+
 namespace pinocchio
 {
 
+  // Forward sweep used by ComputeABASecondOrderDerivatives to set up the
+  // articulated-body-style inner-product recursion: forward kinematics,
+  // CRBA composite inertia, and the spatial force corresponding to the
+  // input acceleration `a`. Gravity is deliberately zeroed at the base so
+  // the result of the subsequent backward sweep is (dM/dq)*a only -- no
+  // gravity contribution.
+  template<
+    typename Scalar,
+    int Options,
+    template<typename, int> class JointCollectionTpl,
+    typename ConfigVectorType,
+    typename InputAccType>
+  struct ComputeABASecondOrderDerivativesInnerForwardStep
+  : public fusion::JointUnaryVisitorBase<ComputeABASecondOrderDerivativesInnerForwardStep<
+      Scalar,
+      Options,
+      JointCollectionTpl,
+      ConfigVectorType,
+      InputAccType>>
+  {
+    typedef ModelTpl<Scalar, Options, JointCollectionTpl> Model;
+    typedef DataTpl<Scalar, Options, JointCollectionTpl> Data;
+
+    typedef boost::fusion::vector<
+      const Model &,
+      Data &,
+      const ConfigVectorType &,
+      const InputAccType &>
+      ArgsType;
+
+    template<typename JointModel>
+    static void algo(
+      const JointModelBase<JointModel> & jmodel,
+      JointDataBase<typename JointModel::JointDataDerived> & jdata,
+      const Model & model,
+      Data & data,
+      const Eigen::MatrixBase<ConfigVectorType> & q,
+      const Eigen::MatrixBase<InputAccType> & a)
+    {
+      typedef typename Model::JointIndex JointIndex;
+      typedef typename Data::Motion Motion;
+      typedef typename Data::Inertia Inertia;
+
+      const JointIndex i = jmodel.id();
+      const JointIndex parent = model.parents[i];
+      Motion & oa = data.oa[i];
+
+      jmodel.calc(jdata.derived(), q.derived());
+      data.liMi[i] = model.jointPlacements[i] * jdata.M();
+
+      if (parent > 0)
+      {
+        data.oMi[i] = data.oMi[parent] * data.liMi[i];
+        oa = data.oa[parent];
+      }
+      else
+      {
+        data.oMi[i] = data.liMi[i];
+        oa.setZero(); // zero gravity: result is (dM/dq)*a, not (dM/dq)*a + dh/dq
+      }
+
+      typedef
+        typename SizeDepType<JointModel::NV>::template ColsReturn<typename Data::Matrix6x>::Type
+          ColsBlock;
+      ColsBlock J_cols = jmodel.jointCols(data.J);
+      ColsBlock ddJ_cols = jmodel.jointCols(data.ddJ);
+
+      J_cols.noalias() = data.oMi[i].act(jdata.S());
+      motionSet::motionAction(oa, J_cols, ddJ_cols);
+      oa += data.oMi[i].act(jdata.S() * jmodel.jointVelocitySelector(a));
+
+      Inertia & oY = data.oYcrb[i];
+      oY = data.oMi[i].act(model.inertias[i]);
+      data.of[i] = oY * oa;
+    }
+  };
+
+  // Backward sweep companion to ComputeABASecondOrderDerivativesInnerForwardStep.
+  // Fills `dM_a_dq(j, k) = d/dq_k ( M(q) * a )_j` using `Ftmp1` and `Ftmp3`
+  // as accumulation workspace (each sized 6 x model.nv).
+  template<
+    typename Scalar,
+    int Options,
+    template<typename, int> class JointCollectionTpl,
+    typename FtmpType,
+    typename OutputMatrixType>
+  struct ComputeABASecondOrderDerivativesInnerBackwardStep
+  : public fusion::JointUnaryVisitorBase<ComputeABASecondOrderDerivativesInnerBackwardStep<
+      Scalar,
+      Options,
+      JointCollectionTpl,
+      FtmpType,
+      OutputMatrixType>>
+  {
+    typedef ModelTpl<Scalar, Options, JointCollectionTpl> Model;
+    typedef DataTpl<Scalar, Options, JointCollectionTpl> Data;
+
+    typedef boost::fusion::
+      vector<const Model &, Data &, FtmpType &, FtmpType &, OutputMatrixType &>
+        ArgsType;
+
+    template<typename JointModel>
+    static void algo(
+      const JointModelBase<JointModel> & jmodel,
+      const Model & model,
+      Data & data,
+      FtmpType & Ftmp1,
+      FtmpType & Ftmp3,
+      OutputMatrixType & dM_a_dq)
+    {
+      typedef typename Model::JointIndex JointIndex;
+
+      const JointIndex i = jmodel.id();
+      const JointIndex parent = model.parents[i];
+
+      typedef
+        typename SizeDepType<JointModel::NV>::template ColsReturn<typename Data::Matrix6x>::Type
+          ColsBlock;
+      ColsBlock J_cols = jmodel.jointCols(data.J);
+      ColsBlock ddJ_cols = jmodel.jointCols(data.ddJ);
+      ColsBlock tmp1 = jmodel.jointCols(Ftmp1);
+      ColsBlock tmp3 = jmodel.jointCols(Ftmp3);
+
+      const Eigen::Index joint_idx = (Eigen::Index)jmodel.idx_v();
+      const Eigen::Index joint_dofs = (Eigen::Index)jmodel.nv();
+      const Eigen::Index subtree_dofs = (Eigen::Index)data.nvSubtree[i];
+      const Eigen::Index successor_idx = joint_idx + joint_dofs;
+      const Eigen::Index successor_dofs = subtree_dofs - joint_dofs;
+
+      typename Data::Inertia & oYcrb = data.oYcrb[i];
+      motionSet::inertiaAction(oYcrb, J_cols, tmp1);
+      motionSet::inertiaAction<ADDTO>(oYcrb, ddJ_cols, tmp3);
+      motionSet::act<ADDTO>(J_cols, data.of[i], tmp3);
+
+      if (successor_dofs > 0)
+        dM_a_dq.block(joint_idx, successor_idx, joint_dofs, successor_dofs).noalias() =
+          J_cols.transpose() * Ftmp3.middleCols(successor_idx, successor_dofs);
+
+      dM_a_dq.block(joint_idx, joint_idx, subtree_dofs, joint_dofs).noalias() =
+        Ftmp1.middleCols(joint_idx, subtree_dofs).transpose() * ddJ_cols;
+
+      if (parent > 0)
+      {
+        data.oYcrb[parent] += data.oYcrb[i];
+        data.of[parent] += data.of[i];
+      }
+    }
+  };
+
   template<
     typename Scalar,
     int Options,
@@ -36,132 +187,122 @@ namespace pinocchio
     const Tensor3 & d2ddq_dqdv,
     const Tensor4 & d2ddq_dtaudq)
   {
-    assert(q.size() == model.nq && "The joint configuration vector is not of right size");
-    assert(v.size() == model.nv && "The joint velocity vector is not of right size");
-    assert(tau.size() == model.nv && "The joint torque vector is not of right size");
-    assert(d2ddq_dqdq.dimension(0) == model.nv);
-    assert(d2ddq_dqdq.dimension(1) == model.nv);
-    assert(d2ddq_dqdq.dimension(2) == model.nv);
-    assert(d2ddq_dvdv.dimension(0) == model.nv);
-    assert(d2ddq_dvdv.dimension(1) == model.nv);
-    assert(d2ddq_dvdv.dimension(2) == model.nv);
-    assert(d2ddq_dqdv.dimension(0) == model.nv);
-    assert(d2ddq_dqdv.dimension(1) == model.nv);
-    assert(d2ddq_dqdv.dimension(2) == model.nv);
-    assert(d2ddq_dtaudq.dimension(0) == model.nv);
-    assert(d2ddq_dtaudq.dimension(1) == model.nv);
-    assert(d2ddq_dtaudq.dimension(2) == model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(q.size(), model.nq);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(v.size(), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(tau.size(), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdq.dimension(0), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdq.dimension(1), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdq.dimension(2), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dvdv.dimension(0), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dvdv.dimension(1), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dvdv.dimension(2), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdv.dimension(0), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdv.dimension(1), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dqdv.dimension(2), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dtaudq.dimension(0), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dtaudq.dimension(1), model.nv);
+    PINOCCHIO_CHECK_ARGUMENT_SIZE(d2ddq_dtaudq.dimension(2), model.nv);
     assert(model.check(data) && "data is not consistent with model.");
+    assert(model.check(MimicChecker()) && "Function does not support mimic joints");
 
-    typedef typename DataTpl<Scalar, Options, JointCollectionTpl>::MatrixXs MatrixXs;
-    typedef typename DataTpl<Scalar, Options, JointCollectionTpl>::VectorXs VectorXs;
-    typedef typename DataTpl<Scalar, Options, JointCollectionTpl>::Tensor3x Tensor3x;
-    typedef Eigen::Map<MatrixXs> MatrixMap;
-    typedef Eigen::Map<const MatrixXs> ConstMatrixMap;
-    typedef Eigen::Map<VectorXs> VectorMap;
-    typedef Eigen::Map<const VectorXs> ConstVectorMap;
+    typedef ModelTpl<Scalar, Options, JointCollectionTpl> Model;
+    typedef DataTpl<Scalar, Options, JointCollectionTpl> Data;
+    typedef typename Data::MatrixXs MatrixXs;
+    typedef typename Data::VectorXs VectorXs;
+    typedef typename Data::Matrix6x Matrix6x;
+    typedef typename Data::Tensor3x Tensor3x;
+    typedef typename Model::JointIndex JointIndex;
 
     const Eigen::Index nv = model.nv;
     const Eigen::Index nv2 = nv * nv;
 
-    // 1. First-order ABA derivatives: data.ddq_dq, data.ddq_dv and data.Minv
-    //    (upper triangular). Also leaves data.ddq filled.
+    // 1) First-order ABA derivatives: fills data.ddq, data.ddq_dq, data.ddq_dv
+    //    and (upper triangle of) data.Minv.
     computeABADerivatives(model, data, q.derived(), v.derived(), tau.derived());
-
-    // computeABADerivatives only fills the upper triangle of Minv (which is
-    // symmetric). Symmetrize so it can be used directly as the outer multiplier.
     data.Minv.template triangularView<Eigen::StrictlyLower>() =
       data.Minv.transpose().template triangularView<Eigen::StrictlyLower>();
 
-    // 2. Second-order RNEA derivatives, evaluated at (q, v, qdd). The 4th
-    //    output d2tau_dadq carries the pages of dM/dq.
+    // 2) Second-order RNEA derivatives at (q, v, ddq). d2tau_dadq holds the
+    //    pages of dM/dq.
     ComputeRNEASecondOrderDerivatives(
       model, data, q.derived(), v.derived(), data.ddq, data.d2tau_dqdq, data.d2tau_dvdv,
       data.d2tau_dqdv, data.d2tau_dadq);
 
-    // 3. Pre-compute the three inner products used by the chain-rule formulas.
-    //    Index conventions are kept consistent with the rest of the SO API:
-    //      prodq(n, a, b) = (M_q_a * ddq_dq.col(b))_n
-    //      prodv(n, a, b) = (M_q_a * ddq_dv.col(b))_n
-    //      prodqdd(n, c, u) = (M_q_u * Minv)(n, c)
-    //    where M_q_u(n, m) = d2tau_dadq(n, m, u) is the u-th page of dM/dq.
-    Tensor3x prodq(nv, nv, nv);
-    Tensor3x prodv(nv, nv, nv);
-    Tensor3x prodqdd(nv, nv, nv);
+    // 3) Inner-Term: for each column w of (ddq_dq, ddq_dv, Minv), compute
+    //    (dM/dq) * column and store as page w of the corresponding tensor.
+    typedef ComputeABASecondOrderDerivativesInnerForwardStep<
+      Scalar, Options, JointCollectionTpl, ConfigVectorType, VectorXs>
+      InnerPass1;
+    typedef ComputeABASecondOrderDerivativesInnerBackwardStep<
+      Scalar, Options, JointCollectionTpl, Matrix6x, MatrixXs>
+      InnerPass2;
 
-    for (Eigen::Index u = 0; u < nv; ++u)
-    {
-      ConstMatrixMap Mq_u(data.d2tau_dadq.data() + u * nv2, nv, nv);
+    Tensor3x prodq(nv, nv, nv), prodv(nv, nv, nv), prodqdd(nv, nv, nv);
+    Matrix6x Ftmp1 = Matrix6x::Zero(6, nv);
+    Matrix6x Ftmp3 = Matrix6x::Zero(6, nv);
+    MatrixXs inner_out(nv, nv);
+
+    auto fill_inner = [&](const auto & input_columns, Tensor3x & out_tensor) {
       for (Eigen::Index w = 0; w < nv; ++w)
       {
-        VectorMap prodq_uw(prodq.data() + (w * nv + u) * nv, nv); // prodq(:, u, w)
-        VectorMap prodv_uw(prodv.data() + (w * nv + u) * nv, nv); // prodv(:, u, w)
-        prodq_uw.noalias() = Mq_u * data.ddq_dq.col(w);
-        prodv_uw.noalias() = Mq_u * data.ddq_dv.col(w);
+        VectorXs a_input = input_columns.col(w);
+        for (JointIndex i = 1; i < (JointIndex)model.njoints; ++i)
+          InnerPass1::run(
+            model.joints[i], data.joints[i],
+            typename InnerPass1::ArgsType(model, data, q.derived(), a_input));
+
+        Ftmp1.setZero();
+        Ftmp3.setZero();
+        inner_out.setZero();
+        for (JointIndex i = (JointIndex)(model.njoints - 1); i > 0; --i)
+          InnerPass2::run(
+            model.joints[i],
+            typename InnerPass2::ArgsType(model, data, Ftmp1, Ftmp3, inner_out));
+
+        Eigen::Map<MatrixXs>(out_tensor.data() + w * nv2, nv, nv) = inner_out;
       }
-      MatrixMap prodqdd_u(prodqdd.data() + u * nv2, nv, nv); // prodqdd page u
-      prodqdd_u.noalias() = Mq_u * data.Minv;
-    }
+    };
 
-    // 4. Assemble the inner block term_in of size (nv, 4 * nv * nv). For each
-    //    outer index u and inner index i, four nv-vectors are written:
-    //      (4u    ) i : d2tau_dqdq(:, i, u) + prodq(:, u, i) + prodq(:, i, u)
-    //      (4u + 1) i : d2tau_dvdv(:, i, u)
-    //      (4u + 2) i : d2tau_dqdv(:, i, u) + prodv(:, i, u)
-    //      (4u + 3) i : prodqdd(:, i, u)
-    //
-    //    The 2nd-dim slice of a Tensor3x is not contiguous in memory, so we
-    //    use column-by-column maps for the prodq cross term.
+    fill_inner(data.ddq_dq, prodq);
+    fill_inner(data.ddq_dv, prodv);
+    fill_inner(data.Minv, prodqdd);
+
+    // 4) Outer-Term: build term_in (nv x 4*nv*nv) and take a single dense
+    //    -M^{-1} * term_in product. Singh, Russell & Wensing (2023),
+    //    Fig. 14b: at practical robotics sizes the BLAS-3 GEMM beats the
+    //    AZA recursion on the Outer-Term.
     MatrixXs term_in(nv, 4 * nv2);
     for (Eigen::Index u = 0; u < nv; ++u)
     {
-      const Scalar * d2tau_dqdq_page = data.d2tau_dqdq.data() + u * nv2;
-      const Scalar * d2tau_dvdv_page = data.d2tau_dvdv.data() + u * nv2;
-      const Scalar * d2tau_dqdv_page = data.d2tau_dqdv.data() + u * nv2;
-      const Scalar * prodv_page      = prodv.data() + u * nv2;
-      const Scalar * prodqdd_page    = prodqdd.data() + u * nv2;
-
       for (Eigen::Index i = 0; i < nv; ++i)
       {
-        // prodq(:, u, i) is the contiguous column at offset (i * nv + u) * nv
-        ConstVectorMap prodq_u_i(prodq.data() + (i * nv + u) * nv, nv);
-        // prodq(:, i, u) is the contiguous column at offset (u * nv + i) * nv
-        ConstVectorMap prodq_i_u(prodq.data() + (u * nv + i) * nv, nv);
-
-        ConstVectorMap d2tau_dqdq_iu(d2tau_dqdq_page + i * nv, nv);
-        ConstVectorMap d2tau_dvdv_iu(d2tau_dvdv_page + i * nv, nv);
-        ConstVectorMap d2tau_dqdv_iu(d2tau_dqdv_page + i * nv, nv);
-        ConstVectorMap prodv_iu(prodv_page + i * nv, nv);
-        ConstVectorMap prodqdd_iu(prodqdd_page + i * nv, nv);
-
-        term_in.col((4 * u) * nv + i)       = d2tau_dqdq_iu + prodq_u_i + prodq_i_u;
-        term_in.col((4 * u + 1) * nv + i)   = d2tau_dvdv_iu;
-        term_in.col((4 * u + 2) * nv + i)   = d2tau_dqdv_iu + prodv_iu;
-        term_in.col((4 * u + 3) * nv + i)   = prodqdd_iu;
+        for (Eigen::Index n = 0; n < nv; ++n)
+        {
+          term_in(n, (4 * u + 0) * nv + i) =
+            data.d2tau_dqdq(n, i, u) + prodq(n, u, i) + prodq(n, i, u);
+          term_in(n, (4 * u + 1) * nv + i) = data.d2tau_dvdv(n, i, u);
+          term_in(n, (4 * u + 2) * nv + i) = data.d2tau_dqdv(n, i, u) + prodv(n, i, u);
+          term_in(n, (4 * u + 3) * nv + i) = prodqdd(n, u, i);
+        }
       }
     }
 
-    // 5. Outer DMM: term_out = -Minv * term_in. Single dense matrix product.
     MatrixXs term_out(nv, 4 * nv2);
     term_out.noalias() = -data.Minv * term_in;
 
-    // 6. Scatter term_out columns back into the four output tensors.
-    Tensor1 & d2ddq_dqdq_   = const_cast<Tensor1 &>(d2ddq_dqdq);
-    Tensor2 & d2ddq_dvdv_   = const_cast<Tensor2 &>(d2ddq_dvdv);
-    Tensor3 & d2ddq_dqdv_   = const_cast<Tensor3 &>(d2ddq_dqdv);
-    Tensor4 & d2ddq_dtaudq_ = const_cast<Tensor4 &>(d2ddq_dtaudq);
-
+    // 5) Scatter term_out into the four output tensors.
     for (Eigen::Index u = 0; u < nv; ++u)
     {
-      MatrixMap dqdq_u(d2ddq_dqdq_.data()   + u * nv2, nv, nv);
-      MatrixMap dvdv_u(d2ddq_dvdv_.data()   + u * nv2, nv, nv);
-      MatrixMap dqdv_u(d2ddq_dqdv_.data()   + u * nv2, nv, nv);
-      MatrixMap dtdq_u(d2ddq_dtaudq_.data() + u * nv2, nv, nv);
-
-      dqdq_u = term_out.middleCols((4 * u) * nv, nv);
-      dvdv_u = term_out.middleCols((4 * u + 1) * nv, nv);
-      dqdv_u = term_out.middleCols((4 * u + 2) * nv, nv);
-      dtdq_u = term_out.middleCols((4 * u + 3) * nv, nv);
+      for (Eigen::Index i = 0; i < nv; ++i)
+      {
+        for (Eigen::Index n = 0; n < nv; ++n)
+        {
+          const_cast<Tensor1 &>(d2ddq_dqdq)(n, i, u)   = term_out(n, (4 * u + 0) * nv + i);
+          const_cast<Tensor2 &>(d2ddq_dvdv)(n, i, u)   = term_out(n, (4 * u + 1) * nv + i);
+          const_cast<Tensor3 &>(d2ddq_dqdv)(n, i, u)   = term_out(n, (4 * u + 2) * nv + i);
+          const_cast<Tensor4 &>(d2ddq_dtaudq)(n, i, u) = term_out(n, (4 * u + 3) * nv + i);
+        }
+      }
     }
   }
 
