From e3a82247c2da9db96e6627c7b8ddcc09e88eaf93 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 1 Sep 2022 13:00:02 -0700 Subject: [PATCH] Add sensor matrices to `mjd_transitionFD`. - Also add missing `mj_jacSubtreeCom` to Python bindings. PiperOrigin-RevId: 471610181 Change-Id: I31410d194527ce6a7bdbd01e176f3a41efa52d84 --- doc/changelog.rst | 1 + include/mujoco/mujoco.h | 11 +++- introspect/functions.py | 14 ++++- python/mujoco/functions.cc | 27 ++++++++- src/engine/engine_derivative.c | 45 ++++++++++---- src/engine/engine_derivative.h | 4 +- test/engine/engine_derivative_test.cc | 86 +++++++++++++++++++++++---- 7 files changed, 157 insertions(+), 31 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 844e5a82..621b8af6 100644 --- a/doc/changelog.rst +++ b/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 ^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 1847cddf..135db7e1 100644 --- a/include/mujoco/mujoco.h +++ b/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, - mjtNum* A, mjtNum* B); +// 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* C, mjtNum* D); + #if defined(__cplusplus) diff --git a/introspect/functions.py b/introspect/functions.py index 6b7a1965..7a3a198c 100644 --- a/introspect/functions.py +++ b/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 )), ]) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 5153a2bb..897d7afe 100644 --- a/python/mujoco/functions.cc +++ b/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( + pymodule, + [](const raw::MjModel* m, raw::MjData* d, + std::optional> 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( 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> A, - std::optional> B) { + std::optional> B, + std::optional> C, + std::optional> 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(pymodule); // Skipped: mju_strncpy (doesn't make sense in Python) diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 29b84fa6..cebc1ef1 100644 --- a/src/engine/engine_derivative.c +++ b/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) { - int nv = m->nv, na = m->na, nu = m->nu; - int ndx = 2*nv+na; // row length of Jacobians + mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D) { + int nv = m->nv, na = m->na, nu = m->nu, ns = m->nsensordata; + 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, 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); - } + mjd_stepFD(m, d, eps, centered, DyDq, DyDv, DyDa, BT, DsDq, DsDv, DsDa, DT); + // 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; } diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index 801a5dc0..0cc92bfe 100644 --- a/src/engine/engine_derivative.h +++ b/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 } diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 066a1dc3..47cd03be 100644 --- a/test/engine/engine_derivative_test.cc +++ b/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_LT(CompareMatrices(A, AFDc, 2*nv, 2*nv, eps), - CompareMatrices(A, AFD, 2*nv, 2*nv, eps)); - EXPECT_LT(CompareMatrices(B, BFDc, 2*nv, nu, eps), - CompareMatrices(B, BFD, 2*nv, nu, eps)); + // expect central derivatives to be equal to forward differences + CompareMatrices(AFD, AFDc, 2*nv, 2*nv, eps); + CompareMatrices(BFD, BFDc, 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"( + + + + + + + + + + + + + + + + + + + )"; + + 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