diff --git a/preference_dataset.py b/preference_dataset.py
new file mode 100644
index 0000000..2032d02
--- /dev/null
+++ b/preference_dataset.py
@@ -0,0 +1,118 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import pickle
+import os
+import numpy as np
+import e2e_pb2
+from tqdm import tqdm
+import random
+
+
+class PreferenceDataset(Dataset):    
+    def __init__(self, index_file, data_dir, cache_file='preference_cache_full.pkl'):
+        self.data_dir = data_dir
+
+        if os.path.exists(cache_file):
+            print(f"Loading filtered samples from cache: {cache_file}")
+            with open(cache_file, 'rb') as f:
+                self.valid_samples = pickle.load(f)
+            print(f"Loaded {len(self.valid_samples)} valid preference trajectory samples from cache")
+            return
+        
+        print("Cache not found. Filtering dataset for valid preference trajectories...")
+        with open(index_file, 'rb') as f:
+            all_indexes = pickle.load(f)
+        
+        self.valid_samples = []
+        
+        # Group indexes by filename
+        file_groups = {}
+        for idx, (filename, start_byte, byte_length) in enumerate(all_indexes):
+            if filename not in file_groups:
+                file_groups[filename] = []
+            file_groups[filename].append((idx, start_byte, byte_length))
+        
+        # Process each file once
+        for filename, frame_list in tqdm(file_groups.items(), desc="Processing files"):
+            file_path = os.path.join(data_dir, filename)
+            
+            # Check if file exists
+            if not os.path.exists(file_path):
+                print(f"WARNING: File not found: {file_path}")
+                continue
+                
+            with open(file_path, 'rb') as f:
+                for original_idx, start_byte, byte_length in frame_list:
+                    # Read and parse frame
+                    f.seek(start_byte)
+                    protobuf = f.read(byte_length)
+                    frame = e2e_pb2.E2EDFrame()
+                    frame.ParseFromString(protobuf)
+                    print("num preference trajectories:", len(frame.preference_trajectories))
+                    
+                    # Check if this frame has valid preference trajectories
+                    for traj_idx, pref_traj in enumerate(frame.preference_trajectories):
+                        if pref_traj.preference_score >= 0:  # Valid score
+                            self.valid_samples.append({
+                                'file_info': (filename, start_byte, byte_length),
+                                'traj_idx': traj_idx,
+                                'original_idx': original_idx
+                            })
+        
+        print(f"Found {len(self.valid_samples)} valid preference trajectory samples")
+        
+        # Save to cache
+        print(f"Saving filtered samples to cache: {cache_file}")
+        with open(cache_file, 'wb') as f:
+            pickle.dump(self.valid_samples, f)
+        print("Cache saved!")
+    
+    def __len__(self):
+        return len(self.valid_samples)
+    
+    def __getitem__(self, idx):
+        sample_info = self.valid_samples[idx]
+        filename, start_byte, byte_length = sample_info['file_info']
+        traj_idx = sample_info['traj_idx']
+        
+        # Read frame
+        file_path = os.path.join(self.data_dir, filename)
+        with open(file_path, 'rb') as f:
+            f.seek(start_byte)
+            protobuf = f.read(byte_length)
+            frame = e2e_pb2.E2EDFrame()
+            frame.ParseFromString(protobuf)
+        
+        # Extract past states (T,6)
+        past = np.stack([
+            frame.past_states.pos_x,
+            frame.past_states.pos_y,
+            frame.past_states.vel_x,
+            frame.past_states.vel_y,
+            frame.past_states.accel_x,
+            frame.past_states.accel_y
+        ], axis=-1).astype(np.float32)
+        
+        # Extract intent (one-hot encode)
+        intent = frame.intent
+        intent_onehot = np.zeros(4, dtype=np.float32)
+        intent_onehot[intent] = 1.0
+        
+        # Extract preference trajectory (only pos_x, pos_y are populated)
+        pref_traj = frame.preference_trajectories[traj_idx]
+        trajectory = np.stack([
+            pref_traj.pos_x,
+            pref_traj.pos_y
+        ], axis=-1).astype(np.float32)
+        
+        # Extract preference score (target)
+        score = pref_traj.preference_score
+        
+        return {
+            'past_states': torch.from_numpy(past.flatten()),
+            'intent': torch.from_numpy(intent_onehot),
+            'trajectory': torch.from_numpy(trajectory.flatten()),
+            'score': torch.tensor(score, dtype=torch.float32)
+        }
diff --git a/src/camera-based-e2e/dataset/export_dataset.py b/src/camera-based-e2e/dataset/export_dataset.py
index 274de66..a9853b0 100644
--- a/src/camera-based-e2e/dataset/export_dataset.py
+++ b/src/camera-based-e2e/dataset/export_dataset.py
@@ -15,17 +15,18 @@
 
 def safe_name(name: str) -> str:
     # Keep filesystem-friendly ASCII.
-    return re.sub(r"[^A-Za-z0-9_.-]+", "_", name)
+    return re.sub(r"[^A-Za-z0-9_.-]+", "_", name) #substitutes any non ascii chars with _
 
 
-def camera_name(name_enum: int) -> str:
+def camera_name(name_enum: int) -> str: #which camera (front, left, etc)
     return dataset_pb2.CameraName.Name.Name(name_enum)
+    #access dataset_pb2 class, then CameraName nested class, then Name enum nested class, then Name constructor method to get string name
 
 
 def intent_name(intent_enum: int) -> str:
     return e2e_pb2.EgoIntent.Intent.Name(intent_enum)
 
-
+#NOT SURE ABOUT THIS FUNCTION
 def transform_list(transform_msg) -> list:
     return list(transform_msg.transform)
 
