add bilinear alternation. (#38)

Use wo_input_limits_trigpoly_demo.py as a demonstration of running
bilinear alternation on systems with equality constraints.

compatible_clf_cbf/clf_cbf.py

@@ -132,7 +132,7 @@ def construct_polynomial(
     b_plus_eps: Optional[List[Degree]]
     state_eq_constraints: Optional[List[Degree]]
 
-    def initialize_lagrangians(
+    def to_lagrangians(
         self, prog: solvers.MathematicalProgram, x: sym.Variables, y: sym.Variables
     ) -> CompatibleLagrangians:
         lambda_y = np.array(
@@ -224,6 +224,41 @@ def get_result(self, result: solvers.MathematicalProgramResult) -> Self:
         )
 
 
+@dataclass
+class UnsafeRegionLagrangianDegrees:
+    cbf: int
+    unsafe_region: List[int]
+    state_eq_constraints: Optional[List[int]]
+
+    def to_lagrangians(
+        self, prog: solvers.MathematicalProgram, x_set: sym.Variables
+    ) -> UnsafeRegionLagrangians:
+        cbf, _ = new_sos_polynomial(prog, x_set, self.cbf)
+
+        unsafe_region = np.array(
+            [
+                new_sos_polynomial(prog, x_set, degree)[0]
+                for degree in self.unsafe_region
+            ]
+        )
+
+        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 UnsafeRegionLagrangians(
+            cbf=cbf,
+            unsafe_region=unsafe_region,
+            state_eq_constraints=state_eq_constraints,
+        )
+
+
 class CompatibleClfCbf:
     """
     Certify and synthesize compatible Control Lyapunov Function (CLF) and
@@ -348,9 +383,7 @@ 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,
+        lagrangian_degrees: UnsafeRegionLagrangianDegrees,
         solver_options: Optional[solvers.SolverOptions] = None,
     ) -> UnsafeRegionLagrangians:
         """
@@ -368,36 +401,10 @@ def certify_cbf_unsafe_region(
             self.unsafe_regions[unsafe_region_index]
           cbf: bᵢ(x) in the documentation above. The CBF function for
             self.unsafe_regions[unsafe_region_index]
-          cbf_lagrangian_degree: The degree of the polynomial ϕᵢ,₀(x).
-          unsafe_region_lagrangian_degrees: unsafe_region_lagrangian_degrees[j]
-          is the degree of the polynomial ϕᵢ,ⱼ₊₁(x)
         """
         prog = solvers.MathematicalProgram()
         prog.AddIndeterminates(self.x_set)
-        cbf_lagrangian, cbf_lagrangian_gram = prog.NewSosPolynomial(
-            self.x_set, cbf_lagrangian_degree
-        )
-        unsafe_lagrangians = np.array(
-            [
-                prog.NewSosPolynomial(self.x_set, deg)[0]
-                for deg in unsafe_region_lagrangian_degrees
-            ]
-        )
-        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,
-        )
+        lagrangians = lagrangian_degrees.to_lagrangians(prog, self.x_set)
         self._add_barrier_safe_constraint(prog, unsafe_region_index, cbf, lagrangians)
         result = solvers.Solve(prog, None, solver_options)
         assert result.is_success()
@@ -435,9 +442,7 @@ def construct_search_compatible_lagrangians(
         """
         prog = solvers.MathematicalProgram()
         prog.AddIndeterminates(self.xy_set)
-        lagrangians = lagrangian_degrees.initialize_lagrangians(
-            prog, self.x_set, self.y_set
-        )
+        lagrangians = lagrangian_degrees.to_lagrangians(prog, self.x_set, self.y_set)
         self._add_compatibility(
             prog=prog,
             V=V,
@@ -635,6 +640,130 @@ def in_compatible_region(
         else:
             return in_b
 
+    def bilinear_alternation(
+        self,
+        V_init: Optional[sym.Polynomial],
+        b_init: np.ndarray,
+        rho_init: Optional[float],
+        compatible_lagrangian_degrees: CompatibleLagrangianDegrees,
+        unsafe_regions_lagrangian_degrees: List[UnsafeRegionLagrangianDegrees],
+        kappa_V: Optional[float],
+        kappa_b: np.ndarray,
+        barrier_eps: np.ndarray,
+        x_equilibrium: np.ndarray,
+        clf_degree: Optional[int],
+        cbf_degrees: List[int],
+        max_iter: int,
+        *,
+        x_inner: np.ndarray,
+        solver_id: Optional[solvers.SolverId] = None,
+        solver_options: Optional[solvers.SolverOptions] = None,
+        lagrangian_coefficient_tol: Optional[float] = None,
+        ellipsoid_trust_region: Optional[float] = 100,
+        binary_search_scale_min: float = 1,
+        binary_search_scale_max: float = 2,
+        binary_search_scale_tol: float = 0.1,
+        find_inner_ellipsoid_max_iter: int = 3,
+    ) -> Tuple[Optional[sym.Polynomial], np.ndarray, Optional[float]]:
+        """
+        Synthesize the compatible CLF and CBF through bilinear alternation. We
+        alternate between
+        1. Fixing the CLF/CBF, searching for Lagrangians.
+        2. Fixing Lagrangians, searching for CLF/CBF.
+
+        Args:
+          x_inner: Each row of x_inner is a state that should be contained in
+            the compatible region.
+          max_iter: The maximal number of bilinear alternation iterations.
+          lagrangian_coefficient_tol: We remove the coefficients whose absolute
+            value is smaller than this tolerance in the Lagrangian polynomials.
+            Use None to preserve all coefficients.
+          ellipsoid_trust_region: when we search for the ellipsoid, we put a
+            trust region constraint. This is the squared radius of that trust
+            region.
+        """
+
+        iteration = 0
+        clf = V_init
+        assert len(b_init) == len(self.unsafe_regions)
+        cbf = b_init
+        rho = rho_init
+
+        compatible_lagrangians = None
+        unsafe_lagrangians: List[UnsafeRegionLagrangians] = [None] * len(
+            self.unsafe_regions
+        )
+
+        for iteration in range(max_iter):
+            print(f"iteration {iteration}")
+            # Search for the Lagrangians.
+            for i in range(len(self.unsafe_regions)):
+                unsafe_lagrangians[i] = self.certify_cbf_unsafe_region(
+                    i, cbf[i], unsafe_regions_lagrangian_degrees[i], solver_options
+                )
+            (
+                prog_compatible,
+                compatible_lagrangians,
+            ) = self.construct_search_compatible_lagrangians(
+                clf,
+                cbf,
+                kappa_V,
+                kappa_b,
+                compatible_lagrangian_degrees,
+                rho,
+                barrier_eps,
+            )
+            result_compatible = solve_with_id(
+                prog_compatible, solver_id, solver_options
+            )
+            assert result_compatible.is_success()
+            compatible_lagrangians_result = compatible_lagrangians.get_result(
+                result_compatible, lagrangian_coefficient_tol
+            )
+
+            # Search for the inner ellipsoid.
+            V_contain_ellipsoid_lagrangian_degree = (
+                self._get_V_contain_ellipsoid_lagrangian_degree(clf)
+            )
+            b_contain_ellipsoid_lagrangian_degree = (
+                self._get_b_contain_ellipsoid_lagrangian_degrees(cbf)
+            )
+            (
+                S_ellipsoid_inner,
+                b_ellipsoid_inner,
+                c_ellipsoid_inner,
+            ) = self._find_max_inner_ellipsoid(
+                clf,
+                cbf,
+                rho,
+                V_contain_ellipsoid_lagrangian_degree,
+                b_contain_ellipsoid_lagrangian_degree,
+                x_inner,
+                solver_id=solver_id,
+                max_iter=find_inner_ellipsoid_max_iter,
+                trust_region=ellipsoid_trust_region,
+            )
+
+            clf, cbf, rho = self.binary_search_clf_cbf(
+                compatible_lagrangians_result,
+                unsafe_lagrangians,
+                clf_degree,
+                cbf_degrees,
+                x_equilibrium,
+                kappa_V,
+                kappa_b,
+                barrier_eps,
+                S_ellipsoid_inner,
+                b_ellipsoid_inner,
+                c_ellipsoid_inner,
+                binary_search_scale_min,
+                binary_search_scale_max,
+                binary_search_scale_tol,
+                solver_id,
+                solver_options,
+            )
+        return clf, cbf, rho
+
     def _calc_xi_Lambda(
         self,
         *,
@@ -1059,3 +1188,37 @@ def _add_ellipsoid_in_compatible_region_constraint(
                 inner_eq_poly=self.state_eq_constraints,
                 outer_poly=-b[i],
             )
+
+    def _get_V_contain_ellipsoid_lagrangian_degree(
+        self, V: Optional[sym.Polynomial]
+    ) -> Optional[ContainmentLagrangianDegree]:
+        if V is None:
+            return None
+        else:
+            return ContainmentLagrangianDegree(
+                inner_ineq=[-1],
+                inner_eq=[]
+                if self.state_eq_constraints is None
+                else [
+                    np.maximum(0, V.TotalDegree() - poly.TotalDegree())
+                    for poly in self.state_eq_constraints
+                ],
+                outer=0,
+            )
+
+    def _get_b_contain_ellipsoid_lagrangian_degrees(
+        self, b: np.ndarray
+    ) -> List[ContainmentLagrangianDegree]:
+        return [
+            ContainmentLagrangianDegree(
+                inner_ineq=[-1],
+                inner_eq=[]
+                if self.state_eq_constraints is None
+                else [
+                    np.maximum(0, b_i.TotalDegree() - poly.TotalDegree())
+                    for poly in self.state_eq_constraints
+                ],
+                outer=0,
+            )
+            for b_i in b
+        ]

examples/linear_toy/linear_toy_demo.py

@@ -49,7 +49,10 @@ def search_compatible_lagrangians(
 def search_barrier_safe_lagrangians(
     dut: clf_cbf.CompatibleClfCbf, b: np.ndarray
 ) -> List[clf_cbf.UnsafeRegionLagrangians]:
-    lagrangians = dut.certify_cbf_unsafe_region(0, b[0], 2, [2])
+    lagrangian_degrees = clf_cbf.UnsafeRegionLagrangianDegrees(
+        cbf=2, unsafe_region=[2], state_eq_constraints=None
+    )
+    lagrangians = dut.certify_cbf_unsafe_region(0, b[0], lagrangian_degrees)
     return [lagrangians]
 
 

examples/nonlinear_toy/wo_input_limits_demo.py

@@ -52,11 +52,14 @@ def main(use_y_squared: bool):
     compatible_result = solvers.Solve(compatible_prog)
     assert compatible_result.is_success()
 
+    unsafe_region_lagrangian_degrees = clf_cbf.UnsafeRegionLagrangianDegrees(
+        cbf=0, unsafe_region=[0], state_eq_constraints=None
+    )
+
     compatible.certify_cbf_unsafe_region(
         unsafe_region_index=0,
         cbf=b_init[0],
-        cbf_lagrangian_degree=0,
-        unsafe_region_lagrangian_degrees=[0],
+        lagrangian_degrees=unsafe_region_lagrangian_degrees,
         solver_options=None,
     )
 

examples/nonlinear_toy/wo_input_limits_trigpoly_demo.py

@@ -5,7 +5,6 @@
 
 import numpy as np
 
-import pydrake.solvers as solvers
 import pydrake.symbolic as sym
 
 import compatible_clf_cbf.clf_cbf as clf_cbf
@@ -30,36 +29,50 @@ def main():
     )
     V_init = sym.Polynomial(x[0] ** 2 + x[1] ** 2 + x[2] ** 2)
     b_init = np.array([sym.Polynomial(0.1 - x[0] ** 2 - x[1] ** 2 - x[2] ** 2)])
+    rho_init = 0.1
 
-    lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees(
+    compatible_lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees(
         lambda_y=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0)],
         xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0),
         y=None,
         rho_minus_V=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2),
         b_plus_eps=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2)],
         state_eq_constraints=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=2)],
     )
-    rho = 0.1
+    unsafe_region_lagrangian_degrees = [
+        clf_cbf.UnsafeRegionLagrangianDegrees(
+            cbf=0, unsafe_region=[0], state_eq_constraints=[0]
+        )
+    ]
     kappa_V = 0.01
     kappa_b = np.array([0.01])
     barrier_eps = np.array([0.001])
-    (
-        compatible_prog,
-        compatible_lagrangians,
-    ) = compatible.construct_search_compatible_lagrangians(
+
+    x_equilibrium = np.array([0, 0.0, 0.0])
+
+    clf_degree = 2
+    cbf_degrees = [2]
+    max_iter = 5
+
+    V, b, rho = compatible.bilinear_alternation(
         V_init,
         b_init,
+        rho_init,
+        compatible_lagrangian_degrees,
+        unsafe_region_lagrangian_degrees,
         kappa_V,
         kappa_b,
-        lagrangian_degrees,
-        rho,
         barrier_eps,
+        x_equilibrium,
+        clf_degree,
+        cbf_degrees,
+        max_iter,
+        x_inner=x_equilibrium,
+        find_inner_ellipsoid_max_iter=1,
     )
-    assert compatible_lagrangians is not None
-    solver_options = solvers.SolverOptions()
-    solver_options.SetOption(solvers.CommonSolverOption.kPrintToConsole, 0)
-    compatible_result = solvers.Solve(compatible_prog, None, solver_options)
-    assert compatible_result.is_success()
+    print(f"V={V}")
+    print(f"rho={rho}")
+    print(f"b={b}")
 
 
 if __name__ == "__main__":