Merge pull request #273 from EveCharbie/CoM_objectives

CoM objectives

bioptim/limits/constraints.py

@@ -286,6 +286,8 @@ class ConstraintFcn(Enum):
     TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,)
     ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,)
     ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,)
+    COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,)
+    COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,)
     CUSTOM = (PenaltyType.CUSTOM,)
     CONTACT_FORCE = (ConstraintFunction.Functions.contact_force,)
     NON_SLIPPING = (ConstraintFunction.Functions.non_slipping,)

bioptim/limits/objective_functions.py

@@ -334,6 +334,8 @@ class Lagrange(Enum):
         TRACK_ALL_CONTROLS = (PenaltyType.TRACK_ALL_CONTROLS,)
         MINIMIZE_CONTACT_FORCES = (PenaltyType.MINIMIZE_CONTACT_FORCES,)
         TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,)
+        MINIMIZE_COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,)
+        MINIMIZE_COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,)
         ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,)
         ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,)
         CUSTOM = (PenaltyType.CUSTOM,)
@@ -368,6 +370,8 @@ class Mayer(Enum):
         MINIMIZE_CONTACT_FORCES = (PenaltyType.MINIMIZE_CONTACT_FORCES,)
         TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,)
         MINIMIZE_PREDICTED_COM_HEIGHT = (PenaltyType.MINIMIZE_PREDICTED_COM_HEIGHT,)
+        MINIMIZE_COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,)
+        MINIMIZE_COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,)
         ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,)
         ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,)
         CUSTOM = (PenaltyType.CUSTOM,)

bioptim/limits/penalty.py

@@ -346,6 +346,61 @@ def minimize_predicted_com_height(penalty, ocp, nlp, t, x, u, p):
                 CoM_height = (CoM_dot[2] * CoM_dot[2]) / (2 * -g) + CoM[2]
                 penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_height, penalty)
 
