From 06557c14899d3bcac7035fed97c3fa4302f42431 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 16 Jan 2023 07:51:36 -0800 Subject: [PATCH] Set `time` correctly in `mjd_transitionFD`, to support time-dependent user code. PiperOrigin-RevId: 502380864 Change-Id: Ief36135a900c091d95e4450992c7c73e86de267c --- src/engine/engine_derivative.c | 32 ++++++++-------- test/engine/engine_derivative_test.cc | 55 +++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 15 deletions(-) diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index cebc1ef1..9da2f846 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -240,10 +240,11 @@ static void getState(const mjModel* m, const mjData* d, mjtNum* state, mjtNum* s // set state=[qpos; qvel; act] and optionally warmstart accelerations -static void setState(const mjModel* m, mjData* d, const mjtNum* state, const mjtNum* ctrl, - const mjtNum* warmstart) { +static void setState(const mjModel* m, mjData* d, mjtNum time, const mjtNum* state, + const mjtNum* ctrl, const mjtNum* warmstart) { int nq = m->nq, nv = m->nv, na = m->na; + d->time = time; mju_copy(d->qpos, state, nq); mju_copy(d->qvel, state+nq, nv); mju_copy(d->act, state+nq+nv, na); @@ -1501,6 +1502,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, mjtNum *ctrl = mj_stackAlloc(d, nu); // save current inputs + mjtNum time = d->time; mju_copy(ctrl, d->ctrl, nu); getState(m, d, state, NULL); if (warmstart) { @@ -1514,7 +1516,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next, sensor); // restore input - setState(m, d, state, ctrl, warmstart); + setState(m, d, time, state, ctrl, warmstart); // finite-difference controls: skip=mjSTAGE_VEL, handle ctrl at range limits if (DyDu || DsDu) { @@ -1531,7 +1533,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_plus, sensor_plus); // reset - setState(m, d, state, ctrl, warmstart); + setState(m, d, time, state, ctrl, warmstart); } // nudge backward, if possible given ctrlrange @@ -1546,7 +1548,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_minus, sensor_minus); // reset - setState(m, d, state, ctrl, warmstart); + setState(m, d, time, state, ctrl, warmstart); } // difference states @@ -1574,7 +1576,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_plus, sensor_plus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); // nudge backward if (centered) { @@ -1586,7 +1588,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_minus, sensor_minus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); } // difference states @@ -1621,7 +1623,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_plus, sensor_plus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); // nudge backward if (centered) { @@ -1633,7 +1635,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_minus, sensor_minus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); } // difference states @@ -1670,7 +1672,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_plus, sensor_plus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); // nudge backward if (centered) { @@ -1684,7 +1686,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, getState(m, d, next_minus, sensor_minus); // reset - setState(m, d, state, NULL, warmstart); + setState(m, d, time, state, NULL, warmstart); } // difference states @@ -1732,10 +1734,10 @@ void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, mjMARKSTACK; // allocate transposed matrices - 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) + 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) { diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 574a3523..2aaf79e1 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -36,6 +36,7 @@ using ::testing::Pointwise; using ::testing::DoubleNear; using ::testing::Eq; using ::testing::Each; +using ::testing::NotNull; using DerivativeTest = MujocoTest; // errors smaller than this are ignored @@ -99,6 +100,7 @@ static const char* const kDampedPendulumPath = "engine/testdata/derivative/damped_pendulum.xml"; 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 @@ -503,5 +505,58 @@ TEST_F(DerivativeTest, SensorDerivatives) { } +// derivatives don't mutate the state +TEST_F(DerivativeTest, NoStateMutation) { + const std::string xml_path = GetTestDataFilePath(kModelPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + ASSERT_THAT(model, NotNull()); + mjData* data0 = mj_makeData(model); + mjData* data = mj_makeData(model); + int nv = model->nv, nu = model->nu, na = model->na, ns = model->nsensordata; + + // set time + data->time = data0->time = 0.5; + + for (int i=0; i < nv; i++) { + data->qpos[i] = data0->qpos[i] = (mjtNum) i+1; + data->qvel[i] = data0->qvel[i] = (mjtNum) i+2; + } + + // set ctrl + for (int i=0; i < nu; i++) { + data->ctrl[i] = data0->ctrl[i] = (mjtNum) i+1; + } + + // set act + for (int i=0; i < na; i++) { + data->act[i] = data0->act[i] = (mjtNum) i+1; + } + + + // allocate Jacobians, call derivatives + int ndx = nv+nv+na; + mjtNum* A = (mjtNum*) mju_malloc(sizeof(mjtNum)*ndx*ndx); + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*ndx*nu); + mjtNum* C = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*ndx); + mjtNum* D = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*nu); + mjtNum eps = 1e-6; + mjd_transitionFD(model, data, eps, /*centered=*/0, A, B, C, D); + + // compare states in data and data0 + EXPECT_EQ(data->time, data0->time); + EXPECT_EQ(AsVector(data->qpos, model->nq), AsVector(data0->qpos, model->nq)); + EXPECT_EQ(AsVector(data->qvel, nv), AsVector(data0->qvel, nv)); + EXPECT_EQ(AsVector(data->act, na), AsVector(data0->act, na)); + EXPECT_EQ(AsVector(data->ctrl, nu), AsVector(data0->ctrl, nu)); + + mju_free(D); + mju_free(C); + mju_free(B); + mju_free(A); + mj_deleteData(data); + mj_deleteData(data0); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco