diff --git a/docs/source/Support/bskReleaseNotesSnippets/mjscene-mass-change-and-perf.rst b/docs/source/Support/bskReleaseNotesSnippets/mjscene-mass-change-and-perf.rst new file mode 100644 index 00000000000..315cd46415e --- /dev/null +++ b/docs/source/Support/bskReleaseNotesSnippets/mjscene-mass-change-and-perf.rst @@ -0,0 +1,2 @@ +- Fixed :ref:`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` 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. diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp index 920be595d28..7885e091338 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp +++ b/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); diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h index 38029906392..4c265df55ca 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h @@ -268,8 +268,10 @@ class MJBody : public MJObject * 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 diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp index 3c1388afaa0..0129893ed90 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/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("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 diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/_UnitTest/test_mass_update.py b/src/simulation/mujocoDynamics/_GeneralModuleFiles/_UnitTest/test_mass_update.py index d1f05fe12d9..5b54765c5b4 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/_UnitTest/test_mass_update.py +++ b/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: