Add sensor matrices to mjd_transitionFD.

- Also add missing `mj_jacSubtreeCom` to Python bindings.

PiperOrigin-RevId: 471610181
Change-Id: I31410d194527ce6a7bdbd01e176f3a41efa52d84

doc/changelog.rst

@@ -61,6 +61,7 @@ General
 - Added a disable flag for sensors.
 - :ref:`mju_mulQuat` and :ref:`mju_mulQuatAxis` support in place computation. For example
   |br| ``mju_mulQuat(a, a, b);`` sets the quaternion ``a`` equal to the product of ``a`` and ``b``.
+- Added sensor matrices to ``mjd_transitionFD`` (note this is an API change).
 
 Deleted/deprecated features
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^

include/mujoco/mujoco.h

@@ -1091,12 +1091,17 @@ MJAPI mjtNum mju_sigmoid(mjtNum x);
 
 //---------------------- Derivatives ---------------------------------------------------------------
 
-// Finite differenced state-transition and control-transition matrices dx(t+h) = A*dx(t) + B*du(t).
+// Finite differenced transition matrices (control theory notation)
+//   d(x_next) = A*dx + B*du
+//   d(sensor) = C*dx + D*du
 //   required output matrix dimensions:
 //      A: (2*nv+na x 2*nv+na)
 //      B: (2*nv+na x nu)
-MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
+//      D: (nsensordata x 2*nv+na)
-                            mjtNum* A, mjtNum* B);
+//      C: (nsensordata x nu)
+void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
+                      mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
+
 
 
 #if defined(__cplusplus)

introspect/functions.py

@@ -6980,7 +6980,19 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
              ),
+             FunctionParameterDecl(
+                 name='C',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='D',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+             ),
          ),
-         doc='Finite differenced state-transition and control-transition matrices dx(t+h) = A*dx(t) + B*du(t).   required output matrix dimensions:      A: (2*nv+na x 2*nv+na)      B: (2*nv+na x nu)',  # pylint: disable=line-too-long
+         doc='Finite differenced transition matrices (control theory notation)   d(x_next) = A*dx + B*du   d(sensor) = C*dx + D*du   required output matrix dimensions:      A: (2*nv+na x 2*nv+na)      B: (2*nv+na x nu)      D: (nsensordata x 2*nv+na)      C: (nsensordata x nu)',  # pylint: disable=line-too-long
      )),
 ])

python/mujoco/functions.cc

@@ -355,6 +355,17 @@ PYBIND11_MODULE(_functions, pymodule) {
             m, d, jacp.has_value() ? jacp->data() : nullptr,
             jacr.has_value() ? jacr->data() : nullptr, body);
       });
+  Def<traits::mj_jacSubtreeCom>(
+      pymodule,
+      [](const raw::MjModel* m, raw::MjData* d,
+         std::optional<Eigen::Ref<EigenArrayXX>> jacp, int body) {
+        if (jacp.has_value() &&
+            (jacp->rows() != 3 || jacp->cols() != m->nv)) {
+          throw py::type_error("jacp should be of shape (3, nv)");
+        }
+        return InterceptMjErrors(::mj_jacSubtreeCom)(
+            m, d, jacp.has_value() ? jacp->data() : nullptr, body);
+      });
   Def<traits::mj_jacGeom>(
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d,
@@ -1033,7 +1044,9 @@ PYBIND11_MODULE(_functions, pymodule) {
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d, mjtNum eps, mjtByte centered,
          std::optional<Eigen::Ref<EigenArrayXX>> A,
-         std::optional<Eigen::Ref<EigenArrayXX>> B) {
+         std::optional<Eigen::Ref<EigenArrayXX>> B,
+         std::optional<Eigen::Ref<EigenArrayXX>> C,
+         std::optional<Eigen::Ref<EigenArrayXX>> D) {
         if (A.has_value() &&
             (A->rows() != 2*m->nv+m->na || A->cols() != 2*m->nv+m->na)) {
           throw py::type_error("A should be of shape (2*nv+na, 2*nv+na)");
@@ -1042,10 +1055,20 @@ PYBIND11_MODULE(_functions, pymodule) {
             (B->rows() != 2*m->nv+m->na || B->cols() != m->nu)) {
           throw py::type_error("B should be of shape (2*nv+na, nu)");
         }
+        if (C.has_value() &&
+            (C->rows() != m->nsensordata || C->cols() != 2*m->nv+m->na)) {
+          throw py::type_error("C should be of shape (nsensordata, 2*nv+na)");
+        }
+        if (D.has_value() &&
+            (D->rows() != m->nsensordata || D->cols() != m->nu)) {
+          throw py::type_error("D should be of shape (nsensordata, nu)");
+        }
         return InterceptMjErrors(::mjd_transitionFD)(
             m, d, eps, centered,
             A.has_value() ? A->data() : nullptr,
-            B.has_value() ? B->data() : nullptr);
+            B.has_value() ? B->data() : nullptr,
+            C.has_value() ? C->data() : nullptr,
+            D.has_value() ? D->data() : nullptr);
       });
   Def<traits::mju_Halton>(pymodule);
   // Skipped: mju_strncpy (doesn't make sense in Python)

src/engine/engine_derivative.c

@@ -1463,7 +1463,7 @@ void mjd_smooth_vel(const mjModel *m, mjData *d) {
 
 // finite differenced Jacobian of  (next_state, sensors) = mj_step(state, control)
 //   all outputs are optional
-//   output dimensions (transposed w.r.t common convention):
+//   output dimensions (transposed w.r.t Control Theory convention):
 //     DyDq: (nv x 2*nv+na)
 //     DyDv: (nv x 2*nv+na)
 //     DyDa: (na x 2*nv+na)
@@ -1712,30 +1712,53 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
 
 
 
-// finite differenced state-transition and control-transition matrices dy = A*dx + B*du
+// finite differenced transition matrices (control theory notation)
+//   d(x_next) = A*dx + B*du
+//   d(sensor) = C*dx + D*du
 //   required output matrix dimensions:
 //      A: (2*nv+na x 2*nv+na)
 //      B: (2*nv+na x nu)
+//      D: (nsensordata x 2*nv+na)
+//      C: (nsensordata x nu)
 void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
-                      mjtNum* A, mjtNum* B) {
+                      mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D) {
-  int nv = m->nv, na = m->na, nu = m->nu;
+  int nv = m->nv, na = m->na, nu = m->nu, ns = m->nsensordata;
-  int ndx = 2*nv+na;  // row length of Jacobians
+  int ndx = 2*nv+na;  // row length of state Jacobians
+
+  // stepFD() offset pointers, initialised to NULL
+  mjtNum *DyDq, *DyDv, *DyDa, *DsDq, *DsDv, *DsDa;
+  DyDq = DyDv = DyDa = DsDq = DsDv = DsDa = NULL;
+
   mjMARKSTACK;
 
   // allocate transposed matrices
-  mjtNum *AT = mj_stackAlloc(d, ndx*ndx);            // state-transition matrix   (transposed)
+  mjtNum *AT = A ? mj_stackAlloc(d, ndx*ndx) : NULL; // state-transition matrix   (transposed)
   mjtNum *BT = B ? mj_stackAlloc(d, nu*ndx) : NULL;  // control-transition matrix (transposed)
+  mjtNum *CT = C ? mj_stackAlloc(d, ndx*ns) : NULL;  // state-observation matrix   (transposed)
+  mjtNum *DT = D ? mj_stackAlloc(d, nu*ns) : NULL;   // control-observation matrix (transposed)
+
+  // set offset pointers
+  if (A) {
+    DyDq = AT;
+    DyDv = AT+ndx*nv;
+    DyDa = AT+ndx*2*nv;
+  }
+
+  if (C) {
+    DsDq = CT;
+    DsDv = CT + ns*nv;
+    DsDa = CT + ns*2*nv;
+  }
 
   // get Jacobians
-  if (A) {
+  mjd_stepFD(m, d, eps, centered, DyDq, DyDv, DyDa, BT, DsDq, DsDv, DsDa, DT);
-    mjd_stepFD(m, d, eps, centered, AT, AT+ndx*nv, AT+ndx*2*nv, BT, NULL, NULL, NULL, NULL);
+
-  } else {
-    mjd_stepFD(m, d, eps, centered, NULL, NULL, NULL, BT, NULL, NULL, NULL, NULL);
-  }
 
   // transpose
   if (A) mju_transpose(A, AT, ndx, ndx);
   if (B) mju_transpose(B, BT, nu, ndx);
+  if (C) mju_transpose(C, CT, ndx, ns);
+  if (D) mju_transpose(D, DT, nu, ns);
 
   mjFREESTACK;
 }

src/engine/engine_derivative.h

@@ -39,9 +39,9 @@ MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* Df
 // advance simulation using control callback, skipstage is mjtStage
 MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
 
-// finite differenced state-transition and control-transition matrices dy = A*dx + B*du
+// finite differenced transition matrices (control theory notation)
 MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
-                            mjtNum* A, mjtNum* B);
+                            mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
 
 #ifdef __cplusplus
 }

test/engine/engine_derivative_test.cc

@@ -352,7 +352,8 @@ TEST_F(DerivativeTest, LinearSystem) {
   mjtNum* AFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv);
   mjtNum* BFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
 
