Fix MJScene body position pinned at reference pose on mass change

docs/source/Support/bskReleaseNotesSnippets/mjscene-mass-change-and-perf.rst

@@ -0,0 +1,2 @@
+- Fixed :ref:`MJScene<MJScene>` reporting a body's position pinned at the model reference pose when its mass changed (for example a depleting fuel tank); the integrator state is now restored after the per-step ``mj_setConst`` refresh.
+- Sped up :ref:`MJScene<MJScene>` equations-of-motion evaluation by caching the ``mjModel``/``mjData`` pointers for the duration of the call instead of re-fetching them through the recompile-checking accessor on every access.

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp

@@ -257,10 +257,8 @@ double MJBody::getMass()
     return this->getSpec().getScene().getMassState()->state(this->getId());
 }
 
-void MJBody::updateMujocoModelFromMassProps()
+void MJBody::updateMujocoModelFromMassProps(mjModel* m)
 {
-    auto m = spec.getMujocoModel();
-
     double oldMass = m->body_mass[this->getId()];
     double newMass = this->getMass();
     auto diff = abs(oldMass - newMass);

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h

@@ -268,8 +268,10 @@ class MJBody : public MJObject<mjsBody>
      * The center of mass remains constant.
      *
      * This method will also mark the kinematics and mujoco model 'const' as stale.
+     *
+     * @param m The compiled MuJoCo model.
      */
-    void updateMujocoModelFromMassProps();
+    void updateMujocoModelFromMassProps(mjModel* m);
 
     /**
      * @brief Updates this body's entry of the bulk mass state derivative with

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp

@@ -210,6 +210,11 @@ MJScene::equationsOfMotion(double t, double timeStep)
     // Make sure the model is compiled
     this->spec.recompileIfNeeded();
 
+    // recompileIfNeeded() above is the only thing that can invalidate these, so cache
+    // them for the rest of the call rather than re-fetching through the accessors.
+    mjModel* model = this->spec.getMujocoModel();
+    mjData* data = this->spec.getMujocoData();
+
     // Copy data from Basilisk state objects to MuJoCo structs
     updateMujocoArraysFromStates();
 
@@ -232,7 +237,7 @@ MJScene::equationsOfMotion(double t, double timeStep)
         // The mass of bodies is stored as a state, which may evolve in time.
         // Mujoco expects mass properties to be stored in mjModel, so we need
         // to update the mjModel with the mass properties of the bodies.
-        body.updateMujocoModelFromMassProps();
+        body.updateMujocoModelFromMassProps(model);
 
         body.writeStateDependentOutputMessages(nanos);
     }
@@ -241,7 +246,12 @@ MJScene::equationsOfMotion(double t, double timeStep)
     // supposed to be constant during mujoco simulations (like body mass). However,
     // we need to alter some of them, in which case we need to update the 'constants'.
     if (areMujocoModelConstStale()) {
-        mj_setConst(this->spec.getMujocoModel(), this->spec.getMujocoData());
+        mj_setConst(model, data);
+
+        // mj_setConst overwrites qpos with the reference pose; restore the integrator
+        // state (and re-flag kinematics stale) before anything downstream reads qpos.
+        updateMujocoArraysFromStates();
+        this->mjModelConstStale = false;
     }
 
     // Execute the dynamics task!
@@ -251,21 +261,17 @@ MJScene::equationsOfMotion(double t, double timeStep)
     // These messages can then be read by the rest of modules.
     this->dynamicsTask.ExecuteTaskList(nanos);
 
-    // TODO: When allowing to prescribe the mass, remember to mj_setConst.
-    // This is similar to how prescribing joint states should be followed
-    // by running forward kinematics.
-
     // If the kinematics became stale while running dynamics modules, refresh
     // MuJoCo's cached position/velocity quantities before actuator, equality,
     // and acceleration calculations read them.
     if (areKinematicsStale()) {
-        mj_fwdPosition(this->spec.getMujocoModel(), this->spec.getMujocoData());
-        mj_fwdVelocity(this->spec.getMujocoModel(), this->spec.getMujocoData());
+        mj_fwdPosition(model, data);
+        mj_fwdVelocity(model, data);
     }
 
     // Update the ctrl array in mjData from the inputs in the actuators
     for (auto&& actuator : this->spec.getActuators()) {
-        actuator->updateCtrl(this->spec.getMujocoData());
+        actuator->updateCtrl(data);
     }
 
     // Update the prescribed joints equalities
@@ -274,13 +280,13 @@ MJScene::equationsOfMotion(double t, double timeStep)
     }
 
     // These methods will compute the accelerations
-    mj_fwdActuation(this->spec.getMujocoModel(), this->spec.getMujocoData());
-    mj_fwdAcceleration(this->spec.getMujocoModel(), this->spec.getMujocoData());
-    mj_fwdConstraint(this->spec.getMujocoModel(), this->spec.getMujocoData());
+    mj_fwdActuation(model, data);
+    mj_fwdAcceleration(model, data);
+    mj_fwdConstraint(model, data);
 
     // Sanity check the produced accelerations
-    auto qacc = this->spec.getMujocoData()->qacc;
-    if (std::any_of(qacc, qacc + this->spec.getMujocoModel()->nv, [](mjtNum v) { return std::isnan(v); })) {
+    auto qacc = data->qacc;
+    if (std::any_of(qacc, qacc + model->nv, [](mjtNum v) { return std::isnan(v); })) {
         logAndThrow<std::runtime_error>("Encountered NaN acceleration at time " + std::to_string(t) +
                                         "s in MJScene with ID: " + std::to_string(moduleID));
     }
@@ -293,13 +299,13 @@ MJScene::equationsOfMotion(double t, double timeStep)
     // The derivative of the bulk velocity is the computed acceleration.
     {
         auto qvelDeriv = this->qvelState->stateDeriv.data();
-        std::copy_n(this->spec.getMujocoData()->qacc, this->spec.getMujocoModel()->nv, qvelDeriv);
+        std::copy_n(data->qacc, model->nv, qvelDeriv);
     }
 
     // Also copy the derivative of the actuator states, if we have them
-    if (this->spec.getMujocoModel()->na > 0) {
+    if (model->na > 0) {
         auto actDeriv = this->actState->stateDeriv.data();
-        std::copy_n(this->spec.getMujocoData()->act_dot, this->spec.getMujocoModel()->na, actDeriv);
+        std::copy_n(data->act_dot, model->na, actDeriv);
     }
 
     // Update the derivative of the body mass property states (into the bulk

src/simulation/mujocoDynamics/_GeneralModuleFiles/_UnitTest/test_mass_update.py

@@ -59,8 +59,13 @@ def test_continuouslyChangingMass(showPlots: bool = False):
 
         v(t) = log(m_0) - log(m_0 - 100t)
 
+    Integrating once more gives the position (with x(0) = 0):
+
+        x(t) = t + (m_0 - 100t)/100 * log((m_0 - 100t)/m_0)
+
     In this function we run this scenario in MuJoCo to check that simulation with
-    variable mass returns the same as the analytical solution of the velocity.
+    variable mass returns the same as the analytical solution of the velocity and
+    position.
     """
 
     dt = 1 # s
@@ -71,6 +76,9 @@ def test_continuouslyChangingMass(showPlots: bool = False):
     dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt)))
 
     scene = mujoco.MJScene.fromFile(XML_PATH_BALL)
+    # With extraEoMCall the final message is written inside equationsOfMotion, so
+    # the recorded position reflects the mid-step mj_setConst qpos state directly.
+    scene.extraEoMCall = True
     scSim.AddModelToTask("test", scene)
 
     actPayload = messaging.SingleActuatorMsgPayload()
@@ -80,8 +88,10 @@ def test_continuouslyChangingMass(showPlots: bool = False):
 
     scene.getSingleActuator("ball").actuatorInMsg.subscribeTo(actMsg)
 
+    massDepletionRate = 100 # kg/s
+
     mDotPayload = messaging.SCMassPropsMsgPayload()
-    mDotPayload.massSC = -100 # kg/s
+    mDotPayload.massSC = -massDepletionRate # kg/s
     mDotMsg = messaging.SCMassPropsMsg()
     mDotMsg.write(mDotPayload)
 
@@ -101,8 +111,11 @@ def test_continuouslyChangingMass(showPlots: bool = False):
     scSim.ConfigureStopTime(macros.sec2nano(tf))
     scSim.ExecuteSimulation()
 
-    m_0 = massRec.massSC[0]
-    expected_v = np.log(m_0) - np.log(m_0 - 100 * macros.NANO2SEC * massRec.times())
+    m_0 = massRec.massSC[0] # kg
+    t = macros.NANO2SEC * massRec.times() # s
+    m = m_0 - massDepletionRate * t # kg
+    expected_v = np.log(m_0) - np.log(m) # m/s
+    expected_x = t + m / massDepletionRate * np.log(m / m_0) # m
 
     if showPlots:
         plt.plot(massRec.times() * macros.NANO2SEC, massRec.massSC)
@@ -119,10 +132,25 @@ def test_continuouslyChangingMass(showPlots: bool = False):
         plt.legend()
         plt.xlabel("Time [s]")
         plt.ylabel("Velocity [m/s]")
+
+        plt.figure()
+        plt.plot(
+            posRec.times() * macros.NANO2SEC, posRec.r_BN_N[:, 0], label="Simulated"
+        )
+        plt.plot(
+            posRec.times() * macros.NANO2SEC, expected_x, "k--", label="Analytical"
+        )
+        plt.legend()
+        plt.xlabel("Time [s]")
+        plt.ylabel("Position [m]")
         plt.show()
 
     assert expected_v[-1] == pytest.approx(posRec.v_BN_N[-1, 0])
 
+    # The position catches mj_setConst leaving qpos at the reference pose
+    # (velocity alone integrates correctly even then).
+    assert expected_x[-1] == pytest.approx(posRec.r_BN_N[-1, 0], rel=1e-3)
+
 
 if __name__ == "__main__":
     if True: