Merge pull request #980 from AVSLab/feature/979-mujoco-stateful-sysmodel

Feature/979 mujoco stateful sysmodel

Author: juan-g-bonilla

docs/source/Support/bskReleaseNotes.rst

@@ -70,7 +70,8 @@ Version |release|
 - Added support for showing ``QuadMap`` quadrilateral surface meshes in Vizard, with scenario :ref:`scenarioQuadMaps` detailing usage. Allows users to draw quads on celestial bodies and spacecraft.
 - Added ``fixedframe2lla()`` function in :ref:`vizSupport` which is useful for computing QuadMap mesh interpolations
 - Added QuadMap mesh support functions (:ref:`quadMapSupport`) for displaying camera FOV boxes as projected on the surface of a reference ellipsoid, and drawing rectangular latitude/longitude defined regions.
-
+- :beta:`Mujoco Support`: Added ``StatefulSysModel`` for models in the dynamics task of ``MJScene`` that need to declare
+  continuous-time states. Modified :ref:`scenarioDeployPanels` to illustrate the use of ``StatefulSysModel``.
 
 Version  2.6.0  (Feb. 21, 2025)
 -------------------------------

examples/mujoco/scenarioDeployPanels.py

@@ -22,13 +22,13 @@
 #. ``examples/mujoco/scenarioArmWithThrusters.py``
 
 This script demonstrates how to simulate a spacecraft with solar panels deployed
-using a Proportional-Derivative (PD) controller. This script uses the MuJoCo-based
-:ref:`DynamicObject<dynamicObject>` :ref:`MJScene<MJScene>`.
+using a Proportional-Integral-Derivative (PID) controller. This script uses the
+MuJoCo-based :ref:`DynamicObject<dynamicObject>` :ref:`MJScene<MJScene>`.
 
 In ``mujoco/scenarioArmWithThrusters.py``, we saw how we can constrain joints to
 follow a specific angle by letting the dynamic engine figure out and apply
 the necessary torques. In this script, we are controlling the joints using a
-PD controller. This is a more adequate simulation setup when you want to simulate
+PID controller. This is a more adequate simulation setup when you want to simulate
 or design the control system for these joints. It is also generally more
 computationally efficient than letting the dynamic engine figure out the torques.
 
@@ -45,11 +45,13 @@
 plots and in the 3D visualization how the panels never get deployed over their
 joint limit.
 
-The deployment of the panels is controlled using a Proportional-Derivative (PD)
-controller. The desired position and velocity profiles for the joints are
+The deployment of the panels is controlled using an analog PID controller.
+The desired position and velocity profiles for the joints are
 generated using a trapezoidal/triangular velocity profile. These profiles are
-then used as inputs to the PD controller, which computes the torque required to achieve
-the desired motion.
+then used as inputs to the PID controller, which computes the torque required to achieve
+the desired motion. Note that the controller class extends ``StatefulSysModel``,
+instead of ``SysModel``, since we need to register the integral error as a
+continuous state.
 
 The simulation is run for 80 minutes and the state of the system is recorded.
 The desired and achieved joint angles, as well as the torque applied to each
@@ -65,7 +67,7 @@
 from Basilisk.simulation import mujoco
 from Basilisk.utilities import SimulationBaseClass
 from Basilisk.utilities import macros
-from Basilisk.architecture import sysModel
+from Basilisk.simulation import StatefulSysModel
 from Basilisk.architecture import messaging
 from Basilisk.simulation import svIntegrators
 
@@ -252,7 +254,7 @@ def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = Fa
         # the measured position and velocity of the joint (in this case the
         # exact values are used, but in a real system these may be the product
         # of a sensor), and the output is the torque to be applied to the joint.
-        pdController = PDController()
+        pdController = PIDController()
         pdController.ModelTag = f"{actuatorName}_controller"
 
         # Connect the interpolators to the PD controller for the desired
@@ -385,13 +387,17 @@ def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = Fa
 
 # The following is an example of a Python-based SysModel that
 # can be added to the dynamics task of a MJScene.
