Add control.py

fsd_model/control.py

@@ -0,0 +1,396 @@
+"""
+Control Module for FSD Model.
+Converts planned trajectory waypoints into actuator commands:
+- Steering angle
+- Throttle (acceleration)
+- Brake
+- Gear (forward/reverse/park)
+
+Uses a combination of:
+1. PID controllers for smooth tracking
+2. Neural network for adaptive control
+3. Stanley controller for lateral control
+4. Bicycle model for vehicle dynamics
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Tuple
+import math
+
+
+class BicycleModel(nn.Module):
+    """
+    Kinematic bicycle model for vehicle dynamics simulation.
+    Used for both prediction and control.
+    State: [x, y, heading, speed]
+    Control: [steering_angle, acceleration]
+    """
+    def __init__(self, wheelbase: float = 2.7, dt: float = 0.1):
+        super().__init__()
+        self.wheelbase = wheelbase
+        self.dt = dt
+    
+    def forward(
+        self, state: torch.Tensor, control: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            state: (B, 4) [x, y, heading, speed]
+            control: (B, 2) [steering_angle, acceleration]
+        Returns:
+            next_state: (B, 4)
+        """
+        x, y, heading, speed = state[:, 0], state[:, 1], state[:, 2], state[:, 3]
+        steer, accel = control[:, 0], control[:, 1]
+        
+        # Kinematic bicycle model equations
+        beta = torch.atan(0.5 * torch.tan(steer))  # slip angle
+        
+        x_new = x + speed * torch.cos(heading + beta) * self.dt
+        y_new = y + speed * torch.sin(heading + beta) * self.dt
+        heading_new = heading + (speed / self.wheelbase) * torch.sin(beta) * self.dt
+        speed_new = speed + accel * self.dt
+        
+        # Clamp speed to be non-negative
+        speed_new = torch.clamp(speed_new, min=0.0)
+        
+        return torch.stack([x_new, y_new, heading_new, speed_new], dim=-1)
+
+
+class PIDController(nn.Module):
+    """
+    Learnable PID controller with neural network gain scheduling.
+    Gains (Kp, Ki, Kd) are predicted based on current state.
+    """
+    def __init__(self, state_dim: int = 6, hidden_dim: int = 64):
+        super().__init__()
+        
+        # Gain predictor network
+        self.gain_net = nn.Sequential(
+            nn.Linear(state_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 6),  # Kp, Ki, Kd for lateral + longitudinal
+            nn.Softplus(),  # Ensure positive gains
+        )
+        
+        # Integral buffer (not a parameter, reset per episode)
+        self.register_buffer('integral_error', torch.zeros(1, 2))
+        self.register_buffer('prev_error', torch.zeros(1, 2))
+    
+    def forward(
+        self,
+        error: torch.Tensor,
+        ego_state: torch.Tensor,
+        dt: float = 0.1,
+    ) -> torch.Tensor:
+        """
+        Args:
+            error: (B, 2) [lateral_error, longitudinal_error]
+            ego_state: (B, 6) current vehicle state
+            dt: time step
+        Returns:
+            control: (B, 2) [steering_correction, accel_correction]
+        """
+        B = error.shape[0]
+        
+        # Predict adaptive gains
+        gains = self.gain_net(ego_state)
+        kp = gains[:, :2]
+        ki = gains[:, 2:4]
+        kd = gains[:, 4:6]
+        
+        # PID computation
+        proportional = kp * error
+        
+        # Integral (with anti-windup) — handle variable batch sizes
+        if self.integral_error.shape[0] != B:
+            self.integral_error = torch.zeros(B, 2, device=error.device)
+        if self.prev_error.shape[0] != B:
+            self.prev_error = torch.zeros(B, 2, device=error.device)
+        
+        self.integral_error = self.integral_error + error * dt
+        self.integral_error = torch.clamp(self.integral_error, -10.0, 10.0)
+        integral = ki * self.integral_error
+        
+        # Derivative
+        derivative = kd * (error - self.prev_error) / dt
+        self.prev_error = error.detach()
+        
+        control = proportional + integral + derivative
+        
+        return control
+    
+    def reset(self):
+        """Reset integral and derivative buffers."""
+        self.integral_error.zero_()
+        self.prev_error.zero_()
+
+
+class StanleyController(nn.Module):
+    """
+    Stanley lateral controller enhanced with learned parameters.
+    Computes steering angle based on:
+    1. Heading error
+    2. Cross-track error
+    """
+    def __init__(self, k_gain: float = 0.5, k_soft: float = 1.0):
+        super().__init__()
+        # Learnable gains
+        self.k_gain = nn.Parameter(torch.tensor(k_gain))
+        self.k_soft = nn.Parameter(torch.tensor(k_soft))
+    
+    def forward(
+        self,
+        heading_error: torch.Tensor,
+        cross_track_error: torch.Tensor,
+        speed: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            heading_error: (B,) heading difference to path
+            cross_track_error: (B,) lateral distance to path
+            speed: (B,) current speed
+        Returns:
+            steering: (B,) desired steering angle (radians)
+        """
+        # Stanley formula
+        cross_track_steer = torch.atan2(
+            self.k_gain * cross_track_error,
+            speed + self.k_soft
+        )
+        steering = heading_error + cross_track_steer
+        
+        # Clamp to max steering angle (~35 degrees)
+        max_steer = math.radians(35)
+        steering = torch.clamp(steering, -max_steer, max_steer)
+        
+        return steering
+
+
+class NeuralController(nn.Module):
+    """
+    End-to-end neural network controller.
+    Takes BEV features + ego state + waypoints and directly outputs
+    steering, throttle, brake commands.
+    Serves as a refinement on top of classical controllers.
+    """
+    def __init__(
+        self,
+        bev_channels: int = 256,
+        waypoint_dim: int = 4,
+        num_waypoints: int = 20,
+        ego_dim: int = 6,
+        hidden_dim: int = 256,
+    ):
+        super().__init__()
+        
+        # BEV feature compression
+        self.bev_encoder = nn.Sequential(
+            nn.AdaptiveAvgPool2d(4),
+            nn.Flatten(),
+            nn.Linear(bev_channels * 16, hidden_dim),
+            nn.ReLU(),
+        )
+        
+        # Waypoint encoder
+        self.waypoint_encoder = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(num_waypoints * waypoint_dim, hidden_dim),
+            nn.ReLU(),
+        )
+        
+        # Ego state encoder
+        self.ego_encoder = nn.Sequential(
+            nn.Linear(ego_dim, hidden_dim // 2),
+            nn.ReLU(),
+        )
+        
+        # Control output
+        self.control_head = nn.Sequential(
+            nn.Linear(hidden_dim * 2 + hidden_dim // 2, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(hidden_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, 3),  # steering, throttle, brake
+        )
+    
+    def forward(
+        self,
+        bev_features: torch.Tensor,
+        waypoints: torch.Tensor,
+        ego_state: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Returns:
+            Dict with steering (-1 to 1), throttle (0 to 1), brake (0 to 1)
+        """
+        bev_feat = self.bev_encoder(bev_features)
+        wp_feat = self.waypoint_encoder(waypoints)
+        ego_feat = self.ego_encoder(ego_state)
+        
+        combined = torch.cat([bev_feat, wp_feat, ego_feat], dim=-1)
+        raw = self.control_head(combined)
+        
+        steering = torch.tanh(raw[:, 0])       # [-1, 1]
+        throttle = torch.sigmoid(raw[:, 1])     # [0, 1]
+        brake = torch.sigmoid(raw[:, 2])        # [0, 1]
+        
+        return {
+            "steering": steering,
+            "throttle": throttle,
+            "brake": brake,
+        }
+
+
+class ControlModule(nn.Module):
+    """
+    Complete control module that combines:
+    1. Neural controller (BEV-aware, end-to-end)
+    2. Stanley controller (geometric lateral control)
+    3. PID controller (error-based correction)
+    4. Bicycle model (physics-based prediction)
+    5. Safety limits enforcement
+    """
+    def __init__(
+        self,
+        bev_channels: int = 256,
+        num_waypoints: int = 20,
+        wheelbase: float = 2.7,
+        max_speed_ms: float = 8.94,
+        max_steering_deg: float = 35.0,
+        max_accel: float = 3.0,
+        max_decel: float = 8.0,
+        dt: float = 0.1,
+    ):
+        super().__init__()
+        self.max_speed_ms = max_speed_ms
+        self.max_steering = math.radians(max_steering_deg)
+        self.max_accel = max_accel
+        self.max_decel = max_decel
+        self.dt = dt
+        
+        # Sub-controllers
+        self.neural_controller = NeuralController(
+            bev_channels=bev_channels,
+            num_waypoints=num_waypoints,
+        )
+        self.stanley_controller = StanleyController()
+        self.pid_controller = PIDController()
+        self.bicycle_model = BicycleModel(wheelbase, dt)
+        
+        # Controller fusion weights (learned)
+        self.fusion_weights = nn.Sequential(
+            nn.Linear(6, 32),  # ego state -> weights
+            nn.ReLU(),
+            nn.Linear(32, 3),  # weights for [neural, stanley, pid]
+            nn.Softmax(dim=-1),
+        )
+    
+    def forward(
+        self,
+        bev_features: torch.Tensor,
+        planned_waypoints: torch.Tensor,
+        ego_state: torch.Tensor,
+        emergency_brake: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            bev_features: (B, C, H, W) BEV features
+            planned_waypoints: (B, T, 4) [x, y, heading, speed]
+            ego_state: (B, 6) [speed, accel, steer, yaw_rate, x, y]
+            emergency_brake: (B, 1) emergency brake probability
+        Returns:
+            Dict with final actuator commands
+        """
+        B = ego_state.shape[0]
+        device = ego_state.device
+        
+        # 1. Neural controller output
+        neural_out = self.neural_controller(bev_features, planned_waypoints, ego_state)
+        
+        # 2. Stanley controller - compute from first waypoint error
+        next_wp = planned_waypoints[:, 0, :]  # next waypoint
+        heading_error = next_wp[:, 2] - ego_state[:, 3]  # yaw_rate as proxy
+        cross_track_error = torch.sqrt(
+            (next_wp[:, 0] - ego_state[:, 4])**2 + 
+            (next_wp[:, 1] - ego_state[:, 5])**2
+        )
+        stanley_steer = self.stanley_controller(
+            heading_error, cross_track_error, ego_state[:, 0]
+        )
+        
+        # 3. PID controller
+        lateral_err = cross_track_error
+        speed_err = next_wp[:, 3] - ego_state[:, 0]
+        pid_error = torch.stack([lateral_err, speed_err], dim=-1)
+        pid_out = self.pid_controller(pid_error, ego_state, self.dt)
+        
+        # 4. Fuse controllers based on driving state
+        weights = self.fusion_weights(ego_state)  # (B, 3)
+        
+        # Neural steering + Stanley steering + PID steering
+        neural_steer = neural_out["steering"] * self.max_steering
+        final_steering = (
+            weights[:, 0] * neural_steer +
+            weights[:, 1] * stanley_steer +
+            weights[:, 2] * torch.clamp(pid_out[:, 0], -self.max_steering, self.max_steering)
+        )
+        
+        # Throttle/brake from neural + PID
+        final_throttle = neural_out["throttle"]
+        final_brake = neural_out["brake"]
+        
+        # PID speed correction
+        pid_accel = pid_out[:, 1]
+        final_throttle = final_throttle + torch.clamp(pid_accel, 0, 1) * weights[:, 2]
+        final_brake = final_brake + torch.clamp(-pid_accel, 0, 1) * weights[:, 2]
+        
+        # 5. Safety overrides
+        if emergency_brake is not None:
+            emergency_mask = (emergency_brake.squeeze(-1) > 0.5).float()
+            final_throttle = final_throttle * (1 - emergency_mask)
+            final_brake = torch.max(final_brake, emergency_mask)
+        
+        # Clamp all outputs
+        final_steering = torch.clamp(final_steering, -self.max_steering, self.max_steering)
+        final_throttle = torch.clamp(final_throttle, 0.0, 1.0)
+        final_brake = torch.clamp(final_brake, 0.0, 1.0)
+        
+        # Mutual exclusion: can't throttle and brake simultaneously
+        # If braking > throttle, zero out throttle
+        brake_dominant = (final_brake > final_throttle).float()
+        final_throttle = final_throttle * (1 - brake_dominant)
+        
+        # Convert to physical units
+        accel_cmd = final_throttle * self.max_accel - final_brake * self.max_decel
+        steer_deg = torch.rad2deg(final_steering)
+        
+        # Predict next state using bicycle model
+        current_state = torch.stack([
+            ego_state[:, 4],  # x
+            ego_state[:, 5],  # y
+            ego_state[:, 3],  # heading (yaw_rate as proxy)
+            ego_state[:, 0],  # speed
+        ], dim=-1)
+        
+        control_input = torch.stack([final_steering, accel_cmd], dim=-1)
+        predicted_next_state = self.bicycle_model(current_state, control_input)
+        
+        return {
+            "steering_rad": final_steering,
+            "steering_deg": steer_deg,
+            "throttle": final_throttle,
+            "brake": final_brake,
+            "acceleration_cmd": accel_cmd,
+            "controller_weights": weights,
+            "predicted_next_state": predicted_next_state,
+        }
+    
+    def reset(self):
+        """Reset controller states (call at start of new episode)."""
+        self.pid_controller.reset()