[RTM after push] Provide two examples of minmax OCP with extra parameters

bioptim/examples/torque_driven_ocp/minimize_maximum_torque_by_extra_parameter.py

@@ -1,9 +1,9 @@
 """
-This example is inspired from the clear pike circle gymnastics skill. It is composed of two pendulums
+This example is inspired from the giant circle gymnastics skill. It is composed of two pendulums
 representing the trunk and legs segments (only the hip flexion is actuated). The objective is to minimize the
-maximum torque (minmax) of the hip flexion while performing the clear pike circle motion. The maximum torque is included to the
+maximum torque (minmax) of the hip flexion while performing the giant circle. The maximum torque is included to the
 problem as a parameter, all torque interval re constrained to be smaller than this parameter, this parameter is the
-minimized.
+minimized. Two options are provided and compared to define initial and final states (0: bounds; 1: constraints)
 """
 
 import numpy as np
@@ -15,6 +15,7 @@
     DynamicsFcn,
     ObjectiveList,
     ConstraintList,
+    ConstraintFcn,
     BoundsList,
     InitialGuessList,
     Node,
@@ -25,123 +26,113 @@
     BiorbdModel,
     PenaltyController,
     ParameterObjectiveList,
-    ParameterConstraintList,
-    ConstraintFcn,
 )
+from matplotlib import pyplot as plt
+
 
+def custom_constraint_max_tau(controller: PenaltyController) -> MX:
+    return controller.parameters["max_tau"].cx - controller.controls["tau"].cx
 
-def custom_constraint_parameters(controller: PenaltyController) -> MX:
-    tau = controller.controls["tau_joints"].cx_start
-    max_param = controller.parameters["max_tau"].cx
-    val = max_param - tau
-    return val
+
+def custom_constraint_min_tau(controller: PenaltyController) -> MX:
+    return controller.parameters["min_tau"].cx - controller.controls["tau"].cx
 
 
 def my_parameter_function(bio_model: biorbd.Model, value: MX):
     return
 
 
 def prepare_ocp(
+    method_bound_states: int = 0,
     bio_model_path: str = "models/double_pendulum.bioMod",
 ) -> OptimalControlProgram:
-    bio_model = (BiorbdModel(bio_model_path), BiorbdModel(bio_model_path))
+    bio_model = BiorbdModel(bio_model_path)
 
     # Problem parameters
-    n_shooting = (40, 40)
-    final_time = (1, 1)
-    tau_min, tau_max, tau_init = -300, 300, 0
-    n_root = bio_model[0].nb_root
-    n_q = bio_model[0].nb_q
-    n_tau = n_q - n_root
+    n_shooting = 50
+    final_time = 1
+    tau_min, tau_max, tau_init = -10, 10, 0
+
+    # Mapping
+    tau_mappings = BiMappingList()
+    tau_mappings.add("tau", to_second=[None, 0], to_first=[1])
 
     # Define the parameter to optimize
     parameters = ParameterList()
+    parameter_init = InitialGuessList()
+    parameter_bounds = BoundsList()
 
     parameters.add(
-        "max_tau",  # The name of the parameter
-        my_parameter_function,  # The function that modifies the biorbd model
+        "max_tau",
+        my_parameter_function,
+        size=1,
+    )
+    parameters.add(
+        "min_tau",
+        my_parameter_function,
         size=1,
     )
 
-    parameter_bounds = BoundsList()
-    parameter_init = InitialGuessList()
+    parameter_init["max_tau"] = 1
+    parameter_init["min_tau"] = -1
 
-    parameter_init["max_tau"] = 0
     parameter_bounds.add("max_tau", min_bound=0, max_bound=tau_max, interpolation=InterpolationType.CONSTANT)
+    parameter_bounds.add("min_tau", min_bound=tau_min, max_bound=0, interpolation=InterpolationType.CONSTANT)
 
-    # Add phase independant objective functions
+    # Add phase independent objective functions
     parameter_objectives = ParameterObjectiveList()
-    parameter_objectives.add(ObjectiveFcn.Parameter.MINIMIZE_PARAMETER, key="max_tau", weight=1000, quadratic=True)
-
-    # Add phase independant constraint functions
-    parameter_constraints = ParameterConstraintList()
-    parameter_constraints.add(ConstraintFcn.TRACK_PARAMETER, min_bound=-100, max_bound=100)
+    parameter_objectives.add(ObjectiveFcn.Parameter.MINIMIZE_PARAMETER, key="max_tau", weight=10, quadratic=True)
+    parameter_objectives.add(ObjectiveFcn.Parameter.MINIMIZE_PARAMETER, key="min_tau", weight=10, quadratic=True)
 
     # Add objective functions
     objective_functions = ObjectiveList()
-    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=1, phase=0, min_bound=0.1, max_bound=3)
-    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=1, phase=1, min_bound=0.1, max_bound=3)
-    objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau_joints", weight=10, phase=0)
-    objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau_joints", weight=10, phase=1)
+    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=1, phase=0, min_bound=1, max_bound=5)
+    objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, phase=0)
 
     # Constraints
     constraints = ConstraintList()
-    constraints.add(custom_constraint_parameters, node=Node.ALL, min_bound=0, max_bound=tau_max)
+
+    # constraints.add(ConstraintFcn.TIME_CONSTRAINT, node=Node.END, target=3.8)
+
+    constraints.add(custom_constraint_max_tau, phase=0, node=Node.ALL_SHOOTING, min_bound=0, max_bound=tau_max)
+    constraints.add(custom_constraint_min_tau, phase=0, node=Node.ALL_SHOOTING, min_bound=tau_min, max_bound=0)
 
     # Dynamics
     dynamics = DynamicsList()
-    dynamics.add(DynamicsFcn.TORQUE_DRIVEN_FREE_FLOATING_BASE, with_contact=False)
-    dynamics.add(DynamicsFcn.TORQUE_DRIVEN_FREE_FLOATING_BASE, with_contact=False)
+    dynamics.add(DynamicsFcn.TORQUE_DRIVEN, with_contact=False)
 
     # Path constraint
     x_bounds = BoundsList()