diff --git a/src/camera-based-e2e/depthLoss.py b/src/camera-based-e2e/depthLoss.py
new file mode 100644
index 0000000..d781f95
--- /dev/null
+++ b/src/camera-based-e2e/depthLoss.py
@@ -0,0 +1,107 @@
+from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+import torch
+import torch.nn.functional as F
+
+DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf"
+
+class DepthLoss:
+    def __init__(self, device):
+        self.device = device
+        self.depth_processor = AutoImageProcessor.from_pretrained(DEPTH_MODEL_ID)
+        self.depth_model = AutoModelForDepthEstimation.from_pretrained(DEPTH_MODEL_ID).to(self.device)
+
+    def get_depth(self, images):
+        """
+        Compute ground truth depth from images
+
+        Args:
+            images (torch.Tensor): Input images of shape (B, C, H, W).           
+        """
+        # Preprocess images
+        inputs = self.depth_processor(images=images, return_tensors="pt").to(self.device)
+
+        # Forward pass through the depth estimation model
+        with torch.no_grad():
+            outputs = self.depth_model(**inputs)
+            predicted_depth = outputs.predicted_depth
+
+        height, width = images.shape[2], images.shape[3]
+        
+        # Interpolate to original size
+        prediction = F.interpolate(
+            predicted_depth.unsqueeze(1),
+            size=(height, width),
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze(1)  # (B, H, W)
+
+        return prediction
+
+    def compute_depth_loss(self, gt_images, pred_depths, loss_fn):
+        pred_depth = self.get_depth(gt_images)
+        depth_loss = loss_fn(pred_depth, pred_depths)
+
+        return depth_loss
+    
+    def __call__(self, gt_images, pred_depths, loss_fn=F.l1_loss):
+        return self.compute_depth_loss(gt_images, pred_depths, loss_fn)
+    
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        print("Must run on GPU")
+
+    import sys
+    from pathlib import Path
+    # Determine the absolute path to the directory containing loader.py
+    # depth_loss.py is in models/losses/
+    current_dir = Path(__file__).resolve().parent
+    project_root = current_dir.parent.parent
+    sys.path.append(str(project_root))
+
+    from loader import WaymoE2E
+    loader = WaymoE2E(indexFile="index_val.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/", images=True)
+    data_iterator = iter(torch.utils.data.DataLoader(loader, batch_size=8, num_workers=4))
+
+    device = torch.device("cuda")
+    depth_loss_fn = DepthLoss(device)
+
+    from matplotlib import pyplot as plt
+    import numpy as np
+
+    for _ in range(6):
+        batch = next(data_iterator)
+    images = batch["IMAGES"][1].to(device)  # front camera
+
+    res = depth_loss_fn.get_depth(images)  # (B, H, W) tensor
+
+    fig, ax = plt.subplots(8, 2, figsize=(16, 48))
+    for i in range(8):
+        # Input image
+        ax[i, 0].set_title(f"Input Image {i+1}")
+        ax[i, 0].imshow(images[i].permute(1, 2, 0).cpu().numpy().astype(np.uint8))
+        ax[i, 0].axis('off')
+        
+        # Predicted depth
+        ax[i, 1].set_title(f"Predicted Depth {i+1}")
+        depth_img = ax[i, 1].imshow(res[i].cpu().numpy(), cmap='plasma')
+        ax[i, 1].axis('off')
+        fig.colorbar(depth_img, ax=ax[i, 1], orientation='vertical', label='Depth')
+    
+    plt.tight_layout()
+    fig.savefig("predicted_depth.png", dpi=150, bbox_inches='tight')
+
+    # throughput test
+    from time import perf_counter
+
+    start_time = perf_counter()
+    times = []
+    for _ in range(100):
+        batch = next(data_iterator)
+        images = batch["IMAGES"][1].to(device)  # front camera
+        t0 = perf_counter()
+        res = depth_loss_fn.get_depth(images)
+        times.append(perf_counter() - t0)
+    end_time = perf_counter()
+    print(f"Throughput: {100 / (end_time - start_time):.2f} batches/sec")
+    print(f"Avg depth inference batch/s: {1 / np.mean(times):.2f} FPS")
diff --git a/src/camera-based-e2e/loader.py b/src/camera-based-e2e/loader.py
index 4cbc140..c97a14f 100644
--- a/src/camera-based-e2e/loader.py
+++ b/src/camera-based-e2e/loader.py
@@ -1,6 +1,6 @@
 import torch
 from torch.utils.data import IterableDataset
-from waymo_open_dataset.protos import end_to_end_driving_data_pb2 as e2e_pb2
+from protos import e2e_pb2
 import torchvision
 import pickle
 import struct
@@ -11,6 +11,10 @@
 import cv2 
 from typing import Optional
 import random
+import tqdm
+import open3d as o3d
+from point_cloud import get_waymo_intrinsics, create_point_cloud
+from depthLoss import DepthLoss
 
 devices = ['cuda:0', 'cuda:1']
 
@@ -55,7 +59,7 @@ def decode_img(self, img):
         gpu_tensors_list = torchvision.io.decode_jpeg(
             img_tensor, 
             mode=torchvision.io.ImageReadMode.UNCHANGED,
-            device= 'cpu' #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2]
+            device= 'cpu', #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2]
         )
         # img_array = np.frombuffer(img, np.uint8)
         return gpu_tensors_list
@@ -96,7 +100,28 @@ def __iter__(self):
             # For submission to waymo evaluation server
             name = frame.frame.context.name
 
-            yield {'PAST': past, 'FUTURE': future, 'IMAGES': [self.decode_img(images.image) for images in frame.frame.images], 'INTENT': frame.intent, 'NAME': name}
+            intrinsics_vec = np.zeros(6, dtype=np.float32)
+            
+            # Find FRONT camera (Enum 1) in the context
+            for calib in frame.frame.context.camera_calibrations:
+                if calib.name == 1: # 1 = FRONT
+                    intrinsics_vec[0] = calib.intrinsic[0] # fx
+                    intrinsics_vec[1] = calib.intrinsic[1] # fy
+                    intrinsics_vec[2] = calib.intrinsic[2] # cx
+                    intrinsics_vec[3] = calib.intrinsic[3] # cy
+                    intrinsics_vec[4] = calib.width
+                    intrinsics_vec[5] = calib.height
+                    break
+            
+            decoded_images = [self.decode_img(img.image) for img in frame.frame.images]
+
+            yield {
+                'PAST': past, 
+                'FUTURE': future, 
+                'IMAGES': decoded_images, 
+                'INTRINSICS': intrinsics_vec, # <--- Passing this to main
+                'NAME': frame.frame.context.name
+            }
 
 if __name__ == "__main__":
 
@@ -104,9 +129,9 @@ def __iter__(self):
     import time
     from tqdm import tqdm
     # NOTE: Replace with your path
-    DATA_DIR = '/scratch/gilbreth/mgagvani/wod/waymo_open_dataset_end_to_end_camera_v_1_0_0/'
+    DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/'
     BATCH_SIZE = 32
-    dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR, images=True)
+    dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR, images=True, n_items= 2)
     loader = DataLoader(
         dataset, 
         batch_size=BATCH_SIZE,
@@ -114,12 +139,62 @@ def __iter__(self):
     )
     
     def main():
-        # start = time.time()
-        for batch_of_frames in tqdm(loader):
-            # print(batch_of_frames["INTENT"])
-            # print(batch_of_frames.keys(), [b.shape for b in batch_of_frames.values() if isinstance(b, torch.Tensor)])
-            pass
-        # print("Total Time:", time.time()-start)
+        device = torch.device("cuda")
+        depth_model = DepthLoss(device)
+        output_dir = "visualizations"
+        
+        for batch_idx, batch_of_frames in enumerate(tqdm(loader)):
+            images = batch_of_frames["IMAGES"][1].to(device)  # Shape: (B, 3, H, W)
+            intrinsics_batch = batch_of_frames["INTRINSICS"]
+            pred_depths = depth_model.get_depth(images)       # Shape: (B, H, W)
+
+            batch_size = images.shape[0] # B
+            
+            for i in range(batch_size):
+                # Image: (H, W, 3) uint8
+
+                img_np = images[i].permute(1, 2, 0).cpu().numpy().copy()  
+
+                if img_np.max() <= 1.0: 
+                    img_np = (img_np * 255).astype(np.uint8)
+                else: 
+                    img_np = img_np.astype(np.uint8)
+                
+                #save img
+                img_filename = f"batch_{batch_idx:04d}_img_{i:02d}.png"
+                img_path = os.path.join(output_dir, img_filename)
+
+                img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+                cv2.imwrite(img_path, img_bgr)
+
+                # Depth: (H, W) float32
+                depth_np = pred_depths[i].cpu().numpy()
+
+                # set camera intrinsics
+                vals = intrinsics_batch[i].numpy()
+                camera_intrinsics = o3d.camera.PinholeCameraIntrinsic()
+                camera_intrinsics.set_intrinsics(
+                    width=int(vals[4]), 
+                    height=int(vals[5]), 
+                    fx=vals[0], fy=vals[1], 
+                    cx=vals[2], cy=vals[3]
+                )
+                
+                # Create Point Cloud
+                pcd = create_point_cloud(
+                    img_np, 
+                    depth_np, 
+                    intrinsics=camera_intrinsics,
+                    depth_scale=1.0
+                )
+
+                # Save to Disk
+                filename = f"batch_{batch_idx:04d}_img_{i:02d}.pcd"
+                file_path = os.path.join(output_dir, filename)
+                
+                # non-blocking save command
+                o3d.io.write_point_cloud(file_path, pcd)
+                
     
     import cProfile
-    main()
+    main()
\ No newline at end of file
diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py
index 9096ea2..b3a547e 100644
--- a/src/camera-based-e2e/models/base_model.py
+++ b/src/camera-based-e2e/models/base_model.py
@@ -2,81 +2,139 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import pytorch_lightning as pl
-
-class BaseModel(nn.Module):
-    def __init__(self, in_dim, out_dim):
-        super(BaseModel, self).__init__()
-        
-        # This is literally just linear regression = y_hat = Wx + b
-        self.nn = nn.Sequential(
-            nn.Linear(in_dim, out_dim)
-        )
-        
-    def forward(self, x: dict) -> torch.Tensor:
-        past, images, intent = x['PAST'], x['IMAGES'], x['INTENT']
-        x = past.reshape(past.size(0), -1)  # Flatten to (B, 16 * 6) = (B, 96)
-        return self.nn(x)
     
+from scipy.stats import spearmanr, kendalltau
+import torch
+
 class LitModel(pl.LightningModule):
-    def __init__(self, model: nn.Module, lr: float):
-        super(LitModel, self).__init__()
+    def __init__(self, model: nn.Module, lr: float, delta: float = 1.0):
+        super().__init__()
         self.model = model
-        self.hparams.lr = lr
-
-        self.example_input_array = ({
-            'PAST': torch.zeros((1, 16, 6)),  # PAST
-            'IMAGES': [torch.zeros((1, 3, 1280, 1920)) for _ in range(6)],  # IMAGES
-            'INTENT': torch.tensor([1.0]),  # INTENT
-        },)
-
-    # ---- Metrics ----
-    def ade_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
-        """
-        Average Displacement Error -> L2 Norm -> Average Euclidean Distance between predicted and ground truth future trajectory
-        """
-        return torch.mean(torch.norm(pred - gt, dim=-1))
-    
-    # ---- optimizers ----
+        self.save_hyperparameters(ignore=["model"])
+
+        self.delta = delta
+        self.val_errors = []
+        self.val_pred = []
+        self.val_true = []
+
+    # ---- optimizer ----
     def configure_optimizers(self):
-        # NOTE: This can be extended and tuned, LR especially will differ and have an impact.
-        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr)
-        return optimizer
-    
-    # ---- forward / step ----
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr)
+
+    # ---- forward ----
+    def forward(self, x):
         return self.model(x)
-    
-    def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
-        past, future, images, intent = batch['PAST'], batch['FUTURE'], batch['IMAGES'], batch['INTENT']
-        
-        # `past` is our input (B, 16, 6) e.g. Batch x Time x (x, y, v_x, v_y, a_x, a_y)
-        # and `future` is our output (B, 20, 2) e.g. Batch x Time x (x, y)
 
-        # create all input data that we are allowed to give to a model
-        model_inputs = {'PAST': past, 'IMAGES': images, 'INTENT': intent}
+    # ---- shared step ----
+    def _shared_step(self, batch, stage: str):
+        past = batch["PAST"]
+        images = batch["IMAGES"]
+        intent = batch["INTENT"]
+        pref_traj = batch["PREF_TRAJ"]
+        pref_score = batch["PREF_SCORE"]
 
-        pred_future = self.forward(model_inputs)  # (B, T*2)
-        loss = self.ade_loss(pred_future.reshape_as(future), future)  # reshape to (B, T, 2
+        pred = self(
+            {
+                "PAST": past,
+                "IMAGES": images,
+                "INTENT": intent,
+                "PREF_TRAJ": pref_traj,
+            }
+        ).squeeze(-1)
 
-        # TODO: improve logging both to disk and to console
-        self.log_dict({
-            f"{stage}_loss": loss,
-        }, prog_bar=True, logger=True)
+        loss = F.smooth_l1_loss(pred, pref_score, beta=self.delta)
+
+        self.log(
+            f"{stage}_loss",
+            loss,
+            batch_size=past.size(0),
+            prog_bar=True,
+            on_epoch=True,
+            on_step=False,
+        )
+
+        if stage == "val":
+            err = (pred.detach() - pref_score.detach()).abs()
+            self.val_errors.append(err.cpu())
+            # Save predictions and ground truth for correlation
+            self.val_pred.append(pred.detach().cpu())
+            self.val_true.append(pref_score.detach().cpu())
 
         return loss
-    
-    def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
+
+    def training_step(self, batch, batch_idx):
         return self._shared_step(batch, "train")
-    
-    def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
+
+    def validation_step(self, batch, batch_idx):
         return self._shared_step(batch, "val")
-    
+
+    # ---- Lightning 2.x hook ----
+    def on_validation_epoch_end(self):
+        if not self.trainer.is_global_zero:
+            return
+
+        import os
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        if len(self.val_errors) == 0:
+            print("[WARN] No validation errors collected")
+            return
+
+        # Concatenate validation results
+        errors = torch.cat(self.val_errors).to(torch.float32).numpy()
+        preds = torch.cat(self.val_pred).to(torch.float32).numpy()
+        trues = torch.cat(self.val_true).to(torch.float32).numpy()
+
+        self.val_errors.clear()
+        self.val_pred.clear()
+        self.val_true.clear()
+
+        # Histogram
+        plt.figure(figsize=(6, 4))
+        plt.hist(errors, bins=50)
+        plt.axvline(np.mean(errors), linestyle="--", label="Mean")
+        plt.xlabel("|pred − gt|")
+        plt.ylabel("Count")
+        plt.title(f"Validation Error Histogram (epoch {self.current_epoch})")
+        plt.legend()
+        out_path = os.path.join("logs", f"val_error_hist_epoch_{self.current_epoch}.png")
+        os.makedirs(os.path.dirname(out_path), exist_ok=True)
+        plt.savefig(out_path)
+        plt.close()
+        print(f"[OK] Saved histogram → {out_path}")
+
+        # Spearman and Kendall correlations
+        spearman_corr = spearmanr(trues, preds).correlation
+        kendall_corr = kendalltau(trues, preds).correlation
+
+        print(f"[INFO] Spearman correlation: {spearman_corr:.4f}")
+        print(f"[INFO] Kendall tau correlation: {kendall_corr:.4f}")
+
+        # Log to Lightning
+        self.log("val_spearman", spearman_corr, prog_bar=True, logger=True)
+        self.log("val_kendall", kendall_corr, prog_bar=True, logger=True)
+
+# helps to batch data with images
 def collate_with_images(batch):
     past = [torch.as_tensor(b["PAST"], dtype=torch.float32) for b in batch]
     future = [torch.as_tensor(b["FUTURE"], dtype=torch.float32) for b in batch]
     intent = torch.as_tensor([b["INTENT"] for b in batch])
     names = [b["NAME"] for b in batch]
 
+    padded_trajs = []
+
+    for b in batch:
+        traj = torch.as_tensor(b["PREF_TRAJ"], dtype=torch.float32)
+        for _ in range(len(traj), 21):
+            traj = F.pad(traj, (0,0,0,1), value=0.0)  # pad to length 20
+        
+        padded_trajs.append(traj)
+
+    # Extra values for preference trajectories 
+    pref_traj = padded_trajs
+    pref_score = torch.as_tensor([b["PREF_SCORE"] for b in batch], dtype=torch.float32)
+
     cams = list(zip(*[b["IMAGES"] for b in batch]))  # per-camera tuples
     images = [torch.stack(cam_imgs, dim=0) for cam_imgs in cams]  # stay on CPU
 
@@ -86,5 +144,6 @@ def collate_with_images(batch):
         "INTENT": intent,
         "IMAGES": images,
         "NAME": names,
+        "PREF_TRAJ": torch.stack(pref_traj, dim=0),
+        "PREF_SCORE": pref_score,
     }
-
diff --git a/src/camera-based-e2e/models/monocular.py b/src/camera-based-e2e/models/monocular.py
index 2edd6cc..59c626d 100644
--- a/src/camera-based-e2e/models/monocular.py
+++ b/src/camera-based-e2e/models/monocular.py
@@ -7,28 +7,36 @@
 import timm
 from math import sqrt
 
-from .base_model import BaseModel, LitModel
+from .base_model import LitModel
 from .blocks import TransformerBlock
 
 class DINOFeatures(nn.Module):
-    def __init__(self, model_name: str = "vit_small_plus_patch16_dinov3.lvd1689m", frozen: bool = True):
+    def __init__(self, model_name: str = "vit_tiny_plus_patch16_dinov3.lvd1689m", frozen: bool = True):
         super(DINOFeatures, self).__init__()
 
+        # load in dino classification model
+        # features_only returns a list of layer outputs
         self.dino_model = timm.create_model(model_name, pretrained=True, features_only=True)
+
+        #sets up input preprocessing transforms
         self.data_config = timm.data.resolve_data_config(model=self.dino_model)
         self.transforms = timm.data.create_transform(**self.data_config, is_training=False)
+
+        # freezes the weights and doesn't continue training the model
         if frozen:
             for param in self.dino_model.parameters():
                 param.requires_grad = False
 
-        self.dims = [384, 384, 384]  # feature dims for each layer
+        self.dims = [384, 384, 384]  # feature the number of channels/features extracted at each layer
+        # intermediate layer outputs are also returned and useful
+
         self.patch_size = 16  # patch size
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         # x: (B, 3, H, W)
         # transforms: resize 256x256, center crop, normalize
-        x_t = self.transforms(x.float()) # preprocess
-        features = self.dino_model(x_t)
+        x_t = self.transforms(x.float()) # preprocess input
+        features = self.dino_model(x_t) # 3 layers all with 384 channels/features and 16x16 patches
         return features # 3 x [B, 384, 16, 16]
     
 class SAMFeatures(nn.Module):
@@ -49,12 +57,13 @@ def __init__(
                 param.requires_grad = False
 
         channels = self.sam_model.feature_info.channels()
-        reductions = self.sam_model.feature_info.reduction()
-        self.feature_stage = feature_stage
+        reductions = self.sam_model.feature_info.reduction() #input is shrunk by cumulative factor of reductions (array) at each stage
+
+        self.feature_stage = feature_stage # chooses which feature stage to output (-1 = last)
         self.dims = [channels[feature_stage]]
-        self.patch_size = reductions[feature_stage]  # effective stride
+        self.patch_size = reductions[feature_stage]  # sets patch size to reduction at that stage
 
-    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]: 
         # x: (B, 3, H, W)
         x_t = self.transforms(x.float())  # preprocess
         feats = self.sam_model(x_t)       # list of feature maps
