Add planning.py

fsd_model/planning.py

@@ -0,0 +1,335 @@
+"""
+Planning Module for FSD Model.
+Handles:
+1. Route Planning (high-level waypoints from navigation)
+2. Behavior Planning (lane changes, turns, stops, yields)
+3. Trajectory Planning (smooth, collision-free path generation)
+4. Safety Verification (collision checking, emergency braking)
+
+Architecture: Transformer-based planner that attends to perception features
+and produces waypoint trajectories. Inspired by UniAD and VAD planners.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Optional, Tuple
+import math
+
+
+class PositionalEncoding2D(nn.Module):
+    """2D sinusoidal positional encoding for BEV features."""
+    def __init__(self, channels: int):
+        super().__init__()
+        self.channels = channels
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = x.shape
+        device = x.device
+        
+        y_pos = torch.arange(H, device=device).float().unsqueeze(1).expand(H, W) / H
+        x_pos = torch.arange(W, device=device).float().unsqueeze(0).expand(H, W) / W
+        
+        dim = torch.arange(0, self.channels // 4, device=device).float()
+        dim = 10000 ** (2 * dim / (self.channels // 2))
+        
+        pe = torch.zeros(self.channels, H, W, device=device)
+        quarter = self.channels // 4
+        pe[0:quarter] = torch.sin(x_pos.unsqueeze(0) / dim.unsqueeze(1).unsqueeze(2))
+        pe[quarter:2*quarter] = torch.cos(x_pos.unsqueeze(0) / dim.unsqueeze(1).unsqueeze(2))
+        pe[2*quarter:3*quarter] = torch.sin(y_pos.unsqueeze(0) / dim.unsqueeze(1).unsqueeze(2))
+        pe[3*quarter:4*quarter] = torch.cos(y_pos.unsqueeze(0) / dim.unsqueeze(1).unsqueeze(2))
+        
+        return x + pe.unsqueeze(0).expand(B, -1, -1, -1)
+
+
+class BehaviorPredictor(nn.Module):
+    """
+    Predicts high-level driving behavior/command.
+    Commands: keep_lane, turn_left, turn_right, lane_change_left, 
+              lane_change_right, stop, yield, park, reverse, emergency_stop
+    """
+    def __init__(self, in_channels: int = 256, num_behaviors: int = 10):
+        super().__init__()
+        self.num_behaviors = num_behaviors
+        
+        self.encoder = nn.Sequential(
+            nn.AdaptiveAvgPool2d(8),
+            nn.Flatten(),
+            nn.Linear(in_channels * 64, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(256, num_behaviors),
+        )
+    
+    def forward(self, bev: torch.Tensor) -> torch.Tensor:
+        """Returns: (B, num_behaviors) logits"""
+        return self.encoder(bev)
+
+
+class TrajectoryTransformer(nn.Module):
+    """
+    Transformer-based trajectory planner.
+    Generates waypoints by attending to BEV features and navigation commands.
+    Uses learnable trajectory queries (similar to DETR object queries).
+    """
+    def __init__(
+        self,
+        bev_channels: int = 256,
+        d_model: int = 256,
+        nhead: int = 8,
+        num_decoder_layers: int = 6,
+        num_waypoints: int = 20,  # planning horizon waypoints
+        dim_feedforward: int = 1024,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.num_waypoints = num_waypoints
+        self.d_model = d_model
+        
+        # BEV feature compression
+        self.bev_compress = nn.Sequential(
+            nn.Conv2d(bev_channels, d_model, 1),
+            nn.BatchNorm2d(d_model),
+            nn.ReLU(),
+        )
+        self.pos_encoding = PositionalEncoding2D(d_model)
+        
+        # Learnable trajectory queries
+        self.trajectory_queries = nn.Parameter(
+            torch.randn(num_waypoints, d_model)
+        )
+        
+        # Navigation command embedding (high-level route)
+        self.command_embed = nn.Embedding(10, d_model)  # 10 possible commands
+        
+        # Ego state embedding (speed, acceleration, steering)
+        self.ego_state_embed = nn.Sequential(
+            nn.Linear(6, d_model),  # speed, accel, steer, yaw_rate, x, y
+            nn.ReLU(),
+        )
+        
+        # Transformer decoder
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.transformer_decoder = nn.TransformerDecoder(
+            decoder_layer, num_layers=num_decoder_layers
+        )
+        
+        # Waypoint prediction heads
+        self.waypoint_head = nn.Sequential(
+            nn.Linear(d_model, 128),
+            nn.ReLU(),
+            nn.Linear(128, 4),  # (x, y, heading, speed)
+        )
+        
+        # Confidence / collision probability per waypoint
+        self.confidence_head = nn.Sequential(
+            nn.Linear(d_model, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1),
+            nn.Sigmoid(),
+        )
+    
+    def forward(
+        self,
+        bev_features: torch.Tensor,
+        ego_state: torch.Tensor,
+        nav_command: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            bev_features: (B, C, H, W) from perception
+            ego_state: (B, 6) current ego state [speed, accel, steer, yaw_rate, x, y]
+            nav_command: (B,) integer navigation command
+        Returns:
+            waypoints: (B, num_waypoints, 4) predicted trajectory
+            confidence: (B, num_waypoints, 1) per-waypoint confidence
+        """
+        B = bev_features.shape[0]
+        device = bev_features.device
+        
+        # Compress and add positional encoding to BEV
+        bev = self.bev_compress(bev_features)
+        bev = self.pos_encoding(bev)
+        
+        # Flatten BEV to sequence: (B, H*W, d_model)
+        bev_seq = bev.flatten(2).permute(0, 2, 1)
+        
+        # Build trajectory queries
+        queries = self.trajectory_queries.unsqueeze(0).expand(B, -1, -1)
+        
+        # Add ego state information to queries
+        ego_feat = self.ego_state_embed(ego_state).unsqueeze(1)
+        queries = queries + ego_feat
+        
+        # Add navigation command if provided
+        if nav_command is not None:
+            cmd_feat = self.command_embed(nav_command).unsqueeze(1)
+            queries = queries + cmd_feat
+        
+        # Transformer decoding
+        decoded = self.transformer_decoder(queries, bev_seq)
+        
+        # Predict waypoints and confidence
+        waypoints = self.waypoint_head(decoded)   # (B, T, 4)
+        confidence = self.confidence_head(decoded)  # (B, T, 1)
+        
+        return {
+            "waypoints": waypoints,
+            "confidence": confidence,
+        }
+
+
+class SafetyChecker(nn.Module):
+    """
+    Verifies planned trajectories against safety constraints.
+    Checks for:
+    - Collision with detected objects
+    - Lane boundary violations
+    - Speed limit violations
+    - Minimum following distance
+    - Emergency stop conditions
+    """
+    def __init__(
+        self,
+        bev_channels: int = 256,
+        max_speed_ms: float = 8.94,  # 20 mph
+        min_following_distance: float = 4.0,  # meters
+        emergency_decel: float = 8.0,  # m/s^2
+    ):
+        super().__init__()
+        self.max_speed_ms = max_speed_ms
+        self.min_following_distance = min_following_distance
+        self.emergency_decel = emergency_decel
+        
+        # Collision risk estimator
+        self.collision_net = nn.Sequential(
+            nn.AdaptiveAvgPool2d(8),
+            nn.Flatten(),
+            nn.Linear(bev_channels * 64, 256),
+            nn.ReLU(),
+            nn.Linear(256, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1),
+            nn.Sigmoid(),
+        )
+        
+        # Emergency brake detector
+        self.emergency_detector = nn.Sequential(
+            nn.AdaptiveAvgPool2d(4),
+            nn.Flatten(),
+            nn.Linear(bev_channels * 16, 128),
+            nn.ReLU(),
+            nn.Linear(128, 2),  # [no_emergency, emergency]
+        )
+    
+    def forward(
+        self,
+        bev: torch.Tensor,
+        planned_waypoints: torch.Tensor,
+        ego_state: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            bev: (B, C, H, W) BEV features with occupancy info
+            planned_waypoints: (B, T, 4) planned trajectory
+            ego_state: (B, 6)
+        Returns:
+            Dict with safety scores and emergency signals
+        """
+        # Collision risk
+        collision_risk = self.collision_net(bev)
+        
+        # Emergency brake
+        emergency_logits = self.emergency_detector(bev)
+        emergency_prob = F.softmax(emergency_logits, dim=-1)[:, 1:]
+        
+        # Speed constraint check
+        planned_speeds = planned_waypoints[:, :, 3]  # speed component
+        speed_violation = (planned_speeds > self.max_speed_ms).float().mean(dim=-1, keepdim=True)
+        
+        # Clamp speeds to max
+        clamped_waypoints = planned_waypoints.clone()
+        clamped_waypoints[:, :, 3] = torch.clamp(
+            planned_waypoints[:, :, 3], 0.0, self.max_speed_ms
+        )
+        
+        return {
+            "collision_risk": collision_risk,
+            "emergency_brake": emergency_prob,
+            "speed_violation": speed_violation,
+            "safe_waypoints": clamped_waypoints,
+        }
+
+
+class PlanningModule(nn.Module):
+    """
+    Complete planning module.
+    Pipeline: BEV → Behavior Prediction → Trajectory Generation → Safety Check
+    """
+    def __init__(
+        self,
+        bev_channels: int = 256,
+        d_model: int = 256,
+        num_waypoints: int = 20,
+        max_speed_ms: float = 8.94,
+        num_behaviors: int = 10,
+    ):
+        super().__init__()
+        
+        self.behavior_predictor = BehaviorPredictor(bev_channels, num_behaviors)
+        self.trajectory_planner = TrajectoryTransformer(
+            bev_channels=bev_channels,
+            d_model=d_model,
+            num_waypoints=num_waypoints,
+        )
+        self.safety_checker = SafetyChecker(
+            bev_channels=bev_channels,
+            max_speed_ms=max_speed_ms,
+        )
+    
+    def forward(
+        self,
+        bev_features: torch.Tensor,
+        ego_state: torch.Tensor,
+        nav_command: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            bev_features: (B, C, H, W)
+            ego_state: (B, 6) [speed, accel, steer, yaw_rate, x, y]
+            nav_command: (B,) high-level navigation command
+        Returns:
+            Complete planning output including safe trajectory
+        """
+        # Predict behavior
+        behavior_logits = self.behavior_predictor(bev_features)
+        
+        # Generate trajectory
+        traj_output = self.trajectory_planner(
+            bev_features, ego_state, nav_command
+        )
+        
+        # Safety verification
+        safety = self.safety_checker(
+            bev_features, traj_output["waypoints"], ego_state
+        )
+        
+        return {
+            "behavior_logits": behavior_logits,
+            "raw_waypoints": traj_output["waypoints"],
+            "waypoint_confidence": traj_output["confidence"],
+            "safe_waypoints": safety["safe_waypoints"],
+            "collision_risk": safety["collision_risk"],
+            "emergency_brake": safety["emergency_brake"],
+            "speed_violation": safety["speed_violation"],
+        }