[RTM] Free floating dynamics

.gitignore

@@ -94,6 +94,7 @@ examples/*/biorbd_optim
 examples/biorbd_optim
 
 # Remove other development artifacts
+examples/sandbox
 */examples/sandbox
 *.npy
 *.bo

README.md

@@ -732,6 +732,14 @@ Please note, this dynamics is expected to be very slow to converge, if it ever d
 One is therefore encourage using TORQUE_DRIVEN instead, and to add the TORQUE_MAX_FROM_ACTUATORS constraint.
 This has been shown to be more efficient and allows defining minimum torque.
 
+#### JOINTS_ACCELERATION_DRIVEN
+The joints acceleration driven defines the states (x) as *q* and *qdot* and the controls (u) as *qddot_joints*. The derivative of *q* is trivially *qdot*.
+The joints' acceleration *qddot_joints* is the acceleration of the actual joints of the `biorb_model` without its root's joints.
+The model's root's joints acceleration *qddot_root* are computed by the `biorbd` function: `qddot_root = boirbd_model.ForwardDynamicsFreeFloatingBase(q, qdot, qddot_joints)`.
+The derivative of *qdot* is the vertical stack of *qddot_root* and *qddot_joints*.
+
+This dynamic is suitable for bodies in free fall.
+
 #### MUSCLE_ACTIVATIONS_DRIVEN
 The torque driven defines the states (x) as *q* and *qdot* and the controls (u) as the muscle activations. 
 The derivative of *q* is trivially *qdot*.
@@ -1769,8 +1777,11 @@ The main goal of this kind of simulation, especially in single shooting (that is
 is to validate the dynamics of multiple shooting. If they both are equal, it usually means that a great confidence
 can be held in the solution. Another goal would be to reload fast a previously saved optimized solution.
 
+### The example_joints_acceleration_driven.py file
+This example shows how to use the joints acceleration dynamic to achieve the same goal as the simple pendulum, but with a double pendulum for which only the angular acceleration of the second pendulum is controled.
+
 ### The pendulum.py file
-This is another way to present the pendulum example of the 'Getting started' section. 
+This is another way to present the pendulum example of the 'Getting started' section.
 
 ## Muscle driven OCP
 In this file, you will find four examples about muscle driven optimal control programs. The two first refer to traking 

bioptim/dynamics/configure_problem.py

@@ -282,6 +282,39 @@ def torque_activations_driven(ocp, nlp, with_contact=False):
             )
         ConfigureProblem.configure_soft_contact_function(ocp, nlp)
 