@@ -73,14 +82,17 @@ def __init__(
 
         # attention 
         self.feature_dim = sum(self.features.dims)  # works for both DINO and SAM
+
+        #converts tokens into key, value spaces
         self.key_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim) # project into "key" space
         self.value_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim)
 
         # condition the query on intent (B,) and past (B, 16, 6)
-        query_input_dim = 3 + 16 * 6  # one hot -- concat -- flattened
+        query_input_dim = 3 + 16 * 6 + 21 * 2  # one hot -- concat -- flattened
         self.query = nn.Sequential(
             nn.Linear(query_input_dim, self.feature_dim),
             nn.LeakyReLU(),
+            nn.Dropout(p=0.1),
             nn.Linear(self.feature_dim, self.feature_dim),
         )
 
@@ -92,6 +104,7 @@ def __init__(
         self.decoder = nn.Sequential(
             nn.Linear(self.feature_dim, self.feature_dim),
             nn.LeakyReLU(),
+            nn.Dropout(p=0.1),
             nn.Linear(self.feature_dim, out_dim),
         )
 
@@ -104,6 +117,8 @@ def __init__(
     def forward(self, x: dict) -> torch.Tensor:
         # past: (B, 16, 6), intent: int
         past, images, intent = x['PAST'], x['IMAGES'], x['INTENT']
+
+        pref_traj = x['PREF_TRAJ']  
         
         # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203]
         front_cam = images[1]
@@ -124,7 +139,9 @@ def forward(self, x: dict) -> torch.Tensor:
 
         intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float()  # (B, 3). minus 1 --> 0, 1, 2
         past_flat = past.view(past.size(0), -1)  # (B, 96)
-        query = self.query(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1)  # (B, 1, 256)
+        pref_traj_flat = pref_traj.view(pref_traj.size(0), -1)  # (B, 42)
+        #print(pref_traj_flat.shape)
+        query = self.query(torch.cat([intent_onehot, past_flat, pref_traj_flat], dim=1)).unsqueeze(1)  
         query = self.query_norm(query)
 
         scores = query @ key.permute((0, 2, 1)) # (B, T, N)
@@ -140,7 +157,7 @@ def __init__(self, feature_extractor, out_dim, n_layers=1):
         self.feature_dim = sum(self.features.dims)
         
         # Initial Query Projection (Intent + Past -> C)
-        query_input_dim = 3 + 16 * 6
+        query_input_dim = 3 + 16 * 6 + 21 * 2
         self.query_init = nn.Linear(query_input_dim, self.feature_dim)
 
         # learnable positional encoding
@@ -163,6 +180,7 @@ def forward(self, x):
         # Copied from MonocularModel
         # past: (B, 16, 6), intent: int
         past, images, intent = x['PAST'], x['IMAGES'], x['INTENT']
+        pref_traj = x['PREF_TRAJ']
         
         # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203]
         front_cam = images[1]
@@ -179,7 +197,8 @@ def forward(self, x):
         # copy procedure to build query_0 from MonocularModel
         intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float()
         past_flat = past.view(past.size(0), -1)
-        query = self.query_init(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1)
+        pref_traj_flat = pref_traj.view(pref_traj.size(0), -1)  # (B, 42)
+        query = self.query_init(torch.cat([intent_onehot, past_flat, pref_traj_flat], dim=1)).unsqueeze(1)
 
         for block in self.blocks:
             query = block(query, tokens)
diff --git a/src/camera-based-e2e/point_cloud.py b/src/camera-based-e2e/point_cloud.py
new file mode 100644
index 0000000..e54e60b
--- /dev/null
+++ b/src/camera-based-e2e/point_cloud.py
@@ -0,0 +1,73 @@
+import open3d as o3d
+import numpy as np
+
+import open3d as o3d
+from protos import dataset_pb2 # Standard waymo import
+
+def get_waymo_intrinsics(frame, camera_name_enum):
+    """
+    Extracts intrinsics for a specific camera from a Waymo Frame proto.
+    
+    Args:
+        frame: The Waymo Open Dataset Frame proto.
+        camera_name_enum: Integer enum for the camera (e.g., 1 for FRONT).
+    
+    Returns:
+        o3d.camera.PinholeCameraIntrinsic object.
+    """
+    # 1. Find the calibration for the requested camera
+    calib = None
+    for c in frame.context.camera_calibrations:
+        if c.name == camera_name_enum:
+            calib = c
+            break
+
+    # Proto definition: [f_u, f_v, c_u, c_v, k1, k2, p1, p2, k3]
+    fx = calib.intrinsic[0]
+    fy = calib.intrinsic[1]
+    cx = calib.intrinsic[2]
+    cy = calib.intrinsic[3]
+    
+    width = calib.width
+    height = calib.height
+
+    # 3. Create Open3D Intrinsic Object
+    intrinsics = o3d.camera.PinholeCameraIntrinsic()
+    intrinsics.set_intrinsics(width, height, fx, fy, cx, cy)
+    
+    return intrinsics
+
+def create_point_cloud(rgb_image, depth_map, intrinsics=None, depth_scale=1.0):
+    """
+    Converts a single RGB image and Depth map into a 3D Point Cloud.
+    
+    Args:
+        rgb_image (np.array): Shape (H, W, 3) - uint8 [0-255]
+        depth_map (np.array): Shape (H, W) - float32
+        intrinsics (o3d.camera.PinholeCameraIntrinsic): Optional custom intrinsics.
+    """
+    height, width = depth_map.shape
+
+    # 1. Create Open3D Image objects
+    o3d_color = o3d.geometry.Image(rgb_image)
+    o3d_depth = o3d.geometry.Image(depth_map.astype(np.float32))
+
+    # 2. Create RGBD Image
+    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
+        o3d_color, 
+        o3d_depth, 
+        depth_scale=1.0, 
+        depth_trunc=10000.0, 
+        convert_rgb_to_intensity=False
+    )
+
+    # 4. Back-project to Point Cloud
+    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
+        rgbd_image, 
+        intrinsics
+    )
+    
+    # Flip the point cloud because Open3D uses Y-down
+    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
+    
+    return pcd
diff --git a/src/camera-based-e2e/preference_cache.pkl b/src/camera-based-e2e/preference_cache.pkl
new file mode 100644
index 0000000..abd77c5
Binary files /dev/null and b/src/camera-based-e2e/preference_cache.pkl differ
diff --git a/src/camera-based-e2e/preference_cache.py b/src/camera-based-e2e/preference_cache.py
new file mode 100644
index 0000000..1b40f63
--- /dev/null
+++ b/src/camera-based-e2e/preference_cache.py
@@ -0,0 +1,55 @@
+import os
+import struct
+import time
+import pickle
+from tqdm import tqdm
+import sys, os
+from protos import e2e_pb2
+sys.path.append(os.getcwd())
+
+DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/'
+
+indexes = []
+
+start = time.time()
+for file in tqdm(os.listdir(DATA_DIR)):
+    offset = 0
+    if '.tfrecord' in file and file.startswith('val') and not file.endswith('.gstmp'):
+        with open(os.path.join(DATA_DIR, file), 'rb') as fp:
+            offset = 0
+            while True:
+                raw_length = fp.read(8)
+                if len(raw_length) < 8:
+                    break
+                record_len = struct.unpack('Q', raw_length)[0]  # unsigned
+
+                fp.read(4)  # skip first CRC
+
+                start_payload_offset = offset + 8 + 4  # where protobuf actually starts
+
+                data = fp.read(record_len)
+                if len(data) != record_len:
+                    print(f"Incomplete record at {file}:{offset}")
+                    break
+
+                fp.read(4)  # skip last CRC
+
+                frame = e2e_pb2.E2EDFrame()
+                try:
+                    frame.ParseFromString(data)
+                except Exception as e:
+                    print(f"Failed to parse frame at {file}:{offset} -> {e}")
+                    offset += 8 + 4 + record_len + 4
+                    print(len(indexes))
+                    continue
+
+                if len(frame.preference_trajectories) and frame.preference_trajectories[0].preference_score != -1:
+                    indexes.append((file, start_payload_offset, record_len))
+
+                # update offset for next record
+                offset += 8 + 4 + record_len + 4
+        print(len(indexes))
+
+#pickle turns objects into bytes
+with open('preference_cache.pkl', 'wb') as f:
+    pickle.dump(indexes, f)
diff --git a/src/camera-based-e2e/preference_loader.py b/src/camera-based-e2e/preference_loader.py
new file mode 100644
index 0000000..6016959
--- /dev/null
+++ b/src/camera-based-e2e/preference_loader.py
@@ -0,0 +1,132 @@
+import torch
+from torch.utils.data import IterableDataset
+from protos import e2e_pb2
+import torchvision
+import pickle
+import struct
+import os
+import numpy as np
+from PIL import Image
+from io import BytesIO 
+import cv2 
+from typing import Optional
+import random
+from models.base_model import collate_with_images
+
+devices = ['cuda:0', 'cuda:1']
+
+random.seed(42) # Deterministic
+
+class WaymoE2E(IterableDataset): 
+    def __init__(
+        self,
+        indexFile = 'preference_cache.pkl',
+        data_dir='./dataset',
+        images = True,
+        n_items: Optional[int] = None,
+        seed: Optional[int] = None,
+    ):
+        self.images = images
+        self.data_dir = data_dir
+        self.seed = seed 
+
+        self.filename = ""
+        self.file = None
+
+        with open(indexFile, 'rb') as f:
+            # NOTE: test does not have reference trajectories
+            # We train on train and validate on val set
+            self.indexes = pickle.load(f)
+
+        if n_items is not None and n_items < len(self.indexes):
+            total = len(self.indexes)
+            # pick a deterministic contiguous block when a seed is provided
+            rng = random.Random(seed) if seed is not None else random
+            start = rng.randint(0, total - n_items)
+            self.indexes = self.indexes[start : start + n_items]
+
+    def decode_img(self, img):
+        if not self.images:
+            return np.array([])
+        
+        img_tensor = torch.from_numpy(np.frombuffer(img, dtype=np.uint8).copy()) 
+        gpu_tensors_list = torchvision.io.decode_jpeg(
+            img_tensor, 
+            mode=torchvision.io.ImageReadMode.UNCHANGED,
+            device= 'cpu' #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2]
+        )
+        # img_array = np.frombuffer(img, np.uint8)
+        return gpu_tensors_list
+    
+    def __len__(self):
+        return len(self.indexes)
+    
+    def __iter__(self):
+        # spawns a new worker, assigns worker.id and num_workers
+        worker = torch.utils.data.get_worker_info()
+        if worker is None:
+            start, step = 0, 1
+        else:
+            start, step = worker.id, worker.num_workers
+
+        # this code makes each worker read a subset of teh data (worker 1 = 0, 16, 32..., worker 2 = 1, 17, 33...)
+        for idx in range(start, len(self.indexes), step):
+            frame = e2e_pb2.E2EDFrame() 
+            filename, start_byte, byte_length = self.indexes[idx]
+
+            if self.filename != filename:
+                if self.file:
+                    self.file.close()
+                    del self.file
+                self.file = open(os.path.join(self.data_dir, filename), 'rb')
+                self.filename = filename
+
+            self.file.seek(start_byte) 
+            protobuf = self.file.read(byte_length) 
+            frame.ParseFromString(protobuf)
+
+            past = np.stack([frame.past_states.pos_x, frame.past_states.pos_y, frame.past_states.vel_x, frame.past_states.vel_y, frame.past_states.accel_x, frame.past_states.accel_y], axis=-1)
+    
+            future = np.stack([frame.future_states.pos_x, frame.future_states.pos_y], axis=-1)
+
+            past = np.array(past, dtype=np.float32) # ensure consistent dtype
+            future = np.array(future, dtype=np.float32)
+
+            # For submission to waymo evaluation server
+            name = frame.frame.context.name
+
+            for traj in frame.preference_trajectories:
+                # stack x, y
+                traj_array = np.stack([traj.pos_x, traj.pos_y], axis=-1)  # shape (T, 2), T = timesteps
+                yield {'PAST': past, 'FUTURE': future, 'IMAGES': [self.decode_img(images.image) for images in frame.frame.images], 'INTENT': frame.intent, 'NAME': name, 'PREF_TRAJ': traj_array, 'PREF_SCORE': traj.preference_score}
+
+if __name__ == "__main__":
+
+    from torch.utils.data import DataLoader
+    import time
+    from tqdm import tqdm
+    # NOTE: Replace with your path
+    DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/'
+    BATCH_SIZE = 32
+    dataset = WaymoE2E(indexFile="preference_cache.pkl", data_dir = DATA_DIR, images=True)
+    loader = DataLoader(
+        dataset, 
+        batch_size=BATCH_SIZE,
+        num_workers=16, 
+        collate_fn=collate_with_images
+    )
+    
+    #now we have this laoder class that I can iterate over like an array
+    #it will return dictionaries with keys PAST, FUTURE, IMAGES, INTENT
+
+    def main():
+        # start = time.time()
+        for batch_of_frames in tqdm(loader):
+            # print(batch_of_frames["INTENT"])
+            # print(batch_of_frames.keys(), [b.shape for b in batch_of_frames.values() if isinstance(b, torch.Tensor)])
+            pass
+        # print("Total Time:", time.time()-start)
+        print(len(loader))
+    
+    import cProfile
+    main()
\ No newline at end of file
diff --git a/src/camera-based-e2e/protos/dataset.proto b/src/camera-based-e2e/protos/dataset.proto
index 0615466..11d8f9e 100644
--- a/src/camera-based-e2e/protos/dataset.proto
+++ b/src/camera-based-e2e/protos/dataset.proto
@@ -58,7 +58,6 @@ message CameraName {
     REAR_LEFT = 6;
     REAR = 7;
     REAR_RIGHT = 8;
-
   }
 }
 
diff --git a/src/camera-based-e2e/protos/e2e.proto b/src/camera-based-e2e/protos/e2e.proto
index 0550d65..f7cbbff 100644
--- a/src/camera-based-e2e/protos/e2e.proto
+++ b/src/camera-based-e2e/protos/e2e.proto
@@ -55,7 +55,7 @@ message E2EDFrame {
   // assigned rater scores of -1 or left empty. Valid scores range from [0, 10].
   repeated EgoTrajectoryStates preference_trajectories = 8;
 
-
+  // Numbers here are field numbers (1,5,6,7,8)
 }
 
 message EgoTrajectoryStates {
diff --git a/src/camera-based-e2e/test.sh b/src/camera-based-e2e/test.sh
new file mode 100755
index 0000000..08f2ed0
--- /dev/null
+++ b/src/camera-based-e2e/test.sh
@@ -0,0 +1,61 @@
+#!/bin/bash
+
+# Configuration
+PARTITIONS=("v100" "a10" "a30")
+ACCOUNT="csso"
+GRES="gpu:1"
+MEM="128G"
+CPUS="16"
+TIME="8:00:00"
+
+echo "Checking estimated wait times for 8h job..."
+echo "--------------------------------------------"
+
+for PART in "${PARTITIONS[@]}"; do
+    NOW_TS=$(date +%s)
+
+    # Run sbatch test-only
+    OUTPUT=$(sbatch --test-only \
+        -A $ACCOUNT \
+        --gres=$GRES \
+        --mem=$MEM \
+        --cpus-per-task=$CPUS \
+        --partition=$PART \
+        -t $TIME \
+        --wrap="hostname" 2>&1)
+
+    if [[ "$OUTPUT" == *"allocation available"* ]] || [[ "$OUTPUT" == *"begins now"* ]]; then
+        echo "${PART}: 0:00:00 delay (Immediate)"
+    else
+        # 1. Capture the text after "start at"
+        EST_TIME_STRING=$(echo "$OUTPUT" | grep -oP 'start at \K.*')
+
+        # 2. FIX: Isolate ONLY the timestamp (take the first "word" of the string)
+        # This removes "using 16 processors on nodes..."
+        CLEAN_TIMESTAMP=$(echo "$EST_TIME_STRING" | awk '{print $1}')
+
+        if [[ -z "$CLEAN_TIMESTAMP" ]]; then
+            echo "${PART}: [Could not parse SLURM output] $OUTPUT"
+        else
+            # 3. Convert clean timestamp to Epoch seconds
+            EST_TS=$(date -d "$CLEAN_TIMESTAMP" +%s 2>/dev/null)
+
+            if [[ -z "$EST_TS" ]]; then
+                 echo "${PART}: [Date parse error] Raw: $CLEAN_TIMESTAMP"
+            else
+                DIFF=$((EST_TS - NOW_TS))
+                if [ $DIFF -lt 0 ]; then DIFF=0; fi
+
+                # Format output
+                FORMATTED_DELAY=$(date -u -d @"$DIFF" +'%H:%M:%S')
+                DAYS=$((DIFF / 86400))
+
+                if [ $DAYS -gt 0 ]; then
+                    FORMATTED_DELAY="${DAYS}d-${FORMATTED_DELAY}"
+                fi
+
+                echo "${PART}: ${FORMATTED_DELAY} delay (Est: $CLEAN_TIMESTAMP)"
+            fi
+        fi
+    fi
+done
\ No newline at end of file
diff --git a/src/camera-based-e2e/testRead.py b/src/camera-based-e2e/testRead.py
new file mode 100644
index 0000000..edfa9d1
--- /dev/null
+++ b/src/camera-based-e2e/testRead.py
@@ -0,0 +1,44 @@
+import os
+import struct
+import time
+import pickle
+from tqdm import tqdm
+import sys, os
+from protos import e2e_pb2
+
+DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/'
+indexes = []
+
+with open(os.path.join(DATA_DIR, 'val_202504211843.tfrecord-00027-of-00093_.gstmp'), 'rb') as fp:
+            offset = 0
+            while True:
+                raw_length = fp.read(8)
+                if len(raw_length) < 8:
+                    break
+                record_len = struct.unpack('Q', raw_length)[0]  # unsigned
+
+                fp.read(4)  # skip first CRC
+
+                start_payload_offset = offset + 8 + 4  # where protobuf actually starts
+
+                data = fp.read(record_len)
+                if len(data) != record_len:
+                    print(f"Incomplete record at {file}:{offset}")
+                    break
+
+                fp.read(4)  # skip last CRC
+
+                frame = e2e_pb2.E2EDFrame()
+                try:
+                    frame.ParseFromString(data)
+                except Exception as e:
+                    print(f"length of indexes: {len(indexes)}")
+                    print(f"Failed to parse frame at {file}:{offset} -> {e}")
+                    offset += 8 + 4 + record_len + 4
+                    continue
+
+                if len(frame.preference_trajectories) and frame.preference_trajectories[0].preference_score != -1:
+                    indexes.append((file, start_payload_offset, record_len))
+
+                # update offset for next record
+                offset += 8 + 4 + record_len + 4
\ No newline at end of file
diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py
index 875c942..028b888 100644
--- a/src/camera-based-e2e/train.py
+++ b/src/camera-based-e2e/train.py
@@ -12,8 +12,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from pathlib import Path
+import pickle
 
-from loader import WaymoE2E
+from preference_loader import WaymoE2E
+from torch.utils.data import random_split, DataLoader
 
 # Replace with your model defined in models/ 
 from models.base_model import LitModel, collate_with_images
@@ -28,17 +30,48 @@
     args = parser.parse_args()
 
     # Data 
-    train_dataset = WaymoE2E(batch_size=args.batch_size, indexFile='index_train.pkl', data_dir=args.data_dir, images=True, n_items=250000)
-    test_dataset = WaymoE2E(batch_size=args.batch_size, indexFile='index_val.pkl', data_dir=args.data_dir, images=True, n_items=50000)
+    # Initialize the full dataset
+    full_dataset = WaymoE2E(
+        indexFile='preference_cache.pkl',
+        data_dir=args.data_dir,
+        images=True,
+        n_items=250000
+    )
+
+    # Decide split sizes
+    train_size = int(0.8 * len(full_dataset))  # 80% train
+    val_size = len(full_dataset) - train_size  # 20% val
+
+    # Random split
+    # Load full dataset indexes
+    with open("preference_cache.pkl", "rb") as f:
+        indexes = pickle.load(f)
+
+    # Shuffle for randomness
+    import random
+    random.seed(42)
+    random.shuffle(indexes)
+
+    # Split 80/20
+    split = int(0.8 * len(indexes))
+    train_indexes = indexes[:split]
+    val_indexes = indexes[split:]
 
+    # Create train / val datasets
+    train_dataset = WaymoE2E(indexFile='preference_cache.pkl', data_dir=args.data_dir, images=True)
+    train_dataset.indexes = train_indexes  # override indexes
+
+    val_dataset = WaymoE2E(indexFile='preference_cache.pkl', data_dir=args.data_dir, images=True)
+    val_dataset.indexes = val_indexes  # override indexes
+    
     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=12, collate_fn=collate_with_images, persistent_workers=False, pin_memory=False)
-    val_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, num_workers=12, collate_fn=collate_with_images, persistent_workers=False, pin_memory=False)
+    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=12, collate_fn=collate_with_images, persistent_workers=False, pin_memory=False)
 
     # Model
-    in_dim = 16 * 6  # Past: (B, 16, 6)
-    out_dim = 20 * 2  # Future: (B, 20, 2)
+    in_dim = 16 * 6 + 3 + 21 * 2 
+    out_dim = 1  
 
-    model = DeepMonocularModel(feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb"), out_dim=out_dim)
+    model = MonocularModel(in_dim=in_dim, out_dim=out_dim, feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb"))
     lit_model = LitModel(model=torch.compile(model, mode="max-autotune"), lr=args.lr)
 
     base_path = Path(args.data_dir).parent.as_posix()
@@ -59,27 +92,6 @@
 
     trainer.fit(lit_model, train_loader, val_loader)
 
-    # Export loss graph to visualizations/
-    try:
-        base_path = Path(base_path)
-        run_dir = sorted((base_path / "logs").glob("camera_e2e_*"))[-1]  # newest run
-        metrics = pd.read_csv(run_dir / "version_0" / "metrics.csv")
-        train = metrics[metrics["train_loss"].notna()]
-        val = metrics[metrics["val_loss"].notna()]
-
-        plt.figure()
-        plt.plot(train["step"], train["train_loss"], label="train_loss")
-        plt.plot(val["step"], val["val_loss"], label="val_loss")
-        plt.xlabel("Step")
-        plt.ylabel("Loss")
-        plt.legend()
-        plt.tight_layout()
-        out = Path("./visualizations")
-        plt.savefig(out / "loss.png", dpi=200)
-    except Exception as e:
-        print(f"Could not save loss plot: {e}")
-
-
 
 
 
diff --git a/src/camera-based-e2e/train.slurm b/src/camera-based-e2e/train.slurm
new file mode 100644
index 0000000..4b1c872
--- /dev/null
+++ b/src/camera-based-e2e/train.slurm
@@ -0,0 +1,20 @@
+#!/bin/bash
+#SBATCH --job-name=waymo_e2e
+#SBATCH --output=logs/%x_%j.out
+#SBATCH --error=logs/%x_%j.err
+#SBATCH --partition=v100
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=20
+#SBATCH --mem=32G
+#SBATCH --time=24:00:00
+#SBATCH --account=csso
+
+# Load modules (edit for your cluster)
+module load cuda/12.1
+module load conda
+
+# Activate your environment
+conda activate waymo_env
+
+# Run training
+python train.py --data_dir /scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/
\ No newline at end of file
diff --git a/src/camera-based-e2e/view_pcd.py b/src/camera-based-e2e/view_pcd.py
new file mode 100644
index 0000000..270f2b1
--- /dev/null
+++ b/src/camera-based-e2e/view_pcd.py
@@ -0,0 +1,11 @@
+import open3d as o3d
+import sys
+
+# Usage: python view_pcd.py filename.pcd
+if len(sys.argv) < 2:
+    print("Please provide a filename.")
+else:
+    filename = sys.argv[1]
+    print(f"Loading {filename}...")
+    pcd = o3d.io.read_point_cloud(filename)
+    o3d.visualization.draw_geometries([pcd])
\ No newline at end of file
diff --git a/src/camera-based-e2e/waymoModel.py b/src/camera-based-e2e/waymoModel.py
new file mode 100644
index 0000000..2067d48
--- /dev/null
+++ b/src/camera-based-e2e/waymoModel.py
@@ -0,0 +1,428 @@
+# Copyright 2025 The Waymo Open Dataset Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""WOD E2E Rater Feedback Score."""
+
+from typing import Dict, List, Tuple
+import numpy as np
+
+
+_THRESHOLD_TIME_SECONDS = np.array([3, 5], dtype=np.int64)
+_BASE_THRESHOLDS = np.array([1.0, 1.8], dtype=np.float64)
+_MINIMUM_SCORE_OUTSIDE_TRUST_REGION = 4.0
+
+
+def get_lat_lng_thresholds(
+    init_speed: np.ndarray,  # [B]
+    lat_lng_threshold_multipliers: Tuple[float, float],
+) -> Tuple[np.ndarray, np.ndarray]:
+  """Get lateral and longitudinal thresholds."""
+  # Set and scale thresholds with the initial velocity
+  lat_threshold_multiplier, lng_threshold_multiplier = (
+      lat_lng_threshold_multipliers
+  )
+  
+  lat_thresholds = _BASE_THRESHOLDS * lat_threshold_multiplier  # [2]
+  lng_thresholds = _BASE_THRESHOLDS * lng_threshold_multiplier  # [2]
+  scale_by_init_speed = np.clip(
+      0.5 + 0.5 * (init_speed - 1.4) / (11 - 1.4), 0.5, 1.0
+  )  # [B]
+  lat_thresholds = scale_by_init_speed[..., None] * lat_thresholds  # [B, 2]
+  lng_thresholds = scale_by_init_speed[..., None] * lng_thresholds  # [B, 2]
+
+  return lat_thresholds, lng_thresholds
+
+
+def process_rater_specified_trajectories(
+    trajectory_batches: List[List[np.ndarray]],
+    trajectory_labels_batches: List[np.ndarray],
+    target_num_waypoints: int,
+    target_num_trajectories_per_batch: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+  """Processes rater-specified trajectories by truncating or padding.
+
+  Args:
+    trajectory_batches: A list where each element is a batch of trajectories. A
+      trajectory is represented as a NumPy array of waypoints.
+    trajectory_labels_batches: A list where each element is a NumPy array of
+      labels corresponding to a batch of trajectories.
+    target_num_waypoints: The fixed number of waypoints each trajectory should
+      have after processing. Trajectories longer than this will be truncated,
+      and shorter ones will be padded by duplicating their last waypoint.
+    target_num_trajectories_per_batch: The fixed number of trajectories each
+      batch should have after processing. Batches with more trajectories will be
+      truncated, and those with fewer will be padded by duplicating their last
+      trajectory (and its label).
+
+  Returns:
+    A tuple containing:
+      - processed_trajectory_batches: The processed trajectory batches after
+        truncation or padding them at both the trajectory and waypoint levels.
+      - processed_labels_batches: The corresponding processed label batches.
+
+  Raises:
+    ValueError: If the number of trajectory batches and label batches do not
+    match, or if within a batch, the number of trajectories and labels do not
+    match.
+  """
+  if len(trajectory_batches) != len(trajectory_labels_batches):
+    raise ValueError(
+        'The number of trajectory batches and label batches must be the same.'
+    )
+
+  processed_trajectory_batches_list = []
+  processed_labels_batches_list = []
+
+  # Iterate over each batch of trajectories and their corresponding labels
+  for i in range(len(trajectory_batches)):
+    current_trajectory_batch = trajectory_batches[i]
+    current_labels_batch = trajectory_labels_batches[i]
+    if len(current_trajectory_batch) != len(current_labels_batch):
+      raise ValueError(
+          'In each batch, the number of trajectories and labels must be the'
+          ' same.'
+      )
+
+    # --- Step 1: Truncate or pad the number of trajectories in the
+    # current batch ---
+    num_trajectories_in_batch = len(current_trajectory_batch)
+
+    if num_trajectories_in_batch > target_num_trajectories_per_batch:
+      # Truncate trajectories and labels
+      processed_batch_trajectories = current_trajectory_batch[
+          :target_num_trajectories_per_batch
+      ]
+      processed_batch_labels = current_labels_batch[
+          :target_num_trajectories_per_batch
+      ]
+    elif num_trajectories_in_batch < target_num_trajectories_per_batch:
+      # Pad trajectories and labels by duplicating the last element
+      num_to_pad = target_num_trajectories_per_batch - num_trajectories_in_batch
+      padding_trajectories = [current_trajectory_batch[-1]] * num_to_pad
+      padding_labels = [current_labels_batch[-1]] * num_to_pad
+
+      processed_batch_trajectories = (
+          current_trajectory_batch + padding_trajectories
+      )
+      processed_batch_labels = current_labels_batch.tolist() + padding_labels
+    else:
+      # Number of trajectories already matches the target
+      processed_batch_trajectories = current_trajectory_batch
+      processed_batch_labels = current_labels_batch
+
+    # --- Step 2: Truncate or pad waypoints for each trajectory
+    # in the processed batch ---
+    final_trajectories_for_current_batch = []
+
+    for trajectory in processed_batch_trajectories:
+      num_waypoints = len(trajectory)
+      if num_waypoints > target_num_waypoints:
+        # Truncate waypoints
+        processed_trajectory = trajectory[:target_num_waypoints]
+      elif num_waypoints < target_num_waypoints:
+        # Pad waypoints by duplicating the last waypoint
+        last_waypoint = trajectory[-1]
+        padding_waypoints = np.array(
+            [last_waypoint] * (target_num_waypoints - num_waypoints)
+        )
+        processed_trajectory = np.concatenate(
+            (trajectory, padding_waypoints), axis=0
+        )
+
+      else:
+        # Number of waypoints already matches the target
+        processed_trajectory = trajectory
+      final_trajectories_for_current_batch.append(processed_trajectory)
+
+    processed_trajectory_batches_list.append(
+        np.array(final_trajectories_for_current_batch)
+    )
+
+    # Labels are already processed at the batch level (number of trajectories)
+    # No per-label processing is typically needed unless labels have internal
+    # structure to pad/truncate
+    if isinstance(processed_batch_labels, np.ndarray):
+      processed_labels_batches_list.append(processed_batch_labels)
+    else:  # if it became a list during padding
+      processed_labels_batches_list.append(np.array(processed_batch_labels))
+
+  # Convert lists of batches to NumPy arrays
+
+  final_processed_trajectories = np.array(processed_trajectory_batches_list)
+  final_processed_labels = np.array(processed_labels_batches_list)
+
+  return final_processed_trajectories, final_processed_labels
+
+
+def get_rater_feedback_score(
+    inference_trajectories: np.ndarray,  # [B, I, T, 2]
+    inference_probs: np.ndarray,  # [B, I]
+    rater_specified_trajectories: List[
+        List[np.ndarray]
+    ],  # [[T1, 2], [T2, 2], ...], ...]
+    rater_feedback_labels: List[np.ndarray],  #  [[P1], [P2], ...]
+    init_speed: np.ndarray,  # [B]
+    lat_lng_threshold_multipliers: Tuple[float, float] = (1.0, 4.0),
+    decay_factor: float = 0.1,
+    frequency: int = 4,
+    length_seconds: int = 5,
+    default_num_of_rater_specified_trajectories: int = 3,
+    output_trust_region_visualization: bool = False,
+    minimum_score_outside_trust_region: float = _MINIMUM_SCORE_OUTSIDE_TRUST_REGION,
+) -> Dict[str, np.ndarray]:
+  """Get rater feedback score (https://waymo.com/open/challenges/2025/e2e-driving/).
+
+  Notations:
+  - B: batch size
+  - I: number of inference trajectories
+  - P: number of rater-specified trajectories
+  - T: number of timesteps
+
+  Args:
+    inference_trajectories: An array of inference trajectories with shape [B, I,
+      T, 2]
+    inference_probs: An array of inference probabilities with shape [B, I]
+    rater_specified_trajectories: An array of rater-specified trajectories with
+      shape [B, P, T, 2]
+    rater_feedback_labels: An array of rater feedback labels (scores between 0
+      and 10, both inclusive) with shape [B, P]
+    init_speed: A batch of initial velocities with shape [B]
+    lat_lng_threshold_multipliers: A tuple of latitude and longitude threshold
+      multipliers with shape [2]
+    decay_factor: A scalar score decay factor outside the trust region
+    frequency: The frequency (Hz) of trajectories to be considered.
+    length_seconds: The length (seconds) of trajectories to be considered.
+    default_num_of_rater_specified_trajectories: The default number of rater
+      specified trajectories to be used.
+    output_trust_region_visualization: Whether to output trust region
+      visualization.
+    minimum_score_outside_trust_region: The minimum score for inference
+      trajectories that are not fully within the trust region.
+
+  Returns:
+    A dictionary of final rater feedback score and output for visualization.
+  """
+  # We first process the rater-specified trajectories and labels by
+  # truncating or padding them to the same length.
+  # After processing, the shape of rater_specified_trajectories is
+  # [B, P, T, 2], and the shape of rater_feedback_labels is [B, P].
+  rater_specified_trajectories, rater_feedback_labels = (
+      process_rater_specified_trajectories(
+          rater_specified_trajectories,
+          rater_feedback_labels,
+          target_num_waypoints=length_seconds * frequency,
+          target_num_trajectories_per_batch=default_num_of_rater_specified_trajectories,
+      )
+  )
+
+  if inference_trajectories.shape[-2] != rater_specified_trajectories.shape[-2]:
+    raise ValueError(
+        'Inference and rater-specified trajectories must have the same number'
+        ' of timesteps.'
+    )
+
+  if (
+      inference_trajectories.shape[-2]
+      < _THRESHOLD_TIME_SECONDS.max() * frequency
+  ):
+    raise ValueError(
+        'Inference trajectories must have at least'
+        f' {_THRESHOLD_TIME_SECONDS.max()} timesteps.'
+    )
+
+  # Make rater-specified trajectories to include the origin
+  padded_rater_specified_trajectories = np.pad(
+      rater_specified_trajectories,
+      ((0, 0), (0, 0), (1, 0), (0, 0)),
+      constant_values=0,
+  )  # [B, P, T + 1, 2]
+
+  # Compute displacement vectors
+  displacement_vectors = (
+      padded_rater_specified_trajectories[..., 1:, :]
+      - padded_rater_specified_trajectories[..., :-1, :]
+  )  # [B, P, T, 2]
+
+  # Get unnormalized directions
+  lng_directions = displacement_vectors  # [B, P, T, 2]
+
+  # When displacement is zero, which means the vehicle did not move, we bring
+  # longitudinal directions from the previous timestep.
+  lng_magnitudes = np.linalg.norm(lng_directions, axis=-1)  # [B, P, T]
+
+  # At the first timestep, we set the longitudinal directions to be (1, 0).
+  # This is because the vehicle coordinate is used.
+  lng_directions[..., 0, 0] = np.where(
+      lng_magnitudes[..., 0] == 0,
+      1,
+      lng_directions[..., 0, 0],
+  )  # x-axis
+  lng_directions[..., 0, 1] = np.where(
+      lng_magnitudes[..., 0] == 0,
+      0,
+      lng_directions[..., 0, 1],
+  )  # y-axis
+
+  # For the rest of the timesteps, we bring the longitudinal directions from the
+  # previous timestep.
+  for t in range(1, lng_directions.shape[2]):
+    lng_directions[..., t, 0] = np.where(
+        lng_magnitudes[..., t] == 0,
+        lng_directions[..., t - 1, 0],
+        lng_directions[..., t, 0],
+    )  # x-axis
+    lng_directions[..., t, 1] = np.where(
+        lng_magnitudes[..., t] == 0,
+        lng_directions[..., t - 1, 1],
+        lng_directions[..., t, 1],
+    )  # y-axis
+
+  # Lateral directions are 90-degree counterclockwise rotation of longitudinal
+  # directions, i.e., (x_new, y_new) = (-y, x)
+  lat_directions = np.stack(
+      [lng_directions[..., 1] * -1, lng_directions[..., 0]], axis=-1
+  )  # [B, P, T, 2]
+
+  # Normalize directions
+  lng_directions = lng_directions / np.linalg.norm(
+      lng_directions, axis=-1, keepdims=True
+  )  # [B, P, T, 2]
+  lat_directions = lat_directions / np.linalg.norm(
+      lat_directions, axis=-1, keepdims=True
+  )  # [B, P, T, 2]
+
+  # Get longitudinal and lateral distances from rater-specified trajectories
+  rater_specified_to_inference_vectors = (
+      inference_trajectories[..., None, :, :, :]
+      - rater_specified_trajectories[..., None, :, :]
+  )  # [B, 1, I, T, 2] - [B, P, 1, T, 2] --> [B, P, I, T, 2]
+  lng_projections = np.sum(
+      lng_directions[..., None, :, :] * rater_specified_to_inference_vectors,
+      axis=-1,
+  )  # [B, P, I, T], directions are broadcasted to the inference trajectories
+  lat_projections = np.sum(
+      lat_directions[..., None, :, :] * rater_specified_to_inference_vectors,
+      axis=-1,
+  )  # [B, P, I, T], directions are broadcasted to the inference trajectories
+  lng_distances = np.abs(lng_projections)  # [B, P, I, T]
+  lat_distances = np.abs(lat_projections)  # [B, P, I, T]
+
+  # Filter distances at 3 and 5 seconds
+  selected_indices = _THRESHOLD_TIME_SECONDS * frequency - 1
+  lng_distances = lng_distances[..., selected_indices]  # [B, P, I, 2]
+  lat_distances = lat_distances[..., selected_indices]  # [B, P, I, 2]
+
+  lat_thresholds, lng_thresholds = get_lat_lng_thresholds(
+      init_speed, lat_lng_threshold_multipliers
+  )
+
+  outputs = {}
+
+  # ---------------------------------------------------------------------------
+  # Visualization
+  # ---------------------------------------------------------------------------
+  if output_trust_region_visualization:
+    center_x = rater_specified_trajectories[..., selected_indices, :][
+        ..., 0
+    ]  # [B, P, T (=2)]
+    center_y = rater_specified_trajectories[..., selected_indices, :][
+        ..., 1
+    ]  # [B, P, T (=2)]
+    width = 2 * lng_thresholds  # [B, 2]
+    height = 2 * lat_thresholds  # [B, 2]
+    angle = np.degrees(
+        np.arctan2(
+            displacement_vectors[..., selected_indices, :][..., 1],
+            displacement_vectors[..., selected_indices, :][..., 0],
+        )
+    )  # [B, P, T (=2)]
+
+    outputs.update({
+        'trust_region_center_x': center_x,  # [B, I, 2]
+        'trust_region_center_y': center_y,  # [B, I, 2]
+        'trust_region_width': width,  # [B, 2]
+        'trust_region_height': height,  # [B, 2]
+        'trust_region_angle': angle,  # [B, I, 2]
+    })
+
+  # ---------------------------------------------------------------------------
+  # Hard matching with decaying
+  # ---------------------------------------------------------------------------
+
+  # Normalize distances with thresholds
+  normalized_lng_distances = (
+      lng_distances / lng_thresholds[..., None, None, :]
+  )  # [B, P, I, 2]
+  normalized_lat_distances = (
+      lat_distances / lat_thresholds[..., None, None, :]
+  )  # [B, P, I, 2]
+
+  # Pick the maximum of the two normalized distances
+  normalized_distances = np.maximum(
+      normalized_lng_distances, normalized_lat_distances
+  )  # [B, P, I, 2]
+
+  # Mask to indicate if the inference trajectory is fully within the trust
+  # region, i.e., distance from trajectory i is near any rated trajectory p.
+  # For inferences not fully within the trust region, scores are clipped to
+  # `minimum_score_outside_trust_region` during score computation below.
+  # [B, P, I, 2] -> [B, I]
+  is_fully_within_trust_region = np.any(
+      np.all(normalized_distances <= 1.0, axis=3), axis=1)
+  outputs['is_fully_within_trust_region'] = (
+      is_fully_within_trust_region  # [B, I]
+  )
+
+  # Make scores flat within the trust region.
+  exponent = np.maximum(normalized_distances - 1.0, 0.0)
+  decay = decay_factor**exponent
+  # Scores between every inference i and rated trajectory p along x,y axes.
+  rater_feedback_scores_per_axis_pairwise = (
+      rater_feedback_labels[..., None, None] * decay
+  )  # [B, P, I, 2]
+
+  # Scores for each inference trajectory along x,y axes.
+  # Each inference trajectory is assigned a score based on its best match with
+  # a rated trajectory.
+  rater_feedback_score_per_axis_per_inference = np.amax(
+      rater_feedback_scores_per_axis_pairwise, axis=1
+  )  # [B, I, 2]
+
+  # Scores for each inference trajectory averaged over x,y axes.
+  rater_feedback_score_per_inference = np.mean(
+      rater_feedback_score_per_axis_per_inference,
+      axis=-1,
+  )  # [B, I]
+
+  # Clip scores for inferences not fully within the trust region.
+  rater_feedback_score_per_inference[~is_fully_within_trust_region] = (
+      np.maximum(
+          minimum_score_outside_trust_region,
+          rater_feedback_score_per_inference[~is_fully_within_trust_region],
+      )
+  )  # [B, I]
+
+  # Weighted sum over scores for each inference trajectory.
+  rater_feedback_score = np.sum(
+      rater_feedback_score_per_inference * inference_probs, axis=-1
+  )  # [B]
+  outputs['rater_feedback_score'] = rater_feedback_score  # [B]
+  # Updated the truncated or padded rater-specified trajectories.
+  # [B, P, T, 2]
+  outputs['rater_specified_trajectories'] = rater_specified_trajectories
+  # Updated the truncated or padded rater feedback labels.
+  # [B, P]
+  outputs['rater_feedback_labels'] = rater_feedback_labels
+
+  return outputs
\ No newline at end of file