cbfpy 0.0.1__py3-none-any.whl → 0.0.4__py3-none-any.whl

cbfpy/cbfs/clf_cbf.py

@@ -1,7 +1,7 @@
 """
 # Control Lyapunov Function / Control Barrier Functions (CLF-CBFs)
 
-Whereas a CBF acts as a safety filter on top of a nominal controller, a CLF-CBF acts as a safe controller itself, 
+Whereas a CBF acts as a safety filter on top of a nominal controller, a CLF-CBF acts as a safe controller itself,
 based on a control objective defined by the CLF and a safety constraint defined by the CBF. Note that the CLF
 objective should be quadratic and positive-definite to fit in this QP framework.
 
@@ -43,17 +43,12 @@ import jax
 import jax.numpy as jnp
 from jax import Array
 from jax.typing import ArrayLike
+import numpy as np
 import qpax
 
 from cbfpy.config.clf_cbf_config import CLFCBFConfig
-from cbfpy.utils.jax_utils import conditional_jit
 from cbfpy.utils.general_utils import print_warning
 
-# Debugging flags to disable jit in specific sections of the code.
-# Note: If any higher-level jits exist, those must also be set to debug (disable jit)
-DEBUG_CONTROLLER = False
-DEBUG_QP_DATA = False
-
 
 @jax.tree_util.register_static
 class CLFCBF:
@@ -87,9 +82,9 @@ class CLFCBF:
         u_min: Optional[tuple],
         u_max: Optional[tuple],
         control_constrained: bool,
-        relax_cbf: bool,
-        cbf_relaxation_penalty: float,
+        relax_qp: bool,
         clf_relaxation_penalty: float,
+        constraint_relaxation_penalties: tuple,
         h_1: Callable[[ArrayLike], Array],
         h_2: Callable[[ArrayLike], Array],
         f: Callable[[ArrayLike], Array],
@@ -111,9 +106,9 @@ class CLFCBF:
         self.u_min = u_min
         self.u_max = u_max
         self.control_constrained = control_constrained
-        self.relax_cbf = relax_cbf
-        self.cbf_relaxation_penalty = cbf_relaxation_penalty
+        self.relax_qp = relax_qp
         self.clf_relaxation_penalty = clf_relaxation_penalty
+        self.constraint_relaxation_penalties = constraint_relaxation_penalties
         self.h_1 = h_1
         self.h_2 = h_2
         self.f = f
@@ -127,10 +122,6 @@ class CLFCBF:
         self.H = H
         self.F = F
         self.solver_tol = solver_tol
-        if relax_cbf:
-            self.qp_solver: Callable = jax.jit(qpax.solve_qp_elastic)
-        else:
-            self.qp_solver: Callable = jax.jit(qpax.solve_qp)
 
     @classmethod
     def from_config(cls, config: CLFCBFConfig) -> "CLFCBF":
@@ -151,9 +142,9 @@ class CLFCBF:
             config.u_min,
             config.u_max,
             config.control_constrained,
-            config.relax_cbf,
-            config.cbf_relaxation_penalty,
+            config.relax_qp,
             config.clf_relaxation_penalty,
+            config.constraint_relaxation_penalties,
             config.h_1,
             config.h_2,
             config.f,
@@ -168,19 +159,16 @@ class CLFCBF:
             config.F,
             config.solver_tol,
         )
-        instance._validate_instance(*config.init_args)
+        instance._validate_instance(*config.init_args, **config.init_kwargs)
         return instance
 
-    def _validate_instance(self, *h_args) -> None:
-        """Checks that the CLF-CBF is valid; warns the user if not
+    def _validate_instance(self, *args, **kwargs) -> None:
+        """Checks that the CLF-CBF is valid; warns the user if not"""
 
-        Args:
-            *h_args: Optional additional arguments for the barrier function.
-        """
-        test_z = jnp.ones(self.n)
+        test_z = np.ones(self.n)
         try:
-            test_lgh = self.Lgh(test_z, *h_args)
-            if jnp.allclose(test_lgh, 0):
+            test_lgh = self.Lgh(test_z, *args, **kwargs)
+            if np.allclose(test_lgh, 0):
                 print_warning(
                     "Lgh is zero. Consider increasing the relative degree or modifying the barrier function."
                 )
@@ -189,37 +177,36 @@ class CLFCBF:
                 "Cannot test Lgh; missing additional arguments.\n"
                 + "Please provide an initial seed for these args in the config's init_args input"
             )
-        test_lgv = self.LgV(test_z)
-        if jnp.allclose(test_lgv, 0):
+        test_lgv = self.LgV(test_z, test_z, *args, **kwargs)
+        if np.allclose(test_lgv, 0):
             print_warning(
                 "LgV is zero. Consider increasing the relative degree or modifying the Lyapunov function."
             )
 
-    @conditional_jit(not DEBUG_CONTROLLER)
-    def controller(self, z: Array, z_des: Array, *h_args) -> Array:
+    @jax.jit
+    def controller(self, z: Array, z_des: Array, *args, **kwargs) -> Array:
         """Compute the CLF-CBF optimal control input, optimizing for the CLF objective while
         satisfying the CBF safety constraint.
 
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: Safe control input, shape (m,)
         """
-        P, q, A, b, G, h = self.qp_data(z, z_des, *h_args)
-        if self.relax_cbf:
-            x_qp, t_qp, s1_qp, s2_qp, z1_qp, z2_qp, converged, iters = self.qp_solver(
+        P, q, A, b, G, h = self.qp_data(z, z_des, *args, **kwargs)
+        if self.relax_qp:
+            x_qp = qpax.solve_qp_elastic_primal(
                 P,
                 q,
                 G,
                 h,
-                self.cbf_relaxation_penalty,
+                penalty=jnp.asarray(self.constraint_relaxation_penalties),
                 solver_tol=self.solver_tol,
             )
         else:
-            x_qp, s_qp, z_qp, y_qp, converged, iters = self.qp_solver(
+            x_qp, s_qp, z_qp, y_qp, converged, iters = qpax.solve_qp(
                 P,
                 q,
                 A,
@@ -228,18 +215,13 @@ class CLFCBF:
                 h,
                 solver_tol=self.solver_tol,
             )
-        if DEBUG_CONTROLLER:
-            print(
-                f"{'Converged' if converged else 'Did not converge'}. Iterations: {iters}"
-            )
         return x_qp[: self.m]
 
-    def h(self, z: ArrayLike, *h_args) -> Array:
+    def h(self, z: ArrayLike, *args, **kwargs) -> Array:
         """Barrier function(s)
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: Barrier function evaluation, shape (num_barr,)
@@ -247,16 +229,16 @@ class CLFCBF:
 
         # Take any relative-degree-2 barrier functions and convert them to relative-degree-1
         def _h_2(state):
-            return self.h_2(state, *h_args)
+            return self.h_2(state, *args, **kwargs)
 
-        h_2, dh_2_dt = jax.jvp(_h_2, (z,), (self.f(z),))
-        h_2_as_rd1 = dh_2_dt + self.alpha_2(h_2)
+        h_2, dh_2_dt = jax.jvp(_h_2, (z,), (self.f(z, *args, **kwargs),))
+        h_2_as_rd1 = dh_2_dt + self.alpha_2(h_2, *args, **kwargs)
 
         # Merge the relative-degree-1 and relative-degree-2 barrier functions
-        return jnp.concatenate([self.h_1(z, *h_args), h_2_as_rd1])
+        return jnp.concatenate([self.h_1(z, *args, **kwargs), h_2_as_rd1])
 
     def h_and_Lfh(  # pylint: disable=invalid-name
-        self, z: ArrayLike, *h_args
+        self, z: ArrayLike, *args, **kwargs
     ) -> Tuple[Array, Array]:
         """Lie derivative of the barrier function(s) wrt the autonomous dynamics `f(z)`
 
@@ -264,7 +246,6 @@ class CLFCBF:
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             h (Array): Barrier function evaluation, shape (num_barr,)
@@ -274,16 +255,17 @@ class CLFCBF:
         # with the bonus benefit of also evaluating the barrier function
 
         def _h(state):
-            return self.h(state, *h_args)
+            return self.h(state, *args, **kwargs)
 
-        return jax.jvp(_h, (z,), (self.f(z),))
+        return jax.jvp(_h, (z,), (self.f(z, *args, **kwargs),))
 
-    def Lgh(self, z: ArrayLike, *h_args) -> Array:  # pylint: disable=invalid-name
+    def Lgh(
+        self, z: ArrayLike, *args, **kwargs
+    ) -> Array:  # pylint: disable=invalid-name
         """Lie derivative of the barrier function(s) wrt the control dynamics `g(z)u`
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: Lgh, shape (num_barr, m)
@@ -291,98 +273,112 @@ class CLFCBF:
         # Note: the below code is just a more efficient way of stating `Lgh = jax.jacobian(self.h)(z) @ self.g(z)`
 
         def _h(state):
-            return self.h(state, *h_args)
+            return self.h(state, *args, **kwargs)
 
         def _jvp(g_column):
             return jax.jvp(_h, (z,), (g_column,))[1]
 
-        return jax.vmap(_jvp, in_axes=1, out_axes=1)(self.g(z))
+        return jax.vmap(_jvp, in_axes=1, out_axes=1)(self.g(z, *args, **kwargs))
 
     ## CLF functions ##
 
-    def V(self, z: ArrayLike) -> Array:
+    def V(self, z: ArrayLike, z_des: ArrayLike, *args, **kwargs) -> Array:
         """Control Lyapunov Function(s)
 
         Args:
             z (ArrayLike): State, shape (n,)
+            z_des (ArrayLike): Desired state, shape (n,)
 
         Returns:
             Array: CLF evaluation, shape (num_clf,)
         """
+
+        def _V_2(state):
+            return self.V_2(state, z_des, *args, **kwargs)
+
         # Take any relative-degree-2 CLFs and convert them to relative-degree-1
         # NOTE: If adding args to the CLF, create a wrapper func like with the barrier function
-        V_2, dV_2_dt = jax.jvp(self.V_2, (z,), (self.f(z),))
-        V2_rd1 = dV_2_dt + self.gamma_2(V_2)
+        V_2, dV_2_dt = jax.jvp(_V_2, (z,), (self.f(z, *args, **kwargs),))
+        V2_rd1 = dV_2_dt + self.gamma_2(V_2, *args, **kwargs)
 
         # Merge the relative-degree-1 and relative-degree-2 CLFs
-        return jnp.concatenate([self.V_1(z), V2_rd1])
+        return jnp.concatenate([self.V_1(z, z_des, *args, **kwargs), V2_rd1])
 
-    def V_and_LfV(self, z: ArrayLike) -> Tuple[Array, Array]:
+    def V_and_LfV(
+        self, z: ArrayLike, z_des: ArrayLike, *args, **kwargs
+    ) -> Tuple[Array, Array]:
         """Lie derivative of the CLF wrt the autonomous dynamics `f(z)`
 
         The evaluation of the CLF is also returned "for free", a byproduct of the jacobian-vector-product
 
         Args:
             z (ArrayLike): State, shape (n,)
+            z_des (ArrayLike): Desired state, shape (n,)
 
         Returns:
             V (Array): CLF evaluation, shape (1,)
             LfV (Array): Lie derivative of `V` w.r.t. `f`, shape (1,)
         """
-        return jax.jvp(self.V, (z,), (self.f(z),))
 
-    def LgV(self, z: ArrayLike) -> Array:
+        def _V(state):
+            return self.V(state, z_des, *args, **kwargs)
+
+        return jax.jvp(_V, (z,), (self.f(z, *args, **kwargs),))
+
+    def LgV(self, z: ArrayLike, z_des: ArrayLike, *args, **kwargs) -> Array:
         """Lie derivative of the CLF wrt the control dynamics `g(z)u`
 
         Args:
             z (ArrayLike): State, shape (n,)
+            z_des (ArrayLike): Desired state, shape (n,)
 
         Returns:
             Array: LgV, shape (m,)
         """
 
+        def _V(state):
+            return self.V(state, z_des, *args, **kwargs)
+
         def _jvp(g_column):
-            return jax.jvp(self.V, (z,), (g_column,))[1]
+            return jax.jvp(_V, (z,), (g_column,))[1]
 
-        return jax.vmap(_jvp, in_axes=1, out_axes=1)(self.g(z))
+        return jax.vmap(_jvp, in_axes=1, out_axes=1)(self.g(z, *args, **kwargs))
 
     ## QP Matrices ##
 
     def P_qp(  # pylint: disable=invalid-name
-        self, z: Array, z_des: Array, *h_args
+        self, z: Array, z_des: Array, *args, **kwargs
     ) -> Array:
         """Quadratic term in the QP objective (`minimize 0.5 * x^T P x + q^T x`)
 
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: P matrix, shape (m, m)
         """
         return jnp.block(
             [
-                [self.H(z), jnp.zeros((self.m, 1))],
+                [self.H(z, *args, **kwargs), jnp.zeros((self.m, 1))],
                 [jnp.zeros((1, self.m)), jnp.atleast_1d(self.clf_relaxation_penalty)],
             ]
         )
 
-    def q_qp(self, z: Array, z_des: Array, *h_args) -> Array:
+    def q_qp(self, z: Array, z_des: Array, *args, **kwargs) -> Array:
         """Linear term in the QP objective (`minimize 0.5 * x^T P x + q^T x`)
 
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: Q vector, shape (m,)
         """
-        return jnp.concatenate([self.F(z), jnp.array([0.0])])
+        return jnp.concatenate([self.F(z, *args, **kwargs), jnp.array([0.0])])
 
     def G_qp(  # pylint: disable=invalid-name
-        self, z: Array, z_des: Array, *h_args
+        self, z: Array, z_des: Array, *args, **kwargs
     ) -> Array:
         """Inequality constraint matrix for the QP (`Gx <= h`)
 
@@ -394,15 +390,17 @@ class CLFCBF:
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: G matrix, shape (num_constraints, m)
         """
         G = jnp.block(
             [
-                [self.LgV(z), -1.0 * jnp.ones((self.num_clf, 1))],
-                [-self.Lgh(z, *h_args), jnp.zeros((self.num_cbf, 1))],
+                [
+                    self.LgV(z, z_des, *args, **kwargs),
+                    -1.0 * jnp.ones((self.num_clf, 1)),
+                ],
+                [-self.Lgh(z, *args, **kwargs), jnp.zeros((self.num_cbf, 1))],
             ]
         )
         if self.control_constrained:
@@ -416,7 +414,7 @@ class CLFCBF:
         else:
             return G
 
-    def h_qp(self, z: Array, z_des: Array, *h_args) -> Array:
+    def h_qp(self, z: Array, z_des: Array, *args, **kwargs) -> Array:
         """Upper bound on constraints for the QP (`Gx <= h`)
 
         Note:
@@ -427,17 +425,16 @@ class CLFCBF:
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: h vector, shape (num_constraints,)
         """
-        hz, lfh = self.h_and_Lfh(z, *h_args)
-        vz, lfv = self.V_and_LfV(z)
+        hz, lfh = self.h_and_Lfh(z, *args, **kwargs)
+        vz, lfv = self.V_and_LfV(z, z_des, *args, **kwargs)
         h = jnp.concatenate(
             [
-                -lfv - self.gamma(vz),
-                self.alpha(hz) + lfh,
+                -lfv - self.gamma(vz, *args, **kwargs),
+                self.alpha(hz, *args, **kwargs) + lfh,
             ]
         )
         if self.control_constrained:
@@ -447,9 +444,8 @@ class CLFCBF:
         else:
             return h
 
-    @conditional_jit(not DEBUG_QP_DATA)
     def qp_data(
-        self, z: Array, z_des: Array, *h_args
+        self, z: Array, z_des: Array, *args, **kwargs
     ) -> Tuple[Array, Array, Array, Array, Array, Array]:
         """Constructs the QP matrices based on the current state and desired control
 
@@ -470,7 +466,6 @@ class CLFCBF:
         Args:
             z (Array): State, shape (n,)
             z_des (Array): Desired state, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             P (Array): Quadratic term in the QP objective, shape (m + 1, m + 1)
@@ -481,10 +476,10 @@ class CLFCBF:
             h (Array): Upper bound on constraints, shape (num_constraints,)
         """
         return (
-            self.P_qp(z, z_des, *h_args),
-            self.q_qp(z, z_des, *h_args),
+            self.P_qp(z, z_des, *args, **kwargs),
+            self.q_qp(z, z_des, *args, **kwargs),
             jnp.zeros((0, self.m + 1)),  # Equality matrix (not used for CLF-CBF)
             jnp.zeros(0),  # Equality vector (not used for CLF-CBF)
-            self.G_qp(z, z_des, *h_args),
-            self.h_qp(z, z_des, *h_args),
+            self.G_qp(z, z_des, *args, **kwargs),
+            self.h_qp(z, z_des, *args, **kwargs),
         )