[RTR] Torque driven free floating base dynamics

bioptim/dynamics/configure_problem.py

@@ -298,6 +298,143 @@ def torque_driven(
                 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_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,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "qdot_joints",
+            name_q_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_q_joints,
+            ocp,
+            nlp,
+            as_states=False,
+            as_controls=False,
+            as_states_dot=True,
+        )
+
+        ConfigureProblem.configure_new_variable(
+            "tau_joints",
+            name_q_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,
@@ -1576,6 +1713,7 @@ class DynamicsFcn(FcnEnum):
     """
 
     TORQUE_DRIVEN = (ConfigureProblem.torque_driven,)
+    TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.torque_driven_free_floating_base,)
     STOCHASTIC_TORQUE_DRIVEN = (ConfigureProblem.stochastic_torque_driven,)
     TORQUE_DERIVATIVE_DRIVEN = (ConfigureProblem.torque_derivative_driven,)
     TORQUE_ACTIVATIONS_DRIVEN = (ConfigureProblem.torque_activations_driven,)

bioptim/dynamics/dynamics_functions.py

@@ -6,6 +6,7 @@
 from ..optimization.non_linear_program import NonLinearProgram
 from .dynamics_evaluation import DynamicsEvaluation
 from ..models.protocols.stochastic_biomodel import StochasticBioModel
+from ..misc.mapping import BiMapping
 
 
 class DynamicsFunctions:
@@ -189,6 +190,81 @@ def torque_driven(
 
         return DynamicsEvaluation(dxdt, defects)
 
+    @staticmethod
+    def torque_driven_free_floating_base(
+        time: MX.sym,
+        states: MX.sym,
+        controls: MX.sym,
+        parameters: MX.sym,
+        stochastic_variables: MX.sym,
+        nlp,
+        with_contact: bool,
+        with_passive_torque: bool,
+        with_ligament: bool,
+        with_friction: bool,
+        external_forces: list = None,
+    ) -> DynamicsEvaluation:
+        """
+        Forward dynamics driven by joint torques without actuation of the free floating base, optional external forces can be declared.
+
+        Parameters
+        ----------
+        time: MX.sym
+            The time of the system
+        states: MX.sym
+            The state of the system
+        controls: MX.sym
+            The controls of the system
+        parameters: MX.sym
+            The parameters of the system
+        stochastic_variables: MX.sym
+            The stochastic_variables of the system
+        nlp: NonLinearProgram
+            The definition of the system
+        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 friction should be used
+        external_forces: list[Any]
+            The external forces
+
+        Returns
+        ----------
+        DynamicsEvaluation
+            The derivative of the states and the defects of the implicit dynamics
+        """
+
+        q_roots = DynamicsFunctions.get(nlp.states["q_roots"], states)
+        q_joints = DynamicsFunctions.get(nlp.states["q_joints"], states)
+        qdot_roots = DynamicsFunctions.get(nlp.states["qdot_roots"], states)
+        qdot_joints = DynamicsFunctions.get(nlp.states["qdot_joints"], states)
+        tau_joints = DynamicsFunctions.get(nlp.controls["tau_joints"], controls)
+
+        q_full = vertcat(q_roots, q_joints)
+        qdot_full = vertcat(qdot_roots, qdot_joints)
+        n_q = q_full.shape[0]
+
+        tau_joints = (
+            tau_joints + nlp.model.passive_joint_torque(q_full, qdot_full) if with_passive_torque else tau_joints
+        )
+        tau_joints = tau_joints + nlp.model.ligament_joint_torque(q_full, qdot_full) if with_ligament else tau_joints
+        tau_joints = tau_joints + nlp.model.friction_coefficients @ qdot_full if with_friction else tau_joints
+
+        tau_full = vertcat(MX.zeros(nlp.model.nb_root), tau_joints)
+
+        ddq = DynamicsFunctions.forward_dynamics(nlp, q_full, qdot_full, tau_full, with_contact, external_forces)
+        dxdt = MX(nlp.states.shape, ddq.shape[1])
+        dxdt[:n_q, :] = horzcat(*[qdot_full for _ in range(ddq.shape[1])])
+        dxdt[n_q:, :] = ddq
+
+        defects = None
+
+        return DynamicsEvaluation(dxdt, defects)
+
     @staticmethod
     def stochastic_torque_driven(
         time: MX.sym,
@@ -500,13 +576,29 @@ def forces_from_torque_driven(
             The contact forces that ensure no acceleration at these contact points
         """
 
-        q_nlp, q_var = (nlp.states["q"], states) if "q" in nlp.states else (nlp.controls["q"], controls)
-        qdot_nlp, qdot_var = (nlp.states["qdot"], states) if "qdot" in nlp.states else (nlp.controls["qdot"], controls)
-        tau_nlp, tau_var = (nlp.states["tau"], states) if "tau" in nlp.states else (nlp.controls["tau"], controls)
+        def get_var_from_states_or_controls(key: str, nlp: NonLinearProgram, states: MX.sym, controls: MX.sym):
+            if key in nlp.states:
+                out_nlp, out_var = (nlp.states[key], states)
+                out = DynamicsFunctions.get(out_nlp, out_var)
+            elif key in nlp.controls:
+                out_nlp, out_var = (nlp.controls[key], controls)
+                out = DynamicsFunctions.get(out_nlp, out_var)
+            elif f"{key}_roots" in nlp.states and f"{key}_joints" in nlp.states:
+                out_roots_nlp, out_roots_var = (nlp.states[f"{key}_roots"], states)
+                out_roots = DynamicsFunctions.get(out_roots_nlp, out_roots_var)
+                out_joints_nlp, out_joints_var = (nlp.states[f"{key}_joints"], states)
+                out_joints = DynamicsFunctions.get(out_joints_nlp, out_joints_var)
+                out = vertcat(out_roots, out_joints)
+            elif f"{key}_joints" in nlp.controls:
+                out_joints_nlp, out_joints_var = (nlp.controls[f"{key}_joints"], controls)
+                out = DynamicsFunctions.get(out_joints_nlp, out_joints_var)
+            else:
+                raise RuntimeError(f"{key} not found in states or controls")
+            return out
 
-        q = DynamicsFunctions.get(q_nlp, q_var)
-        qdot = DynamicsFunctions.get(qdot_nlp, qdot_var)
-        tau = DynamicsFunctions.get(tau_nlp, tau_var)
+        q = get_var_from_states_or_controls("q", nlp, states, controls)
+        qdot = get_var_from_states_or_controls("qdot", nlp, states, controls)
+        tau = get_var_from_states_or_controls("tau", nlp, states, controls)
         tau = tau + nlp.model.passive_joint_torque(q, qdot) if with_passive_torque else tau
         tau = tau + nlp.model.ligament_joint_torque(q, qdot) if with_ligament else tau
 
@@ -936,25 +1028,30 @@ def forward_dynamics(
         -------
         The derivative of qdot
         """
-        qdot_var = nlp.states["qdot"] if "qdot" in nlp.states else nlp.controls["qdot"]
+        if "qdot" in nlp.states:
+            qdot_var_mapping = nlp.states["qdot"].mapping.to_first
+        elif "qdot" in nlp.controls:
+            qdot_var_mapping = nlp.controls["qdot"].mapping.to_first
+        else:
+            qdot_var_mapping = BiMapping([i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])]).to_first
 
         if external_forces is None:
             if with_contact:
                 qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau)
             else:
                 qddot = nlp.model.forward_dynamics(q, qdot, tau)
 
-            return qdot_var.mapping.to_first.map(qddot)
+            return qdot_var_mapping.map(qddot)
         else:
-            dxdt = MX(len(qdot_var.mapping.to_first), nlp.ns)
+            dxdt = MX(len(qdot_var_mapping), nlp.ns)
             # Todo: Should be added to pass f_ext in controls (as a symoblic value)
             #  this would avoid to create multiple equations of motions per node
             for i, f_ext in enumerate(external_forces):
                 if with_contact:
                     qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau, f_ext)
                 else:
                     qddot = nlp.model.forward_dynamics(q, qdot, tau, f_ext)
-                dxdt[:, i] = qdot_var.mapping.to_first.map(qddot)
+                dxdt[:, i] = qdot_var_mapping.map(qddot)
             return dxdt
 
     @staticmethod