-  mjd_transitionFD(model, data, eps, /*centered=*/0, AFD, BFD);
+  mjd_transitionFD(model, data, eps, /*centered=*/0,
+                   AFD, BFD, nullptr, nullptr);
 
   // uncomment for debugging:
   // PrintMatrix(AFD, 2*nv, 2*nv);
@@ -365,13 +366,12 @@ TEST_F(DerivativeTest, LinearSystem) {
   // central differenced A and B
   mjtNum* AFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv);
   mjtNum* BFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
-  mjd_transitionFD(model, data, eps, /*centered=*/1, AFDc, BFDc);
+  mjd_transitionFD(model, data, eps, /*centered=*/1,
+                   AFDc, BFDc, nullptr, nullptr);
 
-  // expect central derivatives to be closer to analytic solution
+  // expect central derivatives to be equal to forward differences
-  EXPECT_LT(CompareMatrices(A, AFDc, 2*nv, 2*nv, eps),
+  CompareMatrices(AFD, AFDc, 2*nv, 2*nv, eps);
-            CompareMatrices(A, AFD, 2*nv, 2*nv, eps));
+  CompareMatrices(BFD, BFDc, 2*nv, nu, eps);
-  EXPECT_LT(CompareMatrices(B, BFDc, 2*nv, nu, eps),
-            CompareMatrices(B, BFD, 2*nv, nu, eps));
 
   mju_free(BFDc);
   mju_free(AFDc);
@@ -400,7 +400,7 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
   // analytic B
   mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
 
-  LinearSystem(model, data, NULL, B);
+  LinearSystem(model, data, nullptr, B);
 
   // forward differenced A and B
   mjtNum eps = 1e-6;
@@ -409,19 +409,22 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
   // set ctrl to the limits, request forward differences
   data->ctrl[0] =  1;
   data->ctrl[1] = -1;
-  mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD);
+  mjd_transitionFD(model, data, eps, /*centered=*/0,
+                   nullptr, BFD, nullptr, nullptr);
   // expect FD and analytic derivatives to be similar to eps precision
   CompareMatrices(B, BFD, 2*nv, nu, eps);
 
   // ctrl remains at limits, request central differences
-  mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD);
+  mjd_transitionFD(model, data, eps, /*centered=*/1,
+                   nullptr, BFD, nullptr, nullptr);
   // expect FD and analytic derivatives to be similar to eps precision
   CompareMatrices(B, BFD, 2*nv, nu, eps);
 
   // set ctrl beyond limits, request forward differences
   data->ctrl[0] =  2;
   data->ctrl[1] = -2;
-  mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD);
+  mjd_transitionFD(model, data, eps, /*centered=*/0,
+                   nullptr, BFD, nullptr, nullptr);
   // expect derivatives to be 0
   EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0)));
 
@@ -430,7 +433,8 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
   EXPECT_EQ(data->ctrl[1], -2.0);
 
   // ctrl remains beyond limits, request centered differences
-  mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD);
+  mjd_transitionFD(model, data, eps, /*centered=*/1,
+                   nullptr, BFD, nullptr, nullptr);
   // expect derivatives to be 0
   EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0)));
 
@@ -440,5 +444,63 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
   mj_deleteModel(model);
 }
 
+// compare FD sensor derivatives to analytic derivatives
+TEST_F(DerivativeTest, SensorDerivatives) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <joint name="joint" type="slide"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+
+    <actuator>
+      <general name="actuator" joint="joint" gainprm="3"/>
+    </actuator>
+
+    <sensor>
+      <jointpos joint="joint"/>
+      <jointvel joint="joint"/>
+      <actuatorfrc actuator="actuator"/>
+    </sensor>
+  </mujoco>
+  )";
+
+  mjModel* model = LoadModelFromString(xml);
+  int nv = model->nv, nu = model->nu, ns = model->nsensordata;
+  mjData* data = mj_makeData(model);
+
+  // finite differenced C and D
+  mjtNum eps = 1e-6;
+  mjtNum* CFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*2*nv);
+  mjtNum* DFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*nu);
+  mjd_transitionFD(model, data, eps, /*centered=*/0,
+                   nullptr, nullptr, CFD, DFD);
+
+  // expected analytic C and D
+  mjtNum C[6] = {
+    1, 0,
+    0, 1,
+    0, 0
+  };
+
+  mjtNum D[3] = {
+    0,
+    0,
+    3,
+  };
+
+  // compare expected and actual values
+  CompareMatrices(CFD, C, ns, 2*nv, eps);
+  CompareMatrices(DFD, D, ns, nu, eps);
+
+  mju_free(DFD);
+  mju_free(CFD);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+
 }  // namespace
 }  // namespace mujoco