Add the linearized function for maximizing ellipsoid volume. (#15)

compatible_clf_cbf/ellipsoid_utils.py

@@ -0,0 +1,67 @@
+"""
+The utility function for manipulating ellipsoids
+"""
+
+import numpy as np
+
+import pydrake.solvers as solvers
+import pydrake.symbolic as sym
+
+
+def add_max_volume_linear_cost(
+    prog: solvers.MathematicalProgram,
+    S: np.ndarray,
+    b: np.ndarray,
+    c: sym.Variable,
+    S_bar: np.ndarray,
+    b_bar: np.ndarray,
+    c_bar: float,
+) -> solvers.Binding[solvers.LinearCost]:
+    """
+    Adds the linear cost as the linearization of an objective which correlates
+    to the volume of the ellipsoid ℰ={x | xᵀSx+bᵀx+c≤0}
+
+    To maximize the ellipsoid voluem, we can maximize
+    log(bᵀS⁻¹b/4-c) - 1/n * log(det(S))
+
+    The linearization of this function at (S̅, b̅, c̅) is
+    max trace(⌈c   bᵀ/2⌉ᵀ  ⌈c̅,   b̅ᵀ/2⌉⁻¹) -(1+1/n)*trace(Sᵀ * S̅⁻¹)
+              ⌊b/2    S⌋   ⌊b̅/2     S̅⌋
+
+    Check doc/maximize_inner_ellipsoid.md for more details.
+
+    Args:
+      prog: The optimization problem to which the cost is added.
+      S: A symmetric matrix of decision variables. S must have been registered
+        in `prog` already.
+      b: A vector of decision variables. b must have been registered in `prog`
+        already.
+      c: A decision variable. c must have been registered in `prog` already.
+      S_bar: A symmetric matrix of floats, where we linearize the objective.
+      b_bar: A vector of floats, where we linearized the objective.
+    """
+
+    n = S.shape[0]
+    assert S.shape == (n, n)
+    assert b.shape == (n,)
+    assert S_bar.shape == (n, n)
+    assert b_bar.shape == (n,)
+    mat = np.empty((n + 1, n + 1), dtype=object)
+    mat[0, 0] = c
+    mat[0, 1:] = b / 2
+    mat[1:, 0] = b / 2
+    mat[1:, 1:] = S
+
+    mat_bar = np.empty((n + 1, n + 1))
+    mat_bar[0, 0] = c_bar
+    mat_bar[0, 1:] = b_bar / 2
+    mat_bar[1:, 0] = b_bar / 2
+    mat_bar[1:, 1:] = S_bar
+    mat_bar_inv = np.linalg.inv(mat_bar)
+    S_bar_inv = np.linalg.inv(S_bar)
+
+    cost_expr = np.trace(mat.T @ mat_bar_inv) - (1 + 1.0 / n) * np.trace(
+        S.T @ S_bar_inv
+    )
+    cost = prog.AddLinearCost(-cost_expr)
+    return cost

doc/maximize_inner_ellipsoid.md

@@ -0,0 +1,120 @@
+# Maximizing ellipsoid volume
+
+Say we have a semi-algebraic set $S = \{x\in\mathbb{R}^n | g(x) < 0 \}$ and we want to find a large inscribed ellipsoid $\mathcal{E}\subset \mathcal{S}$.
+
+## Formulation 1
+
+Consider to parameterize the ellipsoid as $\mathcal{E} = \{x | x^TSx + b^Tx + c \le 0\}$. The condition that $\mathcal{E}\subset\mathcal{S}$ is
+
+$$
+\begin{align}
+-1 - \phi_0(x)g(x) + \phi_1(x)(x^TSx+b^Tx+c) \text{ is sos}\\
+\phi_0(x), \phi_1(x) \text{ are sos}.
+\end{align}
+$$
+
+The volume of the ellipsoid is proportional to
+
+$$
+vol(\mathcal{E})\propto \left(\frac{b^TS^{-1}b/4-c}{\text{det}(S)^{1/n}}\right)^{\frac{n}{2}}
+$$
+
+So to maximize the ellipsoid, we can maximize 
+
+$$
+\begin{align}
+\frac{b^TS^{-1}b/4-c}{\text{det}(S)^{1/n}}
+\end{align}
+$$
+
+I don’t think we can maximize this term directly through convex optimization, so we will seek to maximize a surrogate function instead, which can be done through convex optimization.
+
+Since logarithm function is monotonically increasing, we can maximize the log of Eq 3
+
+$$
+\begin{align}
+\max_{S, b, c} \log(b^TS^{-1}b/4-c) - \frac{1}{n}\log\text{det}(S)
+\end{align}
+$$
+
+This objective (4) is still non-convex. So we consider to linearize this objective and maximize the linearization. To compute the linearized objective, by Schur complement, we know that
+
+$$
+(c - b^TS^{-1}b/4) \bullet \text{det}(S) = \text{det}\left(\begin{bmatrix}c & b^T/2\\b/2 & S\end{bmatrix}\right)
+$$
+
+Therefore we have 
+
+$$
+\log(b^TS^{-1}b/4-c) = \log\left(-\text{det}\begin{bmatrix}c & b^T/2 \\b/2 & S\end{bmatrix}\right) - \log \text{det}(S)
+$$
+
+Hence the objective function (4) can be re-formulated as
+
+$$
+\begin{align}
+\max_{S, b, c} \log \left(-\text{det}\left(\begin{bmatrix}c & b^T/2 \\b/2 & S\end{bmatrix}\right)\right) - (1+\frac{1}{n})\log\text{det}(S)
+\end{align}
+$$
+
+It is easier to compute the linearization of objective (5) than objective (4), by using the following property on the gradient of $\log \text{det}(X)$ for a symmetric matrix $X$:m
+
+$$
+\begin{align}
+\frac{\partial\log\text{det}(X)}{\partial X} = X^{-1}\text{ if } \text{det}(X) > 0\\
+\frac{\partial\log(-\text{det}(X))}{\partial X} = X^{-1}\text{ if } \text{det}(X) < 0
+\end{align}
+$$
+
+So the linearization of the objective in (5) at a point $(S^{(i)}, b^{(i)}, c^{(i)})$ is
+
+$$
+\left<\begin{bmatrix}c & b^T/2\\b/2 & S\end{bmatrix}, \begin{bmatrix}c^{(i)} &(b^{(i)})^T/2\\b^{(i)}/2 & S^{(i)}\end{bmatrix}^{-1}\right> - (1+\frac{1}{n})\left<S, (S^{(i)})^{-1}\right>
+$$
+
+where we use $<X, Y> \equiv \text{trace}(X^TY)$ to denote matrix inner product.
+
+Furthermore, to ensure that the $\{x | x^TSx+b^Tx+c\le 0\}$ really represents an ellipsoid. We need to enforce the constraint $b^TS^{-1}b/4-c \ge 0$. Unfortunately this constraint is non-convex. To remedy this, we consider a sufficient condition, that a given point $\bar{x}$ is inside the ellipsoid, namely
+
+$$
+\bar{x}^TS\bar{x}+b^T\bar{x}+c\le 0
+$$
+
+as a linear constraint on $(S, b, c)$.
+
+So to summarize, we solve a sequence of convex optimization problem
+
+$$
+\begin{align}
+\max_{S, b, c, \phi_0} &\left<\begin{bmatrix}c & b^T/2\\b/2 & S\end{bmatrix}, \begin{bmatrix}c^{(i)} &(b^{(i)})^T/2\\b^{(i)}/2 & S^{(i)}\end{bmatrix}^{-1}\right> - (1+\frac{1}{n})\left<S, (S^{(i)})^{-1}\right>\\
+\text{s.t } &\bar{x}^TS\bar{x} + b^T\bar{x} + c \le 0\\
+&-1 - \phi_0(x)g(x) + \phi_1(x)(x^TSx+b^Tx+c) \text{ is sos}\\
+&\phi_0(x), \phi_1(x) \text{ are sos}.\end{align}
+$$
+
+In the i’th iteration we get a solution $(S^{(i)}, b^{(i)}, c^{(i)})$, we linearize the objective (5) at this solution and maximize the linear objective again.
+
+## Formulation 2
+
+We consider an alternative parameterization of the ellipsoid as $\mathcal{E}=\{x | \Vert Ax+b\Vert_2\le 1\}, A\succ 0$. To maximize the volume of this ellipsoid, we need to minimize $\log \text{det}(A)$. But $\log \text{det}(A)$ is a concave function. To minimize this concave function, we again use the sequential linearization idea as in formulation 1. Namely we linearize this objective at a point $A^{(i)}$, and minimize the linearized objective
+
+$$
+\left<A - A^{(i)}, (A^{(i)})^{-1}\right>
+$$
+
+Moreover, to impose the constraint that $\mathcal{E}\subset\mathcal{S}$, we first use the Schur complement to convert $\Vert Ax+b\Vert_2\le 1$ to a linear constraint on $(A, b)$
+
+$$
+\Vert Ax+b\Vert_2\le 1 \Leftrightarrow \begin{bmatrix} 1 &(Ax+b)^T\\Ax+b & I\end{bmatrix}\succeq 0
+$$
+
+Hence a sufficient condition for $\mathcal{E}\subset\mathcal{S}$ is
+
+$$
+\begin{align}
+g(x)\Phi_0(x) - \phi_1(x)\begin{bmatrix}1 & (Ax+b)^T\\Ax+b & I \end{bmatrix} \text{ is psd}\\
+\Phi_0(x) \text{ is psd}, \phi_1(x) \text{ is sos}
+\end{align}
+$$
+
+Notice that this condition is significantly more complicated than the p-satz condition (Eq (1) and (2)) in Formulation 1, as (12) (13) require matrix-sos conditions.

requirements.txt

@@ -3,6 +3,8 @@ black[jupyter]
 drake
 flake8
 ipykernel
+jax
+jaxlib
 matplotlib
 numpy
 pytest

tests/test_ellipsoid_utils.py

@@ -0,0 +1,104 @@
+import compatible_clf_cbf.ellipsoid_utils as mut
+
+import jax.numpy as jnp
+import jax
+import numpy as np
+import pytest  # noqa
+
+import pydrake.solvers as solvers
+import pydrake.symbolic as sym
+
+
+def eval_max_volume_cost(S: jnp.ndarray, b: jnp.ndarray, c: jnp.ndarray):
+    """
+    Evaluate the cost
+    log(bᵀS⁻¹b/4-c) - 1/n * log(det(S))
+    """
+    n = S.shape[0]
+    return jnp.log(b.dot(jnp.linalg.solve(S, b)) / 4 - c) - 1.0 / n * jnp.log(
+        jnp.linalg.det(S)
+    )
+
+
+def check_add_max_volume_linear_cost(
+    cost: solvers.Binding[solvers.LinearCost],
+    S: np.ndarray,
+    b: np.ndarray,
+    c: sym.Variable,
+    S_val: np.ndarray,
+    b_val: np.ndarray,
+    c_val: float,
+    S_bar: np.ndarray,
+    b_bar: np.ndarray,
+    c_bar: float,
+):
+    """
+    Evaluate the `cost` term at (S, b, c) = (S_val, b_val, c_val). Make sure
+    that the evaluated value is the same as taking the linearization of the
+    nonlinear cost numerically through JAX.
+    """
+    S_bar_jnp = jnp.array(S_bar)
+    b_bar_jnp = jnp.array(b_bar)
+    S_grad, b_grad, c_grad = jax.grad(eval_max_volume_cost, argnums=(0, 1, 2))(
+        S_bar_jnp, b_bar_jnp, c_bar
+    )
+    cost_val_expected = -(
+        np.trace(S_grad.T @ S_val) + b_grad.dot(b_val) + c_grad * c_val
+    )
+    env = dict()
+    for i in range(S.shape[0]):
+        for j in range(i, S.shape[1]):
+            env[S[i, j]] = S_val[i, j]
+    for b_var, b_var_val in zip(b.flat, b_val.flat):
+        env[b_var] = b_var_val
+    env[c] = c_val
+    cost_expr = cost.evaluator().a() @ cost.variables() + cost.evaluator().b()
+    cost_val = cost_expr.Evaluate(env)
+    np.testing.assert_allclose(cost_val, cost_val_expected)
+
+
+def test_add_max_volume_linear_cost1():
+    # Test add_max_volume_linear_cost with an odd dimension n = 3.
+    n = 3
+    prog = solvers.MathematicalProgram()
+    S = prog.NewSymmetricContinuousVariables(n, "S")
+    prog.AddPositiveSemidefiniteConstraint(S)
+    b = prog.NewContinuousVariables(n, "b")
+    c = prog.NewContinuousVariables(1, "c")[0]
+
+    S_bar = np.diag(np.array([1, 2.0, 3.0]))
+    b_bar = np.array([4.0, 5.0, 6.0])
+    c_bar = -10.0
+
+    cost = mut.add_max_volume_linear_cost(prog, S, b, c, S_bar, b_bar, c_bar)
+    S_val = np.array([[6, 2, 3], [2, 10, 4], [3, 4, 10.0]])
+    b_val = np.array([1.0, 2.0, 3.0])
+    c_val = -5.0
+    check_add_max_volume_linear_cost(
+        cost, S, b, c, S_val, b_val, c_val, S_bar, b_bar, c_bar
+    )
+
+
+def test_add_max_volume_linear_cost2():
+    # Test add_max_volume_linear_cost with an even dimension n = 4.
+    n = 4
+    prog = solvers.MathematicalProgram()
+    S = prog.NewSymmetricContinuousVariables(n, "S")
+    prog.AddPositiveSemidefiniteConstraint(S)
+    b = prog.NewContinuousVariables(n, "b")
+    c = prog.NewContinuousVariables(1, "c")[0]
+
+    S_bar = np.diag(np.array([1, 2.0, 3.0, 4.0]))
+    S_bar[3, 0] = 0.1
+    S_bar[0, 3] = 0.1
+    b_bar = np.array([4.0, 5.0, 6.0, 7.0])
+    c_bar = -10.0
+
+    cost = mut.add_max_volume_linear_cost(prog, S, b, c, S_bar, b_bar, c_bar)
+    S_val = np.array([[6, 2, 3, 0], [2, 10, 4, 1], [3, 4, 10.0, 1], [1, 2, 3, 8]])
+    S_val = (S_val + S_val.T) / 2
+    b_val = np.array([1.0, 2.0, 3.0, 4.0])
+    c_val = -5.0
+    check_add_max_volume_linear_cost(
+        cost, S, b, c, S_val, b_val, c_val, S_bar, b_bar, c_bar
+    )