type hints PEP484 style and code formatting

Author: sidtalia

examples/BeamNG/IGHAStarMP.py

@@ -3,13 +3,15 @@
 import time
 import traceback
 from queue import Empty
+from typing import Any, Optional, Tuple, Dict
+
 
 class IGHAStarMP:
     """
     Multiprocessing wrapper for the IGHA* planner. Runs the planner in a separate process and communicates via queues.
     """
 
-    def __init__(self, configs):
+    def __init__(self, configs: Dict[str, Any]) -> None:
         mp.set_start_method("spawn", force=True)  # Safe for CUDA
         self.query_queue = mp.Queue(5)
         self.result_queue = mp.Queue(5)
@@ -23,7 +25,9 @@ def __init__(self, configs):
         self.completed = True
         self.expansion_counter = 0
 
-    def _planner_process(self, configs, query_queue, result_queue):
+    def _planner_process(
+        self, configs: Dict[str, Any], query_queue: Any, result_queue: Any
+    ) -> None:
         """
         Planner process: loads the CUDA/C++ kernel and runs the IGHA* planner in response to queries.
         """
@@ -113,16 +117,16 @@ def _planner_process(self, configs, query_queue, result_queue):
 
     def set_query(
         self,
-        map_center,
-        start_state,
-        goal_,
-        costmap,
-        heightmap,
-        hysteresis,
-        expansion_limit,
-        stop=False,
-        disable=False,
-    ):
+        map_center: np.ndarray,
+        start_state: np.ndarray,
+        goal_: np.ndarray,
+        costmap: np.ndarray,
+        heightmap: np.ndarray,
+        hysteresis: float,
+        expansion_limit: int,
+        stop: bool = False,
+        disable: bool = False,
+    ) -> None:
         """
         Submit a new planning query. Returns immediately; results are available via update().
         """
@@ -148,14 +152,14 @@ def set_query(
         self.completed = False
         self.success = False
 
-    def reset(self):
+    def reset(self) -> None:
         """Reset planner state."""
         self.path = None
         self.success = False
         self.completed = True
         self.expansion_counter = 0
 
-    def update(self):
+    def update(self) -> Tuple[bool, Optional[np.ndarray], int]:
         """
         Call this periodically in your main loop to check for planner results.
         Returns (success, path, expansion_counter).
@@ -176,7 +180,7 @@ def update(self):
         except Empty:
             return False, self.path, self.expansion_counter
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         """Shut down the planner process."""
         self.query_queue.put(None)
         self.process.terminate()

examples/BeamNG/TrackingCost.py

@@ -1,4 +1,5 @@
 import torch
+from typing import Any, Optional, Dict
 
 
 class SimpleCarCost(torch.nn.Module):
@@ -8,11 +9,11 @@ class SimpleCarCost(torch.nn.Module):
 
     def __init__(
         self,
-        Cost_config,
-        Map_config,
-        dtype=torch.float32,
-        device=torch.device("cuda"),
-    ):
+        Cost_config: Dict[str, Any],
+        Map_config: Dict[str, Any],
+        dtype: torch.dtype = torch.float32,
+        device: torch.device = torch.device("cuda"),
+    ) -> None:
         super(SimpleCarCost, self).__init__()
         self.dtype = dtype
         self.d = device
@@ -67,7 +68,12 @@ def __init__(
         self.constraint_violation = False
 
     @torch.jit.export
-    def set_BEV(self, BEVmap_height, BEVmap_normal, BEVmap_cost):
+    def set_BEV(
+        self,
+        BEVmap_height: torch.Tensor,
+        BEVmap_normal: torch.Tensor,
+        BEVmap_cost: torch.Tensor,
+    ) -> None:
         """
         Set BEV (bird's-eye view) map data for cost calculation.
         """
@@ -76,25 +82,25 @@ def set_BEV(self, BEVmap_height, BEVmap_normal, BEVmap_cost):
         self.BEVmap_cost = (255 - BEVmap_cost) / 255
 
     @torch.jit.export
-    def set_goal(self, goal_state):
+    def set_goal(self, goal_state: torch.Tensor) -> None:
         self.goal_state = goal_state[:2]
 
-    def set_path(self, path):
+    def set_path(self, path: torch.Tensor) -> None:
         self.path = torch.tensor(path, dtype=self.dtype, device=self.d)
 
     @torch.jit.export
-    def set_speed_limit(self, speed_lim):
+    def set_speed_limit(self, speed_lim: float) -> None:
         self.speed_target = torch.tensor(speed_lim, dtype=self.dtype, device=self.d)
 
-    def meters_to_px(self, meters):
+    def meters_to_px(self, meters: torch.Tensor) -> torch.Tensor:
         px = ((meters + self.BEVmap_size * 0.5) / self.BEVmap_res).to(
             dtype=torch.long, device=self.d
         )
         px = torch.maximum(px, torch.zeros_like(px))
         px = torch.minimum(px, self.BEVmap_size_px - 1)
         return px
 
-    def forward(self, state, controls):
+    def forward(self, state: torch.Tensor, controls: torch.Tensor) -> torch.Tensor:
         # Unpack state
         x = state[..., 0]
         y = state[..., 1]

examples/BeamNG/example.py

@@ -11,6 +11,8 @@
 import argparse
 import sys
 import pathlib
+import os
+from typing import Any, Optional, Dict, List
 
 BASE_DIR = pathlib.Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(BASE_DIR / "scripts"))
@@ -28,12 +30,12 @@
 
 
 def main(
-    config_path=None,
-    hal_config_path=None,
-    waypoint_folder=None,
-    output_folder=None,
-    args=None,
-):
+    config_path: Optional[str] = None,
+    hal_config_path: Optional[str] = None,
+    waypoint_folder: Optional[str] = None,
+    output_folder: Optional[str] = None,
+    args: Optional[Any] = None,
+) -> None:
     if config_path is None:
         print("no config file provided!")
         exit()

examples/BeamNG/utils.py

@@ -3,11 +3,14 @@
 import time
 import threading
 import torch
+from typing import Any, Optional, List, Tuple, Dict
 
 
 def generate_costmap_from_BEVmap(
-    BEV_lethal, BEV_normal, costmap_cosine_thresh=np.cos(np.radians(60))
-):
+    BEV_lethal: np.ndarray,
+    BEV_normal: np.ndarray,
+    costmap_cosine_thresh: float = np.cos(np.radians(60)),
+) -> np.ndarray:
     """Generate a costmap from BEV lethal and normal maps."""
     dot_product = BEV_normal[:, :, 2]
     angle_cost = np.where(dot_product >= costmap_cosine_thresh, 255, 0).astype(
@@ -20,7 +23,11 @@ def generate_costmap_from_BEVmap(
     return costmap
 
 
-def convert_global_path_to_bng(bng_interface=None, path=None, Map_config=None):
+def convert_global_path_to_bng(
+    bng_interface: Optional[Any] = None,
+    path: Optional[np.ndarray] = None,
+    Map_config: Optional[Dict[str, Any]] = None,
+) -> np.ndarray:
     """Convert a global path to BeamNG waypoints with elevation and quaternion."""
     target_wp = []
     map_res = Map_config["map_res"]
@@ -44,7 +51,14 @@ def convert_global_path_to_bng(bng_interface=None, path=None, Map_config=None):
     return np.array(target_wp)
 
 
-def update_goal(goal, pos, target_WP, current_wp_index, lookahead, wp_radius=1.0):
+def update_goal(
+    goal: Optional[np.ndarray],
+    pos: np.ndarray,
+    target_WP: np.ndarray,
+    current_wp_index: int,
+    lookahead: float,
+    wp_radius: float = 1.0,
+) -> Tuple[np.ndarray, bool, int]:
     """Update the goal position based on lookahead and proximity to final waypoint."""
     final_wp = target_WP[-1]
     dx = final_wp[0] - pos[0]
@@ -65,7 +79,9 @@ def update_goal(goal, pos, target_WP, current_wp_index, lookahead, wp_radius=1.0
     return goal, success, current_wp_index
 
 
-def steering_limiter(steer, state, RPS_config):
+def steering_limiter(
+    steer: float, state: np.ndarray, RPS_config: Dict[str, Any]
+) -> float:
     """Limit steering to prevent rollovers and respect physical constraints."""
     steering_setpoint = steer * RPS_config["steering_max"]
     whspd2 = max(1.0, np.linalg.norm(state[6:8])) ** 2  # speed squared in world frame
@@ -125,7 +141,7 @@ def steering_limiter(steer, state, RPS_config):
 
 
 class PlannerVis:
-    def __init__(self, map_size, resolution_inv):
+    def __init__(self, map_size: int, resolution_inv: float) -> None:
         self.map_size = map_size
         self.resolution_inv = resolution_inv
         self.costmap = None
@@ -140,15 +156,15 @@ def __init__(self, map_size, resolution_inv):
 
     def update_vis(
         self,
-        states,
-        path,
-        costmap,
-        elevation_map,
-        resolution_inv,
-        goal,
-        expansion_counter,
-        hysteresis,
-    ):
+        states: np.ndarray,
+        path: Optional[np.ndarray],
+        costmap: np.ndarray,
+        elevation_map: np.ndarray,
+        resolution_inv: float,
+        goal: np.ndarray,
+        expansion_counter: int,
+        hysteresis: float,
+    ) -> None:
         with self.lock:
             if isinstance(states, torch.Tensor):
                 self.states = states.cpu().numpy()
@@ -162,12 +178,12 @@ def update_vis(
             self.hysteresis = hysteresis
             self.expansion_counter = expansion_counter
 
-    def generate_costmap_from_BEVmap(self, normal):
+    def generate_costmap_from_BEVmap(self, normal: np.ndarray) -> np.ndarray:
         dot_product = normal[:, :, 2]
         costmap = np.where(dot_product >= self.cosine_thresh, 255, 0).astype(np.float32)
         return costmap
 
-    def costmap_vis(self):
+    def costmap_vis(self) -> None:
         while True:
             if (
                 self.states is not None