+        @staticmethod
+        def minimize_com_position(penalty, ocp, nlp, t, x, u, p, axis=None):
+            """
+            Adds the objective that the position of the center of mass of the model should be minimized.
+            If no axis is specified, the squared-norm of the CoM's position is mimimized.
+            Otherwise, the projection of the CoM's position on the specified axis is minimized.
+            """
+
+            target = None
+            if penalty.target is not None:
+                target = PenaltyFunctionAbstract._check_and_fill_tracking_data_size(penalty.target, (1, len(x)))
+
+            PenaltyFunctionAbstract._add_to_casadi_func(nlp, "biorbd_CoM", nlp.model.CoM, nlp.q)
+            for i, v in enumerate(x):
+                q = nlp.mapping["q"].expand.map(v[: nlp.shape["q"]])
+                CoM = nlp.casadi_func["biorbd_CoM"](q)
+
+                if axis == None:
+                    CoM_proj = CoM
+                elif not isinstance(axis, Axe):
+                    raise RuntimeError("axis must be a bioptim.Axe")
+                else:
+                    CoM_proj = CoM[axis]
+
+                penalty.sliced_target = target[:, i] if target is not None else None
+                penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_proj, penalty)
+
+        @staticmethod
+        def minimize_com_velocity(penalty, ocp, nlp, t, x, u, p, axis=None):
+            """
+            Adds the objective that the velocity of the center of mass of the model should be minimized.
+            If no axis is specified, the squared-norm of the CoM's velocity is mimimized.
+            Otherwise, the projection of the CoM's velocity on the specified axis is minimized.
+            """
+
+            target = None
+            if penalty.target is not None:
+                target = PenaltyFunctionAbstract._check_and_fill_tracking_data_size(penalty.target, (1, len(x)))
+
+            PenaltyFunctionAbstract._add_to_casadi_func(nlp, "biorbd_CoMdot", nlp.model.CoMdot, nlp.q, nlp.q_dot)
+            for i, v in enumerate(x):
+                q = nlp.mapping["q"].expand.map(v[: nlp.shape["q"]])
+                q_dot = nlp.mapping["q_dot"].expand.map(v[nlp.shape["q"] :])
+                CoM_dot = nlp.casadi_func["biorbd_CoMdot"](q, q_dot)
+
+                if axis == None:
+                    CoM_dot_proj = CoM_dot[0] ** 2 + CoM_dot[1] ** 2 + CoM_dot[2] ** 2
+                elif not isinstance(axis, Axe):
+                    raise RuntimeError("axis must be a bioptim.Axe")
+                else:
+                    CoM_dot_proj = CoM_dot[axis]
+
+                penalty.sliced_target = target[:, i] if target is not None else None
+                penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_dot_proj, penalty)
+
         @staticmethod
         def minimize_contact_forces(penalty, ocp, nlp, t, x, u, p):
             """
@@ -528,6 +583,8 @@ def _parameter_modifier(penalty_function, parameters):
                 or penalty_function == PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT
                 or penalty_function == PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS
                 or penalty_function == PenaltyType.MINIMIZE_TORQUE_DERIVATIVE
+                or penalty_function == PenaltyType.MINIMIZE_COM_POSITION
+                or penalty_function == PenaltyType.MINIMIZE_COM_VELOCITY
             ):
                 parameters.quadratic = True
             else:
@@ -745,6 +802,8 @@ class PenaltyType(Enum):
     MINIMIZE_CONTACT_FORCES = PenaltyFunctionAbstract.Functions.minimize_contact_forces
     TRACK_CONTACT_FORCES = MINIMIZE_CONTACT_FORCES
     MINIMIZE_PREDICTED_COM_HEIGHT = PenaltyFunctionAbstract.Functions.minimize_predicted_com_height
+    MINIMIZE_COM_POSITION = PenaltyFunctionAbstract.Functions.minimize_com_position
+    MINIMIZE_COM_VELOCITY = PenaltyFunctionAbstract.Functions.minimize_com_velocity
     ALIGN_SEGMENT_WITH_CUSTOM_RT = PenaltyFunctionAbstract.Functions.align_segment_with_custom_rt
     ALIGN_MARKER_WITH_SEGMENT_AXIS = PenaltyFunctionAbstract.Functions.align_marker_with_segment_axis
     CUSTOM = PenaltyFunctionAbstract.Functions.custom

examples/torque_driven_with_contact/maximize_predicted_height_CoM.py

@@ -1,4 +1,5 @@
 import biorbd
+import numpy as np
 
 from bioptim import (
     OptimalControlProgram,
@@ -13,10 +14,22 @@
     InitialGuessList,
     ShowResult,
     OdeSolver,
+    Axe,
+    ConstraintList,
+    ConstraintFcn,
+    Node,
 )
 
 
-def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=False, ode_solver=OdeSolver.RK4):
+def prepare_ocp(
+    model_path,
+    phase_time,
+    number_shooting_points,
+    use_actuators=False,
+    ode_solver=OdeSolver.RK4,
+    objective_name="MINIMIZE_PREDICTED_COM_HEIGHT",
+    com_constraints=False,
+):
     # --- Options --- #
     # Model path
     biorbd_model = biorbd.Model(model_path)
@@ -30,7 +43,12 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa
 
     # Add objective functions
     objective_functions = ObjectiveList()
-    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1)
+    if objective_name == "MINIMIZE_PREDICTED_COM_HEIGHT":
+        objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1)
+    elif objective_name == "MINIMIZE_COM_POSITION":
+        objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_COM_POSITION, axis=Axe.Z, weight=-1)
+    elif objective_name == "MINIMIZE_COM_VELOCITY":
+        objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_COM_VELOCITY, axis=Axe.Z, weight=-1)
     objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_TORQUE, weight=1 / 100)
 
     # Dynamics
@@ -40,6 +58,19 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa
     else:
         dynamics.add(DynamicsFcn.TORQUE_DRIVEN_WITH_CONTACT)
 
+    # Constraints
+    constraints = ConstraintList()
+    if com_constraints:
+        constraints.add(
+            ConstraintFcn.COM_VELOCITY,
+            node=Node.ALL,
+            min_bound=np.array([-100, -100, -100]),
+            max_bound=np.array([100, 100, 100]),
+        )
+        constraints.add(
+            ConstraintFcn.COM_POSITION, node=Node.ALL, min_bound=np.array([-1, -1, -1]), max_bound=np.array([1, 1, 1])
+        )
+
     # Path constraint
     nb_q = biorbd_model.nbQ()
     nb_qdot = nb_q
@@ -72,6 +103,7 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa
         x_bounds,
         u_bounds,
         objective_functions,
+        constraints=constraints,
         tau_mapping=tau_mapping,
         ode_solver=ode_solver,
     )
@@ -81,7 +113,14 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa
     model_path = "2segments_4dof_2contacts.bioMod"
     t = 0.5
     ns = 20
-    ocp = prepare_ocp(model_path=model_path, phase_time=t, number_shooting_points=ns, use_actuators=False)
+    ocp = prepare_ocp(
+        model_path=model_path,
+        phase_time=t,
+        number_shooting_points=ns,
+        use_actuators=False,
+        objective_name="MINIMIZE_COM_VELOCITY",
+        com_constraints=True,
+    )
 
     # --- Solve the program --- #
     sol = ocp.solve(show_online_optim=True)

tests/test_global_torque_driven_with_contact_ocp.py

@@ -16,7 +16,11 @@
 
 
 @pytest.mark.parametrize("ode_solver", [OdeSolver.RK4, OdeSolver.RK8, OdeSolver.IRK])
-def test_maximize_predicted_height_CoM(ode_solver):
+@pytest.mark.parametrize(
+    "objective_name", ["MINIMIZE_PREDICTED_COM_HEIGHT", "MINIMIZE_COM_POSITION", "MINIMIZE_COM_VELOCITY"]
+)
+@pytest.mark.parametrize("com_constraints", [False, True])
+def test_maximize_predicted_height_CoM(ode_solver, objective_name, com_constraints):
     PROJECT_FOLDER = Path(__file__).parent / ".."
     spec = importlib.util.spec_from_file_location(
         "maximize_predicted_height_CoM",
@@ -31,32 +35,64 @@ def test_maximize_predicted_height_CoM(ode_solver):
         number_shooting_points=20,
         use_actuators=False,
         ode_solver=ode_solver,
+        objective_name=objective_name,
+        com_constraints=com_constraints,
     )
     sol = ocp.solve()
 
-    # Check objective function value
     f = np.array(sol["f"])
     np.testing.assert_equal(f.shape, (1, 1))
-    np.testing.assert_almost_equal(f[0, 0], 0.7592028279017864)
 
     # Check constraints
     g = np.array(sol["g"])
-    np.testing.assert_equal(g.shape, (160, 1))
-    np.testing.assert_almost_equal(g, np.zeros((160, 1)))
+    if com_constraints:
+        np.testing.assert_equal(g.shape, (244, 1))
+    else:
+        np.testing.assert_equal(g.shape, (160, 1))
+        np.testing.assert_almost_equal(g, np.zeros((160, 1)))
 
     # Check some of the results
     states, controls = Data.get_data(ocp, sol["x"])
     q, qdot, tau = states["q"], states["q_dot"], controls["tau"]
 
-    # initial and final position
+    # initial position
     np.testing.assert_almost_equal(q[:, 0], np.array((0.0, 0.0, -0.5, 0.5)))
-    np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996)))
-    # initial and final velocities
+    # initial velocities
     np.testing.assert_almost_equal(qdot[:, 0], np.array((0, 0, 0, 0)))
-    np.testing.assert_almost_equal(qdot[:, -1], np.array((1.2636414, -1.3010929, -3.6274687, 3.6274687)))
-    # initial and final controls
-    np.testing.assert_almost_equal(tau[:, 0], np.array((-22.1218282)))
-    np.testing.assert_almost_equal(tau[:, -1], np.array(0.2653957))
+
+    if objective_name == "MINIMIZE_PREDICTED_COM_HEIGHT":
+        # Check objective function value
+        np.testing.assert_almost_equal(f[0, 0], 0.7592028279017864)
+
+        # final position
+        np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996)))
+        # final velocities
+        np.testing.assert_almost_equal(qdot[:, -1], np.array((1.2636414, -1.3010929, -3.6274687, 3.6274687)))
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array((-22.1218282)))
+        np.testing.assert_almost_equal(tau[:, -1], np.array(0.2653957))
+    elif objective_name == "MINIMIZE_COM_POSITION":
+        # Check objective function value
+        np.testing.assert_almost_equal(f[0, 0], 0.458575464873056)
+
+        # final position
+        np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996)))
+        # final velocities
+        np.testing.assert_almost_equal(qdot[:, -1], np.array((1.24525494, -1.28216182, -3.57468814, 3.57468814)))
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array((-21.96213697)))
+        np.testing.assert_almost_equal(tau[:, -1], np.array(-0.22120207))
+    elif objective_name == "MINIMIZE_COM_VELOCITY":
+        # Check objective function value
+        np.testing.assert_almost_equal(f[0, 0], 0.4709888694097001)
+
+        # final position
+        np.testing.assert_almost_equal(q[:, -1], np.array((0.1189652, -0.09043785, -0.79999979, 0.79999979)))
+        # final velocities
+        np.testing.assert_almost_equal(qdot[:, -1], np.array((1.26103572, -1.29841047, -3.61998944, 3.61998944)))
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array((-22.18008227)))
+        np.testing.assert_almost_equal(tau[:, -1], np.array(-0.02280469))
 
     # save and load
     TestUtils.save_and_load(sol, ocp, False)

tests/test_penalty.py

@@ -689,6 +689,39 @@ def test_penalty_minimize_predicted_com_height(value):
     )
 
 
+@pytest.mark.parametrize("penalty_origin", [ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer, ConstraintFcn])
+@pytest.mark.parametrize("value", [0.1, -10])
+def test_penalty_minimize_com_position(value, penalty_origin):
+    ocp = prepare_test_ocp()
+    x = [DM.ones((12, 1)) * value]
+    if "COM_POSITION" in penalty_origin._member_names_:
+        penalty_type = penalty_origin.COM_POSITION
+    else:
+        penalty_type = penalty_origin.MINIMIZE_COM_POSITION
+
+    if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)):
+        penalty = Objective(penalty_type)
+    else:
+        penalty = Constraint(penalty_type)
+
+    penalty_type.value[0](penalty, ocp, ocp.nlp[0], [], x, [], [])
+
+    if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)):
+        res = ocp.nlp[0].J[0][0]["val"]
+    else:
+        res = ocp.nlp[0].g[0][0]["val"]
+
+    expected = np.array([[0.05], [0.05], [0.05]])
+    if value == -10:
+        expected = np.array([[-5], [0.05], [-5]])
+
+    np.testing.assert_almost_equal(res, expected)
+
+    if isinstance(penalty_type, ConstraintFcn):
+        np.testing.assert_almost_equal(ocp.nlp[0].g[0][0]["bounds"].min, np.array(0))
+        np.testing.assert_almost_equal(ocp.nlp[0].g[0][0]["bounds"].max, np.array(0))
+
+
 @pytest.mark.parametrize("penalty_origin", [ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer, ConstraintFcn])
 @pytest.mark.parametrize("value", [0.1, -10])
 def test_penalty_align_segment_with_custom_rt(penalty_origin, value):