cbfpy 0.0.1__py3-none-any.whl

cbfpy/envs/arm_envs.py

@@ -0,0 +1,84 @@
+"""
+# Simulation environents for robot arms
+
+This currently includes a very simple 3-DOF environment which helps demonstrate joint limit avoidance,
+but more will be added in the future
+"""
+
+import numpy as np
+import pybullet
+from pybullet_utils.bullet_client import BulletClient
+
+from cbfpy.utils.general_utils import find_assets_dir, stdout_redirected
+from cbfpy.envs.base_env import BaseEnv
+
+URDF = find_assets_dir() + "three_dof_arm.urdf"
+
+
+class JointLimitsEnv(BaseEnv):
+    """Simulation environment for the 3-DOF arm joint-limit-avoidance demo
+
+    This includes a desired reference trajectory which is unsafe: it will command sinusoidal
+    joint motions (with different frequencies per link) that will exceed the joint limits of the robot
+    """
+
+    def __init__(self):
+        with stdout_redirected(restore=False):
+            self.client: pybullet = BulletClient(pybullet.GUI)
+        self.robot = pybullet.loadURDF(URDF, useFixedBase=True)
+        self.num_joints = self.client.getNumJoints(self.robot)
+        self.q_min = np.array(
+            [self.client.getJointInfo(self.robot, i)[8] for i in range(self.num_joints)]
+        )
+        self.q_max = np.array(
+            [self.client.getJointInfo(self.robot, i)[9] for i in range(self.num_joints)]
+        )
+        self.timestep = self.client.getPhysicsEngineParameters()["fixedTimeStep"]
+        self.t = 0
+
+        # Sinusoids for the desired joint positions
+        # Setting the amplitude to be the full joint range means we will command DOUBLE
+        # the joint range, exceeding our limits
+        self.omegas = 0.1 * np.array([1.0, 2.0, 3.0])
+        self.amps = self.q_max - self.q_min
+        self.offsets = np.zeros(3)
+
+    def step(self):
+        self.client.stepSimulation()
+        self.t += self.timestep
+
+    def get_state(self):
+        states = self.client.getJointStates(self.robot, range(self.num_joints))
+        return np.array([states[i][0] for i in range(self.num_joints)])
+
+    def get_desired_state(self):
+        # Evaluate our unsafe sinusoidal trajectory
+        return self.amps * np.sin(self.omegas * self.t) + self.offsets
+
+    def apply_control(self, u):
+        self.client.setJointMotorControlArray(
+            self.robot,
+            list(range(self.num_joints)),
+            self.client.VELOCITY_CONTROL,
+            targetVelocities=u,
+        )
+
+
+def _test_env():
+    """Test the environment behavior under an **unsafe** nominal controller"""
+
+    def nominal_controller(z, z_des):
+        k = 1.0
+        return k * (z_des - z)
+
+    env = JointLimitsEnv()
+    while True:
+        z = env.get_state()
+        z_des = env.get_desired_state()
+        u = nominal_controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+if __name__ == "__main__":
+    _test_env()

cbfpy/envs/base_env.py

@@ -0,0 +1,69 @@
+"""
+# Base Environment
+
+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, 
+```
+while True:
+    z = get_state()
+    z_des = get_desired_state()
+    u_nom = nominal_controller(z, z_des)
+    u = cbf.safety_filter(z, u_nom)
+    apply_control(u)
+    step() 
+```
+We use this base environment to set up the `get_state`, `get_desired_state`, `apply_control`, and `step` methods in
+any derived environments. 
+"""
+
+from abc import ABC, abstractmethod
+
+from jax.typing import ArrayLike
+
+
+class BaseEnv(ABC):
+    """Simulation environment Abstract Base Class for testing CBFs
+
+    Any environment inheriting from this class should implement the following methods:
+
+    - `step`: Run a single simulation step
+    - `get_state`: Get the current state of the robot
+    - `get_desired_state`: Get the desired state of the robot
+    - `apply_control`: Apply a control input to the robot
+    """
+
+    @abstractmethod
+    def step(self) -> None:
+        """Runs a single simulation step for the environment
+
+        This should update any dynamics and visuals
+        """
+        pass
+
+    @abstractmethod
+    def get_state(self) -> ArrayLike:
+        """Returns the current state of the environment
+
+        Returns:
+            ArrayLike: State, shape (n,)
+        """
+        pass
+
+    @abstractmethod
+    def get_desired_state(self) -> ArrayLike:
+        """Returns the desired state of the environment
+
+        Returns:
+            ArrayLike: Desired state, shape (n,)
+        """
+        pass
+
+    @abstractmethod
+    def apply_control(self, u: ArrayLike) -> None:
+        """Applies the control input to the environment
+
+        Args:
+            u (ArrayLike): Control, shape (m,)
+        """
+        pass

cbfpy/envs/car_env.py

@@ -0,0 +1,332 @@
+"""
+# Adaptive Cruise Control Simulation Environment
+
+Simple simulation of a leader-follower vehicle system for adaptive cruise control testing
+"""
+
+from typing import Tuple
+
+import numpy as np
+import pygame
+
+from cbfpy.envs.base_env import BaseEnv
+
+CAR_ASCII = """
+        _______
+       //  ||\\ \\
+ _____//___||_\\ \\___
+ )  _          _    \\
+ |_/ \\________/ \\___|
+___\\_/________\\_/______
+"""
+
+
+class VehicleEnv(BaseEnv):
+    """Leader/follower vehicle simulation environment for adaptive cruise control testing
+
+    This will bring up an interactive Pygame window where you can control the speed of the leader car,
+    and the follower car will be controlled by a simple adaptive cruise control algorithm.
+
+    Args:
+        controller_name (str): Name of the controller being tested
+        mass (float, optional): Mass of the follower vehicle. Defaults to 1650.0 (kg)
+        drag_coeffs (Tuple[float, float, float]): Coefficients of a simple polynomial friction model, as
+            described in the Ames CBF paper. Defaults to (0.1, 5.0, 0.25)
+        v_des (float, optional): Desired velocity of the follower car. Defaults to 24.0 (m/s)
+        init_leader_pos (float, optional): Initial position of the leader car. Defaults to 0.0 (meters)
+        init_leader_vel (float, optional): Initial velocity of the leader car. Defaults to 14.0 (m/s)
+        init_follower_pos (float, optional): Initial position of the follower car. Defaults to -20.0 (meters)
+        init_follower_vel (float, optional): Initial velocity of the follower car. Defaults to 10.0 (m/s)
+        u_min (float, optional): Minimum control input, in Newtons. Defaults to -np.inf (unconstrained)
+        u_max (float, optional): Maximum control input, in Newtons. Defaults to np.inf (unconstrained)
+    """
+
+    # Constants
+    PIXELS_PER_METER = 15
+    MIN_FOLLOW_DIST_M = 1
+    MIN_FOLLOW_DIST_PX = PIXELS_PER_METER * MIN_FOLLOW_DIST_M
+
+    # Screen dimensions
+    SCREEN_WIDTH = 600
+    SCREEN_HEIGHT = 800
+
+    # Colors
+    WHITE = (255, 255, 255)
+    GRAY = (100, 100, 100)
+    BLACK = (0, 0, 0)
+    BLUE = (0, 0, 255)
+    RED = (255, 0, 0)
+    YELLOW = (255, 255, 0)
+
+    # Road properties
+    DASH_HEIGHT_M = 2
+    DASH_WIDTH_M = 0.3
+    DASH_GAP_M = 1
+    DASH_HEIGHT_PX = PIXELS_PER_METER * DASH_HEIGHT_M
+    DASH_WIDTH_PX = PIXELS_PER_METER * DASH_WIDTH_M
+    DASH_GAP_PX = PIXELS_PER_METER * DASH_GAP_M
+
+    # Car properties
+    CAR_WIDTH_M = 2
+    CAR_HEIGHT_M = 4
+    CAR_WIDTH_PX = PIXELS_PER_METER * CAR_WIDTH_M
+    CAR_HEIGHT_PX = PIXELS_PER_METER * CAR_HEIGHT_M
+
+    def __init__(
+        self,
+        controller_name: str,
+        mass: float = 1650.0,
+        drag_coeffs: Tuple[float, float, float] = (0.1, 5.0, 0.25),
+        v_des: float = 24.0,
+        init_leader_pos: float = 0.0,
+        init_leader_vel: float = 14.0,
+        init_follower_pos: float = -20.0,
+        init_follower_vel: float = 10.0,
+        u_min: float = -np.inf,
+        u_max: float = np.inf,
+    ):
+        self.mass = mass
+        self.v_des = v_des  # Desired velocity of the follower car
+        assert len(drag_coeffs) == 3
+        self.f0, self.f1, self.f2 = drag_coeffs
+        # Initialize Pygame
+        pygame.init()
+        # Set up the display
+        self.screen = pygame.display.set_mode(
+            (VehicleEnv.SCREEN_WIDTH, VehicleEnv.SCREEN_HEIGHT)
+        )
+        pygame.display.set_caption(f"Adaptive Cruise Control: {controller_name}")
+        # Create car sprites
+        self.leader_sprite = pygame.Surface(
+            (VehicleEnv.CAR_WIDTH_PX, VehicleEnv.CAR_HEIGHT_PX)
+        )
+        self.leader_sprite.fill(VehicleEnv.RED)
+        self.follower_sprite = pygame.Surface(
+            (VehicleEnv.CAR_WIDTH_PX, VehicleEnv.CAR_HEIGHT_PX)
+        )
+        self.follower_sprite.fill(VehicleEnv.BLUE)
+        # Positioning the cars in the display
+        self.leader_x = int(VehicleEnv.SCREEN_WIDTH / 4 - VehicleEnv.CAR_WIDTH_PX / 2)
+        self.leader_y = VehicleEnv.SCREEN_HEIGHT // 5
+        self.follower_x = self.leader_x
+        self.follower_y = self.leader_y + VehicleEnv.MIN_FOLLOW_DIST_PX
+
+        # Init dynamics
+        self.leader_pos = init_leader_pos
+        self.leader_vel = init_leader_vel
+        self.follower_pos = init_follower_pos
+        self.follower_vel = init_follower_vel
+
+        self.leader_vel_des = init_leader_vel
+
+        # Initial position of dashed lines
+        self.dash_offset = 0
+
+        self.font = pygame.font.SysFont("Arial", 20)
+
+        self.fps = 60
+        self.dt = 1 / self.fps
+        self.running = True
+
+        self.last_control = 0.0
+
+        self.u_min = u_min
+        self.u_max = u_max
+
+        # Print instructions
+        print(CAR_ASCII)
+        print("Beginning Vehicle Simulation")
+        print("Press UP and DOWN to control the speed of the leader car")
+        print("Press ESC to quit")
+
+    def pixels_to_meters(self, n_pixels: int) -> float:
+        """Helper function: Converts pixels to meters
+
+        Args:
+            n_pixels (int): Number of pixels
+
+        Returns:
+            float: Distance in meters
+        """
+        return n_pixels / self.PIXELS_PER_METER
+
+    def meters_to_pixels(self, n_meters: float) -> float:
+        """Helper function: Converts meters to pixels
+
+        Args:
+            n_meters (float): Distance in meters
+
+        Returns:
+            float: Number of pixels (Note: Not rounded to an integer value)
+        """
+        return n_meters * self.PIXELS_PER_METER
+
+    def friction(self, v: float) -> float:
+        """Computes the drag force on the vehicle with the simple model presented in the Ames CBF paper
+
+        Args:
+            v (float): Car speed, in m/s
+
+        Returns:
+            float: Drag force on the car, in Newtons
+        """
+        return self.f0 + self.f1 * v + self.f2 * v**2
+
+    @property
+    def follow_distance(self):
+        """Distance between the leader and follower cars, in meters"""
+        return self.leader_pos - self.follower_pos - self.CAR_HEIGHT_M
+
+    def get_state(self):
+        return np.array(
+            [
+                self.follower_vel,
+                self.leader_vel,
+                self.follow_distance,
+            ]
+        )
+
+    def get_desired_state(self):
+        leader_speed_kmh = 3.6 * self.leader_vel
+        safe_follow_distance = leader_speed_kmh / 2 + self.MIN_FOLLOW_DIST_M
+        # Note: This desired state is primarily for the simple nominal controller and NOT the CLF-CBF
+        return np.array(
+            [
+                self.v_des,
+                self.leader_vel,
+                safe_follow_distance,
+            ]
+        )
+
+    def apply_control(self, u) -> None:
+        u = np.clip(u, self.u_min, self.u_max).item()
+        if np.isnan(u):
+            print("Infeasible. Using last safe control")
+            u = self.last_control
+        force = u - self.friction(self.follower_vel)
+        follower_accel = force / self.mass
+        self.follower_vel += follower_accel * self.dt
+        self.follower_pos += self.follower_vel * self.dt
+        self.last_control = u
+
+    def leader_controller(self) -> float:
+        """Simple controller for the leader car
+
+        The desired leader velocity will be set by the user, and this controller will try to track that velocity
+
+        Returns:
+            float: Control input to the leader car: Acceleration force, in Newtons
+        """
+        kv = 1000
+        return kv * (self.leader_vel_des - self.leader_vel)
+
+    def step(self):
+        # Handle events
+        # This includes where the speed of the main controlled by the user
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                self.running = False
+                return
+            elif event.type == pygame.KEYDOWN:
+                if event.key == pygame.K_UP:
+                    self.leader_vel_des += 1
+                elif event.key == pygame.K_DOWN:
+                    self.leader_vel_des -= 1
+                elif event.key == pygame.K_ESCAPE:
+                    self.running = False
+                    return
+
+        leader_u = self.leader_controller()
+        leader_force = leader_u - self.friction(self.leader_vel)
+        # Assume both vehicles have the same mass
+        leader_accel = leader_force / self.mass
+        self.leader_vel += leader_accel * self.dt
+        self.leader_pos += self.leader_vel * self.dt
+        follow_dist = self.follow_distance
+
+        # Update locations of the cars in pixel frame
+        self.follower_y = (
+            self.leader_y + self.meters_to_pixels(follow_dist) + self.CAR_HEIGHT_PX
+        )
+
+        # Clear the screen
+        self.screen.fill(VehicleEnv.WHITE)
+
+        # Draw the road
+        pygame.draw.rect(
+            self.screen,
+            VehicleEnv.GRAY,
+            (
+                0,
+                0,
+                VehicleEnv.SCREEN_WIDTH // 2,
+                VehicleEnv.SCREEN_HEIGHT,
+            ),
+        )
+
+        # Draw dashed lines on the road
+        dash_y = self.dash_offset
+        while dash_y < VehicleEnv.SCREEN_HEIGHT:
+            pygame.draw.rect(
+                self.screen,
+                VehicleEnv.YELLOW,
+                (
+                    VehicleEnv.SCREEN_WIDTH // 4 - VehicleEnv.DASH_WIDTH_PX // 2,
+                    dash_y,
+                    VehicleEnv.DASH_WIDTH_PX,
+                    VehicleEnv.DASH_HEIGHT_PX,
+                ),
+            )
+            dash_y += VehicleEnv.DASH_HEIGHT_PX + VehicleEnv.DASH_GAP_PX
+
+        # Move the dashed lines
+        self.dash_offset += self.leader_vel
+        if self.dash_offset >= VehicleEnv.DASH_HEIGHT_PX + VehicleEnv.DASH_GAP_PX:
+            self.dash_offset = 0
+
+        # Draw the cars
+        self.screen.blit(self.leader_sprite, (self.leader_x, self.leader_y))
+        self.screen.blit(self.follower_sprite, (self.follower_x, self.follower_y))
+
+        # Print info to the pygame window
+        info_1 = f"Leader desired speed: {self.leader_vel_des:.2f}"
+        info_2 = f"Leader speed: {self.leader_vel:.2f}"
+        info_3 = f"Follower speed: {self.follower_vel:.2f}"
+        info_4 = f"Follow distance: {follow_dist:.2f}"
+        info_5 = f"Control: {self.last_control:.2f}"
+        text = [info_1, info_2, info_3, info_4, info_5]
+        for i, line in enumerate(text):
+            self.screen.blit(
+                self.font.render(line, True, VehicleEnv.BLACK),
+                (VehicleEnv.SCREEN_WIDTH // 2 + 10, 10 + 30 * i),
+            )
+
+        # Update the display
+        pygame.display.flip()
+
+        # Cap the frame rate
+        pygame.time.Clock().tick(self.fps)
+
+
+def _test_env():
+    """Test the environment behavior under an **unsafe** nominal controller"""
+
+    def nominal_controller(z, z_des):
+        vf, vl, D = z
+        vf_des, vl_des, D_des = z_des
+        kp = 1000
+        kv = 1000
+        return kp * (D - D_des) + kv * (vf_des - vf)
+
+    env = VehicleEnv("Nominal Controller")
+    print("\nDemoing the vehicle environment with an unsafe controller")
+    while env.running:
+        z = env.get_state()
+        z_des = env.get_desired_state()
+        u = nominal_controller(z, z_des)
+        env.apply_control(u)
+        env.step()
+
+
+if __name__ == "__main__":
+    _test_env()

cbfpy/envs/drone_env.py

@@ -0,0 +1,153 @@
+"""
+# Drone Environment 
+
+This is a wrapper around the gym-pybullet-drones environment, using velocity control.
+"""
+
+import time
+import warnings
+
+import numpy as np
+import pybullet
+from pybullet_utils.bullet_client import BulletClient
+from jax.typing import ArrayLike
+from jax import Array
+
+# Note: the original gym_pybullet_drones repo has a lot of dependencies that are not necessary for this demo.
+# Use the fork at https://github.com/danielpmorton/gym-pybullet-drones instead
+try:
+    from gym_pybullet_drones.envs.VelocityAviary import VelocityAviary
+except ImportError as e:
+    raise ImportError(
+        "Please install the forked version of gym-pybullet-drones:"
+        + "\n'pip install 'gym_pybullet_drones @ git+https://github.com/danielpmorton/gym-pybullet-drones.git''"
+    ) from e
+
+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 DroneEnv(BaseEnv):
+    """Drone Environment class
+
+    This provides an environment where the drone is contained to lie within a safe box,
+    and must avoid a movable obstacle.
+
+    Args:
+        xyz_min (ArrayLike): Minimum bounds of the safe region, shape (3,)
+        xyz_max (ArrayLike): Maximum bounds of the safe region, shape (3,)
+    """
+
+    # Constants
+    RADIUS = 0.1  # TODO tune
+
+    def __init__(
+        self,
+        xyz_min: ArrayLike = (-0.5, -0.5, 0.5),
+        xyz_max: ArrayLike = (0.5, 0.5, 1.5),
+    ):
+        # Suppress pybullet output + a Gym warning about float32 precision
+        with stdout_redirected(restore=False):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                self._env = VelocityAviary(
+                    # drone_model="cf2x",
+                    num_drones=1,
+                    initial_xyzs=np.array([[0, 0, 1]]),
+                    initial_rpys=np.array([[0, 0, 0]]),
+                    # physics="pyb",
+                    neighbourhood_radius=np.inf,
+                    pyb_freq=240,
+                    ctrl_freq=48,
+                    gui=True,
+                    record=False,
+                    obstacles=False,
+                    user_debug_gui=False,
+                )
+            # Hack: Create same client interface as other envs
+            self.client: pybullet = object.__new__(BulletClient)
+            self.client._client = self._env.CLIENT
+        self.xyz_min = np.array(xyz_min)
+        self.xyz_max = np.array(xyz_max)
+        assert len(self.xyz_min) == len(self.xyz_max) == 3
+        self.client.setAdditionalSearchPath(find_assets_dir())
+        self.obstacle = self.client.loadURDF("point_robot.urdf", basePosition=(1, 1, 1))
+        self.client.configureDebugVisualizer(self.client.COV_ENABLE_GUI, 0)
+        self.client.resetDebugVisualizerCamera(1.80, 37.60, -25.00, (0.05, 0.03, 0.75))
+        self.robot = self._env.DRONE_IDS[0]
+        self.client.changeVisualShape(self.obstacle, -1, rgbaColor=[1, 0, 0, 1])
+        self.client.changeDynamics(self.obstacle, -1, 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
+
+        self.action = np.array([[0.0, 0.0, 0.0, 0.0]])
+        self.obs, self.reward, self.terminated, self.truncated, self.info = (
+            self._env.step(self.action)
+        )
+
+    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)
+        obstacle_pos = self.client.getBasePositionAndOrientation(self.obstacle)[0]
+        obstacle_vel = self.client.getBaseVelocity(self.obstacle)[0]
+        return np.array([*robot_pos, *robot_vel]), np.array(
+            [*obstacle_pos, *obstacle_vel]
+        )
+
+    def get_desired_state(self) -> Array:
+        return np.array([0, 0, 1, 0, 0, 0])
+
+    def apply_control(self, u: Array) -> None:
+        # Note: the gym-pybullet-drones API has a weird format for the "velocity action" to take
+        self.action = np.array([[*u, np.linalg.norm(u)]])
+
+    def step(self):
+        # Step the gym-pybullet-drones environment using the last stored action
+        self.obs, self.reward, self.terminated, self.truncated, self.info = (
+            self._env.step(self.action)
+        )
+
+
+def _test_env():
+    """Test the environment behavior under an **unsafe** nominal controller"""
+
+    env = DroneEnv()
+
+    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:])
+
+    print("\nDemoing the drone environment with an unsafe controller")
+    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()
+        time.sleep(1 / 300)
+
+
+if __name__ == "__main__":
+    _test_env()