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

cbfpy/cbfs/cbf.py

@@ -30,17 +30,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.cbf_config import CBFConfig
-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_SAFETY_FILTER = False
-DEBUG_QP_DATA = False
-
 
 @jax.tree_util.register_static
 class CBF:
@@ -73,8 +68,8 @@ class CBF:
         u_min: Optional[tuple],
         u_max: Optional[tuple],
         control_constrained: bool,
-        relax_cbf: bool,
-        cbf_relaxation_penalty: float,
+        relax_qp: bool,
+        constraint_relaxation_penalties: tuple,
         h_1: Callable[[ArrayLike], Array],
         h_2: Callable[[ArrayLike], Array],
         f: Callable[[ArrayLike], Array],
@@ -91,8 +86,8 @@ class CBF:
         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.constraint_relaxation_penalties = constraint_relaxation_penalties
         self.h_1 = h_1
         self.h_2 = h_2
         self.f = f
@@ -102,10 +97,6 @@ class CBF:
         self.P_config = P
         self.q_config = q
         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: CBFConfig) -> "CBF":
@@ -124,8 +115,8 @@ class CBF:
             config.u_min,
             config.u_max,
             config.control_constrained,
-            config.relax_cbf,
-            config.cbf_relaxation_penalty,
+            config.relax_qp,
+            config.constraint_relaxation_penalties,
             config.h_1,
             config.h_2,
             config.f,
@@ -136,19 +127,16 @@ class CBF:
             config.q,
             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 CBF is valid; warns the user if not
+    def _validate_instance(self, *args, **kwargs) -> None:
+        """Checks that the CBF is valid; warns the user if not"""
 
-        Args:
-            *h_args: Optional additional arguments for the barrier function.
-        """
         try:
             # TODO: Decide if this should be checked on a row-by-row basis or via the full matrix
-            test_lgh = self.Lgh(jnp.ones(self.n), *h_args)
-            if jnp.allclose(test_lgh, 0):
+            test_lgh = self.Lgh(np.ones(self.n), *args, **kwargs)
+            if np.allclose(test_lgh, 0):
                 print_warning(
                     "Lgh is zero. Consider increasing the relative degree or modifying the barrier function."
                 )
@@ -158,30 +146,29 @@ class CBF:
                 + "Please provide an initial seed for these args in the config's init_args input"
             )
 
-    @conditional_jit(not DEBUG_SAFETY_FILTER)
-    def safety_filter(self, z: Array, u_des: Array, *h_args) -> Array:
+    @jax.jit
+    def safety_filter(self, z: Array, u_des: Array, *args, **kwargs) -> Array:
         """Apply the CBF safety filter to a nominal control
 
         Args:
             z (Array): State, shape (n,)
             u_des (Array): Desired control input, shape (m,)
-            *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, u_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, u_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,
@@ -190,18 +177,13 @@ class CBF:
                 h,
                 solver_tol=self.solver_tol,
             )
-        if DEBUG_SAFETY_FILTER:
-            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,)
@@ -209,16 +191,16 @@ class CBF:
 
         # 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)`
 
@@ -226,7 +208,6 @@ class CBF:
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             h (Array): Barrier function evaluation, shape (num_barr,)
@@ -236,16 +217,17 @@ class CBF:
         # 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)
@@ -253,47 +235,45 @@ class CBF:
         # 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))
 
     ## QP Matrices ##
 
     def P_qp(  # pylint: disable=invalid-name
-        self, z: Array, u_des: Array, *h_args
+        self, z: Array, u_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,)
             u_des (Array): Desired control input, shape (m,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: P matrix, shape (m, m)
         """
         # This is user-modifiable in the config, but defaults to 2 * I for the standard min-norm CBF objective
-        return self.P_config(z, u_des, *h_args)
+        return self.P_config(z, u_des, *args, **kwargs)
 
-    def q_qp(self, z: Array, u_des: Array, *h_args) -> Array:
+    def q_qp(self, z: Array, u_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,)
             u_des (Array): Desired control input, shape (m,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: q vector, shape (m,)
         """
         # This is user-modifiable in the config, but defaults to -2 * u_des for the standard min-norm CBF objective
-        return self.q_config(z, u_des, *h_args)
+        return self.q_config(z, u_des, *args, **kwargs)
 
     def G_qp(  # pylint: disable=invalid-name
-        self, z: Array, u_des: Array, *h_args
+        self, z: Array, u_des: Array, *args, **kwargs
     ) -> Array:
         """Inequality constraint matrix for the QP (`Gx <= h`)
 
@@ -305,18 +285,17 @@ class CBF:
         Args:
             z (Array): State, shape (n,)
             u_des (Array): Desired control input, shape (m,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             Array: G matrix, shape (num_constraints, m)
         """
-        G = -self.Lgh(z, *h_args)
+        G = -self.Lgh(z, *args, **kwargs)
         if self.control_constrained:
             return jnp.block([[G], [jnp.eye(self.m)], [-jnp.eye(self.m)]])
         else:
             return G
 
-    def h_qp(self, z: Array, u_des: Array, *h_args) -> Array:
+    def h_qp(self, z: Array, u_des: Array, *args, **kwargs) -> Array:
         """Upper bound on constraints for the QP (`Gx <= h`)
 
         Note:
@@ -327,13 +306,12 @@ class CBF:
         Args:
             z (Array): State, shape (n,)
             u_des (Array): Desired control input, shape (m,)
-            *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)
-        h = self.alpha(hz) + lfh
+        hz, lfh = self.h_and_Lfh(z, *args, **kwargs)
+        h = self.alpha(hz, *args, **kwargs) + lfh
         if self.control_constrained:
             return jnp.concatenate(
                 [h, jnp.asarray(self.u_max), -jnp.asarray(self.u_min)]
@@ -341,9 +319,8 @@ class CBF:
         else:
             return h
 
-    @conditional_jit(not DEBUG_QP_DATA)
     def qp_data(
-        self, z: Array, u_des: Array, *h_args
+        self, z: Array, u_des: Array, *args, **kwargs
     ) -> Tuple[Array, Array, Array, Array, Array, Array]:
         """Constructs the QP matrices based on the current state and desired control
 
@@ -364,7 +341,6 @@ class CBF:
         Args:
             z (Array): State, shape (n,)
             u_des (Array): Desired control input, shape (m,)
-            *h_args: Optional additional arguments for the barrier function.
 
         Returns:
             P (Array): Quadratic term in the QP objective, shape (m, m)
@@ -375,10 +351,10 @@ class CBF:
             h (Array): Upper bound on constraints, shape (num_constraints,)
         """
         return (
-            self.P_qp(z, u_des, *h_args),
-            self.q_qp(z, u_des, *h_args),
+            self.P_qp(z, u_des, *args, **kwargs),
+            self.q_qp(z, u_des, *args, **kwargs),
             jnp.zeros((0, self.m)),  # Equality matrix (not used for CBF)
             jnp.zeros(0),  # Equality vector (not used for CBF)
-            self.G_qp(z, u_des, *h_args),
-            self.h_qp(z, u_des, *h_args),
+            self.G_qp(z, u_des, *args, **kwargs),
+            self.h_qp(z, u_des, *args, **kwargs),
         )