Nysa/cleaned up ipad

src/camera-based-e2e/create_vocab.py

@@ -13,8 +13,8 @@
 def get_all_trajectories():
     # Instantiate dataloader w/ n_items None to get everything
     dataset = WaymoE2E(indexFile='index_train.pkl', 
-                       data_dir='/anvil/scratch/x-mgagvani/wod/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0', 
-                       n_items=None # set to None eventually,
+                       data_dir='/anvil/scratch/x-mgagvani/wod/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0',
+                       n_items=None  # set to None eventually
                        )
 
     dataloader = DataLoader(dataset, batch_size=512, num_workers=0, collate_fn=collate_with_images, pin_memory=True)

src/camera-based-e2e/loader.py

@@ -21,6 +21,7 @@ def __init__(
     ):
         self.data_dir = data_dir
         self.seed = seed
+        self._index_file = indexFile
 
         self.filename = ""
         self.file = None
@@ -101,7 +102,7 @@ def collate_with_images(batch):
     import time
     from tqdm import tqdm
     # NOTE: Replace with your path
-    DATA_DIR = '/anvil/scratch/x-mgagvani/wod/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+    DATA_DIR = '/scratch/gilbreth/mathur91/waymo/waymo_open_dataset_end_to_end_camera_v_1_0_0'
     BATCH_SIZE = 256
     dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR)
     loader = DataLoader(

src/camera-based-e2e/models/feature_extractors.py

@@ -3,6 +3,31 @@
 import torch.nn as nn
 import timm
 
+
+class ResNetFeatures(nn.Module):
+    """ResNet backbone (resnet50) - widely available in timm, fallback when DINO/SAM unavailable."""
+
+    def __init__(self, model_name: str = "resnet50", frozen: bool = True, feature_stage: int = -1):
+        super().__init__()
+        self.backbone = timm.create_model(model_name, pretrained=True, features_only=True)
+        self.data_config = timm.data.resolve_data_config(model=self.backbone)
+        self.transforms = timm.data.create_transform(**self.data_config, is_training=False)
+        if frozen:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+            self.backbone.eval()
+        channels = self.backbone.feature_info.channels()
+        reductions = self.backbone.feature_info.reduction()
+        self.feature_stage = feature_stage
+        self.dims = [channels[feature_stage]]
+        self.patch_size = reductions[feature_stage]
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        x_t = self.transforms(x.float().div(255.0))
+        feats = self.backbone(x_t)
+        return [feats[self.feature_stage]]
+
+
 class DINOFeatures(nn.Module):
     def __init__(self, model_name: str = "vit_small_plus_patch16_dinov3.lvd1689m", frozen: bool = True):
         super(DINOFeatures, self).__init__()

src/camera-based-e2e/models/proposal_init.py

@@ -0,0 +1,63 @@
+"""
+Proposal initialization (iPad-style).
+
+Per-timestep learnable embeddings for N proposals × T timesteps,
+conditioned on ego status (past trajectory + intent).
+
+Output: bev_feature (B, N*T, C) — the initial BEV proposal queries.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ProposalInit(nn.Module):
+    """
+    Initialize per-timestep BEV proposal features from learnable embeddings
+    plus ego status encoding.  Matches iPad's init_feature + hist_encoding.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_proposals: int = 16,
+        horizon: int = 20,
+        past_dim: int = 16 * 6,
+        intent_classes: int = 3,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.num_proposals = num_proposals
+        self.horizon = horizon
+
+        # Per-(proposal, timestep) learnable embeddings
+        self.proposal_embed = nn.Parameter(
+            torch.zeros(1, num_proposals * horizon, d_model)
+        )
+        nn.init.trunc_normal_(self.proposal_embed, std=0.1)
+
+        ego_dim = past_dim + intent_classes
+        self.ego_enc = nn.Sequential(
+            nn.Linear(ego_dim, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+
+    def forward(self, past: torch.Tensor, intent: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            past: (B, 16, 6)
+            intent: (B,) integer 1/2/3
+        Returns:
+            bev_feature: (B, N*T, d_model)
+        """
+        B = past.size(0)
+        past_flat = past.view(B, -1)
+        intent_onehot = F.one_hot(
+            (intent - 1).long().clamp(0, 2), num_classes=3
+        ).float()
+        ego = torch.cat([intent_onehot, past_flat], dim=1)
+        ego_feat = self.ego_enc(ego)  # (B, d_model)
+
+        bev_feature = self.proposal_embed.expand(B, -1, -1) + ego_feat[:, None, :]
+        return bev_feature

src/camera-based-e2e/models/proposal_planner.py

@@ -0,0 +1,107 @@
+"""
+Proposal-centric E2E planner (iPad-style).
+
+Pipeline:
+  scene encoder → proposal init → iterative refinement → scorer
+                                  (with intermediate proposal supervision)
+"""
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+
+from .scene_encoder import SceneEncoder
+from .proposal_init import ProposalInit
+from .refinement import Refinement
+from .scorer import Scorer
+
+
+@dataclass
+class IPadConfig:
+    """Loss / scorer hyperparameters for the iPad-style ProposalPlanner.
+
+    Bundled here so the LitModel only needs a single config argument instead of
+    a long list of keyword args. Defaults match the original LitModel signature.
+    """
+
+    # Loss weights
+    rfs_weight: float = 0.0
+    smoothness_weight: float = 0.0
+    collision_weight: float = 0.0
+    comfort_weight: float = 0.0
+    diversity_weight: float = 0.0
+
+    # Scorer
+    score_weight: float = 1.0
+    score_warmup_epochs: int = 2
+    score_temperature: float = 5.0
+    score_loss_type: str = "bce"      # bce | ce | bce_pairwise | listnet
+    score_target_type: str = "l1"     # l1 | navsim | rfs
+    score_rank_weight: float = 0.0
+    score_margin: float = 0.2
+    score_topk: int = 0
+
+    # Misc
+    comfort_jerk_threshold: float = 5.0
+    prev_weight: float = 0.1
+    rfs_target_use_comfort: bool = True
+
+
+class ProposalPlanner(nn.Module):
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        d_model: int = 256,
+        num_proposals: int = 16,
+        num_refinement_steps: int = 4,
+        horizon: int = 20,
+        num_cams: int = 8,
+    ):
+        super().__init__()
+        self.horizon = horizon
+        self.n_proposals = num_proposals
+
+        self.scene_encoder = SceneEncoder(backbone, d_model=d_model, num_cams=num_cams)
+        self.proposal_init = ProposalInit(
+            d_model=d_model,
+            num_proposals=num_proposals,
+            horizon=horizon,
+        )
+        self.refinement = Refinement(
+            d_model=d_model,
+            num_steps=num_refinement_steps,
+            num_heads=8,
+            num_proposals=num_proposals,
+            horizon=horizon,
+        )
+        self.scorer = Scorer(
+            d_model=d_model,
+            num_proposals=num_proposals,
+            horizon=horizon,
+        )
+
+    def forward(self, x: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        past = x["PAST"]
+        images: List[torch.Tensor] = x["IMAGES"]
+        intent = x["INTENT"]
+
+        scene_feat = self.scene_encoder(images)
+        bev_feature = self.proposal_init(past, intent)
+        proposals, bev_feature, proposal_list = self.refinement(bev_feature, scene_feat)
+        scores = self.scorer(bev_feature)
+
+        if getattr(self, "_debug", False):
+            self._debug_proposals = proposals.detach()
+            self._debug_scores = scores.detach()
+            self._debug_proposal_list = [p.detach() for p in proposal_list]
+
+        B, K, T, _ = proposals.shape
+        trajectory_flat = proposals.view(B, K * T * 2)
+
+        return {
+            "trajectory": trajectory_flat,
+            "scores": scores,
+            "proposal_list": proposal_list,
+        }

src/camera-based-e2e/models/refinement.py

@@ -0,0 +1,122 @@
+"""
+Iterative refinement (iPad-style predict-attend-refine loop).
+
+A single shared RefinementBlock is applied K times.  At each iteration:
+  1. Decode per-timestep features into full trajectory proposals
+  2. Encode proposals back into feature space (proposal-anchored)
+  3. Cross-attend proposal features to scene tokens
+  4. FFN update
+Returns all intermediate proposals for discounted supervision.
+"""
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+from .blocks import MHA
+
+
+class RefinementBlock(nn.Module):
+    """One iteration: decode proposals, encode them back, cross-attend to scene."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int = 8,
+        num_proposals: int = 16,
+        horizon: int = 20,
+        mlp_ratio: int = 4,
+    ):
+        super().__init__()
+        self.num_proposals = num_proposals
+        self.horizon = horizon
+
+        # Per-timestep feature → (x, y)
+        self.traj_decoder = nn.Sequential(
+            nn.Linear(d_model, d_model * mlp_ratio),
+            nn.GELU(),
+            nn.Linear(d_model * mlp_ratio, 2),
+        )
+
+        # Encode (x, y) back into feature space
+        self.traj_enc = nn.Sequential(
+            nn.Linear(2, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+
+        self.ln1 = nn.LayerNorm(d_model)
+        self.cross_attn = MHA(d_model, num_heads)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model, d_model * mlp_ratio),
+            nn.GELU(),
+            nn.Linear(d_model * mlp_ratio, d_model),
+        )
+
+    def forward(
+        self,
+        bev_feature: torch.Tensor,
+        scene_feat: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            bev_feature: (B, N*T, C) per-timestep proposal features
+            scene_feat:  (B, S, C) visual tokens from scene encoder
+        Returns:
+            proposals:   (B, N, T, 2) fully predicted trajectories
+            bev_feature: (B, N*T, C) refined features
+        """
+        B = bev_feature.size(0)
+        N, T = self.num_proposals, self.horizon
+
+        proposals = self.traj_decoder(bev_feature).view(B, N, T, 2)
+
+        prop_enc = self.traj_enc(proposals.view(B, N * T, 2))
+        bev_feature = bev_feature + prop_enc
+
+        bev_feature = bev_feature + self.cross_attn(
+            self.ln1(bev_feature), context=scene_feat
+        )
+        bev_feature = bev_feature + self.mlp(self.ln2(bev_feature))
+
+        return proposals, bev_feature
+
+
+class Refinement(nn.Module):
+    """Weight-shared iterative refinement applied num_steps times."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_steps: int = 4,
+        num_heads: int = 8,
+        num_proposals: int = 16,
+        horizon: int = 20,
+    ):
+        super().__init__()
+        self.num_steps = num_steps
+        self.block = RefinementBlock(
+            d_model=d_model,
+            num_heads=num_heads,
+            num_proposals=num_proposals,
+            horizon=horizon,
+        )
+
+    def forward(
+        self,
+        bev_feature: torch.Tensor,
+        scene_feat: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:
+        """
+        Returns:
+            proposals:     (B, N, T, 2) final iteration proposals
+            bev_feature:   (B, N*T, C) final features
+            proposal_list: list of (B, N, T, 2) from each iteration
+        """
+        proposal_list: List[torch.Tensor] = []
+        proposals = None
+        for _ in range(self.num_steps):
+            proposals, bev_feature = self.block(bev_feature, scene_feat)
+            proposal_list.append(proposals)
+        return proposals, bev_feature, proposal_list