chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
PiperOrigin-RevId: 502380864
Change-Id: Ief36135a900c091d95e4450992c7c73e86de267c
This commit is contained in:
Yuval Tassa
2023-01-16 07:51:36 -08:00
committed by Copybara-Service
parent b5813c60ca
commit 06557c1489
2 changed files with 72 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_derivative.c
+17 -15
View File
@@ -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) {
test/engine/engine_derivative_test.cc
+55
View File
@@ -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