Set time correctly in mjd_transitionFD, to support time-dependent user code.

PiperOrigin-RevId: 502380864
Change-Id: Ief36135a900c091d95e4450992c7c73e86de267c

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,
+static void setState(const mjModel* m, mjData* d, mjtNum time, const mjtNum* state,
-                     const mjtNum* warmstart) {
+                     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 *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 *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 *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 *DT = D ? mj_stackAlloc(d, nu*ns) : NULL;    // control-observation matrix (transposed)
 
   // set offset pointers
   if (A) {

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