Merge pull request #1 from shaoanlu/refactoring

Refactor folder structure

README.md

@@ -4,12 +4,12 @@
 This Python program demonstrates the use of control barrier functions (CBFs) for safe robot control in a simple 2D environment. The program simulates a controllable green robot navigating through a scenario with stationary and patrolling robots. CBFs are employed to avoid collisions between the green robot and other robots in the environment. When safety filter is activated, it modified the user command to prevent collision while keeping the modification minimal.
 
 
-![](demo_safety_filter_off.gif) ![](demo_safety_filter_on.gif)
+![](assets/demo_safety_filter_off.gif) ![](assets/demo_safety_filter_on.gif)
 
 (Left: w/o safety filter; Right: w/ safety filter)
 
 ## Dependencies
-The program was develop[ed and tested in the following environment.
+The program was developed and tested in the following environment.
 - Python 3.6+
 - `torch==1.8.1+cpu` (optional)
 - `osqp==0.6.2.post8`
@@ -21,8 +21,17 @@ The program was develop[ed and tested in the following environment.
 ### Execution
 Run `demo.py` to start the simulation. A simple version w/ code generation is provided as `demo_codegen_osqp.py`.
 ```bash
-python demo.py  # or python demo_codegen_osqp.py
+python demo.py  # a PyGame window will be spawn
 ```
+
+Unit test can be executed through the following command.
+```bash
+python -m unittest
+# ----------------------------------------------------------------------
+# Ran 11 tests in 0.000s
+# OK
+```
+
 ### Controls
 - `Arrow keys`: Control the green robot's movement.
 - `X`: Toggle CBF ON/OFF.
@@ -34,8 +43,8 @@ python demo.py  # or python demo_codegen_osqp.py
 ## Program Overview
 The program consists of the following classes and files:
 
-- `Robot`: A class in `robot.py` representing a robot in the environment. It includes methods for controlling the robot and detecting collisions.
-- `SimpleRobotDynamics`: A class in `models.py` defining the dynamics of the robot. It includes methods for calculating state transition, control barrier functions, and the derivatives of control barrier funcitons.
+- `robot.py`: Implements the `Robot` class representing a robot in the environment. It includes methods for controlling the robot and detecting collisions.
+- `models.py`: Implements the `SimpleRobotDynamics` class defining the dynamics of the robot. It includes methods for calculating state transition, control barrier functions, and the derivatives of control barrier funcitons.
 - `demo.py`: The entry point of the program that handles user input, robot movement, collision detection, and rendering.
 - `utils.py`: This file contains a helper function to draw the robots on the screen.
 

controllers/i_controller.py

@@ -0,0 +1,14 @@
+from abc import abstractmethod
+
+
+class ControllerInterface:
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def control(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def detect_collision(self):
+        raise NotImplementedError

robot.py → controllers/robot_cbf.py

@@ -3,20 +3,22 @@
 import osqp
 from typing import Optional
 
-from models import SimpleRobotDynamics
+from models.robot_dynamics import SimpleRobotDynamics
+from controllers.i_controller import ControllerInterface
 
 
+# IDs for detecting arrow keypress in PyGame
 K_LEFT = 1073741904
 K_RIGHT = 1073741903
 K_UP = 1073741906
 K_DOWN = 1073741905
 
 
-class Robot:
+class RobotCBF(ControllerInterface):
     def __init__(
         self,
         model: SimpleRobotDynamics,
-        color: tuple = (0, 255, 0),
+        color: tuple = (0, 0, 255),
         vel: float = 3,
         size: int = 30,
     ):
@@ -25,8 +27,8 @@ def __init__(
         self.uy = 0  # actual control input
         self.nominal_ux = 0  # user control input
         self.nominal_uy = 0  # user control input
+        self.color = color  # color of the robot in PyGame GUI
         self.vel = vel  # robot velocity generated by use control input
-        self.color = color
         self.size = size
         self.is_collided = False
         self.prob = None
@@ -80,11 +82,11 @@ def control(
     def _apply_nominal_control(self):
         self.ux, self.uy = self.nominal_ux, self.nominal_uy
 
-    def _update_positions(self, ux, uy):
+    def _update_positions(self, ux: float, uy: float):
         model_control = np.array([ux, uy])
         self.model.forward(model_control)
 
-    def _update_nominal_control(self, key):
+    def _update_nominal_control(self, key: int):
         ux, uy = 0, 0
         if key is not None:
             if key == K_LEFT:
@@ -98,7 +100,13 @@ def _update_nominal_control(self, key):
         self.nominal_ux, self.nominal_uy = ux, uy
 
     def _apply_cbf_safe_control(
-        self, ux, uy, cbf_alpha, penalty_slack, force_direction_unchanged, collision_objects
+        self,
+        ux: float,
+        uy: float,
+        cbf_alpha: float,
+        penalty_slack: float,
+        force_direction_unchanged: bool,
+        collision_objects: list,
     ):
         """
         Calculate the safe command by solveing the following optimization problem
@@ -151,7 +159,7 @@ def _apply_cbf_safe_control(
 
         self.ux, self.uy = ux, uy
 
-    def _calculate_h_and_coeffs_dhdx(self, collision_objects):
+    def _calculate_h_and_coeffs_dhdx(self, collision_objects: list):
         h = []  # barrier values (here, remaining distance to each obstacle)
         coeffs_dhdx = []  # dhdt = dhdx * dxdt = dhdx * u
         for obj in collision_objects:
@@ -167,7 +175,14 @@ def _calculate_h_and_coeffs_dhdx(self, collision_objects):
         return h, coeffs_dhdx
 
     def _define_QP_problem_data(
-        self, ux, uy, cbf_alpha, penalty_slack, h, coeffs_dhdx, force_direction_unchanged
+        self,
+        ux: float,
+        uy: float,
+        cbf_alpha: float,
+        penalty_slack: float,
+        h: np.ndarray,
+        coeffs_dhdx: np.ndarray,
+        force_direction_unchanged: bool,
     ):
         # P: shape (nx, nx)
         # q: shape (nx,)

demo.py

@@ -2,8 +2,8 @@
 import pygame
 import numpy as np
 
-from models import SimpleRobotDynamics
-from robot import Robot
+from models.robot_dynamics import SimpleRobotDynamics
+from controllers.robot_cbf import RobotCBF
 from utils import draw_robot
 
 
@@ -16,10 +16,10 @@ def run():
     pygame.key.set_repeat(10)
 
     # init robots
-    auto_robot = Robot(SimpleRobotDynamics(x0=np.array([50, 350])), (0, 255, 0), vel=3)
-    static_robot = Robot(SimpleRobotDynamics(x0=np.array([120, 200])), (0, 0, 255), vel=0)
-    patrol_robot1 = Robot(SimpleRobotDynamics(x0=np.array([230, 300])), (0, 0, 255), vel=1)
-    patrol_robot2 = Robot(SimpleRobotDynamics(x0=np.array([300, 70])), (0, 0, 255), vel=1)
+    ego_robot = RobotCBF(SimpleRobotDynamics(x0=np.array([50, 350])), (0, 255, 0), vel=3)
+    static_robot = RobotCBF(SimpleRobotDynamics(x0=np.array([120, 200])), (0, 0, 255), vel=0)
+    patrol_robot1 = RobotCBF(SimpleRobotDynamics(x0=np.array([230, 300])), (0, 0, 255), vel=1)
+    patrol_robot2 = RobotCBF(SimpleRobotDynamics(x0=np.array([300, 70])), (0, 0, 255), vel=1)
     collision_objects = [static_robot, patrol_robot1, patrol_robot2]
 
     # init control configs
@@ -71,7 +71,7 @@ def run():
                     cbf_alphas.append(cbf_alphas.pop(0))  # roll left
                 if e.key == pygame.K_r:
                     # reset
-                    auto_robot.x, auto_robot.y = 50, 350
+                    ego_robot.x, ego_robot.y = 50, 350
                     static_robot.x, static_robot.y = 120, 200
                     patrol_robot1.x, patrol_robot1.y = 230, 300
                     patrol_robot2.x, patrol_robot2.y = 300, 70
@@ -81,7 +81,7 @@ def run():
 
         # move robots
         static_robot.control(None)
-        auto_robot.control(
+        ego_robot.control(
             pressed_key,
             use_cbf=use_cbf,
             cbf_alpha=cbf_alphas[0],
@@ -91,32 +91,32 @@ def run():
         patrol_robot1.control(
             direction_patrol_robot1,
             use_cbf=use_cbf_patrol_robots,
-            collision_objects=[auto_robot, static_robot, patrol_robot2]
+            collision_objects=[ego_robot, static_robot, patrol_robot2]
             if use_cbf_patrol_robots
             else [],
             force_direction_unchanged=cbf_force_direction_unchanged,
         )
         patrol_robot2.control(
             direction_patrol_robot2,
             use_cbf=use_cbf_patrol_robots,
-            collision_objects=[auto_robot, static_robot, patrol_robot1]
+            collision_objects=[ego_robot, static_robot, patrol_robot1]
             if use_cbf_patrol_robots
             else [],
             force_direction_unchanged=cbf_force_direction_unchanged,
         )
 
         # detect collision
-        auto_robot.detect_collision(collision_objects=collision_objects)
+        ego_robot.detect_collision(collision_objects=collision_objects)
 
         # draw robots
         screen.fill((0, 0, 0))
         draw_robot(screen, static_robot)
         draw_robot(screen, patrol_robot1, draw_filtered_command=use_cbf_patrol_robots)
         draw_robot(screen, patrol_robot2, draw_filtered_command=use_cbf_patrol_robots)
-        draw_robot(screen, auto_robot, draw_filtered_command=use_cbf)
+        draw_robot(screen, ego_robot, draw_filtered_command=use_cbf)
 
         # draw texts
-        if auto_robot.is_collided:
+        if ego_robot.is_collided:
             screen.blit(text_surface, (230, 350))
         if use_cbf:
             screen.blit(cbf_on_text_surface, (0, 0))