Don't normalize mjData->qpos quaternions in-place.

PiperOrigin-RevId: 647927542
Change-Id: I13b0be55498d1da3af2cdc414cfed4c3908a6fe1

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
 ~~~

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 \

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) {

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);
         }
 

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];

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;

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);
 }
 
 

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<mjtNum> qpos = AsVector(data->qpos, nq);
     std::vector<mjtNum> 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)]

test/engine/engine_forward_test.cc

@@ -27,11 +27,16 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mujoco.h>
+#include <mujoco/mjxmacro.h>
 #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 <sanitizer/msan_interface.h>
+#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"(
+  <mujoco>
+    <option integrator="implicit">
+      <flag warmstart="disable" energy="enable"/>
+    </option>
+    <worldbody>
+      <body name="free">
+        <freejoint/>
+        <geom size="1" pos=".1 .2 .3"/>
+      </body>
+      <body pos="3 0 0">
+        <joint name="ball" type="ball" stiffness="100" range="0 10"/>
+        <geom size="1" pos=".1 .2 .3"/>
+      </body>
+    </worldbody>
+    <sensor>
+      <ballquat joint="ball"/>
+      <framequat objtype="body" objname="free"/>
+    </sensor>
+  </mujoco>
+  )";
+  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

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);