From 4d4b0bb2c3989dae9de70ccc23e91d5a744ac220 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 29 Jun 2024 03:06:07 -0700 Subject: [PATCH] Don't normalize `mjData->qpos` quaternions in-place. PiperOrigin-RevId: 647927542 Change-Id: I13b0be55498d1da3af2cdc414cfed4c3908a6fe1 --- doc/changelog.rst | 4 +- mjx/requirements.txt | 3 +- src/engine/engine_core_constraint.c | 10 ++- src/engine/engine_core_smooth.c | 16 ++-- src/engine/engine_passive.c | 6 +- src/engine/engine_sensor.c | 9 ++- src/engine/engine_util_spatial.c | 2 +- test/engine/engine_derivative_test.cc | 15 +++- test/engine/engine_forward_test.cc | 108 ++++++++++++++++++++++++++ test/engine/engine_sensor_test.cc | 6 +- 10 files changed, 157 insertions(+), 22 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 79d4db8d70..6832af733f 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -12,7 +12,7 @@ General 1. Removed deprecated ``mj_makeEmptyFileVFS`` and ``mj_findFileVFS`` functions. - **Migration:** Use:ref:`mj_addBufferVFS` to copy a buffer into a VFS file directly. + **Migration:** Use :ref:`mj_addBufferVFS` to copy a buffer into a VFS file directly. 2. Calls to:ref:`mj_defaultVFS` may allocate memory inside VFS, and the corresponding :ref:`mj_deleteVFS` must be called to deallocate any internal allocated memory. @@ -44,6 +44,8 @@ General point, as usual. If ``mjUSESINGLE`` is defined, MuJoCo will use single-precision floating point. See :ref:`mjtNum`. Relatedly, fixed various type errors that prevented building with single-precision. +10. Quaternions in ``mjData->qpos`` are no longer normalized in-place by :ref:`mj_kinematics`. Instead they are + normalized when they are used. After the first step, quaternions will be normalized. MJX ~~~ diff --git a/mjx/requirements.txt b/mjx/requirements.txt index 927cd13497..d3ca30bee0 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -36,7 +36,8 @@ jaxlib==0.4.18; python_version >= '3.9' \ --hash=sha256:2b17b3f05b3bbf8e0ddb85fba339525ac03bac21c9f26d0f83dcea1b1654353e \ --hash=sha256:b35ec08984e2aa5e96ba3f3f8b88e90dee0283649e037f213dec8e85638fa17d \ --hash=sha256:0bcc4768d29be80d20fd542aafd3a02510a5b1e47c7953beef8b03a9941fa64d \ - --hash=sha256:1e4b0a0d2cfad3905123af2aa0afc2f6fcfaa9a2f74eb247cd5f140645817ab2 + --hash=sha256:1e4b0a0d2cfad3905123af2aa0afc2f6fcfaa9a2f74eb247cd5f140645817ab2 \ + --hash=sha256:c97469fb0b2be880ae716ecd9254c15c494fa73fa469daa2f31d233fe7a8bb6f pip==23.3.1 \ --hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b pytest==7.4.2 \ diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index e17c2badfa..6cb5c26a7b 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -818,7 +818,10 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { // BALL joint else if (m->jnt_type[i] == mjJNT_BALL) { // convert joint quaternion to axis-angle - mju_quat2Vel(angleAxis, d->qpos+m->jnt_qposadr[i], 1); + int adr = m->jnt_qposadr[i]; + mjtNum quat[4] = {d->qpos[adr], d->qpos[adr+1], d->qpos[adr+2], d->qpos[adr+3]}; + mju_normalize4(quat); + mju_quat2Vel(angleAxis, quat, 1); // get rotation angle, normalize value = mju_normalize3(angleAxis); @@ -1767,7 +1770,10 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { } else if (m->jnt_type[i] == mjJNT_BALL) { mjtNum angleAxis[3]; - mju_quat2Vel(angleAxis, d->qpos+m->jnt_qposadr[i], 1); + int adr = m->jnt_qposadr[i]; + mjtNum quat[4] = {d->qpos[adr], d->qpos[adr+1], d->qpos[adr+2], d->qpos[adr+3]}; + mju_normalize4(quat); + mju_quat2Vel(angleAxis, quat, 1); value = mju_normalize3(angleAxis); dist = mju_max(m->jnt_range[2*i], m->jnt_range[2*i+1]) - value; if (dist < margin) { diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 3dfbb09bf3..8d1e2144be 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -44,9 +44,6 @@ void mj_kinematics(const mjModel* m, mjData* d) { d->xmat[0] = d->xmat[4] = d->xmat[8] = 1; d->ximat[0] = d->ximat[4] = d->ximat[8] = 1; - // normalize all quaternions in qpos - mj_normalizeQuat(m, d->qpos); - // normalize mocap quaternions for (int i=0; i < m->nmocap; i++) { mju_normalize4(d->mocap_quat+4*i); @@ -66,6 +63,7 @@ void mj_kinematics(const mjModel* m, mjData* d) { // copy pos and quat from qpos mju_copy3(xpos, d->qpos+qadr); mju_copy4(xquat, d->qpos+qadr+3); + mju_normalize4(xquat); // assign xanchor and xaxis mju_copy3(d->xanchor+3*jntadr, xpos); @@ -125,6 +123,7 @@ void mj_kinematics(const mjModel* m, mjData* d) { mjtNum qloc[4]; if (jtype == mjJNT_BALL) { mju_copy4(qloc, d->qpos+qadr); + mju_normalize4(qloc); } else { mju_axisAngle2Quat(qloc, m->jnt_axis+3*jid, d->qpos[qadr] - m->qpos0[qadr]); } @@ -890,13 +889,15 @@ void mj_transmission(const mjModel* m, mjData* d) { int j = m->jnt_qposadr[id]; // axis: expmap representation of quaternion - mju_quat2Vel(axis, d->qpos+j, 1); + mju_copy4(quat, d->qpos+j); + mju_normalize4(quat); + mju_quat2Vel(axis, quat, 1); // gearAxis: rotate to parent frame if necessary if (m->actuator_trntype[i] == mjTRN_JOINT) { mju_copy3(gearAxis, gear); } else { - mju_negQuat(quat, d->qpos+j); + mju_negQuat(quat, quat); mju_rotVecQuat(gearAxis, gear, quat); } @@ -923,12 +924,15 @@ void mj_transmission(const mjModel* m, mjData* d) { // axis: expmap representation of quaternion mju_quat2Vel(axis, d->qpos+j+3, 1); + mju_copy4(quat, d->qpos+j+3); + mju_normalize4(quat); + mju_quat2Vel(axis, quat, 1); // gearAxis: rotate to world frame if necessary if (m->actuator_trntype[i] == mjTRN_JOINT) { mju_copy3(gearAxis, gear+3); } else { - mju_negQuat(quat, d->qpos+j+3); + mju_negQuat(quat, quat); mju_rotVecQuat(gearAxis, gear+3, quat); } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 8412f8f52d..c66bb7e43a 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -64,9 +64,11 @@ static void mj_springdamper(const mjModel* m, mjData* d) { case mjJNT_BALL: { - mjtNum dif[3]; // convert quatertion difference into angular "velocity" - mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr); + mjtNum dif[3], quat[4]; + mju_copy4(quat, d->qpos+padr); + mju_normalize4(quat); + mju_subQuat(dif, quat, m->qpos_spring + padr); // apply torque d->qfrc_spring[dadr+0] = -stiffness*dif[0]; diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index 9fe2c2c30d..415472dc48 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -281,6 +281,7 @@ void mj_sensorPos(const mjModel* m, mjData* d) { case mjSENS_BALLQUAT: // ballquat mju_copy4(d->sensordata+adr, d->qpos+m->jnt_qposadr[objid]); + mju_normalize4(d->sensordata+adr); break; case mjSENS_JOINTLIMITPOS: // jointlimitpos @@ -899,7 +900,7 @@ void mj_sensorAcc(const mjModel* m, mjData* d) { // position-dependent energy (potential) void mj_energyPos(const mjModel* m, mjData* d) { int padr; - mjtNum dif[3], stiffness; + mjtNum dif[3], quat[4], stiffness; // disabled: clear and return if (!mjENABLED(mjENBL_ENERGY)) { @@ -923,7 +924,9 @@ void mj_energyPos(const mjModel* m, mjData* d) { switch ((mjtJoint) m->jnt_type[i]) { case mjJNT_FREE: - mju_sub3(dif, d->qpos+padr, m->qpos_spring+padr); + mju_copy4(quat, d->qpos+padr); + mju_normalize4(quat); + mju_sub3(dif, quat, m->qpos_spring+padr); d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif); // continue with rotations @@ -932,6 +935,8 @@ void mj_energyPos(const mjModel* m, mjData* d) { case mjJNT_BALL: // covert quatertion difference into angular "velocity" + mju_copy4(quat, d->qpos+padr); + mju_normalize4(quat); mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr); d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif); break; diff --git a/src/engine/engine_util_spatial.c b/src/engine/engine_util_spatial.c index 60c4af9d74..39cc30b9af 100644 --- a/src/engine/engine_util_spatial.c +++ b/src/engine/engine_util_spatial.c @@ -247,8 +247,8 @@ void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale) { mju_copy3(tmp, vel); angle = scale * mju_normalize3(tmp); mju_axisAngle2Quat(qrot, tmp, angle); - mju_mulQuat(quat, quat, qrot); mju_normalize4(quat); + mju_mulQuat(quat, quat, qrot); } diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 2b9a4702e1..85ecf60ac7 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -223,11 +223,23 @@ TEST_F(DerivativeTest, StepSkip) { mjINT_IMPLICITFAST}) { model->opt.integrator = integrator; - // reset, take 20 steps, save initial state + // reset, take 20 steps mj_resetData(model, data); for (int i=0; i < 20; i++) { mj_step(model, data); } + + // denormalize the quat, just to see that it doesn't make a difference + for (int j=0; j < model->njnt; j++) { + if (model->jnt_type[j] == mjJNT_BALL) { + int adr = model->jnt_qposadr[j]; + for (int k=0; k < 4; k++) { + data->qpos[adr + k] *= 8; + } + } + } + + // save state std::vector qpos = AsVector(data->qpos, nq); std::vector qvel = AsVector(data->qvel, nv); @@ -279,7 +291,6 @@ TEST_F(DerivativeTest, StepSkip) { mj_deleteModel(model); } - // Analytic transition matrices for linear dynamical system xn = A*x + B*u // given modified mass matrix H (`data->qH`) and // Ac = H^-1 [diag(-stiffness) diag(-damping)] diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index 41d58a711f..72f3bd12e8 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -27,11 +27,16 @@ #include #include #include +#include #include "src/cc/array_safety.h" #include "src/engine/engine_callback.h" #include "src/engine/engine_io.h" #include "test/fixture.h" +#ifdef MEMORY_SANITIZER + #include +#endif + namespace mujoco { namespace { @@ -607,6 +612,109 @@ TEST_F(ForwardTest, eq_active) { mj_deleteModel(model); } +// test that normalized and denormalized quats give the same result +TEST_F(ForwardTest, NormalizeQuats) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + )"; + mjModel* model = LoadModelFromString(xml); + ASSERT_THAT(model, NotNull()); + + mjData* data_u = mj_makeData(model); + + // we'll compare all the memory, so unpoison it first + #ifdef MEMORY_SANITIZER + __msan_unpoison(data_u->buffer, data_u->nbuffer); + __msan_unpoison(data_u->arena, data_u->narena); + #endif + + // set quats to denormalized values, non-zero velocities + for (int i = 3; i < model->nq; i++) data_u->qpos[i] = i; + for (int i = 0; i < model->nv; i++) data_u->qvel[i] = 0.1*i; + + // copy data and normalize quats + mjData* data_n = mj_copyData(nullptr, model, data_u); + mj_normalizeQuat(model, data_n->qpos); + + // call forward, expect quats to be untouched + mj_forward(model, data_u); + for (int i = 3; i < model->nq; i++) { + EXPECT_EQ(data_u->qpos[i], (mjtNum)i); + } + + // expect that the ball joint limit is active + EXPECT_EQ(data_u->nl, 1); + + // step both models + mj_step(model, data_u); + mj_step(model, data_n); + + // expect everything to match + MJDATA_POINTERS_PREAMBLE(model) + #define X(type, name, nr, nc) \ + for (int i = 0; i < model->nr; i++) \ + for (int j = 0; j < nc; j++) \ + EXPECT_EQ(data_n->name[i*nc+j], data_u->name[i*nc+j]); + MJDATA_POINTERS; + #undef X + + // repeat the above with RK4 integrator + model->opt.integrator = mjINT_RK4; + + // reset data, unpoison + mj_resetData(model, data_u); + #ifdef MEMORY_SANITIZER + __msan_unpoison(data_u->buffer, data_u->nbuffer); + __msan_unpoison(data_u->arena, data_u->narena); + #endif + + // set quats to un-normalized values, non-zero velocities + for (int i = 3; i < model->nq; i++) data_u->qpos[i] = i; + for (int i = 0; i < model->nv; i++) data_u->qvel[i] = 0.1*i; + + // copy data and normalize quats + mj_copyData(data_n, model, data_u); + mj_normalizeQuat(model, data_n->qpos); + + // step both models + mj_step(model, data_u); + mj_step(model, data_n); + + // expect everything to match + #define X(type, name, nr, nc) \ + for (int i = 0; i < model->nr; i++) \ + for (int j = 0; j < nc; j++) \ + EXPECT_EQ(data_n->name[i*nc+j], data_u->name[i*nc+j]); + MJDATA_POINTERS; + #undef X + + mj_deleteData(data_n); + mj_deleteData(data_u); + mj_deleteModel(model); + +#ifdef STOP_MSAN + #define MEMORY_SANITIZER +#endif +} + // user defined 2nd-order activation dynamics: frequency-controlled oscillator // note that scalar mjcb_act_dyn callbacks are expected to return act_dot, but // since we have a vector output we write into act_dot directly diff --git a/test/engine/engine_sensor_test.cc b/test/engine/engine_sensor_test.cc index 0667c3213f..005136bdae 100644 --- a/test/engine/engine_sensor_test.cc +++ b/test/engine/engine_sensor_test.cc @@ -207,15 +207,11 @@ TEST_F(RelativeFrameSensorTest, ReferencePosMatQuat) { } mj_forward(model, data); - // note that in the loop above the quat is unnormalized, but that's ok, - // quaternions are automatically normalized in place: - EXPECT_NEAR(mju_norm(data->qpos+3, 4), 1.0, tol); - // get values from relative sensors after moving the object std::vector actual_values(data->sensordata+nsensordata/2, data->sensordata+nsensordata); - // object and reference have moved together, we expect values to not unchange + // object and reference have moved together, we expect values to not change EXPECT_THAT(actual_values, Pointwise(DoubleNear(tol), expected_values)); mj_deleteData(data);