Add CompatibleStateOptions to encourage growing the compatible region…

… to cover certain states.

.flake8

@@ -1,2 +1,4 @@
 [flake8]
 max-line-length = 88
+exclude =
+    venv

.gitignore

@@ -1,2 +1,3 @@
 *.pyc
 compatible_clf_cbf.egg-info/*
+venv/*

compatible_clf_cbf/cbf.py

@@ -1,6 +1,7 @@
 """
 Certify and search Control Barrier Function (CBF) through sum-of-squares optimization.
 """
+
 from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 from typing_extensions import Self

compatible_clf_cbf/clf.py

@@ -1,6 +1,7 @@
 """
 Certify and search Control Lyapunov function (CLF) through sum-of-squares optimization.
 """
+
 from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 from typing_extensions import Self

compatible_clf_cbf/clf_cbf.py

@@ -220,10 +220,12 @@ def get_result(
             unsafe_region=get_polynomial_result(
                 result, self.unsafe_region, coefficient_tol
             ),
-            state_eq_constraints=None
-            if self.state_eq_constraints is None
-            else get_polynomial_result(
-                result, self.state_eq_constraints, coefficient_tol
+            state_eq_constraints=(
+                None
+                if self.state_eq_constraints is None
+                else get_polynomial_result(
+                    result, self.state_eq_constraints, coefficient_tol
+                )
             ),
         )
 
@@ -263,6 +265,123 @@ def to_lagrangians(
         )
 
 
+@dataclass
+class CompatibleStatesOptions:
+    """
+    This option is used to encourage the compatible region to include certain
+    candidate states. Namely b(x_candidate) >= 0 and V(x_candidate) <= 1.
+    """
+
+    candidate_compatible_states: np.ndarray
+    # To avoid arbitrarily scaling the CBF, we need to impose the
+    # constraint that
+    # b_anchor_bounds[i][0] <= b[i](anchor_states) <= b_anchor_bounds[i][1]
+    anchor_states: Optional[np.ndarray]
+    b_anchor_bounds: Optional[List[Tuple[np.ndarray, np.ndarray]]]
+
+    # To encourage the compatible region to cover the candidate states, we add
+    # this cost
+    # weight_V * ReLU(V(x_candidates) - 1)
+    #    + weight_b[i] * ReLU(-b[i](x_candidates))
+    weight_V: Optional[float]
+    weight_b: np.ndarray
+
+    def add_cost(
+        self,
+        prog: solvers.MathematicalProgram,
+        x: np.ndarray,
+        V: Optional[sym.Polynomial],
+        b: np.ndarray,
+    ) -> Tuple[solvers.Binding[solvers.LinearCost], Optional[np.ndarray], np.ndarray]:
+        """
+        Adds the cost
+        weight_V * ReLU(V(x_candidates) - 1)
+           + weight_b[i] * ReLU(-b[i](x_candidates))
+        """
+        assert b.shape == self.weight_b.shape
+        num_candidates = self.candidate_compatible_states.shape[0]
+        if V is not None:
+            # Add the slack variable representing ReLU(V(x_candidates)-1)
+            V_relu = prog.NewContinuousVariables(num_candidates, "V_relu")
+            prog.AddBoundingBoxConstraint(0, np.inf, V_relu)
+            # Now evaluate V(x_candidates) as A_v * V_decision_vars + b_v
+            A_v, V_decision_vars, b_v = V.EvaluateWithAffineCoefficients(
+                x, self.candidate_compatible_states.T
+            )
+            # Now impose the constraint V_relu >= V(x_candidates) - 1 as
+            # V_relu - A_v * V_decision_vars >= b_v -1
+            prog.AddLinearConstraint(
+                np.concatenate((-A_v, np.eye(num_candidates)), axis=1),
+                b_v - 1,
+                np.full_like(b_v, np.inf),
+                np.concatenate((V_decision_vars, V_relu)),
+            )
+        else:
+            V_relu = None
+        # Add the slack variable b_relu[i] representing ReLU(-b[i](x_candidates))
+        b_relu = prog.NewContinuousVariables(b.shape[0], num_candidates, "b_relu")
+        prog.AddBoundingBoxConstraint(0, np.inf, b_relu.reshape((-1,)))
+        for i in range(b.shape[0]):
+            A_b, b_decision_vars, b_b = b[i].EvaluateWithAffineCoefficients(
+                x, self.candidate_compatible_states.T
+            )
+            # Now impose the constraint b_relu[i] >= -b[i](x_candidates) as
+            # A_b * b_decision_vars + b_relu[i] >= - b_b
+            prog.AddLinearConstraint(
+                np.concatenate((A_b, np.eye(num_candidates)), axis=1),
+                -b_b,
+                np.full_like(b_b, np.inf),
+                np.concatenate((b_decision_vars, b_relu[i])),
+            )
+
+        cost_coeff = (self.weight_b.reshape((-1, 1)) * np.ones_like(b_relu)).reshape(
+            (-1,)
+        )
+        cost_vars = b_relu.reshape((-1,))
+        if V is not None:
+            assert self.weight_V is not None
+            cost_coeff = np.concatenate(
+                (cost_coeff, self.weight_V * np.ones(num_candidates))
+            )
+            assert V_relu is not None
+            cost_vars = np.concatenate((cost_vars, V_relu))
+        cost = prog.AddLinearCost(cost_coeff, 0.0, cost_vars)
+        return cost, V_relu, b_relu
+
+    def add_constraint(
+        self, prog: solvers.MathematicalProgram, x: np.ndarray, b: np.ndarray
+    ) -> Optional[List[solvers.Binding[solvers.LinearConstraint]]]:
+        """
+        Add the constraint
+        b_anchor_bounds[i][0] <= b[i](anchor_states) <= b_anchor_bounds[i][1]
+        """
+        if self.b_anchor_bounds is not None:
+            assert b.shape == (len(self.b_anchor_bounds),)
+            assert self.anchor_states is not None
+            constraints: List[solvers.Binding[solvers.LinearConstraint]] = [None] * len(
+                self.b_anchor_bounds
+            )
+            for i in range(len(self.b_anchor_bounds)):
+                assert (
+                    self.b_anchor_bounds[i][0].size
+                    == self.b_anchor_bounds[i][1].size
+                    == self.anchor_states.shape[0]
+                )
+                # Evaluate b[i](anchor_states) as A_b * decision_vars_b + b_b
+                A_b, decision_vars_b, b_b = b[i].EvaluateWithAffineCoefficients(
+                    x, self.anchor_states.T
+                )
+                # Adds the constraint
+                constraints[i] = prog.AddLinearConstraint(
+                    A_b,
+                    self.b_anchor_bounds[i][0] - b_b,
+                    self.b_anchor_bounds[i][1] - b_b,
+                    decision_vars_b,
+                )
+            return constraints
+        return None
+
+
 class CompatibleClfCbf:
     """
     Certify and synthesize compatible Control Lyapunov Function (CLF) and
@@ -518,9 +637,7 @@ def search_clf_cbf_given_lagrangian(
         kappa_V: Optional[float],
         kappa_b: np.ndarray,
         barrier_eps: np.ndarray,
-        S_ellipsoid_inner: np.ndarray,
-        b_ellipsoid_inner: np.ndarray,
-        c_ellipsoid_inner: float,
+        ellipsoid_inner: Optional[ellipsoid_utils.Ellipsoid],
         solver_id: Optional[solvers.SolverId] = None,
         solver_options: Optional[solvers.SolverOptions] = None,
     ) -> Tuple[
@@ -548,9 +665,10 @@ def search_clf_cbf_given_lagrangian(
             barrier_eps,
         )
 
-        self._add_ellipsoid_in_compatible_region_constraint(
-            prog, V, b, S_ellipsoid_inner, b_ellipsoid_inner, c_ellipsoid_inner
-        )
+        if ellipsoid_inner is not None:
+            self._add_ellipsoid_in_compatible_region_constraint(
+                prog, V, b, ellipsoid_inner.S, ellipsoid_inner.b, ellipsoid_inner.c
+            )
 
         result = solve_with_id(prog, solver_id, solver_options)
         if result.is_success():
@@ -571,9 +689,7 @@ def binary_search_clf_cbf(
         kappa_V: Optional[float],
         kappa_b: np.ndarray,
         barrier_eps: np.ndarray,
-        S_ellipsoid_inner: np.ndarray,
-        b_ellipsoid_inner: np.ndarray,
-        c_ellipsoid_inner: float,
+        ellipsoid_inner: ellipsoid_utils.Ellipsoid,
         scale_min: float,
         scale_max: float,
         scale_tol: float,
@@ -605,7 +721,7 @@ def search(
             solvers.MathematicalProgramResult,
         ]:
             c_new = ellipsoid_utils.scale_ellipsoid(
-                S_ellipsoid_inner, b_ellipsoid_inner, c_ellipsoid_inner, scale
+                ellipsoid_inner.S, ellipsoid_inner.b, ellipsoid_inner.c, scale
             )
             V, b, result = self.search_clf_cbf_given_lagrangian(
                 compatible_lagrangians,
@@ -616,9 +732,7 @@ def search(
                 kappa_V,
                 kappa_b,
                 barrier_eps,
-                S_ellipsoid_inner,
-                b_ellipsoid_inner,
-                c_new,
+                ellipsoid_utils.Ellipsoid(ellipsoid_inner.S, ellipsoid_inner.b, c_new),
                 solver_id,
                 solver_options,
             )
@@ -789,9 +903,9 @@ def bilinear_alternation(
                 kappa_V,
                 kappa_b,
                 barrier_eps,
-                S_ellipsoid_inner,
-                b_ellipsoid_inner,
-                c_ellipsoid_inner,
+                ellipsoid_utils.Ellipsoid(
+                    S_ellipsoid_inner, b_ellipsoid_inner, c_ellipsoid_inner
+                ),
                 binary_search_scale_min,
                 binary_search_scale_max,
                 binary_search_scale_tol,
@@ -1044,7 +1158,11 @@ def _construct_search_clf_cbf_program(
         kappa_b: np.ndarray,
         barrier_eps: np.ndarray,
         local_clf: bool = True,
-    ) -> Tuple[solvers.MathematicalProgram, Optional[sym.Polynomial], np.ndarray,]:
+    ) -> Tuple[
+        solvers.MathematicalProgram,
+        Optional[sym.Polynomial],
+        np.ndarray,
+    ]:
         """
         Construct a program to search for compatible CLF/CBFs given the Lagrangians.
         Notice that we have not imposed the cost to the program yet.
@@ -1275,12 +1393,14 @@ def _get_V_contain_ellipsoid_lagrangian_degree(
         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
-                ],
+                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,
             )
 
@@ -1290,12 +1410,14 @@ def _get_b_contain_ellipsoid_lagrangian_degrees(
         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
-                ],
+                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

compatible_clf_cbf/ellipsoid_utils.py

@@ -367,3 +367,20 @@ def draw_ellipsoid2d(
         + d.reshape((-1, 1))
     )
     return ax.plot(xy_pts[0], xy_pts[1], **kwargs)
+
+
+class Ellipsoid:
+    """
+    An ellipsoid as
+    {x | xᵀSx+bᵀx+c≤0}
+    """
+
+    def __init__(self, S: np.ndarray, b: np.ndarray, c: float):
+        assert S.shape[0] == S.shape[1]
+        assert b.shape == (S.shape[0],)
+        self.S = (S + S.T) / 2
+        self.b = b
+        self.c = c
+
+    def to_affine_ball(self) -> Tuple[np.ndarray, np.ndarray]:
+        return to_affine_ball(self.S, self.b, self.c)

examples/linear_toy/linear_toy_demo.py

@@ -1,6 +1,7 @@
 """
 We search for compatible CLF and CBF for a 2D linear system.
 """
+
 from typing import List, Tuple
 
 import numpy as np
@@ -27,11 +28,14 @@ def search_compatible_lagrangians(
             clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0) for _ in range(dut.nx)
         ],
         xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0),
-        y=None
-        if dut.use_y_squared
-        else [
-            clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0) for _ in range(y_size)
-        ],
+        y=(
+            None
+            if dut.use_y_squared
+            else [
+                clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0)
+                for _ in range(y_size)
+            ]
+        ),
         rho_minus_V=None,
         b_plus_eps=None,
         state_eq_constraints=None,

examples/nonlinear_toy/wo_input_limits_demo.py

@@ -1,6 +1,7 @@
 """
 Find the compatible CLF/CBF with no input limits.
 """
+
 import numpy as np
 import pydrake.solvers as solvers
 import pydrake.symbolic as sym
@@ -31,12 +32,14 @@ def main(use_y_squared: bool):
     lagrangian_degrees = clf_cbf.CompatibleLagrangianDegrees(
         lambda_y=[clf_cbf.CompatibleLagrangianDegrees.Degree(x=3, y=0)],
         xi_y=clf_cbf.CompatibleLagrangianDegrees.Degree(x=2, y=0),
-        y=None
-        if use_y_squared
-        else [
-            clf_cbf.CompatibleLagrangianDegrees.Degree(x=4, y=0)
-            for _ in range(compatible.y.size)
-        ],
+        y=(
+            None
+            if use_y_squared
+            else [
+                clf_cbf.CompatibleLagrangianDegrees.Degree(x=4, y=0)
+                for _ in range(compatible.y.size)
+            ]
+        ),
         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,

examples/single_integrator/wo_input_limit_demo.py

@@ -3,6 +3,7 @@
 2-dimensional single-integrator system xdot = u. I assume the obstacle is a
 sphere.
 """
+
 from typing import Tuple
 
 import numpy as np