From 2b192cebb6694f9af4b62f8c19a75f2d7daab402 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Mon, 11 Sep 2023 11:08:08 -0700
Subject: [PATCH] PUBLIC: In `mju_normalize4`, don't normalize if norm is less
 than `mjMINVAL` away from 1.

This makes mj_forward calls idempotent if warmstart is disabled and quats in qpos are normalized.

PiperOrigin-RevId: 564442666
Change-Id: I9c43add82bcfb4101c19fbfbfc8c30091b8906d5
---
 src/engine/engine_util_blas.c      |  2 +-
 test/engine/engine_inverse_test.cc |  2 +-
 test/pipeline_test.cc              | 46 ++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c
index bb7d16f5..d5eb960a 100644
--- a/src/engine/engine_util_blas.c
+++ b/src/engine/engine_util_blas.c
@@ -220,7 +220,7 @@ mjtNum mju_normalize4(mjtNum vec[4]) {
     vec[1] = 0;
     vec[2] = 0;
     vec[3] = 0;
-  } else {
+  } else if (mju_abs(norm - 1) > mjMINVAL) {
     mjtNum normInv = 1/norm;
     vec[0] *= normInv;
     vec[1] *= normInv;
diff --git a/test/engine/engine_inverse_test.cc b/test/engine/engine_inverse_test.cc
index a15b9c78..29936b0b 100644
--- a/test/engine/engine_inverse_test.cc
+++ b/test/engine/engine_inverse_test.cc
@@ -104,7 +104,7 @@ TEST_F(InverseTest, DiscreteInverseMatch) {
 
       // depending on mjENBL_INVDISCRETE flag, expect mismatch to be small/large
       if (invdiscrete) {
-        mjtNum epsilon = 1e-10;
+        mjtNum epsilon = 1e-9;
         EXPECT_LT(data->solver_fwdinv[0], epsilon);
         EXPECT_LT(data->solver_fwdinv[1], epsilon);
       } else {
diff --git a/test/pipeline_test.cc b/test/pipeline_test.cc
index 2d9cb28e..fa2ffb9f 100644
--- a/test/pipeline_test.cc
+++ b/test/pipeline_test.cc
@@ -14,6 +14,9 @@
 
 // Tests of the entire pipeline that are not easily associated with one file.
 
+#include <string>
+#include <vector>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
@@ -61,5 +64,48 @@ TEST_F(PipelineTest, SparseDenseEquivalent) {
   mj_deleteModel(model);
 }
 
+// mj_forward should be deterministic when warm starts are disabled
+TEST_F(PipelineTest, DeterministicNoWarmstart) {
+  const std::string xml_path = GetTestDataFilePath(kDefaultModel);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+  mjData* data2 = mj_makeData(model);
+
+  // disable warmstarts
+  model->opt.disableflags |= mjDSBL_WARMSTART;
+
+  int nv = model->nv;
+
+  int kNumSteps = 50;
+
+  for (mjtSolver solver : {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG}) {
+    model->opt.solver = solver;
+    mj_resetData(model, data);
+    mj_resetData(model, data2);
+
+    for (int step = 0; step < kNumSteps; step++) {
+      mj_step(model, data);
+      mj_forward(model, data);
+
+      mj_step(model, data2);
+      mj_forward(model, data2);
+
+      // test determinism: both models steps did the same thing
+      EXPECT_EQ(AsVector(data->qacc, nv), AsVector(data2->qacc, nv));
+
+      // one more mj_forward call on data2
+      mj_forward(model, data2);
+
+      // expect that the extra mj_forward call didn't change anything
+      EXPECT_EQ(AsVector(data->qacc, nv), AsVector(data2->qacc, nv));
+    }
+  }
+
+  mj_deleteData(data2);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+
 }  // namespace
 }  // namespace mujoco