-class PDController(sysModel.SysModel):
+class PIDController(StatefulSysModel.StatefulSysModel):
     """
-    A Proportional-Derivative (PD) Controller class for controlling joint states.
+    A Proportional-Integral-Derivative (PID) Controller class for controlling joint states.
+
+    This models an analog PID controller, which means that its output evolves in continuous
+    time, not discrete time. Thus, it should be used within the dynamics task of ``MJScene``.
 
     Attributes:
-        K (float): Proportional gain.
-        P (float): Derivative gain.
+        K_p (float): Proportional gain.
+        K_d (float): Derivative gain.
+        K_i (float): Integral gain.
 
         measuredInMsg (messaging.ScalarJointStateMsgReader): Reader for the measured joint state.
         desiredInMsg (messaging.ScalarJointStateMsgReader): Reader for the desired joint state.
@@ -404,8 +410,9 @@ class PDController(sysModel.SysModel):
     def __init__(self, *args: Any):
         """Initialize"""
         super().__init__(*args)
-        self.K = 0.1
-        self.P = 0.002
+        self.K_p = 0.1
+        self.K_d = 0.002
+        self.K_i = 0.0001
 
         self.measuredInMsg = messaging.ScalarJointStateMsgReader()
         self.desiredInMsg = messaging.ScalarJointStateMsgReader()
@@ -415,15 +422,23 @@ def __init__(self, *args: Any):
 
         self.outputOutMsg = messaging.SingleActuatorMsg()
 
+    def registerStates(self, registerer: StatefulSysModel.DynParamRegisterer):
+        self.integralErrorState = registerer.registerState(1, 1, "integralError")
+        self.integralErrorState.setState([[0]]) # explicitely zero initialize
+
     def UpdateState(self, CurrentSimNanos: int):
         """Computes the control command from the measured and desired
         joint position and velocity."""
         # Compute the error in the state and its derivative
         stateError = self.desiredInMsg().state - self.measuredInMsg().state
         stateDotError = self.desiredDotInMsg().state - self.measuredDotInMsg().state
+        stateIntegralError = self.integralErrorState.getState()[0][0]
 
         # Compute the control output
-        control_output = self.K * stateError + self.P * stateDotError
+        control_output = self.K_p * stateError + self.K_d * stateDotError + self.K_i * stateIntegralError
+
+        # Set the derivative of the integral error inner state
+        self.integralErrorState.setDerivative([[stateError]])
 
         # Write the control output to the output message
         payload = messaging.SingleActuatorMsgPayload()

src/architecture/_GeneralModuleFiles/py_sys_model.i

@@ -1,5 +1,5 @@
 
-%module(directors="1",threads="1") sysModel
+%module(directors="1",threads="1",package="Basilisk.architecture") sysModel
 %{
    #include "sys_model.h"
 %}

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp

@@ -207,9 +207,9 @@ void MJBody::writeStateDependentOutputMessages(uint64_t CurrentSimNanos)
     }
 }
 
-void MJBody::registerStates(DynParamManager& paramManager)
+void MJBody::registerStates(DynParamRegisterer paramManager)
 {
-    this->massState = paramManager.registerState(1, 1, "mujocoBodyMass_" + this->name);
+    this->massState = paramManager.registerState(1, 1, "mass");
 }
 
 void MJBody::updateMujocoModelFromMassProps()

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h

@@ -36,6 +36,7 @@
 #include "MJJoint.h"
 #include "MJObject.h"
 #include "MJSite.h"
+#include "StatefulSysModel.h"
 
 /// @cond
 /**
@@ -252,13 +253,13 @@ class MJBody : public MJObject<mjsBody>
     void writeStateDependentOutputMessages(uint64_t CurrentSimNanos);
 
     /**
-     * @brief Registers the body's states with a dynamic parameter manager.
+     * @brief Registers the body's states with a dynamic parameter registerer.
      *
      * Currently, only the mass of the body is considered a parameter.
      *
-     * @param paramManager Reference to the dynamic parameter manager.
+     * @param paramManager The dynamic parameter registerer.
      */
-    void registerStates(DynParamManager& paramManager);
+    void registerStates(DynParamRegisterer paramManager);
 
     /**
      * @brief Updates the MuJoCo model from the mass properties of the body.

src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp

@@ -20,6 +20,7 @@
 #include "MJScene.h"
 
 #include "MJFwdKinematics.h"
+#include "StatefulSysModel.h"
 
 #include "simulation/dynamics/_GeneralModuleFiles/svIntegratorRK4.h"
 #include "architecture/utilities/macroDefinitions.h"
@@ -79,7 +80,10 @@ void MJScene::initializeDynamics()
     this->actState = this->dynManager.registerState(1, 1, "mujocoAct");
 
     for (auto&& body : this->spec.getBodies()) {
-        body.registerStates(this->dynManager);
+        body.registerStates(DynParamRegisterer(
+            this->dynManager,
+            "body_" + body.getName() + "_"
+        ));
     }
 
     // Make sure the spec is compiled
@@ -89,6 +93,19 @@ void MJScene::initializeDynamics()
     if (!recompiled) {
         this->spec.configure();
     }
+
+    // Register the states of the models in the dynamics task
+    for (auto[_, sysModelPtr] : this->dynamicsTask.TaskModels)
+    {
+        if (auto statefulSysModelPtr = dynamic_cast<StatefulSysModel*>(sysModelPtr))
+        {
+            statefulSysModelPtr->registerStates(DynParamRegisterer(
+                this->dynManager,
+                sysModelPtr->ModelTag.empty() ? std::string("model") : sysModelPtr->ModelTag
+                + "_" + std::to_string(sysModelPtr->moduleID) + "_"
+            ));
+        }
+    }
 }
 
 void MJScene::UpdateState(uint64_t CurrentSimNanos)

src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h

@@ -0,0 +1,113 @@
+/*
+ ISC License
+
+ Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder
+
+ Permission to use, copy, modify, and/or distribute this software for any
+ purpose with or without fee is hereby granted, provided that the above
+ copyright notice and this permission notice appear in all copies.
+
+ THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+ */
+
+#ifndef STATEFUL_SYS_MODEL_H
+#define STATEFUL_SYS_MODEL_H
+
+#include "simulation/dynamics/_GeneralModuleFiles/dynParamManager.h"
+#include "architecture/_GeneralModuleFiles/sys_model.h"
+
+/**A short-lived class passed to StatefulSysModel for them to register
+ * their states.
+ *
+ * This class serves two purposes. First, it adds a prefix to every state
+ * name before registering it on the actual DynParamManager. This prevents
+ * state name collisions between StatefulSysModel as long as the prefix
+ * are unique. Second, it exponses only the registerState method from
+ * the DynParamManager. This prevents StatefulSysModel from registering
+ * properties or accessing the states of other models, which would allow for
+ * information to flow between models without going through the message system.
+ * If a model needs to access information from another model, it should do so
+ * thorugh a message, not by sharing a state or property.
+ */
+class DynParamRegisterer
+{
+public:
+    /** Constructor */
+    DynParamRegisterer(DynParamManager& manager, std::string stateNamePrefix)
+        : manager(manager)
+        , stateNamePrefix(stateNamePrefix)
+        {}
+
+    /** Creates and returns a new state, which will be managed by the
+     * underlying ``DynParamManager``.
+     *
+     * The state name should be unique: registering two states with the
+     * same name on this class will cause an error. Different StatefulSysModel
+     * are allowed to use the same state name, however.
+     *
+     * This method may optionally be templated to create StateData of
+     * subclasses of StateData.
+     */
+    template <typename StateDataType = StateData,
+              std::enable_if_t<std::is_base_of_v<StateData, StateDataType>, bool> = true>
+    inline StateDataType* registerState(uint32_t nRow, uint32_t nCol, std::string stateName)
+    {
+        return this->manager.registerState<StateDataType>(
+            nRow, nCol, this->stateNamePrefix + stateName
+        );
+    }
+
+protected:
+    DynParamManager& manager; ///< wrapped manager
+    std::string stateNamePrefix; ///< prefix added to all registered state names
+};
+
+/** A SysModel that has continuous-time states.
+ *
+ * StatefulSysModel are added on the dynamics task of an MJScene.
+ * On its UpdateState method, a StatefulSysModel should call each state's
+ * setDerivative method. This value will be used by the integrator to
+ * update the state for the next integrator step.
+ *
+ * The sample code below shows how to get the current value of the state
+ * and how to set its derivative. In this case, ``x`` would follow an
+ * exponential trajectory:
+ * \code{.cpp}
+ * void UpdateState(uint64_t CurrentSimNanos) override {
+ *     auto x = this->xState->getState();
+ *     this->xState->setDerivative( x );
+ * }
+ * \endcode
+ */
+class StatefulSysModel : public SysModel
+{
+public:
+    /** Default constructor */
+    StatefulSysModel() = default;
+
+    /**Used to register states on the given DynParamRegisterer.
+     *
+     * The main purpose of this method is for this class to call
+     * ``registerState`` on the registerer. Note that state names
+     * should not be repeated within the same StatefulSysModel.
+     *
+     * \code{.cpp}
+     * void registerStates(DynParamRegisterer& registerer) override {
+     *     this->posState = registerer.register(3, 1, "pos");
+     *     this->massState = registerer.register(1, 1, "mass");
+     *     // etc.
+     * }
+     * \endcode
+     *
+     */
+    virtual void registerStates(DynParamRegisterer registerer) = 0;
+};
+
+#endif