+    @staticmethod
+    def joints_acceleration_driven(ocp, nlp, implicit_dynamics: bool = False):
+        """
+        Configure the dynamics for a joints acceleration driven program
+        (states are q and qdot, controls are qddot_joints)
+
+        Parameters
+        ----------
+        ocp: OptimalControlProgram
+            A reference to the ocp
+        nlp: NonLinearProgram
+            A reference to the phase
+        implicit_dynamics: bool
+            If the implicit dynamic should be used
+        """
+        if implicit_dynamics:
+            raise NotImplementedError("Implicit dynamics not implemented yet.")
+
+        ConfigureProblem.configure_q(nlp, as_states=True, as_controls=False)
+        ConfigureProblem.configure_qdot(nlp, as_states=True, as_controls=False)
+        # Configure qddot joints
+        nb_root = nlp.model.nbRoot()
+        if not nb_root > 0:
+            raise RuntimeError("Model must have at least one DoF on root.")
+
+        name_qddot_joints = [str(i + nb_root) for i in range(nlp.model.nbQddot() - nb_root)]
+        ConfigureProblem.configure_new_variable(
+            "qddot_joints", name_qddot_joints, nlp, as_states=False, as_controls=True
+        )
+        ConfigureProblem.configure_dynamics_function(
+            ocp, nlp, DynamicsFunctions.joints_acceleration_driven, expand=False
+        )
+
     @staticmethod
     def muscle_driven(
         ocp,
@@ -928,6 +961,7 @@ class DynamicsFcn(Enum):
     TORQUE_DRIVEN = (ConfigureProblem.torque_driven,)
     TORQUE_DERIVATIVE_DRIVEN = (ConfigureProblem.torque_derivative_driven,)
     TORQUE_ACTIVATIONS_DRIVEN = (ConfigureProblem.torque_activations_driven,)
+    JOINTS_ACCELERATION_DRIVEN = (ConfigureProblem.joints_acceleration_driven,)
     MUSCLE_DRIVEN = (ConfigureProblem.muscle_driven,)
     CUSTOM = (ConfigureProblem.custom,)
 

bioptim/dynamics/dynamics_functions.py

@@ -1,6 +1,6 @@
 from typing import Union
 
-from casadi import horzcat, vertcat, MX, SX
+from casadi import horzcat, vertcat, MX, SX, Function
 
 from .fatigue.fatigue_dynamics import FatigueList
 from ..optimization.optimization_variable import OptimizationVariable
@@ -460,6 +460,44 @@ def forces_from_muscle_driven(states: MX.sym, controls: MX.sym, parameters: MX.s
         tau = muscles_tau + residual_tau if residual_tau is not None else muscles_tau
         return DynamicsFunctions.contact_forces(nlp, q, qdot, tau)
 
+    @staticmethod
+    def joints_acceleration_driven(
+        states: MX.sym, controls: MX.sym, parameters: MX.sym, nlp, implicit_dynamics: bool = False
+    ) -> MX:
+        """
+        Forward dynamics driven by joints accelerations of a free floating body.
+
+        Parameters
+        ----------
+        states: MX.sym
+            The state of the system
+        controls: MX.sym
+            The controls of the system
+        parameters: MX.sym
+            The parameters of the system
+        nlp: NonLinearProgram
+            The definition of the system
+        implicit_dynamics: bool
+            If implicit dynamics should be used
+
+        Returns
+        ----------
+        MX.sym
+            The derivative of states
+        """
+        if implicit_dynamics:
+            raise NotImplementedError("Implicit dynamics not implemented yet.")
+
+        DynamicsFunctions.apply_parameters(parameters, nlp)
+        q = DynamicsFunctions.get(nlp.states["q"], states)
+        qdot = DynamicsFunctions.get(nlp.states["qdot"], states)
+        qddot_joints = DynamicsFunctions.get(nlp.controls["qddot_joints"], controls)
+
+        qddot_root = nlp.model.ForwardDynamicsFreeFloatingBase(q, qdot, qddot_joints).to_mx()
+        qddot_root_func = Function("qddot_root_func", [q, qdot, qddot_joints], [qddot_root]).expand()
+
+        return vertcat(qdot, vertcat(qddot_root_func(q, qdot, qddot_joints), qddot_joints))
+
     @staticmethod
     def get(var: OptimizationVariable, cx: Union[MX, SX]):
         """

bioptim/examples/getting_started/example_joints_acceleration_driven.py

@@ -0,0 +1,124 @@
+"""
+A simple optimal control program consisting of a double pendulum starting downward and ending upward while requiring the
+ minimum of generalized forces. The solver is only allowed to apply an angular acceleration the joint linking the second
+ pendulum to the first one.
+
+During the optimization process, the graphs are updated real-time (even though it is a bit too fast and short to really
+appreciate it). Finally, once it finished optimizing, it animates the model using the optimal solution
+"""
+
+import biorbd_casadi as biorbd
+from bioptim import (
+    OptimalControlProgram,
+    DynamicsFcn,
+    Dynamics,
+    Bounds,
+    QAndQDotBounds,
+    InitialGuess,
+    ObjectiveFcn,
+    Objective,
+    OdeSolver,
+    CostType,
+    Solver,
+)
+
+
+def prepare_ocp(
+    biorbd_model_path: str,
+    final_time: float,
+    n_shooting: int,
+    ode_solver: OdeSolver = OdeSolver.RK4(),
+    use_sx: bool = True,
+    n_threads: int = 1,
+) -> OptimalControlProgram:
+    """
+    The initialization of an ocp
+
+    Parameters
+    ----------
+    biorbd_model_path: str
+        The path to the biorbd model
+    final_time: float
+        The time in second required to perform the task
+    n_shooting: int
+        The number of shooting points to define int the direct multiple shooting program
+    ode_solver: OdeSolver = OdeSolver.RK4()
+        Which type of OdeSolver to use
+    use_sx: bool
+        If the SX variable should be used instead of MX (can be extensive on RAM)
+    n_threads: int
+        The number of threads to use in the paralleling (1 = no parallel computing)
+
+    Returns
+    -------
+    The OptimalControlProgram ready to be solved
+    """
+
+    biorbd_model = biorbd.Model(biorbd_model_path)
+
+    # Add objective functions
+    objective_functions = Objective(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="qddot_joints")
+
+    # Dynamics
+    dynamics = Dynamics(DynamicsFcn.JOINTS_ACCELERATION_DRIVEN)
+
+    # Path constraint
+    x_bounds = QAndQDotBounds(biorbd_model)
+    x_bounds[:, [0, -1]] = 0
+    x_bounds[0, -1] = 3.14
+    x_bounds[1, -1] = 0
+
+    # Initial guess
+    n_q = biorbd_model.nbQ()
+    n_qdot = biorbd_model.nbQdot()
+    x_init = InitialGuess([0] * (n_q + n_qdot))
+
+    # Define control path constraint
+    n_qddot_joints = biorbd_model.nbQddot() - biorbd_model.nbRoot()  # 2 - 1 = 1 in this example
+    qddot_joints_min, qddot_joints_max, qddot_joints_init = -100, 100, 0
+    u_bounds = Bounds([qddot_joints_min] * n_qddot_joints, [qddot_joints_max] * n_qddot_joints)
+
+    u_init = InitialGuess([qddot_joints_init] * n_qddot_joints)
+
+    return OptimalControlProgram(
+        biorbd_model,
+        dynamics,
+        n_shooting,
+        final_time,
+        x_init=x_init,
+        u_init=u_init,
+        x_bounds=x_bounds,
+        u_bounds=u_bounds,
+        objective_functions=objective_functions,
+        ode_solver=ode_solver,
+        use_sx=use_sx,
+        n_threads=n_threads,
+    )
+
+
+def main():
+    """
+    If pendulum is run as a script, it will perform the optimization and animates it
+    """
+
+    # --- Prepare the ocp --- #
+    ocp = prepare_ocp(biorbd_model_path="models/double_pendulum.bioMod", final_time=10, n_shooting=100)
+
+    # Custom plots
+    ocp.add_plot_penalty(CostType.ALL)
+
+    # --- Print ocp structure --- #
+    ocp.print(to_console=False, to_graph=False)
+
+    # --- Solve the ocp --- #
+    sol = ocp.solve(Solver.IPOPT(show_online_optim=True))
+    # sol.graphs()
+
+    # --- Show the results in a bioviz animation --- #
+    sol.detailed_cost_values()
+    sol.print_cost()
+    sol.animate(n_frames=100)
+
+
+if __name__ == "__main__":
+    main()

bioptim/examples/getting_started/models/double_pendulum.bioMod

@@ -0,0 +1,54 @@
+version 4
+
+segment Seg1
+    rotations	x
+    ranges
+        -10*pi 10*pi
+    mass 1
+    inertia
+        0.0391  0.0000	0.0000
+	    0.0000	0.0335	-0.0032
+	    0.0000	-0.0032	0.0090
+    com  -0.0005 0.0688 -0.9542
+    meshfile mesh/pendulum.STL
+endsegment
+
+    // Marker 1
+    marker marker_1
+        parent Seg1
+        position 0 0 0
+    endmarker
+
+    // Marker 2
+    marker marker_2
+        parent Seg1
+        position 0 0 -1
+    endmarker
+
+segment Seg2
+    parent Seg1
+    rotations	x
+    rtinmatrix 0
+    rt 0 0 0 xyz 0 0 -1
+    ranges
+        -pi/2 pi/2
+    mass 1
+    inertia
+        0.0391  0.0000	0.0000
+	    0.0000	0.0335	-0.0032
+	    0.0000	-0.0032	0.0090
+    com  -0.0005 0.0688 -0.9542
+    meshfile mesh/pendulum.STL
+endsegment
+
+    // Marker 3
+    marker marker_3
+        parent Seg2
+        position 0 0 0
+    endmarker
+
+    // Marker 4
+    marker marker_4
+        parent Seg2
+        position 0 0 -1
+    endmarker

environment.yml

@@ -5,7 +5,7 @@ channels:
 - default
 dependencies:
 - python >=3.9
-- biorbd >=1.8.3
+- biorbd >=1.8.8
 - bioviz
 - scipy
 - numpy

tests/test_dynamics.py

@@ -1115,6 +1115,35 @@ def test_muscle_driven(with_excitations, with_contact, with_torque, with_externa
                         )
 
 
+@pytest.mark.parametrize("cx", [MX, SX])
+def test_joints_acceleration_driven(cx):
+    # Prepare the program
+    nlp = NonLinearProgram()
+    nlp.model = biorbd.Model(TestUtils.bioptim_folder() + "/examples/getting_started/models/double_pendulum.bioMod")
+    nlp.ns = 5
+    nlp.cx = cx
+
+    nlp.x_bounds = np.zeros((nlp.model.nbQ() * 3, 1))
+    nlp.u_bounds = np.zeros((nlp.model.nbQ(), 1))
+    ocp = OptimalControlProgram(nlp)
+    nlp.control_type = ControlType.CONSTANT
+    NonLinearProgram.add(ocp, "dynamics_type", Dynamics(DynamicsFcn.JOINTS_ACCELERATION_DRIVEN), False)
+
+    np.random.seed(42)
+
+    # Prepare the dynamics
+    ConfigureProblem.initialize(ocp, nlp)
+
+    # Test the results
+    states = np.random.rand(nlp.states.shape, nlp.ns)
+    controls = np.random.rand(nlp.controls.shape, nlp.ns)
+    params = np.random.rand(nlp.parameters.shape, nlp.ns)
+    x_out = np.array(nlp.dynamics_func(states, controls, params))
+
+    # obtained using Ipuch reference implementation. [https://github.com/Ipuch/OnDynamicsForSomersaults]
+    np.testing.assert_almost_equal(x_out[:, 0], [0.02058449, 0.18340451, -2.95556261, 0.61185289])
+
+
 @pytest.mark.parametrize("with_contact", [False, True])
 def test_custom_dynamics(with_contact):
     def custom_dynamic(states, controls, parameters, nlp, with_contact=False) -> tuple: