Merge pull request #817 from mickaelbegon/bound_again

Add state and control bounds using constraints

README.md

@@ -1186,6 +1186,8 @@ constraint_list.add(constraint)
 ### Class: ConstraintFcn
 The `ConstraintFcn` class is the declaration of all the already available constraints in `bioptim`. 
 Since this is an Enum, it is possible to use the tab key on the keyboard to dynamically list them all, depending on the capabilities of your IDE. The existing contraint functions in alphabetical order:
+- **BOUND_STATE**  — Adds bounds on states. Same aim as `bounds["state_name"] = min_bounds, max_bounds` but with a different numerical behaviour.
+- **BOUND_CONTROL**  — Adds bounds on controls. Same aim as `bounds["control_name"] = min_bounds, max_bounds` but with a different numerical behaviour.
 - **NON_SLIPPING**  — Adds a constraint of static friction at contact points constraining for small tangential forces.  
 This constraint assumes that the normal forces is positive (that is having an additional TRACK_CONTACT_FORCES with `max_bound=np.inf`). The extra parameters `tangential_component_idx: int`, `normal_component_idx: int`, and `static_friction_coefficient: float` must be passed to the `Constraint` constructor.
 - **PROPORTIONAL_CONTROL** — Links one control to another, such that `u[first_dof] - first_dof_intercept = coef * (u[second_dof] - second_dof_intercept)`. The extra parameters `first_dof: int` and `second_dof: int` must be passed to the `Constraint` constructor.
@@ -1997,6 +1999,12 @@ can be held in the solution. Another goal would be to reload fast a previously s
 ### The [pendulum.py](./bioptim/examples/getting_started/pendulum.py) file
 This example is another way to present the pendulum example of the 'Getting started' section.
 
+### The [pendulum_constrained_states_controls.py](./bioptim/examples/getting_started/pendulum_constrained_states_controls.py) file 
+This example is a clone of the pendulum.py example with the difference that the
+states and controls are constrained instead of bounded. Sometimes the OCP converges faster with constraints than boundaries. 
+
+It is designed to show how to use `bound_state` and `bound_control`.
+
 ## Torque-driven OCP
 In this section, you will find different examples showing how to implement torque-driven optimal control programs.
 

bioptim/examples/getting_started/pendulum.py

@@ -82,7 +82,7 @@ def prepare_ocp(
     # Dynamics
     dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=expand_dynamics, phase_dynamics=phase_dynamics)
 
-    # Path constraint
+    # Path bounds
     x_bounds = BoundsList()
     x_bounds["q"] = bio_model.bounds_from_ranges("q")
     x_bounds["q"][:, [0, -1]] = 0  # Start and end at 0...
@@ -95,7 +95,7 @@ def prepare_ocp(
     x_init["q"] = [0] * bio_model.nb_q
     x_init["qdot"] = [0] * bio_model.nb_qdot
 
-    # Define control path constraint
+    # Define control path bounds
     n_tau = bio_model.nb_tau
     u_bounds = BoundsList()
     u_bounds["tau"] = [-100] * n_tau, [100] * n_tau  # Limit the strength of the pendulum to (-100 to 100)...

bioptim/examples/getting_started/pendulum_constrained_states_controls.py

@@ -0,0 +1,166 @@
+"""
+A very simple yet meaningful optimal control program consisting in a pendulum starting downward and ending upward
+while requiring the minimum of generalized forces. The solver is only allowed to move the pendulum sideways.
+
+This simple example is a good place to start investigating bioptim as it describes the most common dynamics out there
+(the joint torque driven), it defines an objective function and some constraints on states and controls
+
+Similar to pendulum.py but bounds are replaced by constraints
+
+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 platform
+
+import numpy as np
+
+from bioptim import (
+    OptimalControlProgram,
+    DynamicsFcn,
+    Dynamics,
+    ConstraintList,
+    ConstraintFcn,
+    ObjectiveFcn,
+    Objective,
+    OdeSolver,
+    OdeSolverBase,
+    CostType,
+    Solver,
+    BiorbdModel,
+    ControlType,
+    PhaseDynamics,
+    Node,
+)
+
+
+def prepare_ocp(
+    biorbd_model_path: str,
+    final_time: float,
+    n_shooting: int,
+    ode_solver: OdeSolverBase = OdeSolver.RK4(),
+    use_sx: bool = True,
+    n_threads: int = 1,
+    phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE,
+    expand_dynamics: bool = True,
+    control_type: ControlType = ControlType.CONSTANT,
+) -> 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: OdeSolverBase = 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)
+    phase_dynamics: PhaseDynamics
+        If the dynamics equation within a phase is unique or changes at each node.
+        PhaseDynamics.SHARED_DURING_THE_PHASE is much faster, but lacks the capability to have changing dynamics within
+        a phase. A good example of when PhaseDynamics.ONE_PER_NODE should be used is when different external forces
+        are applied at each node
+    expand_dynamics: bool
+        If the dynamics function should be expanded. Please note, this will solve the problem faster, but will slow down
+        the declaration of the OCP, so it is a trade-off. Also depending on the solver, it may or may not work
+        (for instance IRK is not compatible with expanded dynamics)
+    control_type: ControlType
+        The type of the controls
+
+    Returns
+    -------
+    The OptimalControlProgram ready to be solved
+    """
+
+    bio_model = BiorbdModel(biorbd_model_path)
+
+    # Add objective functions
+    objective_functions = Objective(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau")
+
+    # Dynamics
+    dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=expand_dynamics, phase_dynamics=phase_dynamics)
+
+    # Path state constraints
+    constraints = ConstraintList()
+
+    constraints.add(ConstraintFcn.TRACK_STATE, key="q", node=Node.START, target=[0, 0])  # Initial state
+    constraints.add(ConstraintFcn.TRACK_STATE, key="qdot", node=Node.START, target=[0, 0])
+
+    constraints.add(ConstraintFcn.TRACK_STATE, key="q", index=0, node=Node.END, target=0)  # Final state
+    constraints.add(ConstraintFcn.TRACK_STATE, key="q", index=1, node=Node.END, target=3.14)
+    constraints.add(ConstraintFcn.TRACK_STATE, key="qdot", node=Node.END, target=np.array([0, 0]))
+
+    constraints.add(
+        ConstraintFcn.BOUND_STATE,
+        key="q",
+        node=Node.MID,
+        min_bound=np.array([-1, -2 * np.pi]),
+        max_bound=np.array([5, 2 * np.pi]),
+    )  # In-between
+
+    # Control path constraint
+    constraints.add(
+        ConstraintFcn.BOUND_CONTROL,
+        key="tau",
+        index=0,
+        node=Node.ALL_SHOOTING,
+        min_bound=-100,
+        max_bound=100,
+    )
+    constraints.add(ConstraintFcn.TRACK_CONTROL, key="tau", index=1, node=Node.ALL_SHOOTING, target=0)  # Passive joint
+
+    return OptimalControlProgram(
+        bio_model,
+        dynamics,
+        n_shooting,
+        final_time,
+        constraints=constraints,
+        objective_functions=objective_functions,
+        ode_solver=ode_solver,
+        use_sx=use_sx,
+        n_threads=n_threads,
+        control_type=control_type,
+    )
+
+
+def main():
+    """
+    If pendulum_constrained_states_controls is run as a script, it will perform the optimization and animates it
+    """
+
+    # --- Prepare the ocp --- #
+    ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2)
+
+    # Custom plots
+    ocp.add_plot_penalty(CostType.ALL)
+
+    # --- If one is interested in checking the conditioning of the problem, they can uncomment the following line --- #
+    # ocp.check_conditioning()
+
+    # --- Print ocp structure --- #
+    ocp.print(to_console=False, to_graph=False)
+
+    # --- Solve the ocp. Please note that online graphics only works with the Linux operating system --- #
+    sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux"))
+    sol.print_cost()
+
+    # --- Show the results (graph or animation) --- #
+    # sol.graphs(show_bounds=True)
+    sol.animate(n_frames=100)
+
+    # # --- Save the solution --- #
+    # import pickle
+    # with open("pendulum.pkl", "wb") as file:
+    #     del sol.ocp
+    #     pickle.dump(sol, file)
+
+
+if __name__ == "__main__":
+    main()

bioptim/limits/constraints.py

@@ -346,6 +346,46 @@ def time_constraint(_: Constraint, controller: PenaltyController, **unused_param
 
             return controller.tf
 
+        @staticmethod
+        def bound_state(
+            _: Constraint,
+            controller: PenaltyController,
+            key: str,
+        ):
+            """
+            Bound the state according to key
+            Parameters
+            ----------
+            _: Constraint
+                The actual constraint to declare
+            controller: PenaltyController
+                The penalty node elements
+            key: str
+                The name of the state to constraint
+            """
+
+            return controller.states[key].cx_start
+
+        @staticmethod
+        def bound_control(
+            _: Constraint,
+            controller: PenaltyController,
+            key: str,
+        ):
+            """
+            Bound the control according to key
+            Parameters
+            ----------
+            _: Constraint
+                The actual constraint to declare
+            controller: PenaltyController
+                The penalty node elements
+            key: str
+                The name of the control to constraint
+            """
+
+            return controller.controls[key].cx_start
+
         @staticmethod
         def qddot_equals_forward_dynamics(
             _: Constraint,
@@ -1010,6 +1050,8 @@ def get_type() -> Callable
         Returns the type of the penalty
     """
 
+    BOUND_CONTROL = (ConstraintFunction.Functions.bound_control,)
+    BOUND_STATE = (ConstraintFunction.Functions.bound_state,)
     STATE_CONTINUITY = (PenaltyFunctionAbstract.Functions.state_continuity,)
     FIRST_COLLOCATION_HELPER_EQUALS_STATE = (PenaltyFunctionAbstract.Functions.first_collocation_point_equals_state,)
     CUSTOM = (PenaltyFunctionAbstract.Functions.custom,)

tests/shard1/test_all_examples.py

@@ -4,6 +4,7 @@
 from bioptim import InterpolationType, PhaseDynamics
 
 
+## examples/acados
 def test__acados__cube():
     from bioptim.examples.acados import cube as ocp_module
 
@@ -43,6 +44,7 @@ def test__acados__static_arm():
     )
 
 
+## examples/getting_started
 def test__getting_started__custom_bounds():
     from bioptim.examples.getting_started import custom_bounds as ocp_module
 
@@ -159,6 +161,9 @@ def test__getting_started__custom_plotting():
     )
 
 
+# todo: Add example_continuity_as_objective.py?
+
+
 def test__getting_started__example_cyclic_movement():
     from bioptim.examples.getting_started import example_cyclic_movement as ocp_module
 
@@ -185,6 +190,9 @@ def test__getting_started__example_external_forces():
     )
 
 
+# todo: Add example_external_forces.py?
+
+
 def test__getting_started__example_inequality_constraint():
     from bioptim.examples.getting_started import example_inequality_constraint as ocp_module
 
@@ -202,21 +210,11 @@ def test__getting_started__example_inequality_constraint():
     )
 
 
-def test__getting_started__example_mapping():
-    from bioptim.examples.getting_started import example_mapping as ocp_module
-
-
-def test__getting_started__example_multiphase():
-    from bioptim.examples.getting_started import example_multiphase as ocp_module
+# todo: Add example_joint_acceleratio_driven.py?
 
-    bioptim_folder = os.path.dirname(ocp_module.__file__)
 
-    ocp_module.prepare_ocp(
-        biorbd_model_path=bioptim_folder + "/models/cube.bioMod",
-        long_optim=True,
-        phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE,
-        expand_dynamics=False,
-    )
+def test__getting_started__example_mapping():
+    from bioptim.examples.getting_started import example_mapping as ocp_module
 
 
 def test__getting_started__example_multinode_constraints():
@@ -259,6 +257,22 @@ def test__getting_started__example_multinode_objective():
         )
 
 
+def test__getting_started__example_multiphase():
+    from bioptim.examples.getting_started import example_multiphase as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ocp_module.prepare_ocp(
+        biorbd_model_path=bioptim_folder + "/models/cube.bioMod",
+        long_optim=True,
+        phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE,
+        expand_dynamics=False,
+    )
+
+
+# todo: Add example_multistart.py?
+
+
 def test__getting_started__example_optimal_time():
     from bioptim.examples.getting_started import example_optimal_time as ocp_module
 
@@ -267,6 +281,9 @@ def test__getting_started__example_simulation():
     from bioptim.examples.getting_started import example_optimal_time as ocp_module
 
 
+# todo: Add example_variable_scaling.py?
+
+
 def test__getting_started__pendulum():
     from bioptim.examples.getting_started import pendulum as ocp_module
 
@@ -281,6 +298,21 @@ def test__getting_started__pendulum():
     )
 
 
+def test__getting_started__pendulum_constrained_states_controls():
+    from bioptim.examples.getting_started import pendulum_constrained_states_controls as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ocp_module.prepare_ocp(
+        biorbd_model_path=bioptim_folder + "/models/pendulum.bioMod",
+        final_time=3,
+        n_shooting=100,
+        phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE,
+        expand_dynamics=False,
+    )
+
+
+## examples/moving_horizon_estimation
 def test__moving_horizon_estimation__mhe():
     from bioptim.examples.moving_horizon_estimation import mhe as ocp_module