Implement _add_compatibility.

Add the p-satz condition to prove the emptiness of a set, which certifies the
compatibility between CLF/CBFs according to Farkas lemma.

compatible_clf_cbf/clf_cbf.py

@@ -1,14 +1,60 @@
+from dataclasses import dataclass
 from typing import List, Optional, Tuple
+from typing_extensions import Self
 
 import numpy as np
-import numpy.typing as npt
 
 import pydrake.symbolic as sym
 import pydrake.solvers as solvers
 
 from compatible_clf_cbf.utils import check_array_of_polynomials
 
 
+@dataclass
+class CompatibleLagrangians:
+    """
+    The Lagrangians for proving the compatibility condition, namely set (1) or (2)
+    defined in CompatibleClfCbf class documentation is empty.
+    """
+
+    # An array of symbolic polynomials. The Lagrangian multiplies with Λ(x)ᵀy if
+    # use_y_squared = False, or Λ(x)ᵀy² if use_y_squared = True.
+    # Each entry in this Lagrangian multiplier is a free polynomial.
+    lambda_y: np.ndarray
+    # The Lagrangian polynomial multiplies with ξ(x)ᵀy if use_y_squared = False,
+    # or ξ(x)ᵀy² if use_y_squared = True. This multiplier is a free polynomial.
+    xi_y: sym.Polynomial
+    # The Lagrangian polynomial multiplies with y if use_y_squared = False.
+    # This multiplier is an array of SOS polynomials.
+    y: Optional[np.ndarray]
+    # The Lagrangian polynomial multiplies with ρ − V when with_clf = True,
+    # should be a SOS polynomial.
+    rho_minus_V: Optional[sym.Polynomial]
+    # The Lagrangian polynomials multiplies with b(x)+ε. Should be an array of SOS
+    # polynomials.
+    b_plus_eps: Optional[np.ndarray]
+
+    def get_result(self, result: solvers.MathematicalProgramResult) -> Self:
+        lambda_y_result = result.GetSolution(self.lambda_y)
+        xi_y_result = result.GetSolution(self.xi_y)
+        y_result = result.GetSolution(self.y) if self.y is not None else None
+        rho_minus_V_result = (
+            result.GetSolution(self.rho_minus_V)
+            if self.rho_minus_V is not None
+            else None
+        )
+        b_plus_eps_result = (
+            result.GetSolution(self.b_plus_eps) if self.b_plus_eps 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,
+        )
+
+
 class CompatibleClfCbf:
     """
     Certify and synthesize compatible Control Lyapunov Function (CLF) and
@@ -22,13 +68,15 @@ class CompatibleClfCbf:
     For simplicity, let's first consider that u is un-constrained, namely 𝒰 is
     the entire space.
     By Farkas lemma, this is equivalent to the following set being empty
-    {(x, y) | [y(0)]ᵀ*[-∂b/∂x*g(x)] = 0, [y(0)]ᵀ*[ ∂b/∂x*f(x)+κ_b*b(x)] = -1, y>=0}
+
+    {(x, y) | [y(0)]ᵀ*[-∂b/∂x*g(x)] = 0, [y(0)]ᵀ*[ ∂b/∂x*f(x)+κ_b*b(x)] = -1, y>=0}        (1)
               [y(1)]  [ ∂V/∂x*g(x)]      [y(1)]  [-∂V/∂x*f(x)-κ_V*V(x)]
+
     We can then use Positivstellensatz to certify the emptiness of this set.
 
     The same math applies to multiple CBFs, or when u is constrained within a
     polyhedron.
-    """
+    """  # noqa E501
 
     def __init__(
         self,
@@ -73,13 +121,13 @@ def __init__(
             total degree of the polynomials. Set use_y_squared=True if we use
             y², and we certify the set
 
-            {(x, y) | [y(0)²]ᵀ*[-∂b/∂x*g(x)] = 0, [y(0)²]ᵀ*[ ∂b/∂x*f(x)+κ_b*b(x)] = -1}
+            {(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.
 
           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.
-        """
+        """  # noqa E501
         assert len(f.shape) == 1
         assert len(g.shape) == 2
         self.nx: int = f.shape[0]
@@ -107,12 +155,21 @@ def __init__(
         self.y: np.ndarray = sym.MakeVectorContinuousVariable(y_size, "y")
         self.y_set: sym.Variables = sym.Variables(self.y)
         self.xy_set: sym.Variables = sym.Variables(np.concatenate((self.x, self.y)))
+        # y_poly[i] is just the polynomial y[i]. I wrote it in this more complicated
+        # form to save some computation.
+        self.y_poly = np.array(
+            [sym.Polynomial(sym.Monomial(self.y[i], 1)) for i in range(y_size)]
+        )
+        # y_squared_poly[i] is just the polynomial y[i]**2.
+        self.y_squared_poly = np.array(
+            [sym.Polynomial(sym.Monomial(self.y[i], 2)) for i in range(y_size)]
+        )
 
     def _calc_xi_Lambda(
         self,
         *,
         V: Optional[sym.Polynomial],
-        b: npt.NDArray[sym.Polynomial],
+        b: np.ndarray,
         kappa_V: Optional[float],
         kappa_b: np.ndarray,
     ) -> Tuple[np.ndarray, np.ndarray]:
@@ -149,20 +206,27 @@ def _calc_xi_Lambda(
         lambda_mat[:num_unsafe_regions] = -dbdx @ self.g
         xi = np.empty((xi_rows,), dtype=object)
         xi[:num_unsafe_regions] = dbdx @ self.f + kappa_b * b
+        # TODO(hongkai.dai): support input bounds Au * u <= bu
+        assert self.Au is None and self.bu is None
 
         if self.with_clf:
             lambda_mat[-1] = dVdx @ self.g
-            xi[-1] = -(dVdx @ self.f) - kappa_V * V
+            xi[-1] = -dVdx.dot(self.f) - kappa_V * V
 
         return (xi, lambda_mat)
 
     def _add_compatibility(
         self,
         *,
         prog: solvers.MathematicalProgram,
+        V: Optional[sym.Polynomial],
+        b: np.ndarray,
+        kappa_V: Optional[float],
+        kappa_b: np.ndarray,
+        lagrangians: CompatibleLagrangians,
         rho: Optional[float],
-        barrier_eps: np.ndarray,
-    ):
+        barrier_eps: Optional[np.ndarray],
+    ) -> sym.Polynomial:
         """
         Add the p-satz condition that certifies the following set is empty
         if use_y_squared = False:
@@ -179,11 +243,55 @@ def _add_compatibility(
         ξ(x) = [ ∂b/∂x*f(x)+κ_b*b(x)]
                [-∂V/∂x*f(x)-κ_V*V(x)]
         To certify the emptiness of the set in (1), we can use the sufficient condition
-        s₀(x, y)ᵀ Λ(x)ᵀy + s₁(x, y)(ξ(x)ᵀy+1) + s₂(x, y)ᵀy + s₃(x, y)(ρ − V) + s₄(x, y)ᵀ(b(x)+ε) = -1
+        s₀(x, y)ᵀ Λ(x)ᵀy + s₁(x, y)(ξ(x)ᵀy+1) + s₂(x, y)ᵀy + s₃(x, y)(ρ − V) + s₄(x, y)ᵀ(b(x)+ε) = -1          (3)
         s₂(x, y), s₃(x, y), s₄(x, y) are all sos.
 
         To certify the emptiness of the set in (2), we can use the sufficient condition
-        s₀(x, y)ᵀ Λ(x)ᵀy² + s₁(x, y)(ξ(x)ᵀy²+1) + s₂(x, y)ᵀy + s₃(x, y)(ρ − V) + s₄(x, y)ᵀ(b(x)+ε) = -1
+        s₀(x, y)ᵀ Λ(x)ᵀy² + s₁(x, y)(ξ(x)ᵀy²+1) + s₃(x, y)(ρ − V) + s₄(x, y)ᵀ(b(x)+ε) = -1                     (4)
         s₃(x, y), s₄(x, y) are all sos.
+
+        Note that we do NOT add the constraint
+        s₂(x, y), s₃(x, y), s₄(x, y) are all sos.
+        in this function. The user should add this constraint separately.
+
+        Returns:
+          poly: The polynomial on the left hand side of equation (3) or (4).
         """  # noqa: E501
-        pass
+        xi, lambda_mat = self._calc_xi_Lambda(
+            V=V, b=b, kappa_V=kappa_V, kappa_b=kappa_b
+        )
+        poly = sym.Polynomial()
+        # Compute s₀(x, y)ᵀ Λ(x)ᵀy
+        if self.use_y_squared:
+            lambda_y = lambda_mat.T @ self.y_squared_poly
+        else:
+            lambda_y = lambda_mat.T @ self.y_poly
+        poly += lagrangians.lambda_y.dot(lambda_y)
+
+        # Compute s₁(x, y)(ξ(x)ᵀy+1)
+        # This is just polynomial 1.
+        poly_one = sym.Polynomial(sym.Monomial())
+        if self.use_y_squared:
+            xi_y = xi.dot(self.y_squared_poly) + poly_one
+        else:
+            xi_y = xi.dot(self.y_poly) + poly_one
+        poly += lagrangians.xi_y * xi_y
+
+        # Compute s₂(x, y)ᵀy
+        if not self.use_y_squared:
+            assert lagrangians.y is not None
+            poly += lagrangians.y.dot(self.y_poly)
+
+        # Compute s₃(x, y)(ρ − V)
+        if rho is not None and self.with_clf:
+            assert V is not None
+            assert lagrangians.rho_minus_V is not None
+            poly += lagrangians.rho_minus_V * (rho * poly_one - V)
+
+        # Compute s₄(x, y)ᵀ(b(x)+ε)
+        if barrier_eps is not None:
+            assert lagrangians.b_plus_eps is not None
+            poly += lagrangians.b_plus_eps.dot(barrier_eps + b)
+
+        prog.AddEqualityConstraintBetweenPolynomials(poly, -poly_one)
+        return poly

tests/test_clf_cbf.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest  # noqa
 
+import pydrake.solvers as solvers
 import pydrake.symbolic as sym
 
 import compatible_clf_cbf.utils as utils
@@ -46,7 +47,17 @@ def test_constructor(self):
             with_clf=True,
             use_y_squared=True,
         )
+
+        def check_members(cls: mut.CompatibleClfCbf):
+            assert cls.y_poly.shape == cls.y.shape
+            for y_poly_i, y_i in zip(cls.y_poly.flat, cls.y.flat):
+                assert y_poly_i.EqualTo(sym.Polynomial(y_i))
+            assert cls.y_squared_poly.shape == cls.y.shape
+            for y_squared_poly_i, y_i in zip(cls.y_squared_poly.flat, cls.y.flat):
+                assert y_squared_poly_i.EqualTo(sym.Polynomial(y_i**2))
+
         assert dut.y.shape == (len(self.unsafe_regions) + 1,)
+        check_members(dut)
 
         # Now construct with with_clf=False
         dut = mut.CompatibleClfCbf(
@@ -60,6 +71,7 @@ def test_constructor(self):
             use_y_squared=True,
         )
         assert dut.y.shape == (len(self.unsafe_regions),)
+        check_members(dut)
 
     def test_calc_xi_Lambda_w_clf(self):
         """
@@ -141,3 +153,75 @@ def test_calc_xi_Lambda_wo_clf(self):
         lambda_mat_expected = np.empty((2, self.nu), dtype=object)
         lambda_mat_expected = -dbdx @ self.g
         utils.check_polynomial_arrays_equal(lambda_mat, lambda_mat_expected, 1e-8)
+
+    def test_add_compatibility_w_clf_y_squared(self):
+        """
+        Test _add_compatibility with CLF and use_y_squared=True
+        """
+        dut = mut.CompatibleClfCbf(
+            f=self.f,
+            g=self.g,
+            x=self.x,
+            unsafe_regions=self.unsafe_regions,
+            Au=None,
+            bu=None,
+            with_clf=True,
+            use_y_squared=True,
+        )
+        prog = solvers.MathematicalProgram()
+        prog.AddIndeterminates(dut.xy_set)
+
+        V = prog.NewFreePolynomial(dut.x_set, deg=2)
+        V = sym.Polynomial(dut.x[0] ** 2 + 4 * dut.x[1] ** 2)
+        b = np.array(
+            [
+                sym.Polynomial(1 - dut.x[0] ** 2 - dut.x[1] ** 2),
+                sym.Polynomial(2 - dut.x[0] ** 2 - dut.x[2] ** 2),
+            ]
+        )
+        kappa_V = 0.01
+        kappa_b = np.array([0.02, 0.03])
+
+        # Set up Lagrangians.
+        lambda_y_lagrangian = np.array(
+            [prog.NewFreePolynomial(dut.xy_set, deg=2) for _ in range(dut.nu)]
+        )
+        xi_y_lagrangian = prog.NewFreePolynomial(dut.xy_set, deg=2)
+        y_lagrangian = None
+        rho_minus_V_lagrangian, _ = prog.NewSosPolynomial(dut.xy_set, degree=2)
+        b_plus_eps_lagrangian = np.array(
+            [
+                prog.NewSosPolynomial(dut.xy_set, degree=2)[0]
+                for _ in range(len(dut.unsafe_regions))
+            ]
+        )
+        lagrangians = mut.CompatibleLagrangians(
+            lambda_y=lambda_y_lagrangian,
+            xi_y=xi_y_lagrangian,
+            y=y_lagrangian,
+            rho_minus_V=rho_minus_V_lagrangian,
+            b_plus_eps=b_plus_eps_lagrangian,
+        )
+
+        rho = 0.1
+        barrier_eps = np.array([0.01, 0.02])
+        poly = dut._add_compatibility(
+            prog=prog,
+            V=V,
+            b=b,
+            kappa_V=kappa_V,
+            kappa_b=kappa_b,
+            lagrangians=lagrangians,
+            rho=rho,
+            barrier_eps=barrier_eps,
+        )
+        (xi, lambda_mat) = dut._calc_xi_Lambda(
+            V=V, b=b, kappa_V=kappa_V, kappa_b=kappa_b
+        )
+        poly_expected = (
+            lagrangians.lambda_y.dot(lambda_mat.T @ dut.y_squared_poly)
+            + lagrangians.xi_y * (xi.dot(dut.y_squared_poly) + 1)
+            + lagrangians.rho_minus_V * (rho - V)
+            + lagrangians.b_plus_eps.dot(b + barrier_eps)
+        )
+        assert poly.CoefficientsAlmostEqual(poly_expected, tolerance=1e-5)