Add diffusion dynamics task and related methods to MJScene class

Author: juan-g-bonilla

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp

@@ -32,6 +32,7 @@
 #include <stdexcept>
 #include <vector>
 #include <cmath>
+#include <unordered_set>
 
 MJScene::MJScene(std::string xml, const std::vector<std::string>& files)
     : spec(*this, xml, files)
@@ -59,13 +60,31 @@ void MJScene::AddFwdKinematicsToDynamicsTask(int32_t priority)
     this->AddModelToDynamicsTask(this->ownedSysModel.back().get(), priority);
 }
 
-void MJScene::SelfInit() { this->dynamicsTask.SelfInitTaskList(); }
+void MJScene::AddModelToDiffusionDynamicsTask(SysModel* model, int32_t priority)
+{
+    this->dynamicsDiffusionTask.AddNewObject(model, priority);
+}
+
+void MJScene::AddFwdKinematicsToDiffusionDynamicsTask(int32_t priority)
+{
+    this->ownedSysModel.emplace_back(std::make_unique<MJFwdKinematics>(*this));
+    this->ownedSysModel.back()->ModelTag = "FwdKinematics" + std::to_string(this->ownedSysModel.size()-1);
+    this->AddModelToDiffusionDynamicsTask(this->ownedSysModel.back().get(), priority);
+}
+
+void MJScene::SelfInit()
+{
+    this->dynamicsTask.SelfInitTaskList();
+    this->dynamicsDiffusionTask.SelfInitTaskList();
+}
 
 void MJScene::Reset(uint64_t CurrentSimNanos)
 {
     this->timeBefore = CurrentSimNanos * NANO2SEC;
     this->dynamicsTask.TaskName = "Dynamics:" + this->ModelTag;
     this->dynamicsTask.ResetTaskList(CurrentSimNanos);
+    this->dynamicsDiffusionTask.TaskName = "DiffusionDynamics:" + this->ModelTag;
+    this->dynamicsDiffusionTask.ResetTaskList(CurrentSimNanos);
     this->initializeDynamics();
     this->writeOutputStateMessages(CurrentSimNanos);
 }
@@ -95,16 +114,31 @@ void MJScene::initializeDynamics()
     }
 
     // Register the states of the models in the dynamics task
-    for (auto[_, sysModelPtr] : this->dynamicsTask.TaskModels)
+    std::unordered_set<StatefulSysModel*> alreadyRegisteredModels;
+    auto registerStatesOnSysModel = [this, &alreadyRegisteredModels](SysModel* sysModelPtr)
     {
         if (auto statefulSysModelPtr = dynamic_cast<StatefulSysModel*>(sysModelPtr))
         {
+            // Don't registerStates in a model twice!
+            if (alreadyRegisteredModels.count(statefulSysModelPtr) > 0) return;
+
             statefulSysModelPtr->registerStates(DynParamRegisterer(
                 this->dynManager,
                 sysModelPtr->ModelTag.empty() ? std::string("model") : sysModelPtr->ModelTag
                 + "_" + std::to_string(sysModelPtr->moduleID) + "_"
             ));
+
+            alreadyRegisteredModels.emplace(statefulSysModelPtr);
         }
+    };
+
+    for (auto[_, sysModelPtr] : this->dynamicsTask.TaskModels)
+    {
+        registerStatesOnSysModel(sysModelPtr);
+    }
+    for (auto[_, sysModelPtr] : this->dynamicsDiffusionTask.TaskModels)
+    {
+        registerStatesOnSysModel(sysModelPtr);
     }
 }
 
@@ -206,6 +240,12 @@ void MJScene::equationsOfMotion(double t, double timeStep)
     }
 }
 
+void MJScene::equationsOfMotionDiffusion(double t, double timeStep)
+{
+    auto nanos = static_cast<uint64_t>(t * SEC2NANO);
+    this->dynamicsDiffusionTask.ExecuteTaskList(nanos);
+}
+
 void MJScene::preIntegration(double callTime) { this->timeStep = callTime - this->timeBefore; }
 
 void MJScene::postIntegration(double callTime)

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h

