cbfpy 0.0.1__py3-none-any.whl

cbfpy/envs/point_robot_envs.py

@@ -0,0 +1,209 @@
+"""
+# Simulation environments for point robots
+"""
+
+import numpy as np
+from jax import Array
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+import pybullet
+from pybullet_utils.bullet_client import BulletClient
+
+from cbfpy.utils.visualization import visualize_3D_box
+from cbfpy.envs.base_env import BaseEnv
+from cbfpy.utils.general_utils import find_assets_dir, stdout_redirected
+
+
+class PointRobotEnv(BaseEnv):
+    """Simulation environment for a point robot trying to approach a target position in 3D,
+    while remaining in a safe set defined as a box
+
+    Args:
+        xyz_min (ArrayLike, optional): Minimum bounds of the safe region, shape (3,). Defaults to (-1.0, -1.0, -1.0)
+        xyz_max (ArrayLike, optional): Maximum bounds of the safe region, shape (3,). Defaults to (1.0, 1.0, 1.0)
+    """
+
+    # Constants
+    # Based on the values in the point robot URDF
+    URDF = "point_robot.urdf"
+    RADIUS = 0.25
+    MASS = 1.0
+
+    def __init__(
+        self,
+        xyz_min: ArrayLike = (-1.0, -1.0, -1.0),
+        xyz_max: ArrayLike = (1.0, 1.0, 1.0),
+    ):
+        self.xyz_min = np.array(xyz_min)
+        self.xyz_max = np.array(xyz_max)
+        assert len(self.xyz_min) == len(self.xyz_max) == 3
+        with stdout_redirected(restore=False):
+            self.client: pybullet = BulletClient(pybullet.GUI)
+        self.client.setAdditionalSearchPath(find_assets_dir())
+        self.client.configureDebugVisualizer(self.client.COV_ENABLE_GUI, 0)
+        self.robot = self.client.loadURDF(self.URDF)
+        self.target = self.client.loadURDF(self.URDF, basePosition=(3, 1, 1))
+        self.client.changeVisualShape(self.target, -1, rgbaColor=[1, 0, 0, 1])
+        self.client.changeDynamics(self.target, -1, linearDamping=10, angularDamping=10)
+        self.box = visualize_3D_box(
+            [self.xyz_min - self.RADIUS, self.xyz_max + self.RADIUS],
+            rgba=(0, 1, 0, 0.5),
+        )
+
+        # For color determination
+        self.is_in_box = True
+        self.tol = 1e-3
+
+    def _update_color(self, robot_pos: ArrayLike) -> None:
+        """Update the color of the box depending on if the robot is inside or not (Green inside, red outside)"""
+        robot_pos = np.array(robot_pos)
+        if np.any(robot_pos < self.xyz_min - self.tol) or np.any(
+            robot_pos > self.xyz_max + self.tol
+        ):
+            if self.is_in_box:
+                self.client.changeVisualShape(self.box, -1, rgbaColor=[1, 0, 0, 0.5])
+            self.is_in_box = False
+        else:
+            if not self.is_in_box:
+                self.client.changeVisualShape(self.box, -1, rgbaColor=[0, 1, 0, 0.5])
+            self.is_in_box = True
+
+    def get_state(self) -> Array:
+        robot_pos = self.client.getBasePositionAndOrientation(self.robot)[0]
+        robot_vel = self.client.getBaseVelocity(self.robot)[0]
+        self._update_color(robot_pos)
+        return np.array([*robot_pos, *robot_vel])
+
+    def get_desired_state(self) -> Array:
+        target_pos = self.client.getBasePositionAndOrientation(self.target)[0]
+        target_vel = self.client.getBaseVelocity(self.target)[0]
+        return np.array([*target_pos, *target_vel])
+
+    def apply_control(self, u: Array) -> None:
+        robot_pos = self.client.getBasePositionAndOrientation(self.robot)[0]
+        self.client.applyExternalForce(
+            self.robot, -1, u, robot_pos, self.client.WORLD_FRAME
+        )
+
+    def step(self):
+        self.client.stepSimulation()
+
+
+class PointRobotObstacleEnv(BaseEnv):
+    """Simulation environment for a point robot avoiding a movable obstacle,
+    while remaining in a safe set defined as a box
+
+    Args:
+        robot_pos (ArrayLike, optional): Initial position of the robot. Defaults to (0, 0, 0).
+        robot_vel (ArrayLike, optional): Initial velocity of the robot. Defaults to (0, 0, 0).
+        obstacle_pos (ArrayLike, optional): Initial position of the obstacle. Defaults to (1, 1, 1).
+        obstacle_vel (ArrayLike, optional): Initial velocity of the obstacle. Defaults to (-0.5, -0.6, -0.7).
+        xyz_min (ArrayLike, optional): Minimum bounds of the safe region, shape (3,). Defaults to (-1.0, -1.0, -1.0)
+        xyz_max (ArrayLike, optional): Maximum bounds of the safe region, shape (3,). Defaults to (1.0, 1.0, 1.0)
+    """
+
+    # Constants
+    # Based on the values in the point robot URDF
+    URDF = "point_robot.urdf"
+    RADIUS = 0.25
+    MASS = 1.0
+
+    def __init__(
+        self,
+        robot_pos: ArrayLike = (0, 0, 0),
+        robot_vel: ArrayLike = (0, 0, 0),
+        obstacle_pos: ArrayLike = (1, 1, 1),
+        obstacle_vel: ArrayLike = (-0.5, -0.6, -0.7),
+        xyz_min: ArrayLike = (-1.0, -1.0, -1.0),
+        xyz_max: ArrayLike = (1.0, 1.0, 1.0),
+    ):
+        with stdout_redirected(restore=False):
+            self.client: pybullet = BulletClient(pybullet.GUI)
+        self.client.setAdditionalSearchPath(find_assets_dir())
+        self.client.configureDebugVisualizer(self.client.COV_ENABLE_GUI, 0)
+        self.robot = self.client.loadURDF(self.URDF, basePosition=robot_pos)
+        self.target = self.client.loadURDF(self.URDF, basePosition=obstacle_pos)
+        self.client.changeVisualShape(self.target, -1, rgbaColor=[1, 0, 0, 1])
+        self.client.changeDynamics(self.target, -1, angularDamping=10)
+        self.client.resetBaseVelocity(self.target, obstacle_vel, (0, 0, 0))
+        self.client.resetBaseVelocity(self.robot, robot_vel, (0, 0, 0))
+        self.client.addUserDebugPoints([[0, 0, 0]], [[1, 0, 0]], 10, 0)
+        self.xyz_min = np.array(xyz_min)
+        self.xyz_max = np.array(xyz_max)
+        self.box = visualize_3D_box(
+            [self.xyz_min - self.RADIUS, self.xyz_max + self.RADIUS],
+            rgba=(0, 1, 0, 0.5),
+        )
+
+    def get_state(self):
+        robot_pos = self.client.getBasePositionAndOrientation(self.robot)[0]
+        robot_vel = self.client.getBaseVelocity(self.robot)[0]
+        obstacle_pos = self.client.getBasePositionAndOrientation(self.target)[0]
+        obstacle_vel = self.client.getBaseVelocity(self.target)[0]
+        return np.array([*robot_pos, *robot_vel]), np.array(
+            [*obstacle_pos, *obstacle_vel]
+        )
+
+    def get_desired_state(self):
+        return np.zeros(6)
+
+    def apply_control(self, u):
+        robot_pos = self.client.getBasePositionAndOrientation(self.robot)[0]
+        self.client.applyExternalForce(
+            self.robot, -1, u, robot_pos, self.client.WORLD_FRAME
+        )
+
+    def step(self):
+        self.client.stepSimulation()
+
+
+## Tests ##
+
+
+def _test_standard_env():
+    def nominal_controller(z, z_des):
+        # Use a PD controller to try to "touch" the target robot
+        Kp = 1.0
+        Kd = 1.0
+        pos_diff = z_des[:3] - z[:3]
+        des_pos = z_des[:3] - 2 * PointRobotEnv.RADIUS * pos_diff / jnp.linalg.norm(
+            pos_diff
+        )
+        return -Kp * (z[:3] - des_pos) - Kd * (z[3:] - z_des[3:])
+
+    env = PointRobotEnv()
+    while True:
+        z = env.get_state()
+        z_des = env.get_desired_state()
+        u = nominal_controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+def _test_obstacle_env():
+    def nominal_controller(z, z_des):
+        Kp = 1.0
+        Kd = 1.0
+        return -Kp * (z[:3] - z_des[:3]) - Kd * (z[3:] - z_des[3:])
+
+    env = PointRobotObstacleEnv()
+    while True:
+        z, z_obs = env.get_state()
+        z_des = env.get_desired_state()
+        u = nominal_controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+def _test_env(env_class: BaseEnv):
+    print("\nDemoing the point robot environment with an unsafe controller")
+    if env_class == PointRobotEnv:
+        return _test_standard_env()
+    elif env_class == PointRobotObstacleEnv:
+        return _test_obstacle_env()
+    else:
+        raise ValueError("Invalid env class")
+
+
+if __name__ == "__main__":
+    _test_env(PointRobotEnv)

cbfpy/examples/adaptive_cruise_control_demo.py

@@ -0,0 +1,117 @@
+"""
+# Adaptive Cruise Control CLF-CBF Demo
+
+This will enforce that the follower vehicle maintains a safe distance from the leader vehicle,
+while also tracking a desired velocity.
+
+We define the state z as [v_follower, v_leader, follow_distance] and the control u as the follower's wheel force
+
+The dynamics incorporate a simple drag force model using empirically-derived coefficients
+
+Note: There are a few parameters to tune in this CLF-CBF, such as the weightings between the inputs
+the slack variable in the CLF objective. This is tricky to tune in general and values have been left
+at what has been seen in other references.
+
+Reference:
+
+- "Control Barrier Function Based Quadratic Programs for Safety Critical Systems" - TAC 2017
+- "Control Barrier Function based Quadratic Programs with Application to Adaptive Cruise Control" - CDC 2014
+
+Some parameters are based on Jason Choi's https://github.com/HybridRobotics/CBF-CLF-Helper
+"""
+
+from jax import Array
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+
+from cbfpy.envs.car_env import VehicleEnv
+from cbfpy.config.clf_cbf_config import CLFCBFConfig
+from cbfpy.cbfs.clf_cbf import CLFCBF
+
+
+class ACCConfig(CLFCBFConfig):
+    """Configuration for the Adaptive Cruise Control CLF-CBF demo"""
+
+    def __init__(self):
+        self.gravity = 9.81
+        self.mass = 1650.0
+        self.drag_coeffs = (0.1, 5.0, 0.25)  # Drag coeffs
+        self.v_des = 24.0  # Desired velocity
+        self.T = 1.8  # Lookahead time
+        self.cd = 0.3  # Coefficient of maximum deceleration
+        self.ca = 0.3  # Coefficient of maximum acceleration
+        u_min = -self.cd * self.mass * self.gravity  # Min. control input (max braking)
+        u_max = self.ca * self.mass * self.gravity  # Max. control input (max throttle)
+        super().__init__(
+            n=3,
+            m=1,
+            u_min=u_min,
+            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
+            clf_relaxation_penalty=10.0,
+            cbf_relaxation_penalty=1e6,
+        )
+
+    def drag_force(self, v: float) -> float:
+        """Compute the drag force on the follower car using a simple polynomial model
+
+        Args:
+            v (float): Velocity of the follower vehicle, in m/s
+
+        Returns:
+            float: Drag force, in Newtons
+        """
+        return (
+            self.drag_coeffs[0] + self.drag_coeffs[1] * v + self.drag_coeffs[2] * v**2
+        )
+
+    def f(self, z: ArrayLike) -> Array:
+        v_f, v_l, D = z
+        # Note: We assume that the leader vehicle is at constant velocity here
+        return jnp.array([-self.drag_force(v_f) / self.mass, 0.0, v_l - v_f])
+
+    def g(self, z: ArrayLike) -> Array:
+        return jnp.array([(1 / self.mass), 0.0, 0.0]).reshape(-1, 1)
+
+    def h_1(self, z: ArrayLike) -> Array:
+        v_f, v_l, D = z
+        return jnp.array(
+            [D - self.T * v_f - 0.5 * (v_l - v_f) ** 2 / (self.cd * self.gravity)]
+        )
+
+    def V_1(self, z: ArrayLike) -> float:
+        # CLF: Squared error between the follower velocity and the desired velocity
+        return jnp.array([(z[0] - self.v_des) ** 2])
+
+    def H(self, z: ArrayLike) -> Array:
+        return jnp.eye(self.m) * (2 / self.mass**2)
+
+    def F(self, z: ArrayLike) -> Array:
+        return jnp.array([-2 * self.drag_force(z[0]) / self.mass**2])
+
+
+def main():
+    config = ACCConfig()
+    clf_cbf = CLFCBF.from_config(config)
+    # Ensure that parameters match up between the environment and the CLF-CBF
+    env = VehicleEnv(
+        "CLF-CBF Controller",
+        mass=config.mass,
+        drag_coeffs=config.drag_coeffs,
+        v_des=config.v_des,
+        u_min=config.u_min,
+        u_max=config.u_max,
+    )
+    while env.running:
+        z = env.get_state()
+        z_des = env.get_desired_state()
+        u = clf_cbf.controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+if __name__ == "__main__":
+    main()

cbfpy/examples/drone_demo.py

@@ -0,0 +1,109 @@
+"""
+# Drone Obstacle Avoidance Demo
+
+We use velocity control, which can be easily applied to other drones
+(for instance, PX4-controlled quadrotors can take in velocity commands)
+
+The CBF uses a reduced model of the drone dynamics, while the simulation environment
+does reflect a (somewhat-accurate) model of the true drone dynamics
+
+Here, our state is the position and velocity of the drone: z = [position, velocity]
+and the control input is the velocity of the drone: u = [velocity]
+
+We also incorporate the state of the obstacle as an additional input to the CBF: z_obs = [position, velocity]
+
+Note: since there is gravity in this simulation, the obstacle will fall to the ground initially. This is fine: just
+use the mouse to drag it around near the drone to see the CBF response
+
+See https://danielpmorton.github.io/drone_fencing/ for a demo of this on real hardware, and see the "point robot
+obstacle avoidance" demo in CBFpy for a simplified version of this demo
+"""
+
+import time
+import jax
+import jax.numpy as jnp
+
+from cbfpy import CBF, CBFConfig
+from cbfpy.envs.drone_env import DroneEnv
+
+
+class DroneConfig(CBFConfig):
+    """Config for the drone obstacle avoidance / safe set containment demo"""
+
+    def __init__(self):
+        self.mass = 1.0
+        self.pos_min = jnp.array([-2.0, -2.0, 0.7])
+        self.pos_max = jnp.array([2.0, 2.0, 2.0])
+        init_z_obs = jnp.array([3.0, 1.0, 1.0, -0.1, -0.1, -0.1])
+        super().__init__(
+            n=6,  # State = [position, velocity]
+            m=3,  # Control = [velocity]
+            relax_cbf=True,
+            init_args=(init_z_obs,),
+            cbf_relaxation_penalty=1e6,
+        )
+
+    def f(self, z):
+        # Assume we are directly controlling the robot's velocity
+        return jnp.zeros(self.n)
+
+    def g(self, z):
+        # 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):
+        obstacle_radius = 0.25
+        robot_radius = 0.25
+        padding = 0.1
+        pos_robot = z[:3]
+        vel_robot = z[3:]
+        pos_obs = z_obs[:3]
+        vel_obs = z_obs[3:]
+        dist_between_centers = jnp.linalg.norm(pos_obs - pos_robot)
+        dir_obs_to_robot = (pos_robot - pos_obs) / dist_between_centers
+        collision_velocity_component = (vel_obs - vel_robot).T @ dir_obs_to_robot
+        lookahead_time = 2.0
+        padding = 0.1
+        h_obstacle_avoidance = jnp.array(
+            [
+                dist_between_centers
+                - collision_velocity_component * lookahead_time
+                - obstacle_radius
+                - robot_radius
+                - padding
+            ]
+        )
+        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):
+        return jnp.array([3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) * h
+
+
+def nominal_controller(z, z_des):
+    Kp = 1.0
+    Kv = 1.0
+    return Kp * (z_des[:3] - z[:3]) + Kv * (z_des[3:] - z[3:])
+
+
+def main():
+    config = DroneConfig()
+    cbf = CBF.from_config(config)
+    env = DroneEnv(config.pos_min, config.pos_max)
+
+    @jax.jit
+    def safe_controller(z, z_des, z_obs):
+        u = nominal_controller(z, z_des)
+        return cbf.safety_filter(z, u, z_obs)
+
+    while True:
+        z, z_obs = env.get_state()
+        z_des = env.get_desired_state()
+        u = safe_controller(z, z_des, z_obs)
+        env.apply_control(u)
+        env.step()
+        time.sleep(1 / 300)
+
+
+if __name__ == "__main__":
+    main()

cbfpy/examples/joint_limits_demo.py

@@ -0,0 +1,150 @@
+"""
+# Manipulator Joint Limit Avoidance Demo
+
+We will command a joint position trajectory that exceeds the joint limits, and the CBF will ensure that we stay
+within the limits (+ some margin)
+
+This uses a single-integrator reduced model of the manipulator dynamics.
+We define the state as the joint positions and assume that we can directly control the joint velocities
+i.e. z = [q1, q2, q3] and u = [q1_dot, q2_dot, q3_dot] 
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from jax import Array
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+
+from cbfpy import CBF, CBFConfig
+from cbfpy.envs.arm_envs import JointLimitsEnv
+from cbfpy.utils.general_utils import find_assets_dir
+
+URDF = find_assets_dir() + "three_dof_arm.urdf"
+
+
+class JointLimitsConfig(CBFConfig):
+    """Config for the 3-DOF arm, avoiding its joint limits using velocity control"""
+
+    def __init__(self):
+        self.num_joints = 3
+        # Joint limit values from the URDF
+        self.q_min = -np.pi / 2 * np.ones(self.num_joints)
+        self.q_max = np.pi / 2 * np.ones(self.num_joints)
+        # Pad joint limts (to better evauate CBF performance)
+        self.padding = 0.3
+        super().__init__(n=self.num_joints, m=self.num_joints)
+
+    def f(self, z):
+        return jnp.zeros(self.num_joints)
+
+    def g(self, z):
+        return jnp.eye(self.num_joints)
+
+    def h_1(self, z):
+        q = z
+        return jnp.concatenate(
+            [self.q_max - q - self.padding, q - self.q_min - self.padding]
+        )
+
+
+def nominal_controller(q: Array, q_des: Array) -> Array:
+    """Very simple proportional controller: Commands joint velocities to reduce a position error
+
+    Args:
+        q (Array): Joint positions, shape (num_joints,)
+        q_des (Array): Desired joint positions, shape (num_joints,)
+
+    Returns:
+        Array: Joint velocity command, shape (num_joints,)
+    """
+    k = 1.0
+    return k * (q_des - q)
+
+
+def main():
+    config = JointLimitsConfig()
+    cbf = CBF.from_config(config)
+    env = JointLimitsEnv()
+
+    @jax.jit
+    def safety_filter(q, u):
+        return cbf.safety_filter(q, u)
+
+    q_hist = []
+    q_des_hist = []
+    u_safe_hist = []
+    u_unsafe_hist = []
+    sim_time = 100.0
+    timestep = env.timestep
+    num_timesteps = int(sim_time / timestep)
+    for i in range(num_timesteps):
+        q = env.get_state()
+        q_des = env.get_desired_state()
+        u_unsafe = nominal_controller(q, q_des)
+        u = safety_filter(q, u_unsafe)
+        env.apply_control(u)
+        env.step()
+        q_hist.append(q)
+        q_des_hist.append(q_des)
+        u_unsafe_hist.append(u_unsafe)
+        u_safe_hist.append(u)
+
+    ## Plotting ##
+
+    fig, axs = plt.subplots(2, 3)
+    ts = timestep * np.arange(num_timesteps)
+
+    # On the top row, plot the q and q des for each joint, along with the joint limits indicated
+    for i in range(3):
+        (q_line,) = axs[0, i].plot(ts, np.array(q_hist)[:, i], label="q")
+        (q_des_line,) = axs[0, i].plot(ts, np.array(q_des_hist)[:, i], label="q_des")
+        axs[0, i].plot(
+            ts,
+            env.q_min[i] * np.ones(num_timesteps),
+            ls="--",
+            c="red",
+            label="q_min (True)",
+        )
+        axs[0, i].plot(
+            ts,
+            env.q_max[i] * np.ones(num_timesteps),
+            ls="--",
+            c="red",
+            label="q_max (True)",
+        )
+        axs[0, i].plot(
+            ts,
+            (env.q_min[i] + config.padding) * np.ones(num_timesteps),
+            ls="--",
+            c="blue",
+            label="q_min (CBF)",
+        )
+        axs[0, i].plot(
+            ts,
+            (env.q_max[i] - config.padding) * np.ones(num_timesteps),
+            ls="--",
+            c="blue",
+            label="q_max (CBF)",
+        )
+        legend_elements = [
+            q_line,
+            q_des_line,
+            Line2D([0], [0], color="red", ls="--", label="True limits"),
+            Line2D([0], [0], color="blue", ls="--", label="CBF limits"),
+        ]
+        axs[0, i].legend(handles=legend_elements, loc="lower left")
+        axs[0, i].set_title(f"Joint {i} position")
+
+    # On the bottom row, plot the safe and unsafe velocities for each joint
+    for i in range(3):
+        axs[1, i].plot(ts, np.array(u_safe_hist)[:, i], label="Safe")
+        axs[1, i].plot(ts, np.array(u_unsafe_hist)[:, i], label="Unsafe")
+        axs[1, i].legend(loc="lower left")
+        axs[1, i].set_title(f"Joint {i} velocity")
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

cbfpy/examples/point_robot_demo.py

@@ -0,0 +1,91 @@
+"""
+# Point Robot Safe-Set Containment Demo
+
+Demo of a point robot in 3D constrained to lie within a box via a CBF safety filter on a PD controller
+
+We define the state z as [x, y, z, vx, vy, vz] and the control u as [Fx, Fy, Fz]
+
+The dynamics are that of a simple double integrator:
+z_dot = [vx, vy, vz, 0, 0, 0] + [0, 0, 0, Fx/m, Fy/m, Fz/m]
+
+In matrix form, this can be expressed as z_dot = A z + B u with A and B as implemented in the config.
+
+The safety constraints are set as an upper and lower bound on the position of the robot.
+
+This is a relative-degree-2 system, so we use the RD2 version of the CBF constraints.
+"""
+
+import jax
+from jax import Array
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+
+from cbfpy import CBF, CBFConfig
+from cbfpy.envs.point_robot_envs import PointRobotEnv
+
+
+class PointRobotConfig(CBFConfig):
+    """Configuration for the 3D 'point-robot-in-a-box' example."""
+
+    def __init__(self):
+        self.mass = 1.0
+        self.pos_min = jnp.array([-1.0, -1.0, -1.0])
+        self.pos_max = jnp.array([1.0, 1.0, 1.0])
+        super().__init__(n=6, m=3)
+
+    def f(self, z):
+        A = jnp.block(
+            [[jnp.zeros((3, 3)), jnp.eye(3)], [jnp.zeros((3, 3)), jnp.zeros((3, 3))]]
+        )
+        return A @ z
+
+    def g(self, z):
+        B = jnp.block([[jnp.zeros((3, 3))], [jnp.eye(3) / self.mass]])
+        return B
+
+    def h_2(self, z):
+        return jnp.concatenate([self.pos_max - z[:3], z[:3] - self.pos_min])
+
+
+@jax.jit
+def nominal_controller(z: ArrayLike, z_des: ArrayLike) -> Array:
+    """A simple PD controller for the point robot.
+
+    This is unsafe without the CBF, as there is no guarantee that the robot will stay within the safe region
+
+    Args:
+        z (ArrayLike): The current state of the robot [x, y, z, vx, vy, vz]
+        z_des (ArrayLike): The desired state of the robot [x_des, y_des, z_des, vx_des, vy_des, vz_des]
+    """
+    Kp = 1.0
+    Kd = 1.0
+    radius = 0.25
+    # We assume we are in our Pybullet simulation environment where we have two point robots loaded
+    # One is the robot we are controlling, and the other is the target robot, which is interactive via the GUI
+    # This will attempt to "touch" the target robot
+    pos_diff = z_des[:3] - z[:3]
+    des_pos = z_des[:3] - 2 * radius * pos_diff / jnp.linalg.norm(pos_diff)
+    u = -Kp * (z[:3] - des_pos) - Kd * (z[3:] - z_des[3:])
+    return u
+
+
+def main():
+    config = PointRobotConfig()
+    cbf = CBF.from_config(config)
+    env = PointRobotEnv(config.pos_min, config.pos_max)
+
+    @jax.jit
+    def safe_controller(z, z_des):
+        u = nominal_controller(z, z_des)
+        return cbf.safety_filter(z, u)
+
+    while True:
+        z = env.get_state()
+        z_des = env.get_desired_state()
+        u = safe_controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+if __name__ == "__main__":
+    main()