From 49066e17943fb1494ef91f3c270b945f06274de8 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sun, 7 May 2023 02:20:51 -0700
Subject: [PATCH] equivalent but clearer quat-diff call in mj_differentiatePos

PiperOrigin-RevId: 530075836
Change-Id: I3bbe5dec17ccdf53c372bb81087dddde30a24c5c
---
 src/engine/engine_support.c        | 13 ++----
 test/engine/engine_support_test.cc | 73 ++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 8 deletions(-)

diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index 4815c1b1..fe21f5fa 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -1224,14 +1224,11 @@ void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]
 // compute velocity by finite-differencing two positions
 void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
                          const mjtNum* qpos1, const mjtNum* qpos2) {
-  int padr, vadr;
-  mjtNum neg[4], dif[4];
-
   // loop over joints
   for (int j=0; j<m->njnt; j++) {
     // get addresses in qpos and qvel
-    padr = m->jnt_qposadr[j];
-    vadr = m->jnt_dofadr[j];
+    int padr = m->jnt_qposadr[j];
+    int vadr = m->jnt_dofadr[j];
 
     switch (m->jnt_type[j]) {
     case mjJNT_FREE:
@@ -1245,9 +1242,9 @@ void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
       mjFALLTHROUGH;
 
     case mjJNT_BALL:
-      mju_negQuat(neg, qpos1+padr);           // solve:  qpos1 * dif = qpos2
-      mju_mulQuat(dif, neg, qpos2+padr);
-      mju_quat2Vel(qvel+vadr, dif, dt);
+      // solve:  qpos1 * quat(qvel * dt) = qpos2
+      mju_subQuat(qvel+vadr, qpos2+padr, qpos1+padr);
+      mju_scl3(qvel+vadr, qvel+vadr, 1/dt);
       break;
 
     case mjJNT_HINGE:
diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc
index 5b351edd..b77c7fd9 100644
--- a/test/engine/engine_support_test.cc
+++ b/test/engine/engine_support_test.cc
@@ -14,6 +14,7 @@
 
 // Tests for engine/engine_support.c.
 
+#include <random>
 #include <string>
 
 #include <gmock/gmock.h>
@@ -25,9 +26,14 @@
 namespace mujoco {
 namespace {
 
+std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
+  return std::vector<mjtNum>(array, array + n);
+}
+
 using ::testing::DoubleNear;
 using ::testing::ContainsRegex;
 using ::testing::MatchesRegex;
+using ::testing::Pointwise;
 using JacobianTest = MujocoTest;
 static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
 
@@ -278,5 +284,72 @@ TEST_F(VersionTest, MjVersionString) {
   EXPECT_THAT(std::string(mj_versionString()), regex_matcher);
 }
 
+
+using SupportTest = MujocoTest;
+
+// utility: generate two random quaternions with a given angle difference
+void randomQuatPair(mjtNum qa[4], mjtNum qb[4], mjtNum angle, int seed) {
+  // make distribution using seed
+  std::mt19937_64 rng;
+  rng.seed(seed);
+  std::normal_distribution<double> dist(0, 1);
+
+  // sample qa = qb
+  for (int i=0; i < 4; i++) {
+    qa[i] = qb[i] = dist(rng);
+  }
+  mju_normalize4(qa);
+  mju_normalize4(qb);
+
+  // integrate qb in random direction by angle
+  mjtNum dir[3];
+  for (int i=0; i < 3; i++) {
+    dir[i] = dist(rng);
+  }
+  mju_normalize3(dir);
+  mju_quatIntegrate(qb, dir, angle);
+}
+
+static constexpr char ballJointModel[] = R"(
+<mujoco>
+  <worldbody>
+    <body>
+      <joint type="ball"/>
+      <geom size="1"/>
+    </body>
+  </worldbody>
+</mujoco>
+)";
+
+TEST_F(SupportTest, DifferentiatePosSubQuat) {
+  const mjtNum eps = 1e-12;  // epsilon for float comparison
+
+  mjModel* model = LoadModelFromString(ballJointModel);
+
+  int seed = 1;
+  for (mjtNum angle : {0.0, 1e-5, 1e-2}) {
+    for (mjtNum dt : {1e-6, 1e-3, 1e-1}) {
+      // random quaternion pair with given angle difference
+      mjtNum qpos1[4], qpos2[4];
+      randomQuatPair(qpos1, qpos2, angle, seed++);
+
+      // get velocity given timestep
+      mjtNum qvel[3];
+      mj_differentiatePos(model, qvel, dt, qpos1, qpos2);
+
+      // equivalent computation
+      mjtNum qneg[4], qdif[4], qvel_expect[3];
+      mju_negQuat(qneg, qpos1);
+      mju_mulQuat(qdif, qneg, qpos2);
+      mju_quat2Vel(qvel_expect, qdif, dt);
+
+      // expect numerical equality
+      EXPECT_THAT(AsVector(qvel, 3), Pointwise(DoubleNear(eps), qvel_expect));
+    }
+  }
+
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco