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

cbfpy/config/cbf_config.py

@@ -3,13 +3,13 @@
 
 ## Defining the problem:
 
-CBFs have two primary implementation requirements: the dynamics functions, and the barrier function(s). 
+CBFs have two primary implementation requirements: the dynamics functions, and the barrier function(s).
 These can be specified through the `f`, `g`, and `h` methods, respectively. Note that the main requirements
 for these functions are that (1) the dynamics are control-affine, and (2) the barrier function(s) are "zeroing"
 barriers, as opposed to "reciprocal" barriers. A zeroing barrier is one which is positive in the interior of the
-safe set, and zero on the boundary. 
+safe set, and zero on the boundary.
 
-Depending on the relative degree of your barrier function(s), you should implement the `h_1` method 
+Depending on the relative degree of your barrier function(s), you should implement the `h_1` method
 (for a relative-degree-1 barrier), and/or the `h_2` method (for a relative-degree-2 barrier).
 
 ## Tuning the CBF:
@@ -30,7 +30,7 @@ does require that P is positive semi-definite.
 Depending on the construction of the barrier functions and if control limits are provided, the CBF QP may not always be
 feasible. If allowing for relaxation in the CBFConfig, a slack variable will be introduced to ensure that the
 problem is always feasible, with a high penalty on any infeasibility. This is generally useful for controller
-robustness, but means that safety is not guaranteed. 
+robustness, but means that safety is not guaranteed.
 
 If strict enforcement of the CBF is desired, your higest-level controller should handle the case where the QP
 is infeasible.
@@ -67,13 +67,17 @@ class CBFConfig(ABC):
         m (int): Control dimension
         u_min (ArrayLike, optional): Minimum control input, shape (m,). Defaults to None (Unconstrained).
         u_max (ArrayLike, optional): Maximum control input, shape (m,). Defaults to None (Unconstrained).
-        relax_cbf (bool, optional): Whether to allow for relaxation in the CBF QP. Defaults to True.
-        cbf_relaxation_penalty (float, optional): Penalty on the slack variable in the relaxed CBF QP. Defaults to 1e3.
-            Note: only applies if relax_cbf is True.
+        relax_qp (bool, optional): Whether to allow for relaxation in the CBF QP. Defaults to True.
+            Note: this is required for differentiability through the QP.
+        cbf_relaxation_penalty (float, optional): Penalty on the CBF slack variables in the relaxed QP.
+            Defaults to 1e3. Note: only applies if relax_qp is True.
+        control_relaxation_penalty (float, optional): Penalty on the control constraint slack variables in the
+            relaxed QP. Defaults to 1e5. Note: only applies if relax_qp is True.
         solver_tol (float, optional): Tolerance for the QP solver. Defaults to 1e-3.
-        init_args (tuple, optional): If your barrier function relies on additional arguments other than just the state,
-            include an initial seed for these arguments here. This is to help test the output of the barrier function.
-            Defaults to ().
+        init_args (tuple, optional): If your barriers or dynamics rely on additional (non-differentiable, static shape)
+            args other than just the state, include an initial seed for these args here. Defaults to None.
+        init_kwargs (dict, optional): If your barriers or dynamics rely on additional (non-differentiable, static shape)
+            kwargs other than just the state, include an initial seed for these kwargs here. Defaults to None.
     """
 
     def __init__(
@@ -82,10 +86,12 @@ class CBFConfig(ABC):
         m: int,
         u_min: Optional[ArrayLike] = None,
         u_max: Optional[ArrayLike] = None,
-        relax_cbf: bool = True,
+        relax_qp: bool = True,
         cbf_relaxation_penalty: float = 1e3,
+        control_relaxation_penalty: float = 1e5,
         solver_tol: float = 1e-3,
-        init_args: tuple = (),
+        init_args: Optional[tuple] = None,
+        init_kwargs: Optional[dict] = None,
     ):
         if not (isinstance(n, int) and n > 0):
             raise ValueError(f"n must be a positive integer. Got: {n}")
@@ -95,9 +101,9 @@ class CBFConfig(ABC):
             raise ValueError(f"m must be a positive integer. Got: {m}")
         self.m = m
 
-        if not isinstance(relax_cbf, bool):
-            raise ValueError(f"relax_cbf must be a boolean. Got: {relax_cbf}")
-        self.relax_cbf = relax_cbf
+        if not isinstance(relax_qp, bool):
+            raise ValueError(f"relax_qp must be a boolean. Got: {relax_qp}")
+        self.relax_qp = relax_qp
 
         if not (
             isinstance(cbf_relaxation_penalty, (int, float))
@@ -112,10 +118,33 @@ class CBFConfig(ABC):
             raise ValueError(f"solver_tol must be a positive value. Got: {solver_tol}")
         self.solver_tol = float(solver_tol)
 
-        if not isinstance(init_args, tuple):
+        if self.solver_tol > 1e-2:
+            print(
+                f"WARNING: solver tolerance is quite high ({self.solver_tol}). "
+                + " Solution will likely be poor."
+            )
+
+        if init_args is None:
+            init_args = ()
+        elif not isinstance(init_args, tuple):
             raise ValueError(f"init_args must be a tuple. Got: {init_args}")
         self.init_args = init_args
 
+        if init_kwargs is None:
+            init_kwargs = {}
+        elif not isinstance(init_kwargs, dict):
+            raise ValueError(f"init_kwargs must be a dict. Got: {init_kwargs}")
+        self.init_kwargs = init_kwargs
+
+        if not (
+            isinstance(control_relaxation_penalty, (int, float))
+            and control_relaxation_penalty >= 0
+        ):
+            raise ValueError(
+                f"control_relaxation_penalty must be a non-negative value. Got: {control_relaxation_penalty}"
+            )
+        self.control_relaxation_penalty = float(control_relaxation_penalty)
+
         # Control limits require a bit of extra handling. They can be both None if unconstrained,
         # but we should not have one limit as None and the other as some value
         u_min = np.asarray(u_min, dtype=float).flatten() if u_min is not None else None
@@ -137,11 +166,17 @@ class CBFConfig(ABC):
         self.u_min = u_min
         self.u_max = u_max
 
+        if (
+            self.control_constrained
+            and self.control_relaxation_penalty < self.cbf_relaxation_penalty
+        ):
+            print("WARNING: Control constraints have a lower penalty than the CBFs.")
+
         # Test if the methods are provided and verify their output dimension
-        z_test = jnp.ones(self.n)
-        u_test = jnp.ones(self.m)
-        f_test = self.f(z_test)
-        g_test = self.g(z_test)
+        z_test = np.ones(self.n)
+        u_test = np.ones(self.m)
+        f_test = self.f(z_test, *self.init_args, **self.init_kwargs)
+        g_test = self.g(z_test, *self.init_args, **self.init_kwargs)
         if f_test.shape != (self.n,):
             raise ValueError(
                 f"Invalid shape for f(z). Got {f_test.shape}, expected ({self.n},)"
@@ -151,13 +186,10 @@ class CBFConfig(ABC):
                 f"Invalid shape for g(z). Got {g_test.shape}, expected ({self.n}, {self.m})"
             )
         try:
-            h1_test = self.h_1(z_test, *self.init_args)
-            h2_test = self.h_2(z_test, *self.init_args)
+            h1_test = self.h_1(z_test, *self.init_args, **self.init_kwargs)
+            h2_test = self.h_2(z_test, *self.init_args, **self.init_kwargs)
         except TypeError as e:
-            raise ValueError(
-                "Cannot test the barrier function; likely missing additional arguments.\n"
-                + "Please provide an initial seed for these args in the config's init_args input"
-            ) from e
+            raise ValueError("Cannot test the barrier function") from e
         if h1_test.ndim != 1 or h2_test.ndim != 1:
             raise ValueError("Barrier function(s) must be 1D arrays")
         self.num_rd1_cbf = h1_test.shape[0]
@@ -168,9 +200,9 @@ class CBFConfig(ABC):
                 "No barrier functions provided."
                 + "\nYou can implement this via the h_1 and/or h_2 methods in your config class"
             )
-        h_test = jnp.concatenate([h1_test, h2_test])
-        alpha_test = self.alpha(h_test)
-        alpha_2_test = self.alpha_2(h2_test)
+        h_test = np.concatenate([h1_test, h2_test])
+        alpha_test = self.alpha(h_test, *self.init_args, **self.init_kwargs)
+        alpha_2_test = self.alpha_2(h2_test, *self.init_args, **self.init_kwargs)
         if alpha_test.shape != (self.num_cbf,):
             raise ValueError(
                 f"Invalid shape for alpha(h(z)): {alpha_test.shape}. Expected ({self.num_cbf},)"
@@ -181,15 +213,16 @@ class CBFConfig(ABC):
                 f"Invalid shape for alpha_2(h_2(z)): {alpha_2_test.shape}. Expected ({self.num_rd2_cbf},)"
                 + "\nCheck that the output of the alpha_2() function matches the number of RD2 CBFs"
             )
-        self._check_class_kappa(self.alpha, self.num_cbf)
-        self._check_class_kappa(self.alpha_2, self.num_rd2_cbf)
+        self._check_class_kappa(
+            self.alpha, self.num_cbf, *self.init_args, **self.init_kwargs
+        )
+        self._check_class_kappa(
+            self.alpha_2, self.num_rd2_cbf, *self.init_args, **self.init_kwargs
+        )
         try:
-            P_test = self.P(z_test, u_test, *self.init_args)
+            P_test = self.P(z_test, u_test, *self.init_args, **self.init_kwargs)
         except TypeError as e:
-            raise ValueError(
-                "Cannot test the P matrix; likely missing additional arguments.\n"
-                + "Please provide an initial seed for these args in the config's init_args input"
-            ) from e
+            raise ValueError("Cannot test the P matrix") from e
         if P_test.shape != (self.m, self.m):
             raise ValueError(
                 f"Invalid shape for P(z). Got {P_test.shape}, expected ({self.m}, {self.m})"
@@ -197,10 +230,29 @@ class CBFConfig(ABC):
         if not self._is_symmetric_psd(P_test):
             raise ValueError("P matrix must be symmetric positive semi-definite")
 
+        # Handle QP relaxation penalties, if relaxation is enabled
+        num_qp_constraints = (
+            self.num_cbf if not self.control_constrained else self.num_cbf + 2 * self.m
+        )
+        if self.control_constrained:
+            self.constraint_relaxation_penalties = tuple(
+                np.concatenate(
+                    [
+                        self.cbf_relaxation_penalty * np.ones(self.num_cbf),
+                        self.control_relaxation_penalty * np.ones(2 * self.m),
+                    ]
+                )
+            )
+        else:
+            self.constraint_relaxation_penalties = tuple(
+                self.cbf_relaxation_penalty * np.ones(self.num_cbf)
+            )
+        assert len(self.constraint_relaxation_penalties) == num_qp_constraints
+
     ## Control Affine Dynamics ##
 
     @abstractmethod
-    def f(self, z: ArrayLike) -> Array:
+    def f(self, z: ArrayLike, *args, **kwargs) -> Array:
         """The uncontrolled dynamics function. Possibly nonlinear, and locally Lipschitz
 
         i.e. the function f, such that z_dot = f(z) + g(z) u
@@ -214,7 +266,7 @@ class CBFConfig(ABC):
         pass
 
     @abstractmethod
-    def g(self, z: ArrayLike) -> Array:
+    def g(self, z: ArrayLike, *args, **kwargs) -> Array:
         """The control affine dynamics function. Locally Lipschitz.
 
         i.e. the function g, such that z_dot = f(z) + g(z) u
@@ -229,7 +281,7 @@ class CBFConfig(ABC):
 
     ## Barriers ##
 
-    def h_1(self, z: ArrayLike, *h_args) -> Array:
+    def h_1(self, z: ArrayLike, *args, **kwargs) -> Array:
         """Relative-degree-1 barrier function(s).
 
         A (zeroing) CBF is a continuously-differentiable function h, such that for any state z in the interior of
@@ -244,15 +296,13 @@ class CBFConfig(ABC):
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function. Note: If using additional args with your
-                barrier, these must be a static shape/type, or else this will trigger a recompilation in Jax.
 
         Returns:
             Array: Barrier function(s), shape (num_rd1_barr,)
         """
         return jnp.array([])
 
-    def h_2(self, z: ArrayLike, *h_args) -> Array:
+    def h_2(self, z: ArrayLike, *args, **kwargs) -> Array:
         """Relative-degree-2 (high-order) barrier function(s).
 
         A (zeroing) CBF is a continuously-differentiable function h, such that for any state z in the interior of
@@ -268,8 +318,6 @@ class CBFConfig(ABC):
 
         Args:
             z (ArrayLike): State, shape (n,)
-            *h_args: Optional additional arguments for the barrier function. Note: If using additional args with your
-                barrier, these must be a static shape/type, or else this will trigger a recompilation in Jax.
 
         Returns:
             Array: Barrier function(s), shape (num_rd2_barr,)
@@ -278,7 +326,7 @@ class CBFConfig(ABC):
 
     ## Additional tuning functions ##
 
-    def alpha(self, h: ArrayLike) -> Array:
+    def alpha(self, h: ArrayLike, *args, **kwargs) -> Array:
         """A class Kappa function, dictating the "gain" of the barrier function(s)
 
         For reference, a class Kappa function is a monotonically increasing function which passes through the origin.
@@ -294,7 +342,7 @@ class CBFConfig(ABC):
         """
         return h
 
-    def alpha_2(self, h_2: ArrayLike) -> Array:
+    def alpha_2(self, h_2: ArrayLike, *args, **kwargs) -> Array:
         """A second class Kappa function which dictactes the "gain" associated with the relative-degree-2
         barrier functions
 
@@ -313,7 +361,7 @@ class CBFConfig(ABC):
 
     # Objective function tuning
 
-    def P(self, z: Array, u_des: Array, *h_args) -> Array:
+    def P(self, z: Array, u_des: Array, *args, **kwargs) -> Array:
         """Quadratic term in the CBF QP objective (minimize 0.5 * x^T P x + q^T x)
 
         This defaults to 2 * I, which is the value of P when minimizing the standard CBF objective,
@@ -324,14 +372,13 @@ class CBFConfig(ABC):
         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)
         """
         return 2 * jnp.eye(self.m)
 
-    def q(self, z: Array, u_des: Array, *h_args) -> Array:
+    def q(self, z: Array, u_des: Array, *args, **kwargs) -> Array:
         """Linear term in the CBF QP objective (minimize 0.5 * x^T P x + q^T x)
 
         This defaults to -2 * u_des, which is the value of q when minimizing the standard CBF objective,
@@ -342,7 +389,6 @@ class CBFConfig(ABC):
         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,)
@@ -352,7 +398,7 @@ class CBFConfig(ABC):
     ## Helper functions ##
 
     def _check_class_kappa(
-        self, func: Callable[[ArrayLike], ArrayLike], dim: int
+        self, func: Callable[[ArrayLike], ArrayLike], dim: int, *args, **kwargs
     ) -> None:
         """Checks that the provided function is in class Kappa
 
@@ -361,16 +407,20 @@ class CBFConfig(ABC):
             dim (int): Expected dimension of the output
         """
         assert isinstance(func, Callable)
+
+        def func_wrapper(x):
+            return func(x, *args, **kwargs)
+
         try:
             # Check that func(0) == 0
-            assert jnp.allclose(func(jnp.zeros(dim)), 0.0)
+            assert np.allclose(func_wrapper(np.zeros(dim)), 0.0)
             # Check that func is monotonically increasing
             n_test = 100
-            test_points = jnp.repeat(
-                jnp.linspace(-1e6, 1e6, n_test).reshape(n_test, 1), dim, axis=1
+            test_points = np.repeat(
+                np.linspace(-1e6, 1e6, n_test).reshape(n_test, 1), dim, axis=1
             )
-            a = jax.vmap(func, in_axes=0)(test_points)
-            assert jnp.all(a[:-1, :] < a[1:, :])
+            a = jax.vmap(func_wrapper, in_axes=0)(test_points)
+            assert np.all(a[:-1, :] < a[1:, :])
         except AssertionError as e:
             raise ValueError(
                 f"{func.__name__} does not appear to be a class Kappa function"

cbfpy/config/clf_cbf_config.py

@@ -3,11 +3,11 @@
 
 ## Defining the problem:
 
-As with the CBF, we require implementation of the dynamics functions `f` and `g`, as well as the barrier function(s) 
+As with the CBF, we require implementation of the dynamics functions `f` and `g`, as well as the barrier function(s)
 `h`. Now, with the CLF-CBF, we require the definition of the Control Lyapunov Function (CLF) `V`. This CLF must be a
-positive definite function of the state. 
+positive definite function of the state.
 
-Depending on the relative degree of your barrier function(s), you should implement the `h_1` method 
+Depending on the relative degree of your barrier function(s), you should implement the `h_1` method
 (for a relative-degree-1 barrier), and/or the `h_2` method (for a relative-degree-2 barrier).
 
 Likewise, for the CLF, you should implement the `V_1` method (for a relative-degree-1 CLF), and/or the `V_2` method
@@ -24,7 +24,7 @@ CLF objective. These can be used to adjust the weightings between inputs, for in
 
 ## Relaxation:
 
-If the CBF constraints are not necessarily globally feasible, you can enable further relaxation in the CLFCBFConfig. 
+If the CBF constraints are not necessarily globally feasible, you can enable further relaxation in the CLFCBFConfig.
 However, since the CLF constraint was already relaxed with respect to the CBF constraint, this means that tuning the
 relaxation parameters is critical. In general, the penalty on the CBF relaxation should be much higher than the penalty
 on the CLF relaxation.
@@ -35,6 +35,7 @@ is infeasible.
 
 from typing import Optional
 
+import numpy as np
 import jax.numpy as jnp
 from jax import Array
 from jax.typing import ArrayLike
@@ -66,14 +67,18 @@ class CLFCBFConfig(CBFConfig):
         m (int): Control dimension
         u_min (ArrayLike, optional): Minimum control input, shape (m,). Defaults to None (Unconstrained).
         u_max (ArrayLike, optional): Maximum control input, shape (m,). Defaults to None (Unconstrained).
-        relax_cbf (bool, optional): Whether to allow for relaxation in the CBF QP. Defaults to True.
-        cbf_relaxation_penalty (float, optional): Penalty on the slack variable in the relaxed CBF QP. Defaults to 1e4.
-            Note: only applies if relax_cbf is True.
+        relax_qp (bool, optional): Whether to allow for relaxation in the CLF-CBF QP. Defaults to True.
+            Note: this is required for differentiability through the QP.
+        cbf_relaxation_penalty (float, optional): Penalty on the slack variable in the relaxed QP. Defaults to 1e4.
+            Note: only applies if relax_qp is True.
         clf_relaxation_penalty (float): Penalty on the CLF slack variable when enforcing the CBF. Defaults to 1e2
+        control_relaxation_penalty (float, optional): Penalty on the control constraint slack variables in the
+            relaxed QP. Defaults to 1e5. Note: only applies if relax_qp is True.
         solver_tol (float, optional): Tolerance for the QP solver. Defaults to 1e-3.
-        init_args (tuple, optional): If your barrier function relies on additional arguments other than just the state,
-            include an initial seed for these arguments here. This is to help test the output of the barrier function.
-            Defaults to ().
+        init_args (tuple, optional): If your barriers or dynamics rely on additional (non-differentiable, static shape)
+            args other than just the state, include an initial seed for these args here. Defaults to None.
+        init_kwargs (dict, optional): If your barriers or dynamics rely on additional (non-differentiable, static shape)
+            kwargs other than just the state, include an initial seed for these kwargs here. Defaults to None.
     """
 
     def __init__(
@@ -82,21 +87,25 @@ class CLFCBFConfig(CBFConfig):
         m: int,
         u_min: Optional[ArrayLike] = None,
         u_max: Optional[ArrayLike] = None,
-        relax_cbf: bool = True,
+        relax_qp: bool = True,
         cbf_relaxation_penalty: float = 1e4,
         clf_relaxation_penalty: float = 1e2,
+        control_relaxation_penalty: float = 1e5,
         solver_tol: float = 1e-3,
-        init_args: tuple = (),
+        init_args: Optional[tuple] = None,
+        init_kwargs: Optional[dict] = None,
     ):
         super().__init__(
             n,
             m,
             u_min,
             u_max,
-            relax_cbf,
+            relax_qp,
             cbf_relaxation_penalty,
+            control_relaxation_penalty,
             solver_tol,
             init_args,
+            init_kwargs,
         )
 
         if not (
@@ -108,10 +117,19 @@ class CLFCBFConfig(CBFConfig):
             )
         self.clf_relaxation_penalty = float(clf_relaxation_penalty)
 
+        # If relaxing the QP, need to have CLF penalty < CBF penalty < Control constraint penalty
+        if self.cbf_relaxation_penalty < self.clf_relaxation_penalty:
+            print("WARNING: CBF constraints have a lower penalty than the CLFs")
+        if (
+            self.control_constrained
+            and self.control_relaxation_penalty < self.clf_relaxation_penalty
+        ):
+            print("WARNING: Control constraints have a lower penalty than the CLFs")
+
         # Check on CLF dimension
-        z_test = jnp.ones(self.n)
-        v1_test = self.V_1(jnp.ones(self.n))
-        v2_test = self.V_2(jnp.ones(self.n))
+        z_test = np.ones(self.n)
+        v1_test = self.V_1(z_test, z_test, *self.init_args, **self.init_kwargs)
+        v2_test = self.V_2(z_test, z_test, *self.init_args, **self.init_kwargs)
         if v1_test.ndim != 1 or v2_test.ndim != 1:
             raise ValueError("CLF(s) must output 1D arrays")
         self.num_rd1_clf = v1_test.shape[0]
@@ -122,9 +140,9 @@ class CLFCBFConfig(CBFConfig):
                 "No Lyanpunov functions provided."
                 + "\nYou can implement this via the V_1 and/or V_2 methods in your config class"
             )
-        v_test = jnp.concatenate([v1_test, v2_test])
-        gamma_test = self.gamma(v_test)
-        gamma_2_test = self.gamma_2(v2_test)
+        v_test = np.concatenate([v1_test, v2_test])
+        gamma_test = self.gamma(v_test, *self.init_args, **self.init_kwargs)
+        gamma_2_test = self.gamma_2(v2_test, *self.init_args, **self.init_kwargs)
         if gamma_test.shape != (self.num_clf,):
             raise ValueError(
                 f"Invalid shape for gamma(V(z)): {gamma_test.shape}. Expected ({self.num_clf},)"
@@ -135,18 +153,48 @@ class CLFCBFConfig(CBFConfig):
                 f"Invalid shape for gamma_2(V_2(z)): {gamma_2_test.shape}. Expected ({self.num_rd2_clf},)"
                 + "\nCheck that the output of the gamma_2() function matches the number of RD2 CLFs"
             )
-        self._check_class_kappa(self.gamma, self.num_clf)
-        self._check_class_kappa(self.gamma_2, self.num_rd2_clf)
-        H_test = self.H(z_test)
+        self._check_class_kappa(
+            self.gamma, self.num_clf, *self.init_args, **self.init_kwargs
+        )
+        self._check_class_kappa(
+            self.gamma_2, self.num_rd2_clf, *self.init_args, **self.init_kwargs
+        )
+        H_test = self.H(z_test, *self.init_args, **self.init_kwargs)
         if H_test.shape != (self.m, self.m):
             raise ValueError(
                 f"Invalid shape for H(z): {H_test.shape}. Expected ({self.m}, {self.m})"
             )
         if not self._is_symmetric_psd(H_test):
             raise ValueError("H(z) must be symmetric positive semi-definite")
-        # TODO: add a warning if the CLF relaxation penalty > the QP relaxation penalty?
 
-    def V_1(self, z: ArrayLike) -> Array:
+        # Handle QP relaxation penalties, if relaxation is enabled
+        num_qp_constraints = (
+            self.num_cbf + self.num_clf
+            if not self.control_constrained
+            else self.num_cbf + self.num_clf + 2 * self.m
+        )
+        if self.control_constrained:
+            self.constraint_relaxation_penalties = tuple(
+                np.concatenate(
+                    [
+                        self.clf_relaxation_penalty * np.ones(self.num_clf),
+                        self.cbf_relaxation_penalty * np.ones(self.num_cbf),
+                        self.control_relaxation_penalty * np.ones(2 * self.m),
+                    ]
+                )
+            )
+        else:
+            self.constraint_relaxation_penalties = tuple(
+                np.concatenate(
+                    [
+                        self.clf_relaxation_penalty * np.ones(self.num_clf),
+                        self.cbf_relaxation_penalty * np.ones(self.num_cbf),
+                    ]
+                )
+            )
+        assert len(self.constraint_relaxation_penalties) == num_qp_constraints
+
+    def V_1(self, z: ArrayLike, z_des: ArrayLike, *args, **kwargs) -> Array:
         """Relative-Degree-1 Control Lyapunov Function (CLF)
 
         A CLF is a positive-definite function which evaluates to zero at the equilibrium point, and is
@@ -161,6 +209,7 @@ class CLFCBFConfig(CBFConfig):
 
         Args:
             z (ArrayLike): State, shape (n,)
+            z_des (ArrayLike): Desired state, shape (n,)
 
         Returns:
             Array: V(z): The RD1 CLF evaluation, shape (num_rd1_clf,)
@@ -168,7 +217,7 @@ class CLFCBFConfig(CBFConfig):
         return jnp.array([])
 
     # TODO: Check if the math behind this is actually valid
-    def V_2(self, z: ArrayLike) -> Array:
+    def V_2(self, z: ArrayLike, z_des: ArrayLike, *args, **kwargs) -> Array:
         """Relative-Degree-2 (high-order) Control Lyapunov Function (CLF)
 
         A CLF is a positive-definite function which evaluates to zero at the equilibrium point, and is
@@ -184,13 +233,14 @@ class CLFCBFConfig(CBFConfig):
 
         Args:
             z (ArrayLike): State, shape (n,)
+            z_des (ArrayLike): Desired state, shape (n,)
 
         Returns:
             Array: V(z): The RD2 CLF evaluation, shape (num_rd2_clf,)
         """
         return jnp.array([])
 
-    def gamma(self, v: ArrayLike) -> Array:
+    def gamma(self, v: ArrayLike, *args, **kwargs) -> Array:
         """A class Kappa function, dictating the "gain" of the CLF
 
         For reference, a class Kappa function is a monotonically increasing function which passes through the origin.
@@ -205,7 +255,7 @@ class CLFCBFConfig(CBFConfig):
         """
         return v
 
-    def gamma_2(self, v_2: ArrayLike) -> Array:
+    def gamma_2(self, v_2: ArrayLike, *args, **kwargs) -> Array:
         """A second class Kappa function, dictating the "gain" associated with the derivative of the CLF
 
         For reference, a class Kappa function is a monotonically increasing function which passes through the origin.
@@ -220,7 +270,7 @@ class CLFCBFConfig(CBFConfig):
         """
         return v_2
 
-    def H(self, z: ArrayLike) -> Array:
+    def H(self, z: ArrayLike, *args, **kwargs) -> Array:
         """Matrix defining the quadratic control term in the CLF objective (minimize 0.5 * u^T H u + F^T u)
 
         **Must be PSD!**
@@ -236,7 +286,7 @@ class CLFCBFConfig(CBFConfig):
         """
         return jnp.eye(self.m)
 
-    def F(self, z: ArrayLike) -> Array:
+    def F(self, z: ArrayLike, *args, **kwargs) -> Array:
         """Vector defining the linear term in the CLF objective (minimize 0.5 * u^T H u + F^T u)
 
         The default implementation is a zero vector, but this can be overridden

cbfpy/envs/base_env.py

@@ -3,7 +3,7 @@
 
 This is a convenient structure for building demo environments to test CBFs. However, it is not necessary to use this.
 
-For instance, going back to the CBF usage pseudocode, 
+For instance, going back to the CBF usage pseudocode,
 ```
 while True:
     z = get_state()
@@ -11,10 +11,10 @@ while True:
     u_nom = nominal_controller(z, z_des)
     u = cbf.safety_filter(z, u_nom)
     apply_control(u)
-    step() 
+    step()
 ```
 We use this base environment to set up the `get_state`, `get_desired_state`, `apply_control`, and `step` methods in
-any derived environments. 
+any derived environments.
 """
 
 from abc import ABC, abstractmethod

cbfpy/envs/drone_env.py

@@ -1,5 +1,5 @@
 """
-# Drone Environment 
+# Drone Environment
 
 This is a wrapper around the gym-pybullet-drones environment, using velocity control.
 """

cbfpy/examples/adaptive_cruise_control_demo.py

@@ -49,10 +49,11 @@ class ACCConfig(CLFCBFConfig):
             u_max=u_max,
             # Note: Relaxing the CLF-CBF QP is tricky because there is an additional relaxation
             # parameter already, balancing the CLF and CBF constraints.
-            relax_cbf=False,
-            # If indeed relaxing, ensure that the QP relaxation >> the CLF relaxation
+            relax_qp=False,
+            # If indeed relaxing, ensure that the CBF relaxation >> the CLF relaxation
             clf_relaxation_penalty=10.0,
-            cbf_relaxation_penalty=1e6,
+            cbf_relaxation_penalty=1e5,
+            control_relaxation_penalty=1e6,
         )
 
     def drag_force(self, v: float) -> float:
@@ -82,7 +83,7 @@ class ACCConfig(CLFCBFConfig):
             [D - self.T * v_f - 0.5 * (v_l - v_f) ** 2 / (self.cd * self.gravity)]
         )
 
-    def V_1(self, z: ArrayLike) -> float:
+    def V_1(self, z: ArrayLike, z_des: ArrayLike) -> float:
         # CLF: Squared error between the follower velocity and the desired velocity
         return jnp.array([(z[0] - self.v_des) ** 2])
 

cbfpy/examples/drone_demo.py

@@ -38,20 +38,20 @@ class DroneConfig(CBFConfig):
         super().__init__(
             n=6,  # State = [position, velocity]
             m=3,  # Control = [velocity]
-            relax_cbf=True,
+            relax_qp=True,
             init_args=(init_z_obs,),
             cbf_relaxation_penalty=1e6,
         )
 
-    def f(self, z):
+    def f(self, z, *args, **kwargs):
         # Assume we are directly controlling the robot's velocity
         return jnp.zeros(self.n)
 
-    def g(self, z):
+    def g(self, z, *args, **kwargs):
         # Assume we are directly controlling the robot's velocity
         return jnp.block([[jnp.eye(3)], [jnp.zeros((3, 3))]])
 
-    def h_1(self, z, z_obs):
+    def h_1(self, z, z_obs, **kwargs):
         obstacle_radius = 0.25
         robot_radius = 0.25
         padding = 0.1
@@ -73,10 +73,12 @@ class DroneConfig(CBFConfig):
                 - padding
             ]
         )
-        h_box_containment = jnp.concatenate([self.pos_max - z[:3], z[:3] - self.pos_min])
+        h_box_containment = jnp.concatenate(
+            [self.pos_max - z[:3], z[:3] - self.pos_min]
+        )
         return jnp.concatenate([h_obstacle_avoidance, h_box_containment])
 
-    def alpha(self, h):
+    def alpha(self, h, *args, **kwargs):
         return jnp.array([3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) * h