Autonomous Vehicle AI: 3D Perception, BEVFormer, and End-to-End Learning

# Autonomous driving pipeline abstraction
class AutonomousDrivingStack:
    def __init__(self, perception, prediction, planner, controller):
        self.perception = perception
        self.prediction = prediction
        self.planner = planner
        self.controller = controller

    def step(self, sensor_data):
        # 1. Perceive
        scene = self.perception.process(sensor_data)
        # 2. Predict
        future_states = self.prediction.forecast(scene)
        # 3. Plan
        trajectory = self.planner.plan(scene, future_states)
        # 4. Control
        commands = self.controller.compute(trajectory)
        return commands

import torch
import torch.nn as nn

class SetAbstraction(nn.Module):
    """PointNet++ Set Abstraction layer"""
    def __init__(self, npoint, radius, nsample, in_channel, mlp):
        super().__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel

    def forward(self, xyz, points):
        # xyz: (B, N, 3), points: (B, C, N)
        # Apply MLP after farthest point sampling + ball query
        new_xyz = self.farthest_point_sample(xyz, self.npoint)
        grouped = self.ball_query(xyz, new_xyz, points)
        for conv, bn in zip(self.mlp_convs, self.mlp_bns):
            grouped = torch.relu(bn(conv(grouped)))
        new_points = grouped.max(dim=-1)[0]
        return new_xyz, new_points

    def farthest_point_sample(self, xyz, npoint):
        B, N, _ = xyz.shape
        centroids = torch.zeros(B, npoint, dtype=torch.long)
        distance = torch.ones(B, N) * 1e10
        farthest = torch.randint(0, N, (B,))
        for i in range(npoint):
            centroids[:, i] = farthest
            centroid = xyz[torch.arange(B), farthest, :].unsqueeze(1)
            dist = ((xyz - centroid) ** 2).sum(-1)
            distance = torch.min(distance, dist)
            farthest = distance.max(-1)[1]
        return xyz[torch.arange(B).unsqueeze(1), centroids]

    def ball_query(self, xyz, new_xyz, points):
        B, N, C = points.transpose(1, 2).shape
        return points.unsqueeze(-1).expand(-1, -1, -1, self.nsample)

class VoxelFeatureEncoder(nn.Module):
    """VoxelNet VFE layer"""
    def __init__(self, in_channels=4, out_channels=128):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(in_channels, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Linear(64, out_channels),
            nn.BatchNorm1d(out_channels),
            nn.ReLU(),
        )

    def forward(self, voxel_features, num_points):
        # voxel_features: (N_voxels, max_points, C)
        B, M, C = voxel_features.shape
        x = self.fc(voxel_features.view(-1, C))
        x = x.view(B, M, -1)
        mask = torch.arange(M, device=x.device).unsqueeze(0) < num_points.unsqueeze(1)
        x = x * mask.unsqueeze(-1).float()
        return x.max(dim=1)[0]

import torch
import torch.nn as nn

class BEVFormerLayer(nn.Module):
    """BEVFormer Spatial Cross-Attention layer"""
    def __init__(self, embed_dim=256, num_heads=8, num_cameras=6):
        super().__init__()
        self.num_cameras = num_cameras
        self.embed_dim = embed_dim
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim),
        )

    def forward(self, bev_query, camera_features, lidar2img):
        """
        bev_query: (B, H*W, C) - BEV grid queries
        camera_features: (B, N_cam, H_img, W_img, C) - camera features
        lidar2img: (B, N_cam, 4, 4) - LiDAR-to-image projection matrices
        """
        B, Q, C = bev_query.shape
        ref_3d = self.get_reference_points(B, device=bev_query.device)
        projected = self.project_to_image(ref_3d, lidar2img)

        sampled = self.sample_camera_features(camera_features, projected)
        attn_out, _ = self.attn(bev_query, sampled, sampled)
        bev_query = self.norm1(bev_query + attn_out)
        bev_query = self.norm2(bev_query + self.ffn(bev_query))
        return bev_query

    def get_reference_points(self, B, device, H=200, W=200, Z=8):
        xs = torch.linspace(0, 1, W, device=device)
        ys = torch.linspace(0, 1, H, device=device)
        zs = torch.linspace(0, 1, Z, device=device)
        grid = torch.stack(torch.meshgrid(xs, ys, zs), dim=-1)
        return grid.view(1, -1, 3).expand(B, -1, -1)

    def project_to_image(self, pts_3d, lidar2img):
        return pts_3d[..., :2]

    def sample_camera_features(self, features, coords):
        B, Q, C = features.shape[0], coords.shape[1], features.shape[-1]
        return torch.zeros(B, Q, C, device=features.device)

class CenterPointHead(nn.Module):
    """CenterPoint detection head"""
    def __init__(self, in_channels=256, num_classes=10):
        super().__init__()
        self.heatmap = nn.Conv2d(in_channels, num_classes, 1)
        self.offset = nn.Conv2d(in_channels, 2, 1)   # x, y offset
        self.height = nn.Conv2d(in_channels, 1, 1)   # z
        self.size = nn.Conv2d(in_channels, 3, 1)     # w, l, h
        self.rotation = nn.Conv2d(in_channels, 2, 1) # sin, cos
        self.velocity = nn.Conv2d(in_channels, 2, 1) # vx, vy

    def forward(self, x):
        return {
            'heatmap': self.heatmap(x).sigmoid(),
            'offset': self.offset(x),
            'height': self.height(x),
            'size': self.size(x),
            'rotation': self.rotation(x),
            'velocity': self.velocity(x),
        }

import numpy as np

class KalmanTracker:
    """Kalman filter for 3D object tracking"""
    def __init__(self, initial_state):
        # State vector: [x, y, z, vx, vy, vz, w, l, h, theta]
        self.dt = 0.1
        n = 10
        self.x = np.array(initial_state, dtype=float).reshape(-1, 1)
        self.P = np.eye(n) * 10.0

        # State transition (constant velocity model)
        self.F = np.eye(n)
        for i in range(3):
            self.F[i, i + 3] = self.dt

        # Observation matrix (position + size + heading)
        self.H = np.zeros((7, n))
        for i in range(7):
            self.H[i, i] = 1.0

        self.Q = np.eye(n) * 0.1   # Process noise
        self.R = np.eye(7) * 1.0   # Measurement noise

    def predict(self):
        """Prediction step"""
        self.x = self.F @ self.x
        self.P = self.F @ self.P @ self.F.T + self.Q
        return self.x[:7].flatten()

    def update(self, z):
        """Update step"""
        z = np.array(z).reshape(-1, 1)
        y = z - self.H @ self.x                          # Innovation
        S = self.H @ self.P @ self.H.T + self.R          # Innovation covariance
        K = self.P @ self.H.T @ np.linalg.inv(S)         # Kalman gain
        self.x = self.x + K @ y
        self.P = (np.eye(len(self.x)) - K @ self.H) @ self.P
        return self.x[:7].flatten()


class MOT3D:
    """Hungarian algorithm-based 3D multi-object tracking"""
    def __init__(self, max_age=3, min_hits=3, iou_threshold=0.3):
        self.trackers = []
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        self.frame_count = 0
        self.next_id = 0

    def update(self, detections):
        self.frame_count += 1
        predictions = [t.predict() for t in self.trackers]

        matched, unmatched_dets, unmatched_trks = self.associate(
            detections, predictions
        )

        for d, t in matched:
            self.trackers[t].update(detections[d])

        for i in unmatched_dets:
            trk = KalmanTracker(detections[i])
            trk.id = self.next_id
            self.next_id += 1
            self.trackers.append(trk)

        self.trackers = [
            t for t in self.trackers
            if t.time_since_update <= self.max_age
        ]
        return self.trackers

    def associate(self, detections, predictions):
        if not predictions:
            return [], list(range(len(detections))), []
        return [], list(range(len(detections))), []

class NDTMatching:
    """NDT scan matching for localization"""
    def __init__(self, resolution=1.0):
        self.resolution = resolution
        self.voxel_map = {}

    def build_map(self, points):
        """Convert point cloud to NDT voxel map"""
        for pt in points:
            key = tuple((pt / self.resolution).astype(int))
            if key not in self.voxel_map:
                self.voxel_map[key] = []
            self.voxel_map[key].append(pt)

        for key in self.voxel_map:
            pts = np.array(self.voxel_map[key])
            mean = pts.mean(axis=0)
            cov = np.cov(pts.T) if len(pts) > 3 else np.eye(3)
            self.voxel_map[key] = {'mean': mean, 'cov': cov}

    def score(self, points, transform):
        """Compute NDT score for a given transform"""
        score = 0.0
        R = transform[:3, :3]
        t = transform[:3, 3]
        for pt in points:
            pt_t = R @ pt + t
            key = tuple((pt_t / self.resolution).astype(int))
            if key in self.voxel_map:
                v = self.voxel_map[key]
                diff = pt_t - v['mean']
                score += np.exp(-0.5 * diff @ np.linalg.inv(v['cov']) @ diff)
        return score

import xml.etree.ElementTree as ET

def parse_opendrive(xodr_path):
    """Parse OpenDRIVE HD map format"""
    tree = ET.parse(xodr_path)
    root = tree.getroot()
    roads = []
    for road in root.findall('road'):
        road_data = {
            'id': road.get('id'),
            'length': float(road.get('length')),
            'lanes': [],
        }
        for lane_section in road.findall('.//laneSection'):
            for lane in lane_section.findall('.//lane'):
                lane_data = {
                    'id': int(lane.get('id')),
                    'type': lane.get('type'),
                    'speed': None,
                }
                speed = lane.find('.//speed')
                if speed is not None:
                    lane_data['speed'] = float(speed.get('max', 0))
                road_data['lanes'].append(lane_data)
        roads.append(road_data)
    return roads

class TrajectoryPredictor(nn.Module):
    """Transformer-based multi-agent trajectory prediction"""
    def __init__(self, input_dim=6, embed_dim=128, num_heads=4,
                 num_layers=3, pred_horizon=30, num_modes=6):
        super().__init__()
        self.input_proj = nn.Linear(input_dim, embed_dim)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=num_heads, batch_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.mode_query = nn.Embedding(num_modes, embed_dim)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=embed_dim, nhead=num_heads, batch_first=True
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.traj_head = nn.Linear(embed_dim, pred_horizon * 2)
        self.prob_head = nn.Linear(embed_dim, 1)
        self.num_modes = num_modes

    def forward(self, agent_history, map_features=None):
        """
        agent_history: (B, N_agents, T_hist, 6) - x, y, vx, vy, ax, ay
        Returns: trajectories (B, N_agents, K, T_pred, 2), probs (B, N_agents, K)
        """
        B, N, T, C = agent_history.shape
        x = self.input_proj(agent_history.view(B * N, T, C))
        memory = self.encoder(x)

        queries = self.mode_query.weight.unsqueeze(0).expand(B * N, -1, -1)
        decoded = self.decoder(queries, memory)

        trajs = self.traj_head(decoded)
        probs = self.prob_head(decoded).squeeze(-1).softmax(-1)

        pred_horizon = trajs.shape[-1] // 2
        trajs = trajs.view(B, N, self.num_modes, pred_horizon, 2)
        probs = probs.view(B, N, self.num_modes)
        return trajs, probs

import carla
import numpy as np

def setup_carla_sensors(world, vehicle):
    """Configure CARLA sensors: camera + LiDAR"""
    bp_lib = world.get_blueprint_library()

    # RGB camera
    cam_bp = bp_lib.find('sensor.camera.rgb')
    cam_bp.set_attribute('image_size_x', '1280')
    cam_bp.set_attribute('image_size_y', '720')
    cam_bp.set_attribute('fov', '90')
    cam_transform = carla.Transform(carla.Location(x=1.5, z=2.4))
    camera = world.spawn_actor(cam_bp, cam_transform, attach_to=vehicle)

    # LiDAR
    lidar_bp = bp_lib.find('sensor.lidar.ray_cast')
    lidar_bp.set_attribute('channels', '64')
    lidar_bp.set_attribute('range', '100')
    lidar_bp.set_attribute('points_per_second', '1200000')
    lidar_bp.set_attribute('rotation_frequency', '20')
    lidar_transform = carla.Transform(carla.Location(x=0.0, z=2.0))
    lidar = world.spawn_actor(lidar_bp, lidar_transform, attach_to=vehicle)

    return camera, lidar


def domain_randomization(world):
    """Domain randomization to reduce sim-to-real gap"""
    import random
    weather = carla.WeatherParameters(
        cloudiness=random.uniform(0, 80),
        precipitation=random.uniform(0, 30),
        sun_altitude_angle=random.uniform(15, 90),
        fog_density=random.uniform(0, 20),
        wetness=random.uniform(0, 50),
    )
    world.set_weather(weather)

class BehaviorCloningAgent(nn.Module):
    """Behavior cloning autonomous driving agent"""
    def __init__(self, img_channels=3, cmd_dim=4, action_dim=2):
        super().__init__()
        import torchvision.models as models
        backbone = models.resnet34(pretrained=False)
        self.visual_encoder = nn.Sequential(*list(backbone.children())[:-2])
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.cmd_embed = nn.Embedding(cmd_dim, 64)  # Commands: straight/left/right/u-turn
        self.policy = nn.Sequential(
            nn.Linear(512 + 64, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim),  # steer, throttle
            nn.Tanh(),
        )

    def forward(self, image, command):
        feat = self.pool(self.visual_encoder(image)).flatten(1)
        cmd_feat = self.cmd_embed(command)
        x = torch.cat([feat, cmd_feat], dim=-1)
        return self.policy(x)


def dagger_data_aggregation(expert_data, agent_rollouts):
    """DAgger: mitigate distributional shift via iterative data collection"""
    # Collect expert labels at states visited by the agent
    combined = {
        'obs': torch.cat([expert_data['obs'], agent_rollouts['obs']]),
        'actions': torch.cat([expert_data['actions'], agent_rollouts['expert_actions']]),
    }
    return combined

from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner

class ReactiveIDMPlanner(AbstractPlanner):
    """nuPlan-compatible IDM planner"""
    def __init__(self, target_velocity=15.0, min_gap=2.0):
        self.target_velocity = target_velocity
        self.min_gap = min_gap

    def initialize(self, initialization):
        self.initialization = initialization

    def name(self):
        return 'ReactiveIDMPlanner'

    def observation_type(self):
        from nuplan.planning.simulation.observation.observation_type import DetectionsTracks
        return DetectionsTracks

    def compute_planner_trajectory(self, current_input):
        ego_state = current_input.history.ego_states[-1]
        observations = current_input.observations
        accel = self.idm_accel(ego_state, observations)
        return self.build_trajectory(ego_state, accel)

    def idm_accel(self, ego, obs):
        v = ego.dynamic_car_state.speed
        delta = 4.0
        a_max = 2.0
        s_star = self.min_gap + v * 1.5
        return a_max * (1 - (v / self.target_velocity) ** delta
                        - (s_star / max(self.min_gap, 1.0)) ** 2)

    def build_trajectory(self, ego_state, accel):
        return []

class SafetyMonitor:
    """Runtime safety monitor"""
    def __init__(self, uncertainty_threshold=0.7, min_detection_confidence=0.5):
        self.uncertainty_threshold = uncertainty_threshold
        self.min_conf = min_detection_confidence
        self.fallback_triggered = False

    def check(self, perception_output):
        uncertainty = perception_output.get('uncertainty', 0.0)
        confidence = perception_output.get('confidence', 1.0)

        if uncertainty > self.uncertainty_threshold:
            self.trigger_fallback('high_uncertainty')
            return False
        if confidence < self.min_conf:
            self.trigger_fallback('low_confidence')
            return False
        return True

    def trigger_fallback(self, reason):
        self.fallback_triggered = True
        # MRC (Minimal Risk Condition): initiate safe stop
        print(f'[Safety] Fallback triggered: {reason}')