Support state equality constraints for dynamics. (#34)

compatible_clf_cbf/clf.py

@@ -27,8 +27,10 @@ class ClfWoInputLimitLagrangian:
     dVdx_times_f: sym.Polynomial
     # The array of Lagrangians λ₁(x) in λ₁(x)*∂V/∂x*g(x)
     dVdx_times_g: np.ndarray
-    # The Lagrangian
+    # The Lagrangian for ρ − V(x)
     rho_minus_V: sym.Polynomial
+    # The array of Lagrangians for state equality constraints.
+    state_eq_constraints: Optional[np.ndarray]
 
     def get_result(
         self,
@@ -44,10 +46,18 @@ def get_result(
         rho_minus_V_result = get_polynomial_result(
             result, self.rho_minus_V, coefficient_tol
         )
+        state_eq_constraints_result = (
+            None
+            if self.state_eq_constraints is None
+            else get_polynomial_result(
+                result, self.state_eq_constraints, coefficient_tol
+            )
+        )
         return ClfWoInputLimitLagrangian(
             dVdx_times_f=dVdx_times_f_result,
             dVdx_times_g=dVdx_times_g_result,
             rho_minus_V=rho_minus_V_result,
+            state_eq_constraints=state_eq_constraints_result,
         )
 
 
@@ -56,6 +66,7 @@ class ClfWoInputLimitLagrangianDegrees:
     dVdx_times_f: int
     dVdx_times_g: List[int]
     rho_minus_V: int
+    state_eq_constraints: Optional[List[int]]
 
     def to_lagrangians(
         self,
@@ -79,7 +90,20 @@ def to_lagrangians(
         rho_minus_V, _ = new_sos_polynomial(
             prog, x_set, self.rho_minus_V, bool(np.all(x_equilibrium == 0))
         )
-        return ClfWoInputLimitLagrangian(dVdx_times_f, dVdx_times_g, rho_minus_V)
+
+        state_eq_constraints = (
+            None
+            if self.state_eq_constraints is None
+            else np.array(
+                [
+                    prog.NewFreePolynomial(x_set, degree)
+                    for degree in self.state_eq_constraints
+                ]
+            )
+        )
+        return ClfWoInputLimitLagrangian(
+            dVdx_times_f, dVdx_times_g, rho_minus_V, state_eq_constraints
+        )
 
 
 @dataclass
@@ -92,6 +116,8 @@ class ClfWInputLimitLagrangian:
     V_minus_rho: sym.Polynomial
     # The Lagrangians λᵢ(x) in ∑ᵢ λᵢ(x)*(∂V/∂x*(f(x)+g(x)uᵢ)+κ*V(x))
     Vdot: np.ndarray
+    # The Lagrangians for state equality constraints
+    state_eq_constraints: Optional[np.ndarray]
 
     def get_result(
         self,
@@ -104,15 +130,25 @@ def get_result(
         Vdot_result: np.ndarray = get_polynomial_result(
             result, self.Vdot, coefficient_tol
         )
+        state_eq_constraints_result = (
+            None
+            if self.state_eq_constraints is None
+            else get_polynomial_result(
+                result, self.state_eq_constraints, coefficient_tol
+            )
+        )
         return ClfWInputLimitLagrangian(
-            V_minus_rho=V_minus_rho_result, Vdot=Vdot_result
+            V_minus_rho=V_minus_rho_result,
+            Vdot=Vdot_result,
+            state_eq_constraints=state_eq_constraints_result,
         )
 
 
 @dataclass
 class ClfWInputLimitLagrangianDegrees:
     V_minus_rho: int
     Vdot: List[int]
+    state_eq_constraints: Optional[List[int]]
 
     def to_lagrangians(
         self,
@@ -124,7 +160,17 @@ def to_lagrangians(
         Vdot = np.array(
             [new_sos_polynomial(prog, x_set, degree) for degree in self.Vdot]
         )
-        return ClfWInputLimitLagrangian(V_minus_rho, Vdot)
+        state_eq_constraints = (
+            None
+            if self.state_eq_constraints is None
+            else np.array(
+                [
+                    prog.NewFreePolynomial(x_set, degree)
+                    for degree in self.state_eq_constraints
+                ]
+            )
+        )
+        return ClfWInputLimitLagrangian(V_minus_rho, Vdot, state_eq_constraints)
 
 
 class ControlLyapunov:
@@ -162,6 +208,7 @@ def __init__(
         x: np.ndarray,
         x_equilibrium: np.ndarray,
         u_vertices: Optional[np.ndarray] = None,
+        state_eq_constraints: Optional[np.ndarray] = None,
     ):
         """
         Args:
@@ -178,6 +225,12 @@ def __init__(
             as the equilibrium state.
           u_vertices: The vertices of the input constraint polytope 𝒰. Each row
             is a vertex. If u_vertices=None, then the input is unconstrained.
+          state_eq_constraints: An array of polynomials. Some dynamical systems
+            have equality constraints on its states. For example, when the
+            state include sinθ and cosθ (so that the dynamics is a polynomial
+            function of state), we need to impose the equality constraint
+            sin²θ+cos²θ=1 on the state. state_eq_constraints[i] = 0 is an
+            equality constraint on the state.
 
         """
         assert len(f.shape) == 1
@@ -197,6 +250,9 @@ def __init__(
         if u_vertices is not None:
             assert u_vertices.shape[1] == self.nu
         self.u_vertices = u_vertices
+        if state_eq_constraints is not None:
+            check_array_of_polynomials(state_eq_constraints, self.x_set)
+        self.state_eq_constraints = state_eq_constraints
 
     def search_lagrangian_given_clf(
         self,
@@ -267,12 +323,14 @@ def _add_clf_condition(
                 - dVdx_times_g.squeeze().dot(lagrangians.dVdx_times_g)
                 - lagrangians.rho_minus_V * (rho - V)
             )
-            prog.AddSosConstraint(sos_poly)
         else:
             assert isinstance(lagrangians, ClfWInputLimitLagrangian)
             Vdot = dVdx_times_f + dVdx_times_g @ self.u_vertices
             sos_poly = (1 + lagrangians.V_minus_rho) * (V - rho) * sym.Polynomial(
                 self.x.dot(self.x)
             ) - lagrangians.Vdot.dot(Vdot + kappa * V)
-            prog.AddSosConstraint(sos_poly)
+        if self.state_eq_constraints is not None:
+            assert lagrangians.state_eq_constraints is not None
+            sos_poly -= lagrangians.state_eq_constraints.dot(self.state_eq_constraints)
+        prog.AddSosConstraint(sos_poly)
         return sos_poly

compatible_clf_cbf/clf_cbf.py

@@ -41,39 +41,9 @@ class CompatibleLagrangians:
     # The Lagrangian polynomials multiplies with b(x)+ε. Should be an array of SOS
     # polynomials.
     b_plus_eps: Optional[np.ndarray]
-
-    @classmethod
-    def reserve(
-        cls,
-        nu: int,
-        use_y_squared: bool,
-        y_size,
-        with_rho_minus_V: bool,
-        b_plus_eps_size: Optional[int],
-    ) -> Self:
-        """
-        Reserve the Lagrangian polynomials. Note that the polynomials are
-        initialized to 0, you should properly set their values.
-
-        Args:
-          nu: The dimension of control u.
-          use_y_squared: Check CompatibleClfCbf documentation.
-          y_size: The size of the indeterminates y.
-          with_rho_minus_V: Whether the psatz condition considers ρ - V.
-          b_plus_eps_size: The size of b(x)+ε. If set to None, then we don't consider
-          b(x)+ε in the psatz.
-        """
-        return CompatibleLagrangians(
-            lambda_y=np.array([sym.Polynomial() for _ in range(nu)]),
-            xi_y=sym.Polynomial(),
-            y=None
-            if use_y_squared
-            else np.array([sym.Polynomial() for _ in range(y_size)]),
-            rho_minus_V=sym.Polynomial() if with_rho_minus_V else None,
-            b_plus_eps=None
-            if b_plus_eps_size is None
-            else np.array([sym.Polynomial() for _ in range(b_plus_eps_size)]),
-        )
+    # The free Lagrangian polynomials multiplying the state equality
+    # constraints.
+    state_eq_constraints: Optional[np.ndarray]
 
     def get_result(
         self,
@@ -100,12 +70,18 @@ def get_result(
             if self.b_plus_eps is not None
             else None
         )
+        state_eq_constraints_result = (
+            get_polynomial_result(result, self.state_eq_constraints, coefficient_tol)
+            if self.state_eq_constraints is not None
+            else None
+        )
         return CompatibleLagrangians(
             lambda_y=lambda_y_result,
             xi_y=xi_y_result,
             y=y_result,
             rho_minus_V=rho_minus_V_result,
             b_plus_eps=b_plus_eps_result,
+            state_eq_constraints=state_eq_constraints_result,
         )
 
 
@@ -153,6 +129,7 @@ def construct_polynomial(
     y: Optional[List[Degree]]
     rho_minus_V: Optional[Degree]
     b_plus_eps: Optional[List[Degree]]
+    state_eq_constraints: Optional[List[Degree]]
 
     def initialize_lagrangians(
         self, prog: solvers.MathematicalProgram, x: sym.Variables, y: sym.Variables
@@ -186,12 +163,25 @@ def initialize_lagrangians(
                 ]
             )
         )
+        state_eq_constraints = (
+            None
+            if self.state_eq_constraints is None
+            else np.array(
+                [
+                    state_eq_constraints_i.construct_polynomial(
+                        prog, x, y, is_sos=False
+                    )
+                    for state_eq_constraints_i in self.state_eq_constraints
+                ]
+            )
+        )
         return CompatibleLagrangians(
             lambda_y=lambda_y,
             xi_y=xi_y,
             y=y_lagrangian,
             rho_minus_V=rho_minus_V,
             b_plus_eps=b_plus_eps,
+            state_eq_constraints=state_eq_constraints,
         )
 
 
@@ -215,13 +205,21 @@ class UnsafeRegionLagrangians:
     # An array of sym.Polynomial. The Lagrangians that multiply the unsafe region
     # polynomials. ϕᵢ,ⱼ(x) in the documentation above.
     unsafe_region: np.ndarray
+    # The free Lagrangian that multiplies with the state equality constraints
+    # (such as sin²θ+cos²θ=1)
+    state_eq_constraints: Optional[np.ndarray]
 
     def get_result(self, result: solvers.MathematicalProgramResult) -> Self:
         return UnsafeRegionLagrangians(
             cbf=result.GetSolution(self.cbf),
             unsafe_region=np.array(
                 [result.GetSolution(phi) for phi in self.unsafe_region]
             ),
+            state_eq_constraints=None
+            if self.state_eq_constraints is None
+            else np.array(
+                [result.GetSolution(poly) for poly in self.state_eq_constraints]
+            ),
         )
 
 
@@ -259,6 +257,7 @@ def __init__(
         bu: Optional[np.ndarray] = None,
         with_clf: bool = True,
         use_y_squared: bool = True,
+        state_eq_constraints: Optional[np.ndarray] = None,
     ):
         """
         Args:
@@ -294,6 +293,12 @@ def __init__(
             {(x, y) | [y(0)²]ᵀ*[-∂b/∂x*g(x)] = 0, [y(0)²]ᵀ*[ ∂b/∂x*f(x)+κ_b*b(x)] = -1}       (2)
                       [y(1)²]  [ ∂V/∂x*g(x)]      [y(1)²]  [-∂V/∂x*f(x)-κ_V*V(x)]
             is empty.
+          state_eq_constraints: An array of polynomials. Some dynamical systems
+            have equality constraints on its states. For example, when the
+            state include sinθ and cosθ (so that the dynamics is a polynomial
+            function of state), we need to impose the equality constraint
+            sin²θ+cos²θ=1 on the state. state_eq_constraints[i] = 0 is an
+            equality constraint on the state.
 
           If both Au and bu are None, it means that we don't have input limits.
           They have to be both None or both not None.
@@ -334,13 +339,17 @@ def __init__(
         self.y_squared_poly = np.array(
             [sym.Polynomial(sym.Monomial(self.y[i], 2)) for i in range(y_size)]
         )
+        self.state_eq_constraints = state_eq_constraints
+        if self.state_eq_constraints is not None:
+            check_array_of_polynomials(self.state_eq_constraints, self.x_set)
 
     def certify_cbf_unsafe_region(
         self,
         unsafe_region_index: int,
         cbf: sym.Polynomial,
         cbf_lagrangian_degree: int,
         unsafe_region_lagrangian_degrees: List[int],
+        state_eq_constraints_lagrangian_degrees: Optional[List[int]] = None,
         solver_options: Optional[solvers.SolverOptions] = None,
     ) -> UnsafeRegionLagrangians:
         """
@@ -373,7 +382,21 @@ def certify_cbf_unsafe_region(
                 for deg in unsafe_region_lagrangian_degrees
             ]
         )
-        lagrangians = UnsafeRegionLagrangians(cbf_lagrangian, unsafe_lagrangians)
+        if self.state_eq_constraints is not None:
+            assert state_eq_constraints_lagrangian_degrees is not None
+            state_eq_constraints_lagrangians = np.array(
+                [
+                    prog.NewFreePolynomial(self.x_set, deg)
+                    for deg in state_eq_constraints_lagrangian_degrees
+                ]
+            )
+        else:
+            state_eq_constraints_lagrangians = None
+        lagrangians = UnsafeRegionLagrangians(
+            cbf_lagrangian,
+            unsafe_lagrangians,
+            state_eq_constraints=state_eq_constraints_lagrangians,
+        )
         self._add_barrier_safe_constraint(prog, unsafe_region_index, cbf, lagrangians)
         result = solvers.Solve(prog, None, solver_options)
         assert result.is_success()
@@ -741,6 +764,10 @@ def _add_compatibility(
             assert lagrangians.b_plus_eps is not None
             poly -= lagrangians.b_plus_eps.dot(barrier_eps + b)
 
+        if self.state_eq_constraints is not None:
+            assert lagrangians.state_eq_constraints is not None
+            poly -= lagrangians.state_eq_constraints.dot(self.state_eq_constraints)
+
         prog.AddSosConstraint(poly)
         return poly
 
@@ -781,6 +808,9 @@ def _add_barrier_safe_constraint(
         poly = -(1 + lagrangians.cbf) * b + lagrangians.unsafe_region.dot(
             self.unsafe_regions[unsafe_region_index]
         )
+        if self.state_eq_constraints is not None:
+            assert lagrangians.state_eq_constraints is not None
+            poly -= lagrangians.state_eq_constraints.dot(self.state_eq_constraints)
         prog.AddSosConstraint(poly)
         return poly
 

examples/linear_toy/linear_toy_demo.py

@@ -34,6 +34,7 @@ def search_compatible_lagrangians(
         ],
         rho_minus_V=None,
         b_plus_eps=None,
+        state_eq_constraints=None,
     )
     prog, lagrangians = dut.construct_search_compatible_lagrangians(
         V, b, kappa_V, kappa_b, lagrangian_degrees, rho=None, barrier_eps=None

examples/nonlinear_toy/wo_input_limits_demo.py

@@ -39,6 +39,7 @@ def main(use_y_squared: bool):
         ],
         rho_minus_V=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0),
         b_plus_eps=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0)],
+        state_eq_constraints=None,
     )
     rho = 0.01
     barrier_eps = np.array([0.0001])