[RTM after next push] SOCP: Making the noise dependant on the states + controls

README.md

@@ -2098,7 +2098,7 @@ Let us take a look at the structure of the code. First, tau_min, tau_max, and ta
 to -1, 1 and 0 if the integer `actuator_type` (a parameter of the `prepare_ocp` function) equals 1. 
 In this case, the dynamics function used is `DynamicsFcn.TORQUE_ACTIVATIONS_DRIVEN`. 
 
-### The [trampo_quaternions.py](./bioptim/examples/torque_driven_ocp/trampo_quaternions.py) file
+### The [example_quaternions.py](./bioptim/examples/torque_driven_ocp/example_quaternions.py) file
 This example uses a representation of a human body by a trunk_leg segment and two arms.
 It is designed to show how to use a model that has quaternions in their degrees of freedom.
 

bioptim/__init__.py

@@ -188,6 +188,7 @@
 from .limits.path_conditions import BoundsList, InitialGuessList
 from .limits.fatigue_path_conditions import FatigueBounds, FatigueInitialGuess
 from .limits.penalty_controller import PenaltyController
+from .limits.penalty_helpers import PenaltyHelpers
 from .misc.enums import (
     Axis,
     Node,

bioptim/dynamics/configure_new_variable.py

@@ -418,9 +418,11 @@ def _declare_cx_and_plot(self):
                 )
                 if not self.skip_plot:
                     self.nlp.plot[f"{self.name}_states"] = CustomPlot(
-                        lambda t0, phases_dt, node_idx, x, u, p, a: x[self.nlp.states.key_index(self.name), :]
-                        if x.any()
-                        else np.ndarray((cx[0][0].shape[0], 1)) * np.nan,
+                        lambda t0, phases_dt, node_idx, x, u, p, a: (
+                            x[self.nlp.states.key_index(self.name), :]
+                            if x.any()
+                            else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
+                        ),
                         plot_type=PlotType.INTEGRATED,
                         axes_idx=self.axes_idx,
                         legend=self.legend,
@@ -453,15 +455,19 @@ def _declare_cx_and_plot(self):
                 plot_type = PlotType.PLOT if self.nlp.control_type == ControlType.LINEAR_CONTINUOUS else PlotType.STEP
                 if not self.skip_plot:
                     self.nlp.plot[f"{self.name}_controls"] = CustomPlot(
-                        lambda t0, phases_dt, node_idx, x, u, p, a: u[self.nlp.controls.key_index(self.name), :]
-                        if u.any()
-                        else np.ndarray((cx[0][0].shape[0], 1)) * np.nan,
+                        lambda t0, phases_dt, node_idx, x, u, p, a: (
+                            u[self.nlp.controls.key_index(self.name), :]
+                            if u.any()
+                            else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
+                        ),
                         plot_type=plot_type,
                         axes_idx=self.axes_idx,
                         legend=self.legend,
-                        combine_to=f"{self.name}_states"
-                        if self.as_states and self.combine_state_control_plot
-                        else self.combine_name,
+                        combine_to=(
+                            f"{self.name}_states"
+                            if self.as_states and self.combine_state_control_plot
+                            else self.combine_name
+                        ),
                     )
 
         if self.as_states_dot:
@@ -601,9 +607,9 @@ def _manage_fatigue_to_new_variable(
                 var_names_with_suffix[-1], name_elements, ocp, nlp, as_states, as_controls, skip_plot=True
             )
             nlp.plot[f"{var_names_with_suffix[-1]}_controls"] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key: u[nlp.controls.key_index(key), :]
-                if u.any()
-                else np.ndarray((len(name_elements), 1)) * np.nan,
+                lambda t0, phases_dt, node_idx, x, u, p, a, key: (
+                    u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan
+                ),
                 plot_type=PlotType.STEP,
                 combine_to=control_plot_name,
                 key=var_names_with_suffix[-1],
@@ -612,9 +618,9 @@ def _manage_fatigue_to_new_variable(
         elif i == 0:
             NewVariableConfiguration(f"{name}", name_elements, ocp, nlp, as_states, as_controls, skip_plot=True)
             nlp.plot[f"{name}_controls"] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key: u[nlp.controls.key_index(key), :]
-                if u.any()
-                else np.ndarray((len(name_elements), 1)) * np.nan,
+                lambda t0, phases_dt, node_idx, x, u, p, a, key: (
+                    u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan
+                ),
                 plot_type=PlotType.STEP,
                 combine_to=control_plot_name,
                 key=f"{name}",
@@ -625,9 +631,9 @@ def _manage_fatigue_to_new_variable(
             name_tp = f"{var_names_with_suffix[-1]}_{params}"
             NewVariableConfiguration(name_tp, name_elements, ocp, nlp, True, False, skip_plot=True)
             nlp.plot[name_tp] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key, mod: mod * x[nlp.states.key_index(key), :]
-                if x.any()
-                else np.ndarray((len(name_elements), 1)) * np.nan,
+                lambda t0, phases_dt, node_idx, x, u, p, a, key, mod: (
+                    mod * x[nlp.states.key_index(key), :] if x.any() else np.ndarray((len(name_elements), 1)) * np.nan
+                ),
                 plot_type=PlotType.INTEGRATED,
                 combine_to=fatigue_plot_name,
                 key=name_tp,

bioptim/dynamics/configure_problem.py

@@ -298,18 +298,157 @@
                 phase=nlp.phase_idx,
             )
 
+    @staticmethod
+    def torque_driven_free_floating_base(
+        ocp,
+        nlp,
+        with_contact: bool = False,
+        with_passive_torque: bool = False,
+        with_ligament: bool = False,
+        with_friction: bool = False,
+        external_forces: list = None,
+    ):
+        """
+        Configure the dynamics for a torque driven program with a free floating base.
+        This version of the torque driven dynamics avoids defining a mapping to force the root to generate null forces and torques.
+        (states are q_root, q_joints, qdot_root, and qdot_joints, controls are tau_joints)
+        Please note that it was not meant to be used with quaternions yet.
+
+        Parameters
+        ----------
+        ocp: OptimalControlProgram
+            A reference to the ocp
+        nlp: NonLinearProgram
+            A reference to the phase
+        with_contact: bool
+            If the dynamic with contact should be used
+        with_passive_torque: bool
+            If the dynamic with passive torque should be used
+        with_ligament: bool
+            If the dynamic with ligament should be used
+        with_friction: bool
+            If the dynamic with joint friction should be used (friction = coefficients * qdot)
+        external_forces: list[Any]
+            A list of external forces
+        """
+
+        _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx)
+        _check_external_forces_format(external_forces, nlp.ns, nlp.phase_idx)
+        _check_external_forces_and_phase_dynamics(external_forces, nlp.phase_dynamics, nlp.phase_idx)
+
+        nb_q = nlp.model.nb_q
+        nb_qdot = nlp.model.nb_qdot
+        nb_root = nlp.model.nb_root
+
+        # Declared rigidbody states and controls
+        name_q_roots = [str(i) for i in range(nb_root)]
+        ConfigureProblem.configure_new_variable(
+            "q_roots",
+            name_q_roots,
+            ocp,
+            nlp,
+            as_states=True,
+            as_controls=False,
+            as_states_dot=False,
+        )
+
+        name_q_joints = [str(i) for i in range(nb_root, nb_q)]
+        ConfigureProblem.configure_new_variable(
+            "q_joints",
+            name_q_joints,
+            ocp,
+            nlp,
+            as_states=True,
+            as_controls=False,
+            as_states_dot=False,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "qdot_roots",
+            name_q_roots,
+            ocp,
+            nlp,
+            as_states=True,
+            as_controls=False,
+            as_states_dot=True,
+        )
+
+        name_qdot_joints = [str(i) for i in range(nb_root, nb_qdot)]
+        ConfigureProblem.configure_new_variable(
+            "qdot_joints",
+            name_qdot_joints,
+            ocp,
+            nlp,
+            as_states=True,
+            as_controls=False,
+            as_states_dot=True,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "qddot_roots",
+            name_q_roots,
+            ocp,
+            nlp,
+            as_states=False,
+            as_controls=False,
+            as_states_dot=True,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "qddot_joints",
+            name_qdot_joints,
+            ocp,
+            nlp,
+            as_states=False,
+            as_controls=False,
+            as_states_dot=True,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "tau_joints",
+            name_qdot_joints,
+            ocp,
+            nlp,
+            as_states=False,
+            as_controls=True,
+            as_states_dot=False,
+        )
+
+        # TODO: add implicit constraints + soft contacts + fatigue
+
+        # Configure the actual ODE of the dynamics
+        if nlp.dynamics_type.dynamic_function:
+            ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom)
+        else:
+            ConfigureProblem.configure_dynamics_function(
+                ocp,
+                nlp,
+                DynamicsFunctions.torque_driven_free_floating_base,
+                with_contact=with_contact,
+                with_passive_torque=with_passive_torque,
+                with_ligament=with_ligament,
+                with_friction=with_friction,
+                external_forces=external_forces,
+            )
+
+        # Configure the contact forces
+        if with_contact:
+            ConfigureProblem.configure_contact_function(
+                ocp, nlp, DynamicsFunctions.forces_from_torque_driven, external_forces=external_forces
+            )
+
     @staticmethod
     def stochastic_torque_driven(
         ocp,
         nlp,
         problem_type,
         with_contact: bool = False,
-        with_friction: bool = True,
+        with_friction: bool = False,
         with_cholesky: bool = False,
         initial_matrix: DM = None,
     ):
         """
-        Configure the dynamics for a torque driven program (states are q and qdot, controls are tau)
+        Configure the dynamics for a torque driven stochastic program (states are q and qdot, controls are tau)
 
         Parameters
         ----------
@@ -375,6 +514,80 @@
             with_friction=with_friction,
         )
 
+    @staticmethod
+    def stochastic_torque_driven_free_floating_base(
+        ocp,
+        nlp,
+        problem_type,
+        with_contact: bool = False,
+        with_friction: bool = False,
+        with_cholesky: bool = False,
+        initial_matrix: DM = None,
+    ):
+        """
+        Configure the dynamics for a stochastic torque driven program with a free floating base.
+        (states are q_roots, q_joints, qdot_roots, and qdot_joints, controls are tau_joints)
+
+        Parameters
+        ----------
+        ocp: OptimalControlProgram
+            A reference to the ocp
+        nlp: NonLinearProgram
+            A reference to the phase
+        with_contact: bool
+            If the dynamic with contact should be used
+        with_friction: bool
+            If the dynamic with joint friction should be used (friction = coefficient * qdot)
+        """
+        n_noised_tau = nlp.model.n_noised_controls
+        n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0]
+        n_noised_states = nlp.model.n_noised_states
+
+        # Stochastic variables
+        ConfigureProblem.configure_stochastic_k(
+            ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references
+        )
+        ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references)
+        n_collocation_points = 1
+        if isinstance(problem_type, SocpType.COLLOCATION):
+            n_collocation_points += problem_type.polynomial_degree
+        ConfigureProblem.configure_stochastic_m(
+            ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points
+        )
+
+        if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT):
+            if initial_matrix is None:
+                raise RuntimeError(
+                    "The initial value for the covariance matrix must be provided for TRAPEZOIDAL_EXPLICIT"
+                )
+            ConfigureProblem.configure_stochastic_cov_explicit(
+                ocp, nlp, n_noised_states=n_noised_states, initial_matrix=initial_matrix
+            )
+        else:
+            if with_cholesky:
+                ConfigureProblem.configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states=n_noised_states)
+            else:
+                ConfigureProblem.configure_stochastic_cov_implicit(ocp, nlp, n_noised_states=n_noised_states)
+
+        if isinstance(problem_type, SocpType.TRAPEZOIDAL_IMPLICIT):
+            ConfigureProblem.configure_stochastic_a(ocp, nlp, n_noised_states=n_noised_states)
+            ConfigureProblem.configure_stochastic_c(ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise)
+
+        ConfigureProblem.torque_driven_free_floating_base(
+            ocp=ocp,
+            nlp=nlp,
+            with_contact=with_contact,
+            with_friction=with_friction,
+        )
+
+        ConfigureProblem.configure_dynamics_function(
+            ocp,
+            nlp,
+            DynamicsFunctions.stochastic_torque_driven_free_floating_base,
+            with_contact=with_contact,
+            with_friction=with_friction,
+        )
+
     @staticmethod
     def torque_derivative_driven(
         ocp,
@@ -1576,7 +1789,9 @@
     """
 
     TORQUE_DRIVEN = (ConfigureProblem.torque_driven,)
+    TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.torque_driven_free_floating_base,)
     STOCHASTIC_TORQUE_DRIVEN = (ConfigureProblem.stochastic_torque_driven,)
+    STOCHASTIC_TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.stochastic_torque_driven_free_floating_base,)
     TORQUE_DERIVATIVE_DRIVEN = (ConfigureProblem.torque_derivative_driven,)
     TORQUE_ACTIVATIONS_DRIVEN = (ConfigureProblem.torque_activations_driven,)
     JOINTS_ACCELERATION_DRIVEN = (ConfigureProblem.joints_acceleration_driven,)