-    q_roots_min = bio_model[0].bounds_from_ranges("q").min[:n_root, :]
-    q_roots_max = bio_model[0].bounds_from_ranges("q").max[:n_root, :]
-    q_joints_min = bio_model[0].bounds_from_ranges("q").min[n_root:, :]
-    q_joints_max = bio_model[0].bounds_from_ranges("q").max[n_root:, :]
-    qdot_roots_min = bio_model[0].bounds_from_ranges("qdot").min[:n_root, :]
-    qdot_roots_max = bio_model[0].bounds_from_ranges("qdot").max[:n_root, :]
-    qdot_joints_min = bio_model[0].bounds_from_ranges("qdot").min[n_root:, :]
-    qdot_joints_max = bio_model[0].bounds_from_ranges("qdot").max[n_root:, :]
-    x_bounds.add("q_roots", min_bound=q_roots_min, max_bound=q_roots_max, phase=0)
-    x_bounds.add("q_joints", min_bound=q_joints_min, max_bound=q_joints_max, phase=0)
-    x_bounds.add("qdot_roots", min_bound=qdot_roots_min, max_bound=qdot_roots_max, phase=0)
-    x_bounds.add("qdot_joints", min_bound=qdot_joints_min, max_bound=qdot_joints_max, phase=0)
-    x_bounds.add("q_roots", min_bound=q_roots_min, max_bound=q_roots_max, phase=1)
-    x_bounds.add("q_joints", min_bound=q_joints_min, max_bound=q_joints_max, phase=1)
-    x_bounds.add("qdot_roots", min_bound=qdot_roots_min, max_bound=qdot_roots_max, phase=1)
-    x_bounds.add("qdot_joints", min_bound=qdot_joints_min, max_bound=qdot_joints_max, phase=1)
-
-    # change model bound for -pi, pi
-    for i in range(len(bio_model)):
-        x_bounds[i]["q_joints"].min[0, :] = -np.pi
-        x_bounds[i]["q_joints"].max[0, :] = np.pi
-
-    # Phase 0
-    x_bounds[0]["q_roots"][0, 0] = np.pi
-    x_bounds[0]["q_joints"][0, 0] = 0
-    x_bounds[0]["q_joints"].min[0, -1] = 6 * np.pi / 8 - 0.1
-    x_bounds[0]["q_joints"].max[0, -1] = 6 * np.pi / 8 + 0.1
-
-    # Phase 1
-    x_bounds[1]["q_roots"][0, -1] = 3 * np.pi
-    x_bounds[1]["q_joints"][0, -1] = 0
-
-    # Define control path constraint
-    u_bounds = BoundsList()
-    u_bounds.add(key="tau_joints", min_bound=[tau_min] * n_tau, max_bound=[tau_max] * n_tau, phase=0)
-    u_bounds.add(key="tau_joints", min_bound=[tau_min] * n_tau, max_bound=[tau_max] * n_tau, phase=1)
+    x_bounds.add(key="q", bounds=bio_model.bounds_from_ranges("q"))
+    x_bounds.add(key="qdot", bounds=bio_model.bounds_from_ranges("qdot"))
+
+    x_bounds["q"].min[0, :] = 0
+    x_bounds["q"].max[0, :] = 3 * np.pi
+
+    x_bounds["q"].min[1, :] = -np.pi / 3
+    x_bounds["q"].max[1, :] = np.pi / 5
+
+    if method_bound_states == 0:
+        x_bounds["q"][0, 0] = np.pi + 0.01
+        x_bounds["q"][1, 0] = 0
+        x_bounds["qdot"][0, 0] = 0
+        x_bounds["qdot"][1, 0] = 0
+
+        x_bounds["q"][0, -1] = 3 * np.pi
+        x_bounds["q"][1, -1] = 0
+    else:
+        constraints.add(ConstraintFcn.TRACK_STATE, node=Node.START, key="q", target=[np.pi + 0.01, 0])
+        constraints.add(ConstraintFcn.TRACK_STATE, node=Node.START, key="qdot", target=[0, 0])
+        constraints.add(ConstraintFcn.TRACK_STATE, node=Node.END, key="q", target=[3 * np.pi, 0])
 
     return OptimalControlProgram(
         bio_model,
         dynamics,
         n_shooting,
         final_time,
         x_bounds=x_bounds,
-        u_bounds=u_bounds,
         objective_functions=objective_functions,
         parameter_objectives=parameter_objectives,
-        parameter_constraints=parameter_constraints,
         parameter_bounds=parameter_bounds,
         parameter_init=parameter_init,
         constraints=constraints,
@@ -151,14 +142,73 @@
 
 def main():
     # --- Prepare the ocp --- #
-    ocp = prepare_ocp()
-
-    # --- Solve the ocp --- #
-    sol = ocp.solve()
-
-    # --- Show results --- #
-    sol.animate()
-    sol.graphs(show_bounds=True)
+    fig, axs = plt.subplots(1, 3)
+    axs[0].set_title("Joint coordinates")
+    axs[0].set_ylabel("q [°]")
+    axs[1].set_title("Joint velocities")
+    axs[1].set_ylabel("qdot [°/s]")
+    axs[2].set_title("Generalized forces")
+    axs[2].set_ylabel("tau [Nm]")
+    for ax in [0, 1, 2]:
+        axs[ax].set_xlabel("Time [s]")
+        axs[ax].grid(True)
+
+    linestyles = ["solid", "dashed"]
+
+    for i, linestyle in enumerate(linestyles):
+        # --- Prepare the ocp --- #
+        ocp = prepare_ocp(method_bound_states=i)
+
+        # --- Solve the ocp --- #
+        sol = ocp.solve()
+        # sol.print_cost()
+
+        # --- Show results --- #
+        # sol.animate()
+        #    sol.graphs(show_bounds=True)
+
+        states = sol.decision_states()
+        controls = sol.decision_controls()
+
+        q = np.array([item.flatten() for item in states["q"]])
+        qdot = np.array([item.flatten() for item in states["qdot"]])
+        tau = np.vstack([np.array([item.flatten() for item in controls["tau"]]), np.array([[np.nan]])])
+        time = np.array([item.full().flatten()[0] for item in sol.stepwise_time()])
+
+        print("Duration: ", time[-1])
+        print("sum tau**2 dt = ", np.nansum(tau**2 * time[1]))
+        print("min-max tau: ", np.nanmin(tau), np.nanmax(tau))
+
+        # Plotting q solutions for both DOFs
+        axs[0].plot(
+            time,
+            q * 180 / np.pi,
+            linestyle=linestyle,
+            label=[f"q_0 (method {i})", f"q_1 (method {i})"],
+        )
+        axs[1].plot(
+            time,
+            qdot * 180 / np.pi,
+            linestyle=linestyle,
+            label=[f"qdot_0 (method {i})", f"qdot_1 (method {i})"],
+        )
+        axs[2].step(time, tau, linestyle=linestyle, label=f"tau (method {i})", where="post")
+        xlim = np.asarray(axs[2].get_xlim())
+        xlim = np.hstack([xlim, np.nan, xlim])
+        min_max = np.hstack(
+            [np.ones([2]) * sol.parameters["min_tau"], np.nan, np.ones([2]) * sol.parameters["max_tau"]]
+        )
+        axs[2].plot(xlim, min_max, linestyle=(0, (5, 10)), label=f"params - bounds (method {i})")
+
+    plt.tight_layout()
+    plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=None, hspace=0.3)
+
+    # Display the plot
+    axs[0].legend()
+    axs[1].legend()
+    axs[2].legend()
+
+    plt.show()
 
 
 if __name__ == "__main__":