@@ -290,6 +290,12 @@ class MJScene : public DynamicObject
      * task typically compute and apply forces and torques on the system
      * that depend on the simulation time or the state of the multi-body.
      *
+     * Models in this task should contribute to computing ``f`` in:
+     *
+     *     dx = f(t,x) dt
+     *
+     * where ``x`` represents the state of this dynamic object.
+     *
      * @warning `SysModel` added to the dynamics task should be memoryless.
      * That is, any output message computed should depend strictly on
      * input messages and the current time/states. It should save values
@@ -303,6 +309,23 @@ class MJScene : public DynamicObject
      */
     void AddModelToDynamicsTask(SysModel* model, int32_t priority = -1);
 
+    /** @brief Adds a model to the diffusion dynamics task.
+     *
+     * This task and function are only relevant when the dynamics of
+     * the system are stochastic.
+     *
+     * Similar to ``AddModelToDynamicsTask``, except that models in this
+     * task contribute to computing ``g`` in:
+     *
+     *     dx = f(t,x)dt + g(t,x)dW
+     *
+     * where ``x`` is the state of the system, ``f`` represents the
+     * drift term of the dynamics (classic time derivative), and
+     * ``g`` is the diffusion term of the dynamics (which evaluates
+     * the impact of the random 'noise' ``W`` on the dynamics).
+     */
+    void AddModelToDiffusionDynamicsTask(SysModel* model, int32_t priority = -1);
+
     /**
      * @brief Adds forward kinematics to the dynamics task.
      *
@@ -334,6 +357,17 @@ class MJScene : public DynamicObject
      */
     void AddFwdKinematicsToDynamicsTask(int32_t priority);
 
+    /**
+     * @brief Adds forward kinematics to the diffusion dynamics task.
+     *
+     * See ``AddModelToDiffusionDynamicsTask`` and ``AddFwdKinematicsToDynamicsTask``.
+     *
+     * By default, the diffusion dynamics task has no forward kinematics
+     * model, so one must be added if the diffusion term depends on the
+     * forward kinematics of the multibody.
+     */
+    void AddFwdKinematicsToDiffusionDynamicsTask(int32_t priority);
+
     /**
      * @brief Calls `SelfInit` on all system models in the dynamics task.
      */
@@ -367,11 +401,30 @@ class MJScene : public DynamicObject
      * joints... These information is translated, through MuJoCo, into
      * first derivatives of the joint states and the mass of bodies.
      *
+     * Computes ``f`` in:
+     *
+     *     dx = f(t,x) dt
+     *
+     * where ``x`` represents the current state of the dynamics.
+     *
      * @param t The current simulation time in seconds.
      * @param timeStep The integration time step.
      */
     void equationsOfMotion(double t, double timeStep) override;
 
+    /**
+     * @brief Computes the diffusion of the dynamics of the system.
+     *
+     * Only relevant for systems with stochastic dynamics.
+     *
+     * Computes ``g`` in:
+     *
+     *     dx = f(t,x) dt + g(t,x)dW
+     *
+     * where ``x`` represents the current state of the dynamics.
+     */
+    void equationsOfMotionDiffusion(double t, double timeStep) override;
+
     /**
      * @brief Prepares the system before actual integration.
      *
@@ -558,6 +611,7 @@ class MJScene : public DynamicObject
     bool forwardKinematicsStale = true; ///< Flag indicating stale forward kinematics.
 
     SysModelTask dynamicsTask; ///< Task managing models involved in the dynamics of this scene.
+    SysModelTask dynamicsDiffusionTask; ///< Task managing models involved in the diffusion stochastic dynamics of this scene.
     std::vector<std::unique_ptr<SysModel>> ownedSysModel; ///< System models that should be cleared on this scene destruction.
 
     MJQPosStateData* qposState; ///< Position state data.