From 4d942410ac2ccb07c8df2e035ab57701c1d5b10f Mon Sep 17 00:00:00 2001 From: Pratyush Date: Wed, 1 Apr 2026 12:49:52 -0400 Subject: [PATCH 1/6] to anvil --- src/camera-based-e2e/CHANGES.md | 530 ++++++++++++++++++ src/camera-based-e2e/debug_viz.py | 369 ++++++++++++ src/camera-based-e2e/loader.py | 14 +- src/camera-based-e2e/models/base_model.py | 384 +++++++++++-- .../models/feature_extractors.py | 25 + src/camera-based-e2e/models/proposal_init.py | 63 +++ .../models/proposal_planner.py | 76 +++ src/camera-based-e2e/models/refinement.py | 122 ++++ src/camera-based-e2e/models/scene_encoder.py | 60 ++ src/camera-based-e2e/models/scorer.py | 41 ++ src/camera-based-e2e/score_experiment_viz.py | 299 ++++++++++ .../scripts/run_train_proposal_login_gpu.sh | 47 ++ src/camera-based-e2e/train.py | 172 +++++- 13 files changed, 2133 insertions(+), 69 deletions(-) create mode 100644 src/camera-based-e2e/CHANGES.md create mode 100644 src/camera-based-e2e/debug_viz.py create mode 100644 src/camera-based-e2e/models/proposal_init.py create mode 100644 src/camera-based-e2e/models/proposal_planner.py create mode 100644 src/camera-based-e2e/models/refinement.py create mode 100644 src/camera-based-e2e/models/scene_encoder.py create mode 100644 src/camera-based-e2e/models/scorer.py create mode 100644 src/camera-based-e2e/score_experiment_viz.py create mode 100755 src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh diff --git a/src/camera-based-e2e/CHANGES.md b/src/camera-based-e2e/CHANGES.md new file mode 100644 index 0000000..b301ef2 --- /dev/null +++ b/src/camera-based-e2e/CHANGES.md @@ -0,0 +1,530 @@ +# Proposal Planner — Debug & Fix Changelog + +## Phase 1: Debug Instrumentation + +### Problem +The proposal planner had terrible performance after 10 epochs of training: +- Train ADE ~16.8m, Val ADE ~14.3m (extremely high for 5-second horizon) +- All 16 proposals collapsed to the same trajectory (zero regret, zero diversity) +- Score loss dominated total loss by ~25x + +### What was added + +**`models/debug_callbacks.py`** — `GradientDebugCallback` (Lightning Callback) +- Per-module gradient L2 norms: `scene_encoder`, `proposal_init`, `refinement` (per block), `scorer` +- Per-sublayer gradient norms within refinement (`cross_attn`, `mlp`, `traj_residual`, `traj_enc`) +- Per-sublayer gradient norms within scorer (`geom_proj`, `score_mlp`) and proposal_init (`ego_enc`, `traj_decoder`, `proposal_embed`) +- Parameter norms per module (detect weight explosion/vanishing) +- Gradient dominance ratio (max_module / total — detect when one module dominates) +- Activation statistics via forward hooks: mean, std, abs_max, dead fraction, saturated fraction +- Proposal diagnostics: endpoint spread, pairwise diversity, trajectory lengths, score distribution, refinement delta + +**`debug_viz.py`** — Post-training diagnostic plot generator +- 10 individual plots + 1 summary dashboard +- Reads `metrics.csv` from Lightning CSVLogger +- Can be run standalone: `python debug_viz.py --log_dir /path/to/logs` + +**`train.py`** changes: +- `--debug` flag enables the callback and auto-generates plots after training +- `--debug_log_every N` controls logging frequency + +### Diagnosis from debug plots + +| Signal | Observation | Implication | +|--------|-------------|-------------| +| `loss_score=427` vs `loss_ade=16.8` | Score loss 25x larger | Score loss dominates optimization | +| `ade_pred ≈ ade_oracle`, regret ≈ 0 | All proposals identical | Complete mode collapse | +| `grad_norm/proposal_init` ~1000 | 10x higher than everything else | 90% of gradient signal goes through proposal_init | +| `grad_norm/scene_encoder` ~100 | Order of magnitude lower | Scene encoder barely learns | +| `proposals/pairwise_dist` decreasing | Proposals converging | Diversity loss too weak | +| `proposals/score_range` near 0 | Scorer can't discriminate | All proposals look the same to scorer | +| `act/scene_enc_out_mean` ≈ 0 | Scene features near zero | Visual info not reaching proposals | +| `grad_norm/propinit_embed` very low | Proposal seeds barely updated | Can't maintain diversity | + +### Root cause chain +1. Score loss (MSE to ADE, weight=1.0) produces loss ~400, dwarfing ADE loss ~15 +2. Score loss gradient flows through scorer + proposal_init, bypassing refinement/scene_encoder +3. Without gradient pressure to diversify, all proposals collapse to the mean trajectory +4. Once collapsed, diversity loss = 0, refinement has nothing to differentiate, cross-attention degenerates +5. Scene encoder never gets gradient signal because its only path (refinement cross-attn) is starved + +--- + +## Phase 2: Architecture & Loss Fixes + +### Fix 1: Scene-conditioned ProposalInit (`models/proposal_init.py`) + +**Problem:** ProposalInit only saw ego trajectory + intent. Camera images had no direct path +to the ADE loss — their only route was through the refinement cross-attention bottleneck, +which received 10x less gradient signal than proposal_init. + +**Fix:** Added cross-attention from proposal embeddings to scene feature tokens inside +ProposalInit. This gives the scene encoder a direct gradient path: +`ADE loss → proposals → traj_decoder → scene_cross_attn → scene_encoder.proj` + +Also increased `proposal_embed` init std from 0.02 to 0.1 to provide more initial diversity +among proposal seeds (wider starting spread resists early collapse). + +### Fix 2: Score loss downweight + warmup (`models/base_model.py`) + +**Problem:** Score loss (MSE between predicted scores and actual ADEs) was weighted 1.0x, +producing loss ~400 vs ADE loss ~15. This overwhelmed the trajectory quality signal and +drove all proposals toward identical outputs. + +**Fix:** +- New `score_weight` parameter (default 0.1 instead of implicit 1.0) +- New `score_warmup_epochs` parameter (default 2): score loss is 0 for the first N epochs, + then ramps linearly to `score_weight` over 1 epoch. This lets proposals diversify before + the scorer starts pulling them toward similar predictions. + +### Fix 3: Detach proposals from score loss gradient (`models/base_model.py`) + +**Problem:** The scorer's MSE loss backpropagated through both the scorer weights AND the +proposal generation pipeline. Since all proposals had similar ADE, the scorer gradient +effectively pushed proposals to produce outputs that are easy to score (i.e., identical). + +**Fix:** `pred_scores.detach().argmin(dim=1)` when selecting the best proposal for ADE +computation. The scorer still learns to rank, but its gradients only update scorer weights, +not the proposal generator. The proposal generator is driven only by ADE + diversity losses. + +### Fix 4: Diversity weight increase (`train.py`, `sbatch`) + +**Problem:** Diversity weight was 0.1, but the diversity loss magnitude (~0.3) was negligible +compared to score loss (~400). Even after downweighting score loss, diversity needs to be +strong enough to prevent collapse. + +**Fix:** Default `diversity_weight` changed from 0.1 to 1.0. + +### Fix 5: Gradient clipping (`train.py`) + +**Problem:** Gradient norms reached ~1000 with high variance (spikes visible in plots), +causing unstable updates especially for the scorer. + +**Fix:** Added `gradient_clip_val=1.0` with norm clipping to the Lightning Trainer. +New `--grad_clip` CLI arg (0 to disable). + +### Summary of parameter changes + +| Parameter | Before | After | Reason | +|-----------|--------|-------|--------| +| `score_weight` | 1.0 (implicit) | 0.1 | Was 25x larger than ADE, dominated optimization | +| `score_warmup_epochs` | 0 (N/A) | 2 | Let proposals diversify before scorer activates | +| `diversity_weight` | 0.1 | 1.0 | Too weak to prevent collapse against score loss | +| `proposal_embed` std | 0.02 | 0.1 | Wider initial spread resists early collapse | +| `gradient_clip_val` | None | 1.0 | Stabilize training, prevent gradient spikes | +| ProposalInit inputs | past + intent | past + intent + scene_feat | Give scene encoder a direct gradient path | + +### Files changed + +| File | Change | +|------|--------| +| `models/proposal_init.py` | Added scene cross-attention, increased embed init std | +| `models/proposal_planner.py` | Pass `scene_feat` to `proposal_init()` | +| `models/base_model.py` | Added `score_weight`, `score_warmup_epochs`, detached scorer gradients | +| `train.py` | New CLI args: `--score_weight`, `--score_warmup_epochs`, `--grad_clip`; updated defaults | +| `scripts/run_train_proposal.sbatch` | Updated args for new defaults | +| `models/debug_callbacks.py` | Created (Phase 1) | +| `debug_viz.py` | Created (Phase 1) | +| `CHANGES.md` | This file | + +--- + +## Phase 3: iPad-inspired Architectural Changes + +Reference: [iPad — Iterative Proposal-centric End-to-End Autonomous Driving](https://arxiv.org/abs/2505.15111) + +Phase 2 fixed the optimization dynamics (score loss no longer dominates, proposals no +longer collapse), but raw ADE stayed at ~14.3m throughout training. The model +architecture itself had structural limitations that prevented learning to generate +accurate trajectories. Phase 3 addresses three key gaps relative to the iPad paper. + +### Fix 6: BCE scorer loss with quality target (`models/base_model.py`, `models/scorer.py`) + +**Problem:** The old MSE scorer trained on raw ADE values (range 0–50+). This +produced unbounded loss magnitudes, required aggressive downweighting, and gave the +scorer no natural notion of "good" vs "bad" — an ADE of 12 vs 13 had the same loss +gradient magnitude as 1 vs 2, even though the latter distinction is far more important. + +**Fix:** Replaced `MSE(logit, ADE)` with `BCE(sigmoid(logit), exp(-ADE/τ))`: +- Quality target `exp(-ADE/τ)` maps ADE into [0, 1] — lower ADE → higher quality +- BCE loss is naturally bounded (~0.0 to ~0.7), no need for aggressive downweighting +- The exponential mapping concentrates learning on distinguishing good proposals + (low ADE) rather than wasting capacity on large-ADE differences +- Temperature τ (default 5.0) controls sensitivity: lower τ → sharper discrimination +- Scorer output semantics flipped: **higher score = better** (was: lower = better) +- All `argmin` for score-based selection → `argmax` throughout the codebase +- Default `score_weight` raised from 0.1 to 1.0 (BCE is well-scaled, safe at 1.0) + +New CLI arg: `--score_temperature` (default 5.0) + +### Fix 7: Weight-shared iterative refinement (`models/refinement.py`) + +**Problem:** The old refinement used `num_steps` separate `RefinementBlock` instances, +each with independent weights. This meant: +- Each block only saw gradients from a single position in the chain +- Parameter count scaled linearly with `num_steps` (more steps = more parameters) +- No iterative self-correction: each block learned a fixed transformation, not a + general "look at the scene and improve proposals" operation + +**Fix:** iPad-style weight sharing — a **single** `RefinementBlock` applied +`num_steps` times in a loop: +- Same weights at every iteration → the block learns a general refinement operation +- Increasing `num_steps` (e.g., 2 → 4) adds compute but zero extra parameters +- Gradients flow through all iterations (like an RNN), giving the shared weights + stronger learning signal from multiple applications + +**Practical benefit:** `--num_refinement_steps 4` now uses the same parameter count +as the old 2-step model, but applies 4 iterations of refinement. + +### Fix 8: Full trajectory re-prediction (`models/refinement.py`) + +**Problem:** The old `RefinementBlock` predicted **residual** waypoint deltas: +`proposals_new = proposals + traj_residual(feat)`. This constrained each iteration +to small perturbations around the previous trajectory, preventing the model from +fundamentally correcting bad initial proposals. + +**Fix:** Replaced residual prediction with **full re-prediction**: +`proposals_new = traj_decoder(feat)`. At each iteration the model predicts complete +new trajectories from scratch based on the current proposal features. The features +still accumulate information across iterations (via residual connections in the +feature update path), but the trajectories are free to change substantially. + +This matches iPad's approach where proposals are predicted fresh at each iteration +of the ProFormer loop. + +### Summary of Phase 3 parameter changes + +| Parameter | Phase 2 | Phase 3 | Reason | +|-----------|---------|---------|--------| +| Score loss | MSE | BCE with `exp(-ADE/τ)` target | Bounded loss, [0,1] quality scale | +| `score_weight` | 0.1 | 1.0 | BCE is naturally well-scaled | +| `score_temperature` | N/A | 5.0 (new) | Controls quality target sensitivity | +| Score semantics | lower = better | higher = better | Matches BCE probability interpretation | +| Refinement weights | Independent per step | Shared single block | Stronger gradient signal, free to add iterations | +| `num_refinement_steps` | 2 | 4 (sbatch) | More iterations at zero parameter cost | +| Trajectory prediction | Residual (`proposals + delta`) | Full re-prediction | Can correct bad proposals, not just perturb | + +### Phase 3 files changed + +| File | Change | +|------|--------| +| `models/refinement.py` | Single shared `RefinementBlock`, `traj_residual` → `traj_decoder`, full re-prediction | +| `models/scorer.py` | Updated docstring: higher = better, BCE-trained | +| `models/base_model.py` | BCE loss with quality target, `argmin` → `argmax`, `score_temperature` param | +| `models/debug_callbacks.py` | Updated for shared block (`.block` not `.blocks`), flipped score argmin/argmax | +| `train.py` | New `--score_temperature` arg, `score_weight` default 0.1 → 1.0 | +| `scripts/run_train_proposal.sbatch` | `num_refinement_steps` 2 → 4, `score_weight` 0.1 → 1.0, added `score_temperature` | +| `CHANGES.md` | This section | + +--- + +## Phase 4: Full iPad-faithful Rewrite + +Reference: [iPad source code](https://github.com/Kguo-cs/iPad) and [paper appendix](https://arxiv.org/html/2505.15111v1) + +Phase 3 changes did not move the needle: ADE stayed at 14.3m with complete mode +collapse (all proposals identical, zero regret, zero diversity, score loss stuck at +ln(2)=0.693). After studying the iPad source code in detail, several fundamental +architectural and loss differences were identified. Phase 4 is a faithful rewrite +copying the core iPad design patterns. + +### Root cause: Why Phases 2-3 failed + +The model was still stuck because: +1. **Top-5 WTA L2 loss** averaged the best 5 proposals' L2 errors — this pulled ALL + proposals toward the mean trajectory, actively fighting diversity. +2. **One feature per proposal** (B, K, C) was too limited to represent a 20-timestep + trajectory. Each proposal had a single 256-dim vector to encode 40 coordinates. +3. **No intermediate supervision** — only the final proposals were supervised. The + shared refinement block had no gradient signal for early iterations. +4. **Diversity loss as a band-aid** — negative mean pairwise distance is a weak signal. + iPad doesn't use diversity loss at all; MoN naturally handles it. + +### Fix 9: Per-timestep proposal features (`models/proposal_init.py`) + +**Problem:** Each proposal had one feature vector (B, K, C). A single 256-dim vector +must encode all 20 timesteps' (x,y) positions, severely limiting expressiveness. + +**Fix:** Learnable embeddings of shape (N×T, C) = (16×20, C) = (320, C). Each +proposal gets T=20 separate feature vectors, one per timestep. This matches iPad's +`init_feature = nn.Embedding(poses_num * proposal_num, tf_d_model)`. + +The ego status (past trajectory + intent) is encoded to a single vector and broadcast- +added to all N×T tokens, exactly matching iPad's pattern: +`bev_feature = ego_feature + init_feature.weight[None]` + +Removed: scene cross-attention from ProposalInit (iPad doesn't have it — scene +features are incorporated only through refinement cross-attention). + +### Fix 10: Per-timestep trajectory decoding (`models/refinement.py`) + +**Problem:** The old decoder mapped a single (B, K, C) feature to a full (B, K, T×2) +trajectory via one MLP. This forced a single vector to predict 40 output dimensions. + +**Fix:** Each timestep's feature independently predicts that timestep's (x,y): +`proposals = traj_decoder(bev_feature) # (B, N*T, C) → (B, N*T, 2) → (B,N,T,2)` + +The decoder is a 2-layer MLP with only 2 output dimensions per timestep (matching +iPad's `MLP(d_model, d_ffn, state_size)` where state_size=2 for us). + +### Fix 11: Predict-then-refine loop with intermediate proposals (`models/refinement.py`) + +**Problem:** Old refinement encoded proposals back and then predicted residuals. +No intermediate proposals were returned for supervision. + +**Fix:** iPad-faithful predict→encode→attend→refine loop: +1. `proposals = traj_decoder(bev_feature)` — predict from current features +2. `bev_feature += traj_enc(proposals)` — encode predictions back (proposal-anchoring) +3. `bev_feature += cross_attn(bev_feature, scene_feat)` — attend to scene +4. `bev_feature += mlp(bev_feature)` — FFN update +5. Return both proposals and updated features +6. Collect `proposal_list` across all K iterations for intermediate supervision + +### Fix 12: MoN L1 loss with discounted intermediate supervision (`models/base_model.py`) + +**Problem:** Top-5 WTA L2 ADE loss averaged the best 5 modes, pulling all proposals +toward the mean and causing mode collapse. Only final proposals were supervised. + +**Fix:** Faithful copy of iPad's loss: +- **MoN L1**: `min_n mean_t (|Δx| + |Δy|)` — Minimum over N proposals of mean L1 + displacement. Only the SINGLE closest proposal to GT is optimized, leaving others + free to explore different modes. L1 is more robust to outliers than L2. +- **Discounted intermediate supervision**: `L = Σ_k λ^(K-1-k) × MoN_L1(P_k)` + with λ=0.1 (iPad's `prev_weight`). Later iterations get stronger supervision, + earlier ones are relaxed. This gives gradient signal at every refinement step. +- **Removed diversity loss** (set to weight 0). MoN naturally encourages diversity + because only one proposal is pulled toward GT per sample — the others receive no + gradient and can remain spread out. iPad uses `inter_weight=0`. + +### Fix 13: Simplified scorer (`models/scorer.py`) + +**Problem:** Old scorer fused learned proposal features with handcrafted geometric +features (velocity norms, acceleration norms, flat waypoints). This was complex and +the geometric features were redundant given per-timestep features. + +**Fix:** Match iPad's scorer exactly: +1. Reshape per-timestep features: (B, N×T, C) → (B, N, T, C) +2. Max-pool over temporal dimension: (B, N, T, C) → (B, N, C) +3. MLP → scalar logit: (B, N, C) → (B, N) + +This is simpler and matches iPad's `proposal_feature = bev_feature.amax(-2)` followed +by `pred_score(proposal_feature)`. + +### Fix 14: Score quality target uses L1 (`models/base_model.py`) + +The scorer's BCE quality target now uses L1 displacement (consistent with the +trajectory loss) instead of L2 ADE: `quality = exp(-L1_per_mode / τ)`. + +### Summary of Phase 4 parameter changes + +| Parameter | Phase 3 | Phase 4 | Reason | +|-----------|---------|---------|--------| +| Trajectory loss | Top-5 WTA L2 ADE | MoN L1 | Only best proposal supervised; L1 more robust | +| Intermediate supervision | None | Discounted λ=0.1 | Gradient signal at every refinement step | +| `prev_weight` | N/A | 0.1 (new) | iPad discount factor for earlier iterations | +| Proposal features | (B, K, C) | (B, K×T, C) | Per-timestep features, 20× more expressive | +| Trajectory decoder | MLP(C → T×2) per proposal | MLP(C → 2) per timestep | Each timestep independently decoded | +| Scorer | Geometric + learned features | Max-pool temporal + MLP | Simpler, matches iPad | +| `diversity_weight` | 1.0 | 0.0 | MoN handles diversity; explicit loss fights it | +| `smoothness_weight` | 0.01 | 0.0 | Simplify loss to match iPad | +| `comfort_weight` | 0.01 | 0.0 | Simplify loss to match iPad | +| `max_epochs` | 10 | 20 | iPad trains for 20 epochs | + +### Phase 4 files changed + +| File | Change | +|------|--------| +| `models/proposal_init.py` | Rewritten: per-timestep embeddings (N×T, C), removed scene cross-attn | +| `models/refinement.py` | Rewritten: predict-then-refine loop, returns proposal_list | +| `models/scorer.py` | Rewritten: max-pool temporal + MLP, removed geometric features | +| `models/proposal_planner.py` | Updated: new calling pattern, passes proposal_list through | +| `models/base_model.py` | MoN L1 loss, discounted intermediate supervision, removed diversity loss | +| `models/debug_callbacks.py` | Updated for new architecture (scorer, propinit, proposal_list) | +| `train.py` | New `--prev_weight` arg, `diversity_weight` default 0, `num_refinement_steps` default 4 | +| `scripts/run_train_proposal.sbatch` | Updated all params, 20 epochs | +| `CHANGES.md` | This section | + +--- + +## Phase 5: Scorer Improvements (Ranking-focused) + +Phase 4 achieved strong multimodal proposals (val `ade_oracle` ~0.57 m) but deployed +trajectory quality was bottlenecked by scorer ranking: val `ade_pred` ~1.66 m with +`ade_regret` ~1.09 m (66% of total error). The scorer couldn't reliably pick the +best proposal from the diverse set MoN produced. + +### Fix 15: Cross-entropy ranking loss (`models/base_model.py`) + +**Problem:** Per-proposal BCE with soft quality targets `exp(-L1/τ)` does not directly +optimize for correct ranking. The targets for good vs mediocre proposals are close +together (e.g., 0.90 vs 0.82 with τ=5), giving weak ranking signal. A scorer can +achieve low BCE loss while still ranking proposals incorrectly. + +**Fix:** Replaced BCE with cross-entropy classification: +`loss_score = CrossEntropy(logits, argmin(L1_per_mode))`. This directly optimizes +for `argmax(scores) == argmin(L1)`, exactly matching how proposals are selected at +inference. The `score_temperature` parameter is no longer used by the loss. + +### Fix 16: Detached scorer inputs (`models/scorer.py`) + +**Problem:** Score loss gradients flowed back through the proposal generation pipeline +via `bev_feature`, conflicting with MoN trajectory loss. MoN wants only one proposal +to move toward GT; score BCE pushed all proposal features toward their individual +quality targets through shared representations. + +**Fix:** `bev_feature.detach()` before pooling in the scorer forward pass. Score loss +now only updates scorer weights, making it a pure discriminator. MoN has exclusive +control over proposal generation — no gradient conflict. + +### Fix 17: Trajectory-aware scorer (`models/scorer.py`) + +**Problem:** The scorer only saw abstract learned features (`bev_feature`). After +cross-attention and MLP updates, features for different proposals can be very similar +even when decoded trajectories differ substantially. The scorer was ranking based on +features that didn't fully capture trajectory-level differences. + +**Fix:** Concatenate projected trajectory coordinates with pooled BEV features: +- `traj_proj`: MLP mapping (T*2) to 64-dim trajectory embedding per proposal +- `score_mlp` input: `[max_pooled_feat | traj_feat]` with LayerNorm before the MLP +- Both `bev_feature` and `proposals` are detached so score loss is fully isolated + +This gives the scorer explicit spatial information (velocity profiles, endpoint +locations, curvature) for more discriminative ranking. + +### Summary of Phase 5 changes + +| Parameter | Phase 4 | Phase 5 | Reason | +|-----------|---------|---------|--------| +| Score loss | BCE with `exp(-L1/tau)` target | Cross-entropy on `argmin(L1)` | Directly optimizes ranking for `argmax` selection | +| Scorer input gradient | Flows back to proposals | Detached | Eliminates gradient conflict with MoN | +| Scorer input features | Max-pooled BEV only | BEV + projected trajectory coords | Explicit spatial discrimination | +| Scorer MLP input | `d_model` | `d_model + traj_dim` with LayerNorm | Richer, normalized input | +| `score_temperature` | Used (tau=5.0) | Unused (kept as CLI arg) | CE loss doesn't need soft targets | + +### Phase 5 files changed + +| File | Change | +|------|--------| +| `models/scorer.py` | Rewritten: detached inputs, traj_proj, LayerNorm, updated forward signature | +| `models/proposal_planner.py` | Pass `proposals` to `scorer(bev_feature, proposals)` | +| `models/base_model.py` | BCE to cross-entropy with `argmin(L1)` label | +| `models/debug_callbacks.py` | Added `scorer_traj_proj` gradient norm logging | +| `CHANGES.md` | This section | + +--- + +## Phase 6: Multi-loss scorer experiments + +To enable parallel scorer-loss ablations from one codebase, scorer loss is now +configurable via CLI: + +- `bce`: iPad-faithful BCE with soft quality target `exp(-L1/tau)` +- `ce`: hard top-1 cross-entropy on `argmin(L1)` +- `bce_pairwise`: BCE + pairwise margin-ranking auxiliary +- `listnet`: listwise KL objective with `softmax(-L1/tau)` target distribution + +New args in `train.py`: +- `--score_loss_type` +- `--score_rank_weight` +- `--score_margin` +- `--score_topk` + +Additional scorer diagnostics are logged to metrics: +- `*_score_top1_acc`: `argmax(score) == argmin(L1)` accuracy +- `*_score_gap_best_second`: average logit gap between oracle-best and second-best mode + +--- + +## Phase 7: NAVSIM-style quality target + +### Motivation (iPad paper Section 3.3) + +The iPad scorer uses a multi-factor ground-truth score from log-replay simulation +(NAVSIM Eq. 5): `S = NC * DAC * (5*EP + 5*TTC + 2*Comf) / 12`, covering safety, +efficiency, and comfort — not just trajectory closeness. Our `exp(-L1/tau)` target +only captures geometric accuracy, which may explain why the scorer struggles to +discriminate between proposals that are close in L1 but differ in driving quality. + +### Fix 18: Geometric NAVSIM approximation (`models/base_model.py`) + +Without agent trajectories or road boundaries in the current data pipeline, we +approximate the NAVSIM sub-metrics from trajectory geometry alone: + +| Sub-metric | iPad weight | Our approximation | +|---|---|---| +| EP (Ego Progress) | 5/12 | `clamp(dot(prop_disp, gt_dir) / gt_dist, 0, 1)` | +| TTC (Time-to-Collision) | 5/12 | 1.0 (no agent data) | +| Comf (Comfort) | 2/12 | fraction of timesteps with jerk < threshold | +| NC (No at-fault Collision) | gate | 1.0 (no agent data) | +| DAC (Drivable Area Compliance) | gate | 1.0 (no map data loaded) | + +Combined: `quality = (5*EP + 5*1.0 + 2*Comf) / 12`, clamped to [0, 1]. + +New `_compute_quality_target()` method dispatches between `l1` (existing) and +`navsim` (new) based on `--score_target_type` CLI arg. + +### New CLI args + +- `--score_target_type`: `l1` (default, backward-compatible) or `navsim` +- `--comfort_jerk_threshold`: jerk threshold in m/s^3 for comfort metric (default 5.0) + +### Files changed + +| File | Change | +|------|--------| +| `models/base_model.py` | Added `_compute_navsim_score()`, `_compute_quality_target()`, new hparams | +| `train.py` | New `--score_target_type`, `--comfort_jerk_threshold` args | +| `scripts/run_train_proposal.sbatch` | Added new default args | +| `score_experiment_viz.py` | Tracks `score_target_type` in approach labels and summary | +| `CHANGES.md` | This section | + +--- + +## Phase 8: RFS-based scorer quality target + +### Motivation + +For datasets where full NAVSIM-style supervision is unavailable, we want the scorer +to learn proposal ranking from an RFS-style quality signal (longitudinal/lateral +deviation with speed-aware thresholds) instead of only `exp(-L1/tau)`. + +### Fix 19: Add per-proposal RFS quality target (`models/base_model.py`) + +Added a new scorer target path that computes `(B, K)` quality directly from RFS +logic for **all proposals** (not only top-1): + +- New method: `_compute_rfs_quality(proposals, reference, past)` + - Uses the same directional decomposition as existing `rfs_loss`: + - longitudinal and lateral error components + - time thresholds at 3s / 5s (indices 11, 19) + - speed scaling from the existing RFS formula + - Converts deviation to score with the same piecewise definition: + - `1.0` when within threshold + - `0.1 ** (deviation - 1)` beyond threshold + - Averages over selected timesteps to get per-proposal quality in `[0,1]` + - Optional comfort multiplier from jerk-based comfort term +- Extended `_compute_quality_target(...)`: + - Supports `score_target_type == "rfs"` + - Requires `past` for speed scaling +- Updated scorer call site in `_shared_step(...)` to pass `past` into + `_compute_quality_target(...)`. + +### New/updated hyperparameters and CLI + +- `train.py` + - `--score_target_type`: now supports `l1 | navsim | rfs` + - `--no_rfs_target_comfort`: disables jerk comfort multiplier for RFS targets +- `LitModel` hparams + - `rfs_target_use_comfort: bool = True` + +### Training script default + +- `scripts/run_train_proposal.sbatch` + - Default scorer target changed from `l1` to `rfs` + - Added inline note that this can be overridden via `EXTRA_ARGS` + +### Notes + +- Existing standalone trajectory-side `loss_rfs` path remains intact. +- If using `--score_target_type rfs`, keep `rfs_weight=0` unless you explicitly + want both scorer-target RFS and trajectory regularization RFS active together. diff --git a/src/camera-based-e2e/debug_viz.py b/src/camera-based-e2e/debug_viz.py new file mode 100644 index 0000000..b3f9204 --- /dev/null +++ b/src/camera-based-e2e/debug_viz.py @@ -0,0 +1,369 @@ +""" +Generate debug diagnostic plots from training logs. + +Reads metrics.csv produced by PyTorch Lightning's CSVLogger and renders: + 1. Gradient norms per module group (scene_encoder, proposal_init, refinement, scorer) + 2. Gradient norms for refinement internals (cross_attn, mlp, traj_residual per block) + 3. Gradient dominance ratio over time + 4. Activation statistics (mean, std, dead fraction) per module + 5. Proposal diagnostics (spread, diversity, score distribution) + 6. Loss component breakdown (stacked area) + +Usage: + python debug_viz.py --log_dir /path/to/logs/camera_e2e_YYYYMMDD_HHMM/version_0 + python debug_viz.py --log_dir /path/to/logs # auto-picks newest run +""" +import argparse +from pathlib import Path + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.gridspec as gridspec + + +def find_metrics_csv(log_dir: Path) -> Path: + if (log_dir / "metrics.csv").exists(): + return log_dir / "metrics.csv" + candidates = sorted(log_dir.rglob("metrics.csv")) + if not candidates: + raise FileNotFoundError(f"No metrics.csv found under {log_dir}") + return candidates[-1] + + +def load_and_merge(csv_path: Path) -> pd.DataFrame: + """ + Lightning CSVLogger writes one row per log call, with NaN in columns + not logged in that call. Forward-fill within each step so we can + correlate gradient norms with losses at the same step. + """ + df = pd.read_csv(csv_path) + if "step" not in df.columns: + raise KeyError("metrics.csv must contain a 'step' column") + return df + + +def _get_cols(df: pd.DataFrame, prefix: str) -> list: + return sorted([c for c in df.columns if c.startswith(prefix)]) + + +def _plot_series(ax, df, columns, title, ylabel, legend_strip="", logy=False, ewm_span=20): + for col in columns: + series = df.set_index("step")[col].dropna() + if series.empty: + continue + label = col.replace(legend_strip, "").strip("/") + smoothed = series.ewm(span=ewm_span, min_periods=1).mean() + ax.plot(smoothed.index, smoothed.values, label=label, linewidth=1.2) + ax.fill_between(series.index, series.values, alpha=0.08) + ax.set_title(title, fontsize=11, fontweight="bold") + ax.set_ylabel(ylabel) + ax.set_xlabel("Step") + if logy: + ax.set_yscale("log") + ax.legend(fontsize=7, ncol=2) + ax.grid(True, alpha=0.3) + + +# ====================================================================== +# Individual plot functions +# ====================================================================== + +def plot_gradient_norms_by_module(df: pd.DataFrame, out_dir: Path): + """Bar-style time series of gradient norms for the four main modules.""" + cols = [c for c in _get_cols(df, "grad_norm/") if c in ( + "grad_norm/scene_encoder", "grad_norm/proposal_init", + "grad_norm/scorer", "grad_norm/total", + ) or c.startswith("grad_norm/refine_block_")] + if not cols: + return + + fig, ax = plt.subplots(figsize=(14, 5)) + _plot_series(ax, df, cols, "Gradient Norms by Module", "L2 Norm", "grad_norm/", logy=True) + fig.tight_layout() + fig.savefig(out_dir / "grad_norms_modules.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'grad_norms_modules.png'}") + + +def plot_gradient_norms_refinement(df: pd.DataFrame, out_dir: Path): + """Detailed gradient norms inside refinement blocks.""" + cols = [c for c in _get_cols(df, "grad_norm/refine_") if "block" not in c] + if not cols: + return + + fig, ax = plt.subplots(figsize=(14, 5)) + _plot_series(ax, df, cols, "Refinement Internals — Gradient Norms", "L2 Norm", "grad_norm/", logy=True) + fig.tight_layout() + fig.savefig(out_dir / "grad_norms_refinement_detail.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'grad_norms_refinement_detail.png'}") + + +def plot_gradient_norms_scorer_propinit(df: pd.DataFrame, out_dir: Path): + """Scorer and ProposalInit sub-module gradient norms.""" + cols = _get_cols(df, "grad_norm/scorer_") + _get_cols(df, "grad_norm/propinit_") + if not cols: + return + + fig, ax = plt.subplots(figsize=(14, 5)) + _plot_series(ax, df, cols, "Scorer & ProposalInit Internals — Gradient Norms", "L2 Norm", "grad_norm/", logy=True) + fig.tight_layout() + fig.savefig(out_dir / "grad_norms_scorer_propinit.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'grad_norms_scorer_propinit.png'}") + + +def plot_gradient_dominance(df: pd.DataFrame, out_dir: Path): + """Which module dominates the gradient signal?""" + col = "grad_norm/dominant_ratio" + if col not in df.columns: + return + + fig, ax = plt.subplots(figsize=(14, 4)) + series = df.set_index("step")[col].dropna() + smoothed = series.ewm(span=20, min_periods=1).mean() + ax.plot(smoothed.index, smoothed.values, color="crimson", linewidth=1.5) + ax.fill_between(series.index, series.values, alpha=0.15, color="crimson") + ax.axhline(0.5, color="gray", linestyle="--", alpha=0.5, label="50% dominance") + ax.set_title("Gradient Dominance Ratio (max module / total)", fontsize=11, fontweight="bold") + ax.set_ylabel("Ratio") + ax.set_xlabel("Step") + ax.set_ylim(0, 1.05) + ax.legend() + ax.grid(True, alpha=0.3) + fig.tight_layout() + fig.savefig(out_dir / "grad_dominance.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'grad_dominance.png'}") + + +def plot_activation_stats(df: pd.DataFrame, out_dir: Path): + """Activation mean/std/dead fraction for each hooked module output.""" + mean_cols = [c for c in _get_cols(df, "act/") if c.endswith("_mean")] + std_cols = [c for c in _get_cols(df, "act/") if c.endswith("_std")] + dead_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_dead")] + sat_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_saturated")] + + if not mean_cols: + return + + fig, axes = plt.subplots(2, 2, figsize=(16, 10)) + + _plot_series(axes[0, 0], df, mean_cols, "Activation Mean", "Mean", "act/") + _plot_series(axes[0, 1], df, std_cols, "Activation Std", "Std", "act/") + _plot_series(axes[1, 0], df, dead_cols, "Dead Neuron Fraction (|x| < 1e-6)", "Fraction", "act/") + _plot_series(axes[1, 1], df, sat_cols, "Saturated Fraction (|x| > 10)", "Fraction", "act/") + + fig.suptitle("Activation Statistics by Module Output", fontsize=13, fontweight="bold") + fig.tight_layout(rect=[0, 0, 1, 0.96]) + fig.savefig(out_dir / "activation_stats.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'activation_stats.png'}") + + +def plot_proposal_diagnostics(df: pd.DataFrame, out_dir: Path): + """Proposal spread, diversity, trajectory lengths, and score stats.""" + prop_cols = _get_cols(df, "proposals/") + if not prop_cols: + return + + spread_cols = [c for c in prop_cols if "std" in c or "dist" in c] + length_cols = [c for c in prop_cols if "length" in c] + score_cols = [c for c in prop_cols if "score" in c] + + n_panels = sum(1 for g in [spread_cols, length_cols, score_cols] if g) + if n_panels == 0: + return + + fig, axes = plt.subplots(1, n_panels, figsize=(6 * n_panels, 5)) + if n_panels == 1: + axes = [axes] + + idx = 0 + if spread_cols: + _plot_series(axes[idx], df, spread_cols, "Proposal Spread & Diversity", "Value", "proposals/") + idx += 1 + if length_cols: + _plot_series(axes[idx], df, length_cols, "Trajectory Lengths", "Length (m)", "proposals/") + idx += 1 + if score_cols: + _plot_series(axes[idx], df, score_cols, "Score Distribution", "Score", "proposals/") + idx += 1 + + fig.suptitle("Proposal Diagnostics", fontsize=13, fontweight="bold") + fig.tight_layout(rect=[0, 0, 1, 0.95]) + fig.savefig(out_dir / "proposal_diagnostics.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'proposal_diagnostics.png'}") + + +def plot_loss_breakdown(df: pd.DataFrame, out_dir: Path): + """Stacked area chart of all loss components.""" + loss_cols = [c for c in df.columns if c.startswith("train_loss_") and c != "train_loss"] + if not loss_cols: + return + + fig, axes = plt.subplots(1, 2, figsize=(16, 5)) + + # Individual loss curves + _plot_series(axes[0], df, loss_cols, "Training Loss Components", "Loss", "train_loss_", logy=True) + + # Stacked area (absolute contribution) + sub = df[["step"] + loss_cols].dropna(subset=loss_cols, how="all").copy() + if sub.empty: + plt.close(fig) + return + sub = sub.set_index("step").interpolate().fillna(0).clip(lower=0) + sub_smooth = sub.ewm(span=30, min_periods=1).mean() + axes[1].stackplot(sub_smooth.index, *[sub_smooth[c].values for c in loss_cols], + labels=[c.replace("train_loss_", "") for c in loss_cols], alpha=0.7) + axes[1].set_title("Loss Component Breakdown (stacked)", fontsize=11, fontweight="bold") + axes[1].set_ylabel("Loss") + axes[1].set_xlabel("Step") + axes[1].legend(fontsize=7, loc="upper right") + axes[1].grid(True, alpha=0.3) + + fig.tight_layout() + fig.savefig(out_dir / "loss_breakdown.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'loss_breakdown.png'}") + + +def plot_ade_regret(df: pd.DataFrame, out_dir: Path): + """Oracle ADE vs predicted-best ADE, and regret over time.""" + cols = ["train_ade_pred", "train_ade_oracle", "train_ade_regret"] + present = [c for c in cols if c in df.columns] + if not present: + return + + fig, ax = plt.subplots(figsize=(14, 5)) + _plot_series(ax, df, present, "ADE: Oracle vs Predicted-Best & Regret", "ADE (m)", "train_") + fig.tight_layout() + fig.savefig(out_dir / "ade_regret.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'ade_regret.png'}") + + +def plot_param_norms(df: pd.DataFrame, out_dir: Path): + """Parameter norms to detect weight explosion / vanishing.""" + cols = _get_cols(df, "param_norm/") + if not cols: + return + + fig, ax = plt.subplots(figsize=(14, 5)) + _plot_series(ax, df, cols, "Parameter Norms by Module", "L2 Norm", "param_norm/", logy=True) + fig.tight_layout() + fig.savefig(out_dir / "param_norms.png", dpi=180) + plt.close(fig) + print(f" -> {out_dir / 'param_norms.png'}") + + +# ====================================================================== +# Snapshot: single-image summary for quick glance +# ====================================================================== + +def plot_summary_dashboard(df: pd.DataFrame, out_dir: Path): + """Single 2x3 dashboard with the most important signals.""" + fig = plt.figure(figsize=(20, 12)) + gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.3) + + # 1. Module gradient norms + ax1 = fig.add_subplot(gs[0, 0]) + cols = [c for c in _get_cols(df, "grad_norm/") if c in ( + "grad_norm/scene_encoder", "grad_norm/proposal_init", + "grad_norm/scorer", "grad_norm/total", + ) or c.startswith("grad_norm/refine_block_")] + if cols: + _plot_series(ax1, df, cols, "Gradient Norms (modules)", "L2 Norm", "grad_norm/", logy=True, ewm_span=30) + else: + ax1.text(0.5, 0.5, "No grad_norm data", ha="center", va="center", transform=ax1.transAxes) + + # 2. Loss breakdown + ax2 = fig.add_subplot(gs[0, 1]) + loss_cols = [c for c in df.columns if c.startswith("train_loss_") and c != "train_loss"] + if loss_cols: + _plot_series(ax2, df, loss_cols, "Loss Components", "Loss", "train_loss_", logy=True, ewm_span=30) + else: + ax2.text(0.5, 0.5, "No loss data", ha="center", va="center", transform=ax2.transAxes) + + # 3. ADE regret + ax3 = fig.add_subplot(gs[0, 2]) + ade_cols = [c for c in ("train_ade_pred", "train_ade_oracle", "train_ade_regret") if c in df.columns] + if ade_cols: + _plot_series(ax3, df, ade_cols, "ADE: Oracle vs Predicted", "ADE (m)", "train_", ewm_span=30) + else: + ax3.text(0.5, 0.5, "No ADE data", ha="center", va="center", transform=ax3.transAxes) + + # 4. Activation dead fraction + ax4 = fig.add_subplot(gs[1, 0]) + dead_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_dead")] + if dead_cols: + _plot_series(ax4, df, dead_cols, "Dead Neuron Fraction", "Fraction", "act/", ewm_span=30) + else: + ax4.text(0.5, 0.5, "No activation data", ha="center", va="center", transform=ax4.transAxes) + + # 5. Proposal diversity + ax5 = fig.add_subplot(gs[1, 1]) + prop_cols = [c for c in _get_cols(df, "proposals/") if "dist" in c or "std" in c] + if prop_cols: + _plot_series(ax5, df, prop_cols, "Proposal Diversity", "Value", "proposals/", ewm_span=30) + else: + ax5.text(0.5, 0.5, "No proposal data", ha="center", va="center", transform=ax5.transAxes) + + # 6. Gradient dominance + ax6 = fig.add_subplot(gs[1, 2]) + dom_col = "grad_norm/dominant_ratio" + if dom_col in df.columns: + series = df.set_index("step")[dom_col].dropna() + if not series.empty: + smoothed = series.ewm(span=30, min_periods=1).mean() + ax6.plot(smoothed.index, smoothed.values, color="crimson", linewidth=1.5) + ax6.fill_between(series.index, series.values, alpha=0.12, color="crimson") + ax6.axhline(0.5, color="gray", linestyle="--", alpha=0.5) + ax6.set_title("Gradient Dominance", fontsize=11, fontweight="bold") + ax6.set_ylabel("Ratio (max module / total)") + ax6.set_xlabel("Step") + ax6.set_ylim(0, 1.05) + ax6.grid(True, alpha=0.3) + + fig.suptitle("Training Debug Dashboard", fontsize=15, fontweight="bold", y=0.98) + fig.savefig(out_dir / "debug_dashboard.png", dpi=200) + plt.close(fig) + print(f" -> {out_dir / 'debug_dashboard.png'}") + + +def main(): + parser = argparse.ArgumentParser(description="Generate debug diagnostic plots from training logs") + parser.add_argument("--log_dir", type=str, required=True, + help="Path to log directory (version_0/) or parent (auto-picks newest)") + parser.add_argument("--out_dir", type=str, default=None, + help="Output directory for plots (default: /debug_plots/)") + args = parser.parse_args() + + log_dir = Path(args.log_dir) + csv_path = find_metrics_csv(log_dir) + print(f"Reading {csv_path}") + df = load_and_merge(csv_path) + print(f" {len(df)} rows, {len(df.columns)} columns") + + out_dir = Path(args.out_dir) if args.out_dir else csv_path.parent / "debug_plots" + out_dir.mkdir(parents=True, exist_ok=True) + print(f"Output: {out_dir}\n") + + plot_gradient_norms_by_module(df, out_dir) + plot_gradient_norms_refinement(df, out_dir) + plot_gradient_norms_scorer_propinit(df, out_dir) + plot_gradient_dominance(df, out_dir) + plot_activation_stats(df, out_dir) + plot_proposal_diagnostics(df, out_dir) + plot_loss_breakdown(df, out_dir) + plot_ade_regret(df, out_dir) + plot_param_norms(df, out_dir) + plot_summary_dashboard(df, out_dir) + + print(f"\nDone. {len(list(out_dir.glob('*.png')))} plots saved to {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/src/camera-based-e2e/loader.py b/src/camera-based-e2e/loader.py index 743e61a..d9e088b 100644 --- a/src/camera-based-e2e/loader.py +++ b/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 @@ -66,7 +67,16 @@ def __iter__(self): if self.file: self.file.close() del self.file - self.file = open(os.path.join(self.data_dir, filename), 'rb') + # #region agent log + _full = os.path.join(self.data_dir, filename) + _logpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "debug-aec3fc.log") + try: + _line = '{"sessionId":"aec3fc","hypothesisId":"H3,H4,H5","location":"loader.py:open","message":"file_open_attempt","data":{"data_dir":"%s","filename":"%s","full_path":"%s","exists":%s,"index_file":"%s"}}\n' % (str(self.data_dir), str(filename), str(_full), str(os.path.exists(_full)), str(getattr(self, "_index_file", "?"))) + open(_logpath, "a").write(_line) + except Exception: + pass + # #endregion + self.file = open(_full, 'rb') self.filename = filename self.file.seek(start_byte) # type: ignore @@ -100,7 +110,7 @@ def __iter__(self): 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, images=True) loader = DataLoader( diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py index 27df45c..90be2b4 100644 --- a/src/camera-based-e2e/models/base_model.py +++ b/src/camera-based-e2e/models/base_model.py @@ -1,3 +1,4 @@ +from typing import List, Optional import torch import torch.nn as nn import torch.nn.functional as F @@ -21,12 +22,48 @@ def forward(self, x: dict) -> torch.Tensor: return self.nn(x) class LitModel(pl.LightningModule): - def __init__(self, model: nn.Module, lr: float, lr_vision: float | None = None, rfs_weight: float = 0.0): + def __init__( + self, + model: nn.Module, + lr: float, + lr_vision: Optional[float] = None, + 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, + score_weight: float = 1.0, + score_warmup_epochs: int = 2, + score_temperature: float = 5.0, + score_loss_type: str = "bce", + score_target_type: str = "l1", + score_rank_weight: float = 0.0, + score_margin: float = 0.2, + score_topk: int = 0, + comfort_jerk_threshold: float = 5.0, + prev_weight: float = 0.1, + rfs_target_use_comfort: bool = True, + ): super(LitModel, self).__init__() self.model = model self.hparams.lr = lr self.hparams.lr_vision = lr_vision self.hparams.rfs_weight = rfs_weight + self.hparams.smoothness_weight = smoothness_weight + self.hparams.collision_weight = collision_weight + self.hparams.comfort_weight = comfort_weight + self.hparams.diversity_weight = diversity_weight + self.hparams.score_weight = score_weight + self.hparams.score_warmup_epochs = score_warmup_epochs + self.hparams.score_temperature = score_temperature + self.hparams.score_loss_type = score_loss_type + self.hparams.score_target_type = score_target_type + self.hparams.score_rank_weight = score_rank_weight + self.hparams.score_margin = score_margin + self.hparams.score_topk = score_topk + self.hparams.comfort_jerk_threshold = comfort_jerk_threshold + self.hparams.prev_weight = prev_weight + self.hparams.rfs_target_use_comfort = rfs_target_use_comfort self.example_input_array = ({ 'PAST': torch.zeros((1, 16, 6)), # PAST @@ -46,7 +83,7 @@ def transfer_batch_to_device(self, batch, device, dataloader_idx): return out - def decode_batch_jpeg(self, images_jpeg: list[list[torch.Tensor]]) -> list[torch.Tensor]: + def decode_batch_jpeg(self, images_jpeg: List[List[torch.Tensor]]) -> List[torch.Tensor]: # Flatten cameras flat_encoded, cam_sizes = [], [] for cam in images_jpeg: @@ -148,6 +185,145 @@ def _prepare_rfs_inputs(self, past, future, pred_future): t_idx = torch.tensor([3.0, 5.0], device=future.device).unsqueeze(0).expand(future.size(0), -1) return pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx + # ---- NAVSIM-style quality target ---- + @torch.no_grad() + def _compute_navsim_score( + self, + proposals: torch.Tensor, + future: torch.Tensor, + ) -> torch.Tensor: + """ + Approximate NAVSIM Eq. 5: S = NC * DAC * (5*EP + 5*TTC + 2*Comf) / 12 + + Without agent / map data we set NC=1, DAC=1, TTC=1 and compute EP and + Comf from trajectory geometry alone. + + Args: + proposals: (B, K, T, 2) predicted trajectories + future: (B, T, 2) ground-truth trajectory + Returns: + quality: (B, K) in [0, 1] + """ + B, K, T, _ = proposals.shape + gt = future[:, None, :, :] # (B, 1, T, 2) + + # --- Ego Progress (EP) --- + gt_disp = future[:, -1] - future[:, 0] # (B, 2) + gt_dist = gt_disp.norm(dim=-1, keepdim=True).clamp(min=1e-3) # (B, 1) + gt_dir = gt_disp / gt_dist # (B, 2) + + prop_disp = proposals[:, :, -1] - proposals[:, :, 0] # (B, K, 2) + progress = (prop_disp * gt_dir.unsqueeze(1)).sum(dim=-1) # (B, K) + ep = (progress / gt_dist).clamp(0.0, 1.0) # (B, K) + + # --- Comfort (Comf) --- + dt = 0.25 # 4 Hz + vel = (proposals[:, :, 1:] - proposals[:, :, :-1]) / dt # (B,K,T-1,2) + acc = (vel[:, :, 1:] - vel[:, :, :-1]) / dt # (B,K,T-2,2) + jerk = (acc[:, :, 1:] - acc[:, :, :-1]) / dt # (B,K,T-3,2) + jerk_mag = jerk.norm(dim=-1) # (B,K,T-3) + + jerk_thresh = getattr(self.hparams, "comfort_jerk_threshold", 5.0) + comf = (jerk_mag < jerk_thresh).float().mean(dim=-1) # (B,K) + + nc = 1.0 + dac = 1.0 + ttc = 1.0 + quality = nc * dac * (5.0 * ep + 5.0 * ttc + 2.0 * comf) / 12.0 # (B,K) + + return quality.clamp(0.0, 1.0) + + @torch.no_grad() + def _compute_rfs_quality( + self, + proposals: torch.Tensor, + reference: torch.Tensor, + past: torch.Tensor, + ) -> torch.Tensor: + """ + Per-proposal RFS-style quality in [0, 1] for BCE scorer targets (iPad Eq. 4). + + Same longitudinal/lateral deviation + speed scaling as ``rfs_loss``, evaluated at + 3 s and 5 s (indices 11, 19 at 4 Hz). ``reference`` is typically batch ``FUTURE`` + (expert trajectory); it can be swapped for a route or rollout proxy when GT is absent. + + Optionally multiplies by a jerk comfort factor (same spirit as NAVSIM Comf in Eq. 5). + """ + device = proposals.device + indices = [11, 19] + if reference.shape[1] <= max(indices): + raise ValueError( + f"reference horizon {reference.shape[1]} must exceed max RFS index {max(indices)}" + ) + + speed = torch.norm(past[..., 2:4], dim=-1)[:, -1] + full_lng_dir, full_lat_dir = self.compute_direction(reference) + + ref_slice = reference[:, indices, :] + lng_slice = full_lng_dir[:, indices, :] + lat_slice = full_lat_dir[:, indices, :] + + prop_slice = proposals[:, :, indices, :] + delta = prop_slice - ref_slice.unsqueeze(1) + + delta_lng = (delta * lng_slice.unsqueeze(1)).sum(dim=-1).abs() + delta_lat = (delta * lat_slice.unsqueeze(1)).sum(dim=-1).abs() + + t_idx = torch.tensor([3.0, 5.0], device=device).unsqueeze(0).expand(reference.size(0), -1) + tau_lat_raw, tau_lng_raw = self.time_thresholds(t_idx) + scale = self.speed_scale(speed) + if scale.dim() == 1: + scale = scale.unsqueeze(1) + tau_lat = tau_lat_raw * scale + tau_lng = tau_lng_raw * scale + + deviation = torch.max( + delta_lat / tau_lat[:, None, :], + delta_lng / tau_lng[:, None, :], + ) + score = torch.where( + deviation <= 1, + torch.ones_like(deviation), + torch.pow(0.1, deviation - 1), + ) + rfs_quality = score.mean(dim=-1) + + use_comf = getattr(self.hparams, "rfs_target_use_comfort", True) + if use_comf: + dt = 0.25 + vel = (proposals[:, :, 1:] - proposals[:, :, :-1]) / dt + acc = (vel[:, :, 1:] - vel[:, :, :-1]) / dt + jerk = (acc[:, :, 1:] - acc[:, :, :-1]) / dt + jerk_mag = jerk.norm(dim=-1) + jerk_thresh = getattr(self.hparams, "comfort_jerk_threshold", 5.0) + comf = (jerk_mag < jerk_thresh).float().mean(dim=-1) + rfs_quality = rfs_quality * comf + + return rfs_quality.clamp(0.0, 1.0) + + @torch.no_grad() + def _compute_quality_target( + self, + pred: torch.Tensor, + gt_expanded: torch.Tensor, + future: torch.Tensor, + tau: float, + past: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Returns (B, K) quality target in [0, 1] based on score_target_type. + """ + target_type = getattr(self.hparams, "score_target_type", "l1") + if target_type == "navsim": + return self._compute_navsim_score(pred, future) + if target_type == "rfs": + if past is None: + raise ValueError("score_target_type='rfs' requires past trajectory") + return self._compute_rfs_quality(pred, future, past) + + l1_target = (pred - gt_expanded).abs().sum(dim=-1).mean(dim=-1) # (B, K) + return torch.exp(-l1_target / max(tau, 1e-6)) + # ---- optimizers ---- def configure_optimizers(self): # NOTE: This can be extended and tuned, LR especially will differ and have an impact. @@ -186,18 +362,20 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: images = batch["IMAGES"] else: raise KeyError("Batch must contain either 'IMAGES_JPEG' or 'IMAGES' key.") - - # `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} - pred_future = self.forward(model_inputs) # (B, T*2) + raw_output = self.forward(model_inputs) pred_depth = None - pred_scores: torch.Tensor = None - if isinstance(pred_future, dict): - pred_future, pred_depth, pred_scores = pred_future["trajectory"], pred_future.get("depth", None), pred_future.get("scores", None) + pred_scores = None + proposal_list = None + if isinstance(raw_output, dict): + proposal_list = raw_output.get("proposal_list", None) + pred_scores = raw_output.get("scores", None) + pred_depth = raw_output.get("depth", None) + pred_future = raw_output["trajectory"] + else: + pred_future = raw_output pred = pred_future t_steps = future.shape[1] @@ -205,84 +383,192 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: k_modes = self.model.n_proposals if hasattr(self.model, "n_proposals") else 1 if pred.ndim != 2: - raise ValueError(f"Unexpected pred shape {pred.shape}; expected (B, T*2) or (B, {k_modes}*T*2).") + raise ValueError(f"Unexpected pred shape {pred.shape}") if pred.shape[1] == t2: pred = pred.view(pred.size(0), 1, t_steps, 2) elif pred.shape[1] == k_modes * t2: pred = pred.view(pred.size(0), k_modes, t_steps, 2) else: - raise ValueError(f"Unexpected pred shape {pred.shape}; expected (B, T*2) or (B, {k_modes}*T*2).") - + raise ValueError(f"Unexpected pred shape {pred.shape}") + + if not torch.isfinite(pred).all(): + pred = torch.nan_to_num(pred, nan=0.0, posinf=1e3, neginf=-1e3) + + # ---- MoN L1 trajectory loss (iPad-style) ---- + # min over N proposals of mean L1 displacement per timestep + gt = future[:, None, :, :] # (B, 1, T, 2) + + if proposal_list is not None and len(proposal_list) > 0: + # Discounted intermediate supervision: L = sum_k λ^(K-1-k) * MoN_L1(P_k) + prev_w = getattr(self.hparams, "prev_weight", 0.1) + loss_traj = torch.tensor(0.0, device=self.device) + for proposals_k in proposal_list: + if not torch.isfinite(proposals_k).all(): + proposals_k = torch.nan_to_num(proposals_k, nan=0.0, posinf=1e3, neginf=-1e3) + l1_per_mode = (proposals_k - gt).abs().sum(dim=-1).mean(dim=-1) # (B, K) + mon_l1 = l1_per_mode.amin(dim=1).mean() + loss_traj = prev_w * loss_traj + mon_l1 + else: + l1_per_mode = (pred - gt).abs().sum(dim=-1).mean(dim=-1) # (B, K) + loss_traj = l1_per_mode.amin(dim=1).mean() + + # ---- Metrics (for logging, not loss) ---- + refine_oracle_ades: List[torch.Tensor] = [] + with torch.no_grad(): + dist = torch.norm(pred - gt, dim=-1) # (B, K, T) L2 + ade_per_mode = dist.mean(dim=-1) # (B, K) + oracle_ade = ade_per_mode.min(dim=1).values.mean() + ade_pred = None + if pred_scores is not None and k_modes > 1: + pred_idx = pred_scores.detach().argmax(dim=1) + ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + elif k_modes == 1: + ade_pred = ade_per_mode.squeeze(1).mean() + regret = (ade_pred - oracle_ade) if ade_pred is not None else None + + if proposal_list is not None and len(proposal_list) > 0: + for proposals_k in proposal_list: + pk = proposals_k + if not torch.isfinite(pk).all(): + pk = torch.nan_to_num(pk, nan=0.0, posinf=1e3, neginf=-1e3) + dist_r = torch.norm(pk - gt, dim=-1) + ade_pm = dist_r.mean(dim=-1) + refine_oracle_ades.append(ade_pm.min(dim=1).values.mean()) + + # ---- RFS loss ---- if pred_scores is not None and pred.size(1) > 1: - rfs_pred_idx = pred_scores.argmin(dim=1) + rfs_pred_idx = pred_scores.detach().argmax(dim=1) else: rfs_pred_idx = torch.zeros(pred.size(0), dtype=torch.long, device=pred.device) pred_for_rfs = pred[torch.arange(pred.size(0), device=pred.device), rfs_pred_idx] pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx = self._prepare_rfs_inputs( - past, - future, - pred_for_rfs, + past, future, pred_for_rfs, ) rfs_unweighted = self.rfs_loss(pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx) rfs_weight = getattr(self.hparams, "rfs_weight", 0.0) loss_rfs = rfs_weight * rfs_unweighted - # ADE per mode: (B, K) - dist = torch.norm(pred - future[:, torch.newaxis, :, :], dim=-1) # (B, K, T) - ade_per_mode = dist.mean(dim=-1) - - # Top-M WTA for trajectory loss. Here, we have an "oracle" that picks the best mode - # so, our loss is calculated on the mean of the top 5 trajectories. - top_m = min(5, ade_per_mode.size(1)) - loss_ade = ade_per_mode.topk(top_m, largest=False, dim=1).values.mean() - - # oracle ade is best of all proposals, since we have the GT data during training - oracle_ade = ade_per_mode.min(dim=1).values.mean() - ade_pred = None - # pred_scores is now the predicted ADE of each trajectory - if pred_scores is not None and k_modes > 1: - pred_idx = pred_scores.argmin(dim=1) - ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() - elif k_modes == 1: - ade_pred = ade_per_mode.squeeze(1).mean() - regret = (ade_pred - oracle_ade) if ade_pred is not None else None - - # Scorer Losses -> encourage ranking of predicted scores to match true ranking of ades that are generated + # ---- Score loss (configurable) ---- if k_modes > 1 and pred_scores is not None: - ade = ade_per_mode.detach() # (B, K) - loss_score = F.mse_loss(pred_scores, ade) - pred_idx = pred_scores.argmin(dim=1) - ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + if not torch.isfinite(pred_scores).all(): + pred_scores = torch.nan_to_num(pred_scores, nan=0.0, posinf=1e3, neginf=-1e3) + + score_loss_type = getattr(self.hparams, "score_loss_type", "bce") + tau = getattr(self.hparams, "score_temperature", 5.0) + + with torch.no_grad(): + quality_target = self._compute_quality_target(pred, gt, future, tau, past=past) # (B, K) + best_idx = quality_target.argmax(dim=1) # (B,) + + if score_loss_type == "ce": + loss_score = F.cross_entropy(pred_scores, best_idx) + elif score_loss_type == "listnet": + target_probs = F.softmax(quality_target / max(tau * 0.1, 1e-6), dim=1) + loss_score = F.kl_div( + F.log_softmax(pred_scores, dim=1), + target_probs, + reduction="batchmean", + ) + else: + bce_loss = F.binary_cross_entropy_with_logits(pred_scores, quality_target) + loss_score = bce_loss + + if score_loss_type == "bce_pairwise": + margin = getattr(self.hparams, "score_margin", 0.2) + rank_weight = getattr(self.hparams, "score_rank_weight", 0.2) + topk = int(getattr(self.hparams, "score_topk", 0)) + + best_scores = pred_scores.gather(1, best_idx.unsqueeze(1)) # (B,1) + pairwise_margin = margin - (best_scores - pred_scores) # (B,K) + pairwise_margin.scatter_(1, best_idx.unsqueeze(1), 0.0) + + if topk > 0: + k_eff = min(topk, pairwise_margin.size(1) - 1) + hardest = pairwise_margin.topk(k_eff, dim=1).values + rank_loss = F.relu(hardest).mean() + else: + rank_loss = F.relu(pairwise_margin).mean() + loss_score = bce_loss + rank_weight * rank_loss else: loss_score = torch.tensor(0.0, device=self.device) # Depth Loss if pred_depth is not None: - front_img = images[1] # front camera + front_img = images[1] depth_in = F.interpolate(front_img, size=(128, 128), mode='nearest') loss_depth = self.depth_loss(depth_in, pred_depth, loss_fn=F.l1_loss) else: loss_depth = torch.tensor(0.0, device=self.device) + loss_depth *= 0.1 + + # Score loss warmup + sw_score = getattr(self.hparams, "score_weight", 1.0) + warmup_epochs = getattr(self.hparams, "score_warmup_epochs", 2) + current_epoch = self.current_epoch if hasattr(self, "current_epoch") else 0 + if current_epoch < warmup_epochs: + effective_score_weight = 0.0 + elif current_epoch < warmup_epochs + 1: + effective_score_weight = sw_score * (current_epoch - warmup_epochs) + else: + effective_score_weight = sw_score + + loss_score *= effective_score_weight + total_loss = loss_traj + loss_depth + loss_score + loss_rfs + + # Smoothness / comfort losses on best proposal + loss_smooth = torch.tensor(0.0, device=self.device) + loss_collision = torch.tensor(0.0, device=self.device) + loss_comfort = torch.tensor(0.0, device=self.device) + sw = getattr(self.hparams, "smoothness_weight", 0.0) + cw = getattr(self.hparams, "collision_weight", 0.0) + cfw = getattr(self.hparams, "comfort_weight", 0.0) + if (sw > 0 or cw > 0 or cfw > 0) and pred_for_rfs.shape[1] >= 3: + vel = pred_for_rfs[:, 1:] - pred_for_rfs[:, :-1] + acc = vel[:, 1:] - vel[:, :-1] + jerk = acc[:, 1:] - acc[:, :-1] + if sw > 0: + loss_smooth = (jerk ** 2).mean() + if cfw > 0: + v_mid = vel[:, :-1] + a_mid = acc + v_speed = (v_mid ** 2).sum(dim=-1).clamp(min=0.25).sqrt() + a_mag = (a_mid ** 2).sum(dim=-1).sqrt() + curv = a_mag / (v_speed ** 2) + loss_comfort = curv.mean() + total_loss = total_loss + sw * loss_smooth + cw * loss_collision + cfw * loss_comfort - loss_depth *= 0.1 # slightly enabled - loss_ade *= 1.0 # TODO: tune loss terms - loss_score *= 1.0 - total_loss = loss_ade + loss_depth + loss_score + loss_rfs - # TODO: improve logging both to disk and to console log_payload = { - f"{stage}_loss_ade": loss_ade, + f"{stage}_loss_traj": loss_traj, f"{stage}_loss_score": loss_score, f"{stage}_loss_depth": loss_depth, f"{stage}_loss_rfs": loss_rfs, f"{stage}_rfs_unweighted": rfs_unweighted, f"{stage}_loss": total_loss, } + if sw > 0 or cw > 0 or cfw > 0: + log_payload[f"{stage}_loss_smooth"] = loss_smooth + log_payload[f"{stage}_loss_collision"] = loss_collision + log_payload[f"{stage}_loss_comfort"] = loss_comfort if ade_pred is not None: log_payload[f"{stage}_ade_pred"] = ade_pred log_payload[f"{stage}_ade_oracle"] = oracle_ade log_payload[f"{stage}_ade_regret"] = regret + for ri, oade in enumerate(refine_oracle_ades): + log_payload[f"{stage}_ade_oracle_refine_{ri}"] = oade + if pred_scores is not None and k_modes > 1: + with torch.no_grad(): + mode_l1 = (pred - gt).abs().sum(dim=-1).mean(dim=-1) # (B,K) + best_idx = mode_l1.argmin(dim=1) + pred_idx = pred_scores.argmax(dim=1) + top1_acc = (pred_idx == best_idx).float().mean() + s_best = pred_scores.gather(1, best_idx.unsqueeze(1)).squeeze(1) + masked = pred_scores.clone() + masked.scatter_(1, best_idx.unsqueeze(1), float("-inf")) + s_second = masked.max(dim=1).values + log_payload[f"{stage}_score_top1_acc"] = top1_acc + log_payload[f"{stage}_score_gap_best_second"] = (s_best - s_second).mean() self.log_dict(log_payload, prog_bar=True, logger=True) return total_loss diff --git a/src/camera-based-e2e/models/feature_extractors.py b/src/camera-based-e2e/models/feature_extractors.py index 0864b85..9096d09 100644 --- a/src/camera-based-e2e/models/feature_extractors.py +++ b/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__() diff --git a/src/camera-based-e2e/models/proposal_init.py b/src/camera-based-e2e/models/proposal_init.py new file mode 100644 index 0000000..a0f141d --- /dev/null +++ b/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 diff --git a/src/camera-based-e2e/models/proposal_planner.py b/src/camera-based-e2e/models/proposal_planner.py new file mode 100644 index 0000000..8c2b4a4 --- /dev/null +++ b/src/camera-based-e2e/models/proposal_planner.py @@ -0,0 +1,76 @@ +""" +Proposal-centric E2E planner (iPad-style). + +Pipeline: + scene encoder → proposal init → iterative refinement → scorer + (with intermediate proposal supervision) +""" +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 + + +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.num_proposals = num_proposals + 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, + } diff --git a/src/camera-based-e2e/models/refinement.py b/src/camera-based-e2e/models/refinement.py new file mode 100644 index 0000000..2bd7561 --- /dev/null +++ b/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 diff --git a/src/camera-based-e2e/models/scene_encoder.py b/src/camera-based-e2e/models/scene_encoder.py new file mode 100644 index 0000000..144d2f5 --- /dev/null +++ b/src/camera-based-e2e/models/scene_encoder.py @@ -0,0 +1,60 @@ +""" +Multi-camera scene encoder for proposal-centric E2E planner. +Encodes all 8 Waymo cameras with a shared backbone and fuses via concatenation + camera embeddings. +""" +from typing import List + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class SceneEncoder(nn.Module): + """ + Encode multi-camera images into a single scene feature sequence. + Uses a shared backbone per camera, then concatenates tokens and adds camera embeddings. + """ + + def __init__( + self, + backbone: nn.Module, + d_model: int = 256, + num_cams: int = 8, + ): + super().__init__() + self.backbone = backbone + self.num_cams = num_cams + # Backbone output dim (from feature extractor .dims or .feature_dim) + if hasattr(backbone, "dims"): + self.backbone_dim = sum(backbone.dims) + else: + self.backbone_dim = getattr(backbone, "feature_dim", 384) + + self.proj = nn.Linear(self.backbone_dim, d_model) + # Per-camera embedding (1, num_cams, d_model) added to all tokens of that camera + self.cam_embed = nn.Parameter(torch.zeros(1, num_cams, d_model)) + nn.init.trunc_normal_(self.cam_embed, std=0.02) + self.d_model = d_model + self.ln = nn.LayerNorm(d_model) + + def forward(self, images: List[torch.Tensor]) -> torch.Tensor: + """ + Args: + images: List of (B, C, H, W) tensors, one per camera. len(images) == num_cams. + Returns: + scene_feat: (B, N, d_model) with N = num_cams * n_tokens_per_cam + """ + B = images[0].size(0) + all_tokens = [] + for c, img in enumerate(images): + with torch.no_grad(): + feats = self.backbone(img) + if isinstance(feats, (list, tuple)): + feats = feats[0] # (B, C, H, W) + # (B, C, H, W) -> (B, C, H*W) -> (B, H*W, C) + tokens = feats.flatten(2).permute(0, 2, 1) # (B, n_tokens, backbone_dim) + tokens = self.proj(tokens) + self.cam_embed[:, c : c + 1, :] # (B, n_tokens, d_model) + all_tokens.append(tokens) + # (B, num_cams * n_tokens_per_cam, d_model) + scene = torch.cat(all_tokens, dim=1) + return self.ln(scene) diff --git a/src/camera-based-e2e/models/scorer.py b/src/camera-based-e2e/models/scorer.py new file mode 100644 index 0000000..2ee9cbd --- /dev/null +++ b/src/camera-based-e2e/models/scorer.py @@ -0,0 +1,41 @@ +""" +Scorer (paper-faithful): max-pool per-timestep BEV features -> MLP -> score. +""" +import torch +import torch.nn as nn + + +class Scorer(nn.Module): + """ + Score each proposal from BEV proposal features. + Output (B, K) raw logits — higher = better. + """ + + def __init__( + self, + d_model: int, + num_proposals: int = 16, + horizon: int = 20, + hidden_dim: int = 1024, + ): + super().__init__() + self.num_proposals = num_proposals + self.horizon = horizon + self.score_mlp = nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, 1), + ) + + def forward(self, bev_feature: torch.Tensor) -> torch.Tensor: + """ + Args: + bev_feature: (B, N*T, C) per-timestep features + Returns: + scores: (B, N) raw logits, higher is better + """ + B = bev_feature.size(0) + N, T = self.num_proposals, self.horizon + feat = bev_feature.view(B, N, T, -1).amax(dim=2) # (B, N, C) + scores = self.score_mlp(feat).squeeze(-1) # (B, N) + return scores diff --git a/src/camera-based-e2e/score_experiment_viz.py b/src/camera-based-e2e/score_experiment_viz.py new file mode 100644 index 0000000..a536c00 --- /dev/null +++ b/src/camera-based-e2e/score_experiment_viz.py @@ -0,0 +1,299 @@ +""" +Compare scorer-loss experiments across runs and auto-generate rankings. + +This script scans Lightning CSV logs under a logs root, extracts final validation +metrics per run, groups by scorer-loss configuration, and writes: + - summary CSV + - ranking bar chart + - oracle-vs-pred scatter (diagnose ranking bottleneck) + - regret-vs-top1 scatter (diagnose scorer discrimination) + - markdown report with "what went right/wrong" + +Usage: + python score_experiment_viz.py --logs_root /scratch/.../waymo/logs + python score_experiment_viz.py --logs_root /scratch/.../waymo/logs --latest_n 8 +""" + +import argparse +from pathlib import Path +from typing import Dict, List, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + + +def _read_hparams_kv(hparams_path: Path) -> Dict[str, str]: + """ + Parse a few scalar keys from hparams.yaml without requiring strict YAML parse. + This is intentionally tolerant because some generated hparams files can contain + very large content blocks. + """ + keys = { + "score_loss_type", + "score_target_type", + "score_weight", + "score_temperature", + "score_rank_weight", + "score_margin", + "score_topk", + "comfort_jerk_threshold", + "model_type", + } + out: Dict[str, str] = {} + if not hparams_path.exists(): + return out + + try: + with hparams_path.open("r", errors="replace") as f: + for line in f: + line = line.strip() + if ":" not in line: + continue + k, v = line.split(":", 1) + k = k.strip() + if k in keys: + out[k] = v.strip().strip("'\"") + except Exception: + return out + return out + + +def _to_float(v: Optional[str], default: float) -> float: + if v is None or v == "": + return default + try: + return float(v) + except Exception: + return default + + +def _to_int(v: Optional[str], default: int) -> int: + if v is None or v == "": + return default + try: + return int(float(v)) + except Exception: + return default + + +def _find_runs(logs_root: Path) -> List[Path]: + runs = sorted(logs_root.glob("camera_e2e_*/version_0/metrics.csv")) + return runs + + +def _final_metric(df: pd.DataFrame, col: str) -> float: + if col not in df.columns: + return np.nan + s = df[col].dropna() + if s.empty: + return np.nan + return float(s.iloc[-1]) + + +def collect_run_table(logs_root: Path, latest_n: int = 0) -> pd.DataFrame: + rows = [] + metrics_files = _find_runs(logs_root) + if latest_n > 0: + metrics_files = metrics_files[-latest_n:] + + for mpath in metrics_files: + run_dir = mpath.parent + run_name = run_dir.parent.name + try: + df = pd.read_csv(mpath) + except Exception: + continue + + hp = _read_hparams_kv(run_dir / "hparams.yaml") + score_loss_type = hp.get("score_loss_type", "bce") + + row = { + "run_name": run_name, + "run_dir": str(run_dir), + "score_loss_type": score_loss_type, + "score_target_type": hp.get("score_target_type", "l1"), + "score_weight": _to_float(hp.get("score_weight"), 1.0), + "score_temperature": _to_float(hp.get("score_temperature"), 5.0), + "score_rank_weight": _to_float(hp.get("score_rank_weight"), 0.0), + "score_margin": _to_float(hp.get("score_margin"), 0.2), + "score_topk": _to_int(hp.get("score_topk"), 0), + "val_ade_pred": _final_metric(df, "val_ade_pred"), + "val_ade_oracle": _final_metric(df, "val_ade_oracle"), + "val_ade_regret": _final_metric(df, "val_ade_regret"), + "val_loss": _final_metric(df, "val_loss"), + "val_loss_score": _final_metric(df, "val_loss_score"), + "val_score_top1_acc": _final_metric(df, "val_score_top1_acc"), + "val_score_gap_best_second": _final_metric(df, "val_score_gap_best_second"), + } + rows.append(row) + + if not rows: + return pd.DataFrame() + out = pd.DataFrame(rows) + out = out.sort_values("run_name").reset_index(drop=True) + return out + + +def _approach_label(r: pd.Series) -> str: + t = r["score_loss_type"] + tgt = r.get("score_target_type", "l1") + prefix = f"{t}+{tgt}" if tgt != "l1" else t + if t == "bce_pairwise": + return f"{prefix}(w={r['score_rank_weight']:.2f},m={r['score_margin']:.2f},k={int(r['score_topk'])})" + if t in ("bce", "listnet"): + return f"{prefix}(tau={r['score_temperature']:.1f})" + return prefix + + +def _diagnose_row(r: pd.Series) -> str: + pred = r["val_ade_pred"] + oracle = r["val_ade_oracle"] + regret = r["val_ade_regret"] + top1 = r.get("val_score_top1_acc", np.nan) + + if np.isnan(pred) or np.isnan(oracle): + return "incomplete metrics" + if oracle < 0.8 and regret > 1.0: + return "strong proposals, weak ranking (scorer bottleneck)" + if oracle >= 0.8: + return "trajectory quality bottleneck (oracle not strong enough)" + if not np.isnan(top1) and top1 < 0.3: + return "low top1 match; scorer ordering not learning" + if regret < 0.7: + return "ranking improved" + return "mixed behavior" + + +def _plot_rank_bar(df: pd.DataFrame, out_dir: Path): + g = df.groupby("approach", as_index=False)["val_ade_pred"].min().sort_values("val_ade_pred") + if g.empty: + return + fig, ax = plt.subplots(figsize=(11, 5)) + x = np.arange(len(g)) + ax.bar(x, g["val_ade_pred"], color="#2f6db0") + ax.set_ylabel("Best val_ade_pred (m)") + ax.set_title("Approach Ranking (lower is better)") + ax.set_xticks(x) + ax.set_xticklabels(g["approach"], rotation=25, ha="right") + ax.grid(True, axis="y", alpha=0.3) + fig.tight_layout() + fig.savefig(out_dir / "ranking_val_ade_pred.png", dpi=180) + plt.close(fig) + + +def _plot_oracle_vs_pred(df: pd.DataFrame, out_dir: Path): + fig, ax = plt.subplots(figsize=(7, 6)) + for name, sub in df.groupby("approach"): + ax.scatter(sub["val_ade_oracle"], sub["val_ade_pred"], label=name, alpha=0.85) + lim_lo = np.nanmin([df["val_ade_oracle"].min(), df["val_ade_pred"].min()]) * 0.9 + lim_hi = np.nanmax([df["val_ade_oracle"].max(), df["val_ade_pred"].max()]) * 1.1 + ax.plot([lim_lo, lim_hi], [lim_lo, lim_hi], "--", color="gray", alpha=0.7, label="pred=oracle") + ax.set_xlabel("val_ade_oracle (m)") + ax.set_ylabel("val_ade_pred (m)") + ax.set_title("Oracle vs Selected ADE") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + fig.tight_layout() + fig.savefig(out_dir / "oracle_vs_pred_scatter.png", dpi=180) + plt.close(fig) + + +def _plot_regret_vs_top1(df: pd.DataFrame, out_dir: Path): + if "val_score_top1_acc" not in df.columns or df["val_score_top1_acc"].isna().all(): + return + fig, ax = plt.subplots(figsize=(7, 6)) + for name, sub in df.groupby("approach"): + ax.scatter(sub["val_score_top1_acc"], sub["val_ade_regret"], label=name, alpha=0.85) + ax.set_xlabel("val_score_top1_acc") + ax.set_ylabel("val_ade_regret (m)") + ax.set_title("Ranking Accuracy vs Regret") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=8) + fig.tight_layout() + fig.savefig(out_dir / "regret_vs_top1_scatter.png", dpi=180) + plt.close(fig) + + +def _write_report(df: pd.DataFrame, out_dir: Path): + ranked = df.sort_values("val_ade_pred").reset_index(drop=True) + by_approach = ( + df.groupby("approach", as_index=False)[ + ["val_ade_pred", "val_ade_oracle", "val_ade_regret", "val_loss_score", "val_score_top1_acc"] + ] + .agg("mean") + .sort_values("val_ade_pred") + ) + + lines: List[str] = [] + lines.append("# Scorer Experiment Report") + lines.append("") + lines.append("## Overall ranking (by val_ade_pred)") + lines.append("") + for i, r in ranked.head(10).iterrows(): + lines.append( + f"{i+1}. `{r['run_name']}` | approach `{r['approach']}` | " + f"val_ade_pred={r['val_ade_pred']:.3f}, oracle={r['val_ade_oracle']:.3f}, " + f"regret={r['val_ade_regret']:.3f} | diagnosis: {r['diagnosis']}" + ) + + lines.append("") + lines.append("## Approach-level means") + lines.append("") + lines.append("| approach | val_ade_pred | val_ade_oracle | val_ade_regret | val_loss_score | val_score_top1_acc |") + lines.append("|---|---:|---:|---:|---:|---:|") + for _, r in by_approach.iterrows(): + lines.append( + f"| {r['approach']} | {r['val_ade_pred']:.3f} | {r['val_ade_oracle']:.3f} | " + f"{r['val_ade_regret']:.3f} | {r['val_loss_score']:.3f} | {r['val_score_top1_acc']:.3f} |" + ) + + lines.append("") + lines.append("## What went wrong / right") + lines.append("") + scorer_bad = ranked[(ranked["val_ade_oracle"] < 0.8) & (ranked["val_ade_regret"] > 1.0)] + if not scorer_bad.empty: + lines.append("- **Scorer bottleneck persists** in runs where oracle is strong but regret stays high.") + trajectory_bad = ranked[ranked["val_ade_oracle"] >= 0.8] + if not trajectory_bad.empty: + lines.append("- **Trajectory generation bottleneck** appears in some runs (high oracle ADE).") + improved = ranked[ranked["val_ade_regret"] < 0.7] + if not improved.empty: + lines.append("- **Ranking improved** for runs with regret below 0.7.") + if scorer_bad.empty and trajectory_bad.empty and improved.empty: + lines.append("- Mixed outcomes; inspect scatter plots for separation patterns.") + + (out_dir / "score_experiment_report.md").write_text("\n".join(lines)) + + +def main(): + parser = argparse.ArgumentParser(description="Rank scorer-loss experiments and generate diagnostics") + parser.add_argument("--logs_root", type=str, required=True, help="Root containing camera_e2e_*/version_0") + parser.add_argument("--out_dir", type=str, default=None, help="Output dir (default: /score_experiments)") + parser.add_argument("--latest_n", type=int, default=0, help="Only process latest N runs (0 = all)") + args = parser.parse_args() + + logs_root = Path(args.logs_root) + out_dir = Path(args.out_dir) if args.out_dir else logs_root / "score_experiments" + out_dir.mkdir(parents=True, exist_ok=True) + + df = collect_run_table(logs_root, latest_n=args.latest_n) + if df.empty: + raise RuntimeError(f"No runs with metrics found under {logs_root}") + + df["approach"] = df.apply(_approach_label, axis=1) + df["diagnosis"] = df.apply(_diagnose_row, axis=1) + df.to_csv(out_dir / "score_experiment_summary.csv", index=False) + + _plot_rank_bar(df, out_dir) + _plot_oracle_vs_pred(df, out_dir) + _plot_regret_vs_top1(df, out_dir) + _write_report(df, out_dir) + + print(f"Wrote summary CSV and plots to: {out_dir}") + print(f"Top run by val_ade_pred: {df.sort_values('val_ade_pred').iloc[0]['run_name']}") + + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh b/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh new file mode 100755 index 0000000..7d7019d --- /dev/null +++ b/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# Run proposal training on the current machine (e.g. login) with visible GPU. +# Same args as run_train_proposal.sbatch, without Slurm. Does not pip-install PL. + +set -euo pipefail + +module load conda 2>/dev/null || true +CONDA_BASE=$(conda info --base 2>/dev/null) +if [ -n "$CONDA_BASE" ] && [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then + source "$CONDA_BASE/etc/profile.d/conda.sh" +else + eval "$(conda shell.bash hook)" +fi +conda activate robotvision + +DATA_DIR="/scratch/gilbreth/mathur91/waymo/waymo_open_dataset_end_to_end_camera_v_1_0_0" +PROJECT_DIR="/home/mathur91/robotvision/src/camera-based-e2e" + +mkdir -p "$PROJECT_DIR/logs" +cd "$PROJECT_DIR" + +exec python train.py \ + --data_dir "$DATA_DIR" \ + --model_type proposal \ + --backbone resnet \ + --no_wandb \ + --batch_size 16 \ + --lr 1e-4 \ + --max_epochs 20 \ + --num_proposals 16 \ + --num_refinement_steps 4 \ + --smoothness_weight 0.0 \ + --comfort_weight 0.0 \ + --diversity_weight 0.0 \ + --score_weight 1.0 \ + --score_warmup_epochs 2 \ + --score_temperature 5.0 \ + --score_loss_type bce \ + --score_target_type l1 \ + --score_rank_weight 0.2 \ + --score_margin 0.2 \ + --score_topk 0 \ + --comfort_jerk_threshold 5.0 \ + --prev_weight 0.1 \ + --grad_clip 1.0 \ + --log_every_n_steps 100 \ + ${EXTRA_ARGS:-} diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py index e5cd30a..048579c 100644 --- a/src/camera-based-e2e/train.py +++ b/src/camera-based-e2e/train.py @@ -1,4 +1,5 @@ -import argparse +import argparse +import os from datetime import datetime import pytorch_lightning as pl @@ -17,19 +18,74 @@ # Replace with your model defined in models/ from models.base_model import LitModel, collate_with_images from models.monocular import DeepMonocularModel -from models.feature_extractors import SAMFeatures - +from models.proposal_planner import ProposalPlanner +from models.feature_extractors import SAMFeatures, DINOFeatures, ResNetFeatures +from models.debug_callbacks import GradientDebugCallback + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--data_dir', type=str, required=True, help='Path to Waymo E2E data directory') + parser.add_argument('--model_type', type=str, default='deep_monocular', choices=['deep_monocular', 'proposal'], + help='Model type: deep_monocular or proposal (proposal-centric planner)') parser.add_argument('--batch_size', type=int, default=16, help='Batch size for training') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') parser.add_argument('--max_epochs', type=int, default=10, help='Number of epochs to train') + parser.add_argument('--num_proposals', type=int, default=16, help='Number of proposals (proposal model only)') + parser.add_argument('--num_refinement_steps', type=int, default=4, help='Refinement iterations (weight-shared, iPad default=4)') + parser.add_argument('--smoothness_weight', type=float, default=0.01, help='Smoothness (jerk) loss weight (proposal model)') + parser.add_argument('--collision_weight', type=float, default=0.0, help='Collision penalty weight (proposal model)') + parser.add_argument('--comfort_weight', type=float, default=0.01, help='Comfort (curvature) loss weight (proposal model)') + parser.add_argument('--rfs_weight', type=float, default=0.0, help='RFS loss weight (proposal model)') + parser.add_argument('--diversity_weight', type=float, default=0.0, help='Diversity weight (0 = off, MoN loss handles diversity)') + parser.add_argument('--prev_weight', type=float, default=0.1, help='Discount λ for intermediate proposal losses') + parser.add_argument('--score_weight', type=float, default=1.0, help='Score loss weight (BCE is well-scaled, safe at 1.0)') + parser.add_argument('--score_warmup_epochs', type=int, default=2, help='Epochs before score loss activates (prevents early mode collapse)') + parser.add_argument('--score_temperature', type=float, default=5.0, help='Temperature τ for quality target exp(-ADE/τ)') + parser.add_argument( + '--score_loss_type', + type=str, + default='bce', + choices=['bce', 'ce', 'bce_pairwise', 'listnet'], + help='Scorer objective: bce (iPad-faithful), ce, bce_pairwise, or listnet', + ) + parser.add_argument( + '--score_target_type', + type=str, + default='l1', + choices=['l1', 'navsim', 'rfs'], + help='Quality target: l1 (exp(-L1/tau)), navsim (EP+Comf), or rfs (long/lat RFS at 3s/5s × optional jerk)', + ) + parser.add_argument( + '--no_rfs_target_comfort', + action='store_true', + help='For score_target_type=rfs: use pure RFS mean only (no jerk comfort multiplier)', + ) + parser.add_argument('--score_rank_weight', type=float, default=0.2, help='Aux weight for pairwise ranking term (bce_pairwise only)') + parser.add_argument('--score_margin', type=float, default=0.2, help='Pairwise ranking margin (bce_pairwise only)') + parser.add_argument('--score_topk', type=int, default=0, help='Hard negative top-k for pairwise ranking (0=all negatives)') + parser.add_argument('--comfort_jerk_threshold', type=float, default=5.0, help='Jerk threshold (m/s^3) for NAVSIM comfort sub-metric') + parser.add_argument('--grad_clip', type=float, default=1.0, help='Gradient clipping max norm (0 to disable)') + parser.add_argument('--log_every_n_steps', type=int, default=100, help='How often (steps) to emit trainer logs') + parser.add_argument('--backbone', type=str, default='resnet', choices=['resnet', 'dino', 'sam'], + help='Backbone: resnet (default, widely available), dino, or sam') + parser.add_argument('--no_wandb', action='store_true', help='Disable wandb logging (use CSV only)') parser.add_argument('--compile', action='store_true', help='Whether to compile the model with torch.compile') parser.add_argument('--profile', action='store_true', help='Whether to run the profiler') + parser.add_argument('--debug', action='store_true', help='Enable debug visualizations (gradient norms, activation stats, proposal diagnostics)') + parser.add_argument('--debug_log_every', type=int, default=10, help='How often (steps) to log debug metrics') args = parser.parse_args() # Data + # #region agent log + try: + _logpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "debug-aec3fc.log") + open(_logpath, "a").write( + '{"sessionId":"aec3fc","hypothesisId":"H1,H2","location":"train.py:datasets","message":"train_val_config","data":{"data_dir":"%s","train_index":"index_train.pkl","val_index":"index_val.pkl","has_val_loader":true}}\n' + % (str(args.data_dir),) + ) + except Exception: + pass + # #endregion train_dataset = WaymoE2E(indexFile='index_train.pkl', data_dir=args.data_dir, n_items=250_000) test_dataset = WaymoE2E(indexFile='index_val.pkl', data_dir=args.data_dir, n_items=25_000) @@ -37,36 +93,86 @@ val_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, num_workers=0, 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) - model = DeepMonocularModel(feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_small_patch16_512.fb", frozen=True), out_dim=out_dim, n_blocks=8) + if args.model_type == 'proposal': + backbone = ( + ResNetFeatures(frozen=True) if args.backbone == 'resnet' + else DINOFeatures(frozen=True) if args.backbone == 'dino' + else SAMFeatures(frozen=True) + ) + model = ProposalPlanner( + backbone=backbone, + d_model=256, + num_proposals=args.num_proposals, + num_refinement_steps=args.num_refinement_steps, + horizon=20, + num_cams=8, + ) + else: + backbone = ( + ResNetFeatures(frozen=True) if args.backbone == 'resnet' + else DINOFeatures(frozen=True) if args.backbone == 'dino' + else SAMFeatures(frozen=True) + ) + model = DeepMonocularModel(feature_extractor=backbone, out_dim=out_dim, n_blocks=8) + if args.debug: + model._debug = True + if args.compile: model = torch.compile(model, mode="max-autotune") - lit_model = LitModel(model=model, lr=args.lr) + lit_model = LitModel( + model=model, + lr=args.lr, + smoothness_weight=args.smoothness_weight if args.model_type == 'proposal' else 0.0, + collision_weight=args.collision_weight if args.model_type == 'proposal' else 0.0, + comfort_weight=args.comfort_weight if args.model_type == 'proposal' else 0.0, + rfs_weight=args.rfs_weight if args.model_type == 'proposal' else 0.0, + diversity_weight=args.diversity_weight if args.model_type == 'proposal' else 0.0, + score_weight=args.score_weight if args.model_type == 'proposal' else 0.0, + score_warmup_epochs=args.score_warmup_epochs if args.model_type == 'proposal' else 0, + score_temperature=args.score_temperature if args.model_type == 'proposal' else 5.0, + score_loss_type=args.score_loss_type if args.model_type == 'proposal' else 'bce', + score_target_type=args.score_target_type if args.model_type == 'proposal' else 'l1', + score_rank_weight=args.score_rank_weight if args.model_type == 'proposal' else 0.0, + score_margin=args.score_margin if args.model_type == 'proposal' else 0.2, + score_topk=args.score_topk if args.model_type == 'proposal' else 0, + comfort_jerk_threshold=args.comfort_jerk_threshold if args.model_type == 'proposal' else 5.0, + prev_weight=args.prev_weight if args.model_type == 'proposal' else 0.1, + rfs_target_use_comfort=not args.no_rfs_target_comfort if args.model_type == 'proposal' else True, + ) # We don't want to save logs or checkpoints in the home directory - it'll fill up fast base_path = Path(args.data_dir).parent.as_posix() timestamp = f"camera_e2e_{datetime.now().strftime('%Y%m%d_%H%M')}" - wandb_logger = WandbLogger(name=timestamp, save_dir=base_path + "/logs", project="robotvision", log_model=True) - wandb_logger.watch(lit_model, log="all") + loggers = [CSVLogger(base_path + "/logs", name=timestamp)] + if not args.no_wandb: + wandb_logger = WandbLogger(name=timestamp, save_dir=base_path + "/logs", project="robotvision", log_model=True) + wandb_logger.watch(lit_model, log="all") + loggers.append(wandb_logger) strategy = "ddp_find_unused_parameters_true" if torch.cuda.device_count() > 1 else "auto" + callbacks = [ + ModelCheckpoint(monitor='val_loss', + mode='min', + save_top_k=1, + dirpath=base_path + '/checkpoints', + filename='camera-e2e-{epoch:02d}-{val_loss:.2f}' + ), + ] + if args.debug: + callbacks.append(GradientDebugCallback(log_every=args.debug_log_every)) + trainer = pl.Trainer( max_epochs=args.max_epochs, - logger=[CSVLogger(base_path + "/logs", name=timestamp), wandb_logger], + logger=loggers, strategy=strategy, precision="bf16-mixed" if torch.cuda.is_bf16_supported() else 16, - log_every_n_steps=10, + log_every_n_steps=args.log_every_n_steps, + gradient_clip_val=args.grad_clip if args.grad_clip > 0 else None, + gradient_clip_algorithm="norm", profiler=SimpleProfiler(extended=True) if args.profile else None, - callbacks=[ - ModelCheckpoint(monitor='val_loss', - mode='min', - save_top_k=1, - dirpath=base_path + '/checkpoints', - filename='camera-e2e-{epoch:02d}-{val_loss:.2f}' - ), - ], + callbacks=callbacks, ) trainer.fit(lit_model, train_loader, val_loader) @@ -87,10 +193,40 @@ plt.legend() plt.tight_layout() out = Path("./visualizations") + out.mkdir(parents=True, exist_ok=True) plt.savefig(out / "loss.png", dpi=200) except Exception as e: print(f"Could not save loss plot: {e}") + if args.debug: + try: + from debug_viz import find_metrics_csv, load_and_merge + from debug_viz import ( + plot_gradient_norms_by_module, plot_gradient_norms_refinement, + plot_gradient_norms_scorer_propinit, plot_gradient_dominance, + plot_activation_stats, plot_proposal_diagnostics, + plot_loss_breakdown, plot_ade_regret, plot_param_norms, + plot_summary_dashboard, + ) + csv_path = find_metrics_csv(run_dir / "version_0") + df = load_and_merge(csv_path) + debug_out = base_path / "debug_plots" + debug_out.mkdir(parents=True, exist_ok=True) + print(f"\nGenerating debug plots to {debug_out} ...") + plot_gradient_norms_by_module(df, debug_out) + plot_gradient_norms_refinement(df, debug_out) + plot_gradient_norms_scorer_propinit(df, debug_out) + plot_gradient_dominance(df, debug_out) + plot_activation_stats(df, debug_out) + plot_proposal_diagnostics(df, debug_out) + plot_loss_breakdown(df, debug_out) + plot_ade_regret(df, debug_out) + plot_param_norms(df, debug_out) + plot_summary_dashboard(df, debug_out) + print(f"Debug plots saved to {debug_out}") + except Exception as e: + print(f"Could not generate debug plots: {e}") + From ca08746023fa3f0a486194de8f1632260e7a646a Mon Sep 17 00:00:00 2001 From: Pratyush Date: Wed, 1 Apr 2026 13:01:34 -0400 Subject: [PATCH 2/6] env --- environment.yml | 111 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) create mode 100644 environment.yml diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..287a065 --- /dev/null +++ b/environment.yml @@ -0,0 +1,111 @@ +name: robotvision +channels: + - conda-forge +dependencies: + - _libgcc_mutex=0.1 + - _openmp_mutex=4.5 + - bzip2=1.0.8 + - ca-certificates=2025.10.5 + - icu=75.1 + - ld_impl_linux-64=2.44 + - libexpat=2.7.1 + - libffi=3.5.2 + - libgcc=15.2.0 + - libgcc-ng=15.2.0 + - libgomp=15.2.0 + - liblzma=5.8.1 + - libnsl=2.0.1 + - libsqlite=3.51.0 + - libstdcxx=15.2.0 + - libstdcxx-ng=15.2.0 + - libuuid=2.41.2 + - libxcrypt=4.4.36 + - libzlib=1.3.1 + - ncurses=6.5 + - openssl=3.5.4 + - python=3.9.23 + - readline=8.2 + - tk=8.6.13 + - tzdata=2025b + - zstd=1.5.7 + - pip: + - accelerate==1.10.1 + - aiohappyeyeballs==2.6.1 + - aiohttp==3.13.2 + - aiosignal==1.4.0 + - annotated-types==0.7.0 + - anyio==4.12.0 + - async-timeout==5.0.1 + - attrs==25.4.0 + - blinker==1.9.0 + - certifi==2025.10.5 + - charset-normalizer==3.4.4 + - docker==7.1.0 + - exceptiongroup==1.3.1 + - flask==3.1.2 + - frozenlist==1.8.0 + - gitdb==4.0.12 + - graphql-core==3.2.7 + - graphql-relay==3.2.0 + - greenlet==3.2.4 + - gunicorn==23.0.0 + - h11==0.16.0 + - hf-xet==1.2.0 + - httpcore==1.0.9 + - httpx==0.28.1 + - huggingface-hub==0.36.0 + - idna==3.11 + - itsdangerous==2.2.0 + - lightning-utilities==0.15.2 + - mako==1.3.10 + - matplotlib==3.9.4 + - multidict==6.7.0 + - numpy==2.0.2 + - nvidia-cublas-cu12==12.4.5.8 + - nvidia-cuda-cupti-cu12==12.4.127 + - nvidia-cuda-nvrtc-cu12==12.4.127 + - nvidia-cuda-runtime-cu12==12.4.127 + - nvidia-cudnn-cu12==9.1.0.70 + - nvidia-cufft-cu12==11.2.1.3 + - nvidia-curand-cu12==10.3.5.147 + - nvidia-cusparse-cu12==12.3.1.170 + - nvidia-cusparselt-cu12==0.6.2 + - nvidia-ml-py==12.570.172 + - nvidia-nccl-cu12==2.21.5 + - nvidia-nvjitlink-cu12==12.4.127 + - nvidia-nvtx-cu12==12.4.127 + - nvitop==1.4.2 + - opentelemetry-api==1.39.1 + - opentelemetry-semantic-conventions==0.60b1 + - packaging==24.2 + - propcache==0.4.1 + - protobuf==3.20.3 + - pydantic==2.12.5 + - pydantic-core==2.41.5 + - python-dateutil==2.9.0.post0 + - pytorch-lightning==2.5.6 + - pyyaml==6.0.3 + - regex==2026.1.15 + - requests==2.32.5 + - safetensors==0.7.0 + - scipy==1.13.1 + - seaborn==0.13.2 + - shellingham==1.5.4 + - six==1.17.0 + - smmap==5.0.2 + - sqlalchemy==2.0.46 + - sqlparse==0.5.5 + - sympy==1.13.1 + - timm==1.0.12 + - tokenizers==0.22.2 + - tomli==2.4.0 + - torcheval==0.0.7 + - torchmetrics==1.8.2 + - transformers==4.57.6 + - triton==3.2.0 + - typer-slim==0.20.0 + - typing-inspection==0.4.2 + - urllib3==2.5.0 + - wrapt==2.0.1 + - yarl==1.22.0 +prefix: /scratch/gilbreth/mathur91/conda/robotvision From 734dc8fa326147c3465affd5d56320bd338bfdc7 Mon Sep 17 00:00:00 2001 From: Pratyush Date: Wed, 1 Apr 2026 18:10:45 -0400 Subject: [PATCH 3/6] training script --- .../scripts/run_train_proposal.sbatch | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 src/camera-based-e2e/scripts/run_train_proposal.sbatch diff --git a/src/camera-based-e2e/scripts/run_train_proposal.sbatch b/src/camera-based-e2e/scripts/run_train_proposal.sbatch new file mode 100644 index 0000000..4f45005 --- /dev/null +++ b/src/camera-based-e2e/scripts/run_train_proposal.sbatch @@ -0,0 +1,84 @@ +#!/bin/bash +#SBATCH --job-name=ipad +#SBATCH --output=logs/slurm-%j.out +#SBATCH --error=logs/slurm-%j.err +#SBATCH --time=56:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-node=1 +#SBATCH --cpus-per-task=20 +#SBATCH --mem=64G +#SBATCH --account=csso +#SBATCH --partition=a10 + +module load conda +# Sourcing conda.sh makes activate work in non-interactive (sbatch) shells +CONDA_BASE=$(conda info --base 2>/dev/null) +if [ -n "$CONDA_BASE" ] && [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then + source "$CONDA_BASE/etc/profile.d/conda.sh" +else + eval "$(conda shell.bash hook)" +fi +conda activate robotvision +pip install pytorch-lightning +# Use the python that activation puts on PATH. +# Some clusters don't mount the original conda env prefix on compute nodes. +PYTHON="$(command -v python || true)" +echo "CONDA_PREFIX=${CONDA_PREFIX}" +echo "PYTHON=${PYTHON}" +if [ -z "$PYTHON" ] || [ ! -x "$PYTHON" ]; then + echo "ERROR: python not found after conda activate." + echo "If this is a shared-filesystem issue, recreate the env in a path mounted on compute nodes (e.g. under \$HOME)." + conda info --envs || true + exit 2 +fi + +# Ensure pytorch_lightning is available (required by train.py). +if ! "$PYTHON" -c "import pytorch_lightning as pl; print('pytorch_lightning', pl.__version__)" 2>/dev/null; then + echo "ERROR: pytorch_lightning not found in env: ${CONDA_PREFIX:-unknown}" + echo "Install it in the same env you use for this job, e.g.:" + echo " conda activate robotvision" + echo " pip install pytorch-lightning" + exit 3 +fi +# Waymo E2E data; index_train.pkl and index_val.pkl must be in PROJECT_DIR +DATA_DIR="/scratch/gilbreth/mathur91/waymo/waymo_open_dataset_end_to_end_camera_v_1_0_0" +PROJECT_DIR="/home/mathur91/robotvision/src/camera-based-e2e" + +mkdir -p "$PROJECT_DIR/logs" +cd "$PROJECT_DIR" + +# Scorer targets: rfs = per-proposal RFS quality at 3s/5s (× jerk comfort). Override with EXTRA_ARGS, e.g. --score_target_type l1 +"$PYTHON" train.py \ + --data_dir "$DATA_DIR" \ + --model_type proposal \ + --backbone resnet \ + --no_wandb \ + --batch_size 16 \ + --lr 1e-4 \ + --max_epochs 20 \ + --num_proposals 16 \ + --num_refinement_steps 4 \ + --smoothness_weight 0.0 \ + --comfort_weight 0.0 \ + --rfs_weight 0.2 \ + --diversity_weight 0.0 \ + --score_weight 1.0 \ + --score_warmup_epochs 2 \ + --score_temperature 5.0 \ + --score_loss_type bce \ + --score_target_type rfs \ + --score_rank_weight 0.2 \ + --score_margin 0.2 \ + --score_topk 0 \ + --comfort_jerk_threshold 5.0 \ + --prev_weight 0.1 \ + --grad_clip 1.0 \ + --log_every_n_steps 100 \ + ${EXTRA_ARGS:-} + +# Auto-refresh cross-run scorer experiment ranking plots/report. +LOGS_ROOT="$(dirname "$DATA_DIR")/logs" +if [ -d "$LOGS_ROOT" ]; then + "$PYTHON" score_experiment_viz.py --logs_root "$LOGS_ROOT" --latest_n 24 || true +fi \ No newline at end of file From a0f1ebfb225b02d190bca282dd49ce34cff03e1e Mon Sep 17 00:00:00 2001 From: nysapk Date: Tue, 21 Apr 2026 18:44:51 -0400 Subject: [PATCH 4/6] cleaned up some dead code --- src/camera-based-e2e/create_vocab.py | 3 +- src/camera-based-e2e/loader.py | 13 +- .../models/debug_callbacks.py | 5 + src/camera-based-e2e/models/monocular.py | 13 +- .../models/proposal_planner.py | 1 - src/camera-based-e2e/models/transfuser | 1 - src/camera-based-e2e/score_experiment_viz.py | 1 - .../scripts/run_train_proposal.sbatch | 2 +- src/camera-based-e2e/train.py | 10 - src/camera-based-e2e/viz.py | 337 ------------------ 10 files changed, 13 insertions(+), 373 deletions(-) create mode 100644 src/camera-based-e2e/models/debug_callbacks.py delete mode 160000 src/camera-based-e2e/models/transfuser delete mode 100644 src/camera-based-e2e/viz.py diff --git a/src/camera-based-e2e/create_vocab.py b/src/camera-based-e2e/create_vocab.py index 4ddf3a0..08c6ed5 100644 --- a/src/camera-based-e2e/create_vocab.py +++ b/src/camera-based-e2e/create_vocab.py @@ -13,8 +13,7 @@ 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', - images=False, + 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 ) diff --git a/src/camera-based-e2e/loader.py b/src/camera-based-e2e/loader.py index d9e088b..a7ca4a6 100644 --- a/src/camera-based-e2e/loader.py +++ b/src/camera-based-e2e/loader.py @@ -67,16 +67,7 @@ def __iter__(self): if self.file: self.file.close() del self.file - # #region agent log - _full = os.path.join(self.data_dir, filename) - _logpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "debug-aec3fc.log") - try: - _line = '{"sessionId":"aec3fc","hypothesisId":"H3,H4,H5","location":"loader.py:open","message":"file_open_attempt","data":{"data_dir":"%s","filename":"%s","full_path":"%s","exists":%s,"index_file":"%s"}}\n' % (str(self.data_dir), str(filename), str(_full), str(os.path.exists(_full)), str(getattr(self, "_index_file", "?"))) - open(_logpath, "a").write(_line) - except Exception: - pass - # #endregion - self.file = open(_full, 'rb') + self.file = open(os.path.join(self.data_dir, filename), 'rb') self.filename = filename self.file.seek(start_byte) # type: ignore @@ -112,7 +103,7 @@ def __iter__(self): # NOTE: Replace with your path 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, images=True) + dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR) loader = DataLoader( dataset, batch_size=BATCH_SIZE, diff --git a/src/camera-based-e2e/models/debug_callbacks.py b/src/camera-based-e2e/models/debug_callbacks.py new file mode 100644 index 0000000..a9921a0 --- /dev/null +++ b/src/camera-based-e2e/models/debug_callbacks.py @@ -0,0 +1,5 @@ +import pytorch_lightning as pl + +class GradientDebugCallback(pl.Callback): + def __init__(self, log_every=10): + self.log_every = log_every diff --git a/src/camera-based-e2e/models/monocular.py b/src/camera-based-e2e/models/monocular.py index 6c2035f..289e21a 100644 --- a/src/camera-based-e2e/models/monocular.py +++ b/src/camera-based-e2e/models/monocular.py @@ -163,21 +163,16 @@ def forward(self, x): # Doesn't need no_grad b/c DINO/SAMFeatures will freeze if needed feats_vit = self.features(front_cam) # list or tensor - if len(feats_vit) == 1 and isinstance(feats_vit, list): + if isinstance(feats_vit, (list, tuple)): + if len(feats_vit) != 1: + raise ValueError(f"Expected single feature map, got {len(feats_vit)}") feats_vit = feats_vit[0] - feats = self.visual_adapter(feats_vit) # (B, C, H, W) # Depth Supervision output_depth = F.softplus(self.depth_gen(feats).squeeze(1)) # (B, 128, 128) - # tokens: handle list of features or single tensor - # TODO: is this made redundant by if statement above? - if isinstance(feats, (list, tuple)): - tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) - else: - tokens = feats.flatten(2) # (B, C, N) - tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) + tokens = feats.flatten(2).permute(0, 2, 1) + self.positional_encoding # (B, N, C_total) # copy procedure to build query_0 from MonocularModel intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() diff --git a/src/camera-based-e2e/models/proposal_planner.py b/src/camera-based-e2e/models/proposal_planner.py index 8c2b4a4..2194dec 100644 --- a/src/camera-based-e2e/models/proposal_planner.py +++ b/src/camera-based-e2e/models/proposal_planner.py @@ -28,7 +28,6 @@ def __init__( num_cams: int = 8, ): super().__init__() - self.num_proposals = num_proposals self.horizon = horizon self.n_proposals = num_proposals diff --git a/src/camera-based-e2e/models/transfuser b/src/camera-based-e2e/models/transfuser deleted file mode 160000 index 9d413b2..0000000 --- a/src/camera-based-e2e/models/transfuser +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 9d413b2ad2d2d56c112b34a4a799be081800d77f diff --git a/src/camera-based-e2e/score_experiment_viz.py b/src/camera-based-e2e/score_experiment_viz.py index a536c00..d0139e5 100644 --- a/src/camera-based-e2e/score_experiment_viz.py +++ b/src/camera-based-e2e/score_experiment_viz.py @@ -37,7 +37,6 @@ def _read_hparams_kv(hparams_path: Path) -> Dict[str, str]: "score_rank_weight", "score_margin", "score_topk", - "comfort_jerk_threshold", "model_type", } out: Dict[str, str] = {} diff --git a/src/camera-based-e2e/scripts/run_train_proposal.sbatch b/src/camera-based-e2e/scripts/run_train_proposal.sbatch index 4f45005..2a7b429 100644 --- a/src/camera-based-e2e/scripts/run_train_proposal.sbatch +++ b/src/camera-based-e2e/scripts/run_train_proposal.sbatch @@ -81,4 +81,4 @@ cd "$PROJECT_DIR" LOGS_ROOT="$(dirname "$DATA_DIR")/logs" if [ -d "$LOGS_ROOT" ]; then "$PYTHON" score_experiment_viz.py --logs_root "$LOGS_ROOT" --latest_n 24 || true -fi \ No newline at end of file +fi diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py index 048579c..7754d94 100644 --- a/src/camera-based-e2e/train.py +++ b/src/camera-based-e2e/train.py @@ -76,16 +76,6 @@ args = parser.parse_args() # Data - # #region agent log - try: - _logpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "debug-aec3fc.log") - open(_logpath, "a").write( - '{"sessionId":"aec3fc","hypothesisId":"H1,H2","location":"train.py:datasets","message":"train_val_config","data":{"data_dir":"%s","train_index":"index_train.pkl","val_index":"index_val.pkl","has_val_loader":true}}\n' - % (str(args.data_dir),) - ) - except Exception: - pass - # #endregion train_dataset = WaymoE2E(indexFile='index_train.pkl', data_dir=args.data_dir, n_items=250_000) test_dataset = WaymoE2E(indexFile='index_val.pkl', data_dir=args.data_dir, n_items=25_000) diff --git a/src/camera-based-e2e/viz.py b/src/camera-based-e2e/viz.py deleted file mode 100644 index 78b8c2f..0000000 --- a/src/camera-based-e2e/viz.py +++ /dev/null @@ -1,337 +0,0 @@ -import argparse, os, pathlib -import torch -import matplotlib.pyplot as plt -import matplotlib.animation as animation -import numpy as np -from typing import List, Tuple, Optional -from tqdm import tqdm - -from loader import WaymoE2E -from models.monocular import MonocularModel, SAMFeatures -from torch import nn - -def gen_viz_data(model: nn.Module, data_root: str, num_samples: int, device: torch.device): - """Generate the gt and pred trajectories along with past states""" - dataset = WaymoE2E( - batch_size=1, - indexFile="index_val.pkl", - data_dir=data_root, - images=True, - n_items=num_samples, - ) - loader = torch.utils.data.DataLoader( - dataset, - batch_size=1, - num_workers=4, - pin_memory=True, - persistent_workers=False, - ) - model.eval() - - pred_trajectories = [] - gt_trajectories = [] - past_trajectories = [] - - with torch.no_grad(): - for batch in tqdm(loader, desc="Generating trajectory data"): - past, future, images, intent = batch['PAST'], batch['FUTURE'], batch['IMAGES'], batch['INTENT'] - past = past.to(device, non_blocking=True) - future = future.to(device, non_blocking=True) - intent = intent.to(device, non_blocking=True) - images = [img.to(device, non_blocking=True) for img in images] - B, T, F = past.shape - - # Extract past positions (x, y) from the first 2 features - past_positions = past[:, :, :2].squeeze(0).cpu() # (T, 2) - past_trajectories.append(past_positions) - - pred_future = model({"PAST": past, "IMAGES": images, "INTENT": intent}) - pred_future = pred_future.view(B, -1, 2) - pred_trajectories.append(pred_future.squeeze(0).cpu()) - gt_trajectories.append(future.squeeze(0).cpu()) - - return past_trajectories, pred_trajectories, gt_trajectories - - -def create_animated_trajectory_plot( - past_trajectories: List[torch.Tensor], - pred_trajectories: List[torch.Tensor], - gt_trajectories: List[torch.Tensor], - save_path: Optional[str] = None, - interval: int = 200, -): - """ - Args: - past_trajectories: List of past position tensors [(T_past, 2), ...] - pred_trajectories: List of predicted future position tensors [(T_future, 2), ...] - gt_trajectories: List of ground truth future position tensors [(T_future, 2), ...] - save_path: path to save - interval: 1/framerate (ms) - """ - fig, ax = plt.subplots(figsize=(12, 8)) - - # Calculate global bounds once from all trajectories to prevent jumping - all_x_coords = [] - all_y_coords = [] - for past, pred, gt in zip(past_trajectories, pred_trajectories, gt_trajectories): - all_x_coords.extend([past[:, 0].numpy(), pred[:, 0].numpy(), gt[:, 0].numpy()]) - all_y_coords.extend([past[:, 1].numpy(), pred[:, 1].numpy(), gt[:, 1].numpy()]) - - global_x = np.concatenate(all_x_coords) - global_y = np.concatenate(all_y_coords) - margin = 5.0 - x_min, x_max = global_x.min() - margin, global_x.max() + margin - y_min, y_max = global_y.min() - margin, global_y.max() + margin - - past_color = "blue" - pred_color = "red" - gt_color = "green" - - (past_line,) = ax.plot( - [], - [], - "o-", - color=past_color, - linewidth=2, - markersize=4, - label="Past Trajectory", - alpha=0.8, - ) - (pred_line,) = ax.plot( - [], - [], - "o-", - color=pred_color, - linewidth=2, - markersize=4, - label="Predicted Future", - alpha=0.8, - ) - (gt_line,) = ax.plot( - [], - [], - "o-", - color=gt_color, - linewidth=2, - markersize=4, - label="Ground Truth Future", - alpha=0.8, - ) - - (current_pos,) = ax.plot([], [], "ko", markersize=8, label="Current Position") - - text = ax.text( - 0.02, - 0.98, - "", - transform=ax.transAxes, - fontsize=12, - verticalalignment="top", - bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8), - ) - - def init(): - ax.clear() - ax.set_xlabel("X Position (m)") - ax.set_ylabel("Y Position (m)") - ax.set_title("Vehicle Trajectory Animation: Past → Present → Future") - ax.grid(True, alpha=0.3) - ax.legend() - return past_line, pred_line, gt_line, current_pos, text - - def animate(frame): - # Update progress bar for frame generation - if not hasattr(animate, 'pbar'): - animate.pbar = tqdm(total=len(past_trajectories), desc="Generating animation frames", unit="frame", position=1, leave=False) - animate.pbar.update(1) - - ax.clear() - ax.set_xlabel("X Position (m)") - ax.set_ylabel("Y Position (m)") - ax.set_title("Vehicle Trajectory Animation: Past → Present → Future") - ax.grid(True, alpha=0.3) - - sample_idx = frame % len(past_trajectories) - - past = past_trajectories[sample_idx].numpy() - pred = pred_trajectories[sample_idx].numpy() - gt = gt_trajectories[sample_idx].numpy() - - ax.plot( - past[:, 0], - past[:, 1], - "o-", - color=past_color, - linewidth=2, - markersize=4, - label="Past Trajectory", - alpha=0.8, - ) - ax.plot( - pred[:, 0], - pred[:, 1], - "o-", - color=pred_color, - linewidth=2, - markersize=4, - label="Predicted Future", - alpha=0.8, - ) - ax.plot( - gt[:, 0], - gt[:, 1], - "o-", - color=gt_color, - linewidth=2, - markersize=4, - label="Ground Truth Future", - alpha=0.8, - ) - - current_x, current_y = past[-1, 0], past[-1, 1] - ax.plot(current_x, current_y, "ko", markersize=8, label="Current Position") - - # Use static bounds - ax.set_xlim(x_min, x_max) - ax.set_ylim(y_min, y_max) - - ax.text( - 0.02, - 0.98, - f"Sample {sample_idx + 1}/{len(past_trajectories)}", - transform=ax.transAxes, - fontsize=12, - verticalalignment="top", - bbox=dict(boxstyle="round", facecolor="lightblue", alpha=0.8), - ) - - ax.legend() - ax.set_aspect("equal", adjustable="box") - - return [] - - # Create animation - anim = animation.FuncAnimation( - fig, - animate, - init_func=init, - frames=len(past_trajectories), - interval=interval, - blit=False, - repeat=True, - ) - - if save_path: - anim.save(save_path, writer="pillow", fps=3) - - return anim - - -def calculate_metrics( - pred_trajectories: List[torch.Tensor], gt_trajectories: List[torch.Tensor] -) -> dict: - """Calculate trajectory prediction metrics""" - ade_scores = [] # Average Displacement Error - fde_scores = [] # Final Displacement Error - - for pred, gt in zip(pred_trajectories, gt_trajectories): - # ADE: Average Euclidean distance over all time steps - distances = torch.norm(pred - gt, dim=1) - ade = torch.mean(distances).item() - ade_scores.append(ade) - - # FDE: Euclidean distance at final time step - fde = torch.norm(pred[-1] - gt[-1]).item() - fde_scores.append(fde) - - return { - "ADE_mean": np.mean(ade_scores), - "ADE_std": np.std(ade_scores), - "ADE_p1": np.percentile(ade_scores, 1), - "ADE_p50": np.percentile(ade_scores, 50), - "ADE_p99": np.percentile(ade_scores, 99), - "FDE_mean": np.mean(fde_scores), - "FDE_std": np.std(fde_scores), - "FDE_p1": np.percentile(fde_scores, 1), - "FDE_p50": np.percentile(fde_scores, 50), - "FDE_p99": np.percentile(fde_scores, 99), - } - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--data_root", - type=str, - required=True, - help="Path to the root directory of the Waymo E2E dataset.", - ) - parser.add_argument( - "--model_path", - type=str, - required=True, - help="Path to the trained model checkpoint.", - ) - parser.add_argument( - "--num_samples", type=int, default=100, help="Number of samples to animate." - ) - parser.add_argument( - "--output_dir", - type=str, - default="./visualizations", - help="Directory to save visualizations.", - ) - parser.add_argument( - "--animation_interval", - type=int, - default=1000, - help="Animation interval in milliseconds.", - ) - args = parser.parse_args() - - os.makedirs(args.output_dir, exist_ok=True) - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - in_dim = 16 * 6 # Past: (B, 16, 6) - out_dim = 20 * 2 # Future: (B, 20, 2) - model = MonocularModel(in_dim=in_dim, out_dim=out_dim, feature_extractor=SAMFeatures()) - - # HACK: fix for loading torch.compile()d models - ckpt = torch.load(args.model_path, map_location="cpu") - state = ckpt["state_dict"] if "state_dict" in ckpt else ckpt - mapped = {} - for k, v in state.items(): - k = k.replace("model._orig_mod.", "").replace("model.", "") - if k.startswith("features.sam_pos_embed"): - k = k.replace("features.sam_pos_embed", "features.sam_model.model.pos_embed") - elif k.startswith("features.sam_pos_embed_window"): - k = k.replace("features.sam_pos_embed_window", "features.sam_model.model.pos_embed_window") - elif k.startswith("features.sam_patch_embed"): - k = k.replace("features.sam_patch_embed", "features.sam_model.model.patch_embed") - elif k.startswith("features.sam_blocks"): - k = k.replace("features.sam_blocks", "features.sam_model.model.blocks") - mapped[k] = v - - model.load_state_dict(mapped, strict=True) - model.to(device) - model.eval() - - - past_trajectories, pred_trajectories, gt_trajectories = gen_viz_data( - model, args.data_root, args.num_samples, device - ) - - metrics = calculate_metrics(pred_trajectories, gt_trajectories) - print(f"ADE @ 5s: {metrics['ADE_mean']:.3f} ± {metrics['ADE_std']:.3f} meters") - print(f"ADE percentiles - 1%: {metrics['ADE_p1']:.3f}, 50%: {metrics['ADE_p50']:.3f}, 99%: {metrics['ADE_p99']:.3f}") - print(f"FDE: {metrics['FDE_mean']:.3f} ± {metrics['FDE_std']:.3f} m") - print(f"FDE percentiles - 1%: {metrics['FDE_p1']:.3f}, 50%: {metrics['FDE_p50']:.3f}, 99%: {metrics['FDE_p99']:.3f}") - - anim = create_animated_trajectory_plot( - past_trajectories, - pred_trajectories, - gt_trajectories, - save_path=os.path.join(args.output_dir, "trajectory_animation.gif"), - interval=args.animation_interval, - ) From 4770d884940e3da312f47a1b3eae5d0c3f6ac220 Mon Sep 17 00:00:00 2001 From: Pratyush Date: Thu, 30 Apr 2026 16:02:43 -0400 Subject: [PATCH 5/6] Address PR review comments: remove debug/script files, restore deleted code - Remove debug_callbacks.py, debug_viz.py (debugging artifacts not for main) - Remove CHANGES.md, environment.yml (reviewer requested removal) - Remove run_train_proposal.sbatch and run_train_proposal_login_gpu.sh (no sbatch/scripts in main) - Restore monocular.py, viz.py, and models/transfuser to match main (no reason for diff) - Restore pred_traj_flat, query_for_score variables and dict extraction in base_model.py - Add back type annotations and descriptive ValueError messages in base_model.py Made-with: Cursor --- environment.yml | 111 ---- src/camera-based-e2e/CHANGES.md | 530 ------------------ src/camera-based-e2e/debug_viz.py | 369 ------------ src/camera-based-e2e/models/base_model.py | 19 +- .../models/debug_callbacks.py | 5 - src/camera-based-e2e/models/monocular.py | 13 +- src/camera-based-e2e/models/transfuser | 1 + .../scripts/run_train_proposal.sbatch | 84 --- .../scripts/run_train_proposal_login_gpu.sh | 47 -- src/camera-based-e2e/viz.py | 337 +++++++++++ 10 files changed, 359 insertions(+), 1157 deletions(-) delete mode 100644 environment.yml delete mode 100644 src/camera-based-e2e/CHANGES.md delete mode 100644 src/camera-based-e2e/debug_viz.py delete mode 100644 src/camera-based-e2e/models/debug_callbacks.py create mode 160000 src/camera-based-e2e/models/transfuser delete mode 100644 src/camera-based-e2e/scripts/run_train_proposal.sbatch delete mode 100755 src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh create mode 100644 src/camera-based-e2e/viz.py diff --git a/environment.yml b/environment.yml deleted file mode 100644 index 287a065..0000000 --- a/environment.yml +++ /dev/null @@ -1,111 +0,0 @@ -name: robotvision -channels: - - conda-forge -dependencies: - - _libgcc_mutex=0.1 - - _openmp_mutex=4.5 - - bzip2=1.0.8 - - ca-certificates=2025.10.5 - - icu=75.1 - - ld_impl_linux-64=2.44 - - libexpat=2.7.1 - - libffi=3.5.2 - - libgcc=15.2.0 - - libgcc-ng=15.2.0 - - libgomp=15.2.0 - - liblzma=5.8.1 - - libnsl=2.0.1 - - libsqlite=3.51.0 - - libstdcxx=15.2.0 - - libstdcxx-ng=15.2.0 - - libuuid=2.41.2 - - libxcrypt=4.4.36 - - libzlib=1.3.1 - - ncurses=6.5 - - openssl=3.5.4 - - python=3.9.23 - - readline=8.2 - - tk=8.6.13 - - tzdata=2025b - - zstd=1.5.7 - - pip: - - accelerate==1.10.1 - - aiohappyeyeballs==2.6.1 - - aiohttp==3.13.2 - - aiosignal==1.4.0 - - annotated-types==0.7.0 - - anyio==4.12.0 - - async-timeout==5.0.1 - - attrs==25.4.0 - - blinker==1.9.0 - - certifi==2025.10.5 - - charset-normalizer==3.4.4 - - docker==7.1.0 - - exceptiongroup==1.3.1 - - flask==3.1.2 - - frozenlist==1.8.0 - - gitdb==4.0.12 - - graphql-core==3.2.7 - - graphql-relay==3.2.0 - - greenlet==3.2.4 - - gunicorn==23.0.0 - - h11==0.16.0 - - hf-xet==1.2.0 - - httpcore==1.0.9 - - httpx==0.28.1 - - huggingface-hub==0.36.0 - - idna==3.11 - - itsdangerous==2.2.0 - - lightning-utilities==0.15.2 - - mako==1.3.10 - - matplotlib==3.9.4 - - multidict==6.7.0 - - numpy==2.0.2 - - nvidia-cublas-cu12==12.4.5.8 - - nvidia-cuda-cupti-cu12==12.4.127 - - nvidia-cuda-nvrtc-cu12==12.4.127 - - nvidia-cuda-runtime-cu12==12.4.127 - - nvidia-cudnn-cu12==9.1.0.70 - - nvidia-cufft-cu12==11.2.1.3 - - nvidia-curand-cu12==10.3.5.147 - - nvidia-cusparse-cu12==12.3.1.170 - - nvidia-cusparselt-cu12==0.6.2 - - nvidia-ml-py==12.570.172 - - nvidia-nccl-cu12==2.21.5 - - nvidia-nvjitlink-cu12==12.4.127 - - nvidia-nvtx-cu12==12.4.127 - - nvitop==1.4.2 - - opentelemetry-api==1.39.1 - - opentelemetry-semantic-conventions==0.60b1 - - packaging==24.2 - - propcache==0.4.1 - - protobuf==3.20.3 - - pydantic==2.12.5 - - pydantic-core==2.41.5 - - python-dateutil==2.9.0.post0 - - pytorch-lightning==2.5.6 - - pyyaml==6.0.3 - - regex==2026.1.15 - - requests==2.32.5 - - safetensors==0.7.0 - - scipy==1.13.1 - - seaborn==0.13.2 - - shellingham==1.5.4 - - six==1.17.0 - - smmap==5.0.2 - - sqlalchemy==2.0.46 - - sqlparse==0.5.5 - - sympy==1.13.1 - - timm==1.0.12 - - tokenizers==0.22.2 - - tomli==2.4.0 - - torcheval==0.0.7 - - torchmetrics==1.8.2 - - transformers==4.57.6 - - triton==3.2.0 - - typer-slim==0.20.0 - - typing-inspection==0.4.2 - - urllib3==2.5.0 - - wrapt==2.0.1 - - yarl==1.22.0 -prefix: /scratch/gilbreth/mathur91/conda/robotvision diff --git a/src/camera-based-e2e/CHANGES.md b/src/camera-based-e2e/CHANGES.md deleted file mode 100644 index b301ef2..0000000 --- a/src/camera-based-e2e/CHANGES.md +++ /dev/null @@ -1,530 +0,0 @@ -# Proposal Planner — Debug & Fix Changelog - -## Phase 1: Debug Instrumentation - -### Problem -The proposal planner had terrible performance after 10 epochs of training: -- Train ADE ~16.8m, Val ADE ~14.3m (extremely high for 5-second horizon) -- All 16 proposals collapsed to the same trajectory (zero regret, zero diversity) -- Score loss dominated total loss by ~25x - -### What was added - -**`models/debug_callbacks.py`** — `GradientDebugCallback` (Lightning Callback) -- Per-module gradient L2 norms: `scene_encoder`, `proposal_init`, `refinement` (per block), `scorer` -- Per-sublayer gradient norms within refinement (`cross_attn`, `mlp`, `traj_residual`, `traj_enc`) -- Per-sublayer gradient norms within scorer (`geom_proj`, `score_mlp`) and proposal_init (`ego_enc`, `traj_decoder`, `proposal_embed`) -- Parameter norms per module (detect weight explosion/vanishing) -- Gradient dominance ratio (max_module / total — detect when one module dominates) -- Activation statistics via forward hooks: mean, std, abs_max, dead fraction, saturated fraction -- Proposal diagnostics: endpoint spread, pairwise diversity, trajectory lengths, score distribution, refinement delta - -**`debug_viz.py`** — Post-training diagnostic plot generator -- 10 individual plots + 1 summary dashboard -- Reads `metrics.csv` from Lightning CSVLogger -- Can be run standalone: `python debug_viz.py --log_dir /path/to/logs` - -**`train.py`** changes: -- `--debug` flag enables the callback and auto-generates plots after training -- `--debug_log_every N` controls logging frequency - -### Diagnosis from debug plots - -| Signal | Observation | Implication | -|--------|-------------|-------------| -| `loss_score=427` vs `loss_ade=16.8` | Score loss 25x larger | Score loss dominates optimization | -| `ade_pred ≈ ade_oracle`, regret ≈ 0 | All proposals identical | Complete mode collapse | -| `grad_norm/proposal_init` ~1000 | 10x higher than everything else | 90% of gradient signal goes through proposal_init | -| `grad_norm/scene_encoder` ~100 | Order of magnitude lower | Scene encoder barely learns | -| `proposals/pairwise_dist` decreasing | Proposals converging | Diversity loss too weak | -| `proposals/score_range` near 0 | Scorer can't discriminate | All proposals look the same to scorer | -| `act/scene_enc_out_mean` ≈ 0 | Scene features near zero | Visual info not reaching proposals | -| `grad_norm/propinit_embed` very low | Proposal seeds barely updated | Can't maintain diversity | - -### Root cause chain -1. Score loss (MSE to ADE, weight=1.0) produces loss ~400, dwarfing ADE loss ~15 -2. Score loss gradient flows through scorer + proposal_init, bypassing refinement/scene_encoder -3. Without gradient pressure to diversify, all proposals collapse to the mean trajectory -4. Once collapsed, diversity loss = 0, refinement has nothing to differentiate, cross-attention degenerates -5. Scene encoder never gets gradient signal because its only path (refinement cross-attn) is starved - ---- - -## Phase 2: Architecture & Loss Fixes - -### Fix 1: Scene-conditioned ProposalInit (`models/proposal_init.py`) - -**Problem:** ProposalInit only saw ego trajectory + intent. Camera images had no direct path -to the ADE loss — their only route was through the refinement cross-attention bottleneck, -which received 10x less gradient signal than proposal_init. - -**Fix:** Added cross-attention from proposal embeddings to scene feature tokens inside -ProposalInit. This gives the scene encoder a direct gradient path: -`ADE loss → proposals → traj_decoder → scene_cross_attn → scene_encoder.proj` - -Also increased `proposal_embed` init std from 0.02 to 0.1 to provide more initial diversity -among proposal seeds (wider starting spread resists early collapse). - -### Fix 2: Score loss downweight + warmup (`models/base_model.py`) - -**Problem:** Score loss (MSE between predicted scores and actual ADEs) was weighted 1.0x, -producing loss ~400 vs ADE loss ~15. This overwhelmed the trajectory quality signal and -drove all proposals toward identical outputs. - -**Fix:** -- New `score_weight` parameter (default 0.1 instead of implicit 1.0) -- New `score_warmup_epochs` parameter (default 2): score loss is 0 for the first N epochs, - then ramps linearly to `score_weight` over 1 epoch. This lets proposals diversify before - the scorer starts pulling them toward similar predictions. - -### Fix 3: Detach proposals from score loss gradient (`models/base_model.py`) - -**Problem:** The scorer's MSE loss backpropagated through both the scorer weights AND the -proposal generation pipeline. Since all proposals had similar ADE, the scorer gradient -effectively pushed proposals to produce outputs that are easy to score (i.e., identical). - -**Fix:** `pred_scores.detach().argmin(dim=1)` when selecting the best proposal for ADE -computation. The scorer still learns to rank, but its gradients only update scorer weights, -not the proposal generator. The proposal generator is driven only by ADE + diversity losses. - -### Fix 4: Diversity weight increase (`train.py`, `sbatch`) - -**Problem:** Diversity weight was 0.1, but the diversity loss magnitude (~0.3) was negligible -compared to score loss (~400). Even after downweighting score loss, diversity needs to be -strong enough to prevent collapse. - -**Fix:** Default `diversity_weight` changed from 0.1 to 1.0. - -### Fix 5: Gradient clipping (`train.py`) - -**Problem:** Gradient norms reached ~1000 with high variance (spikes visible in plots), -causing unstable updates especially for the scorer. - -**Fix:** Added `gradient_clip_val=1.0` with norm clipping to the Lightning Trainer. -New `--grad_clip` CLI arg (0 to disable). - -### Summary of parameter changes - -| Parameter | Before | After | Reason | -|-----------|--------|-------|--------| -| `score_weight` | 1.0 (implicit) | 0.1 | Was 25x larger than ADE, dominated optimization | -| `score_warmup_epochs` | 0 (N/A) | 2 | Let proposals diversify before scorer activates | -| `diversity_weight` | 0.1 | 1.0 | Too weak to prevent collapse against score loss | -| `proposal_embed` std | 0.02 | 0.1 | Wider initial spread resists early collapse | -| `gradient_clip_val` | None | 1.0 | Stabilize training, prevent gradient spikes | -| ProposalInit inputs | past + intent | past + intent + scene_feat | Give scene encoder a direct gradient path | - -### Files changed - -| File | Change | -|------|--------| -| `models/proposal_init.py` | Added scene cross-attention, increased embed init std | -| `models/proposal_planner.py` | Pass `scene_feat` to `proposal_init()` | -| `models/base_model.py` | Added `score_weight`, `score_warmup_epochs`, detached scorer gradients | -| `train.py` | New CLI args: `--score_weight`, `--score_warmup_epochs`, `--grad_clip`; updated defaults | -| `scripts/run_train_proposal.sbatch` | Updated args for new defaults | -| `models/debug_callbacks.py` | Created (Phase 1) | -| `debug_viz.py` | Created (Phase 1) | -| `CHANGES.md` | This file | - ---- - -## Phase 3: iPad-inspired Architectural Changes - -Reference: [iPad — Iterative Proposal-centric End-to-End Autonomous Driving](https://arxiv.org/abs/2505.15111) - -Phase 2 fixed the optimization dynamics (score loss no longer dominates, proposals no -longer collapse), but raw ADE stayed at ~14.3m throughout training. The model -architecture itself had structural limitations that prevented learning to generate -accurate trajectories. Phase 3 addresses three key gaps relative to the iPad paper. - -### Fix 6: BCE scorer loss with quality target (`models/base_model.py`, `models/scorer.py`) - -**Problem:** The old MSE scorer trained on raw ADE values (range 0–50+). This -produced unbounded loss magnitudes, required aggressive downweighting, and gave the -scorer no natural notion of "good" vs "bad" — an ADE of 12 vs 13 had the same loss -gradient magnitude as 1 vs 2, even though the latter distinction is far more important. - -**Fix:** Replaced `MSE(logit, ADE)` with `BCE(sigmoid(logit), exp(-ADE/τ))`: -- Quality target `exp(-ADE/τ)` maps ADE into [0, 1] — lower ADE → higher quality -- BCE loss is naturally bounded (~0.0 to ~0.7), no need for aggressive downweighting -- The exponential mapping concentrates learning on distinguishing good proposals - (low ADE) rather than wasting capacity on large-ADE differences -- Temperature τ (default 5.0) controls sensitivity: lower τ → sharper discrimination -- Scorer output semantics flipped: **higher score = better** (was: lower = better) -- All `argmin` for score-based selection → `argmax` throughout the codebase -- Default `score_weight` raised from 0.1 to 1.0 (BCE is well-scaled, safe at 1.0) - -New CLI arg: `--score_temperature` (default 5.0) - -### Fix 7: Weight-shared iterative refinement (`models/refinement.py`) - -**Problem:** The old refinement used `num_steps` separate `RefinementBlock` instances, -each with independent weights. This meant: -- Each block only saw gradients from a single position in the chain -- Parameter count scaled linearly with `num_steps` (more steps = more parameters) -- No iterative self-correction: each block learned a fixed transformation, not a - general "look at the scene and improve proposals" operation - -**Fix:** iPad-style weight sharing — a **single** `RefinementBlock` applied -`num_steps` times in a loop: -- Same weights at every iteration → the block learns a general refinement operation -- Increasing `num_steps` (e.g., 2 → 4) adds compute but zero extra parameters -- Gradients flow through all iterations (like an RNN), giving the shared weights - stronger learning signal from multiple applications - -**Practical benefit:** `--num_refinement_steps 4` now uses the same parameter count -as the old 2-step model, but applies 4 iterations of refinement. - -### Fix 8: Full trajectory re-prediction (`models/refinement.py`) - -**Problem:** The old `RefinementBlock` predicted **residual** waypoint deltas: -`proposals_new = proposals + traj_residual(feat)`. This constrained each iteration -to small perturbations around the previous trajectory, preventing the model from -fundamentally correcting bad initial proposals. - -**Fix:** Replaced residual prediction with **full re-prediction**: -`proposals_new = traj_decoder(feat)`. At each iteration the model predicts complete -new trajectories from scratch based on the current proposal features. The features -still accumulate information across iterations (via residual connections in the -feature update path), but the trajectories are free to change substantially. - -This matches iPad's approach where proposals are predicted fresh at each iteration -of the ProFormer loop. - -### Summary of Phase 3 parameter changes - -| Parameter | Phase 2 | Phase 3 | Reason | -|-----------|---------|---------|--------| -| Score loss | MSE | BCE with `exp(-ADE/τ)` target | Bounded loss, [0,1] quality scale | -| `score_weight` | 0.1 | 1.0 | BCE is naturally well-scaled | -| `score_temperature` | N/A | 5.0 (new) | Controls quality target sensitivity | -| Score semantics | lower = better | higher = better | Matches BCE probability interpretation | -| Refinement weights | Independent per step | Shared single block | Stronger gradient signal, free to add iterations | -| `num_refinement_steps` | 2 | 4 (sbatch) | More iterations at zero parameter cost | -| Trajectory prediction | Residual (`proposals + delta`) | Full re-prediction | Can correct bad proposals, not just perturb | - -### Phase 3 files changed - -| File | Change | -|------|--------| -| `models/refinement.py` | Single shared `RefinementBlock`, `traj_residual` → `traj_decoder`, full re-prediction | -| `models/scorer.py` | Updated docstring: higher = better, BCE-trained | -| `models/base_model.py` | BCE loss with quality target, `argmin` → `argmax`, `score_temperature` param | -| `models/debug_callbacks.py` | Updated for shared block (`.block` not `.blocks`), flipped score argmin/argmax | -| `train.py` | New `--score_temperature` arg, `score_weight` default 0.1 → 1.0 | -| `scripts/run_train_proposal.sbatch` | `num_refinement_steps` 2 → 4, `score_weight` 0.1 → 1.0, added `score_temperature` | -| `CHANGES.md` | This section | - ---- - -## Phase 4: Full iPad-faithful Rewrite - -Reference: [iPad source code](https://github.com/Kguo-cs/iPad) and [paper appendix](https://arxiv.org/html/2505.15111v1) - -Phase 3 changes did not move the needle: ADE stayed at 14.3m with complete mode -collapse (all proposals identical, zero regret, zero diversity, score loss stuck at -ln(2)=0.693). After studying the iPad source code in detail, several fundamental -architectural and loss differences were identified. Phase 4 is a faithful rewrite -copying the core iPad design patterns. - -### Root cause: Why Phases 2-3 failed - -The model was still stuck because: -1. **Top-5 WTA L2 loss** averaged the best 5 proposals' L2 errors — this pulled ALL - proposals toward the mean trajectory, actively fighting diversity. -2. **One feature per proposal** (B, K, C) was too limited to represent a 20-timestep - trajectory. Each proposal had a single 256-dim vector to encode 40 coordinates. -3. **No intermediate supervision** — only the final proposals were supervised. The - shared refinement block had no gradient signal for early iterations. -4. **Diversity loss as a band-aid** — negative mean pairwise distance is a weak signal. - iPad doesn't use diversity loss at all; MoN naturally handles it. - -### Fix 9: Per-timestep proposal features (`models/proposal_init.py`) - -**Problem:** Each proposal had one feature vector (B, K, C). A single 256-dim vector -must encode all 20 timesteps' (x,y) positions, severely limiting expressiveness. - -**Fix:** Learnable embeddings of shape (N×T, C) = (16×20, C) = (320, C). Each -proposal gets T=20 separate feature vectors, one per timestep. This matches iPad's -`init_feature = nn.Embedding(poses_num * proposal_num, tf_d_model)`. - -The ego status (past trajectory + intent) is encoded to a single vector and broadcast- -added to all N×T tokens, exactly matching iPad's pattern: -`bev_feature = ego_feature + init_feature.weight[None]` - -Removed: scene cross-attention from ProposalInit (iPad doesn't have it — scene -features are incorporated only through refinement cross-attention). - -### Fix 10: Per-timestep trajectory decoding (`models/refinement.py`) - -**Problem:** The old decoder mapped a single (B, K, C) feature to a full (B, K, T×2) -trajectory via one MLP. This forced a single vector to predict 40 output dimensions. - -**Fix:** Each timestep's feature independently predicts that timestep's (x,y): -`proposals = traj_decoder(bev_feature) # (B, N*T, C) → (B, N*T, 2) → (B,N,T,2)` - -The decoder is a 2-layer MLP with only 2 output dimensions per timestep (matching -iPad's `MLP(d_model, d_ffn, state_size)` where state_size=2 for us). - -### Fix 11: Predict-then-refine loop with intermediate proposals (`models/refinement.py`) - -**Problem:** Old refinement encoded proposals back and then predicted residuals. -No intermediate proposals were returned for supervision. - -**Fix:** iPad-faithful predict→encode→attend→refine loop: -1. `proposals = traj_decoder(bev_feature)` — predict from current features -2. `bev_feature += traj_enc(proposals)` — encode predictions back (proposal-anchoring) -3. `bev_feature += cross_attn(bev_feature, scene_feat)` — attend to scene -4. `bev_feature += mlp(bev_feature)` — FFN update -5. Return both proposals and updated features -6. Collect `proposal_list` across all K iterations for intermediate supervision - -### Fix 12: MoN L1 loss with discounted intermediate supervision (`models/base_model.py`) - -**Problem:** Top-5 WTA L2 ADE loss averaged the best 5 modes, pulling all proposals -toward the mean and causing mode collapse. Only final proposals were supervised. - -**Fix:** Faithful copy of iPad's loss: -- **MoN L1**: `min_n mean_t (|Δx| + |Δy|)` — Minimum over N proposals of mean L1 - displacement. Only the SINGLE closest proposal to GT is optimized, leaving others - free to explore different modes. L1 is more robust to outliers than L2. -- **Discounted intermediate supervision**: `L = Σ_k λ^(K-1-k) × MoN_L1(P_k)` - with λ=0.1 (iPad's `prev_weight`). Later iterations get stronger supervision, - earlier ones are relaxed. This gives gradient signal at every refinement step. -- **Removed diversity loss** (set to weight 0). MoN naturally encourages diversity - because only one proposal is pulled toward GT per sample — the others receive no - gradient and can remain spread out. iPad uses `inter_weight=0`. - -### Fix 13: Simplified scorer (`models/scorer.py`) - -**Problem:** Old scorer fused learned proposal features with handcrafted geometric -features (velocity norms, acceleration norms, flat waypoints). This was complex and -the geometric features were redundant given per-timestep features. - -**Fix:** Match iPad's scorer exactly: -1. Reshape per-timestep features: (B, N×T, C) → (B, N, T, C) -2. Max-pool over temporal dimension: (B, N, T, C) → (B, N, C) -3. MLP → scalar logit: (B, N, C) → (B, N) - -This is simpler and matches iPad's `proposal_feature = bev_feature.amax(-2)` followed -by `pred_score(proposal_feature)`. - -### Fix 14: Score quality target uses L1 (`models/base_model.py`) - -The scorer's BCE quality target now uses L1 displacement (consistent with the -trajectory loss) instead of L2 ADE: `quality = exp(-L1_per_mode / τ)`. - -### Summary of Phase 4 parameter changes - -| Parameter | Phase 3 | Phase 4 | Reason | -|-----------|---------|---------|--------| -| Trajectory loss | Top-5 WTA L2 ADE | MoN L1 | Only best proposal supervised; L1 more robust | -| Intermediate supervision | None | Discounted λ=0.1 | Gradient signal at every refinement step | -| `prev_weight` | N/A | 0.1 (new) | iPad discount factor for earlier iterations | -| Proposal features | (B, K, C) | (B, K×T, C) | Per-timestep features, 20× more expressive | -| Trajectory decoder | MLP(C → T×2) per proposal | MLP(C → 2) per timestep | Each timestep independently decoded | -| Scorer | Geometric + learned features | Max-pool temporal + MLP | Simpler, matches iPad | -| `diversity_weight` | 1.0 | 0.0 | MoN handles diversity; explicit loss fights it | -| `smoothness_weight` | 0.01 | 0.0 | Simplify loss to match iPad | -| `comfort_weight` | 0.01 | 0.0 | Simplify loss to match iPad | -| `max_epochs` | 10 | 20 | iPad trains for 20 epochs | - -### Phase 4 files changed - -| File | Change | -|------|--------| -| `models/proposal_init.py` | Rewritten: per-timestep embeddings (N×T, C), removed scene cross-attn | -| `models/refinement.py` | Rewritten: predict-then-refine loop, returns proposal_list | -| `models/scorer.py` | Rewritten: max-pool temporal + MLP, removed geometric features | -| `models/proposal_planner.py` | Updated: new calling pattern, passes proposal_list through | -| `models/base_model.py` | MoN L1 loss, discounted intermediate supervision, removed diversity loss | -| `models/debug_callbacks.py` | Updated for new architecture (scorer, propinit, proposal_list) | -| `train.py` | New `--prev_weight` arg, `diversity_weight` default 0, `num_refinement_steps` default 4 | -| `scripts/run_train_proposal.sbatch` | Updated all params, 20 epochs | -| `CHANGES.md` | This section | - ---- - -## Phase 5: Scorer Improvements (Ranking-focused) - -Phase 4 achieved strong multimodal proposals (val `ade_oracle` ~0.57 m) but deployed -trajectory quality was bottlenecked by scorer ranking: val `ade_pred` ~1.66 m with -`ade_regret` ~1.09 m (66% of total error). The scorer couldn't reliably pick the -best proposal from the diverse set MoN produced. - -### Fix 15: Cross-entropy ranking loss (`models/base_model.py`) - -**Problem:** Per-proposal BCE with soft quality targets `exp(-L1/τ)` does not directly -optimize for correct ranking. The targets for good vs mediocre proposals are close -together (e.g., 0.90 vs 0.82 with τ=5), giving weak ranking signal. A scorer can -achieve low BCE loss while still ranking proposals incorrectly. - -**Fix:** Replaced BCE with cross-entropy classification: -`loss_score = CrossEntropy(logits, argmin(L1_per_mode))`. This directly optimizes -for `argmax(scores) == argmin(L1)`, exactly matching how proposals are selected at -inference. The `score_temperature` parameter is no longer used by the loss. - -### Fix 16: Detached scorer inputs (`models/scorer.py`) - -**Problem:** Score loss gradients flowed back through the proposal generation pipeline -via `bev_feature`, conflicting with MoN trajectory loss. MoN wants only one proposal -to move toward GT; score BCE pushed all proposal features toward their individual -quality targets through shared representations. - -**Fix:** `bev_feature.detach()` before pooling in the scorer forward pass. Score loss -now only updates scorer weights, making it a pure discriminator. MoN has exclusive -control over proposal generation — no gradient conflict. - -### Fix 17: Trajectory-aware scorer (`models/scorer.py`) - -**Problem:** The scorer only saw abstract learned features (`bev_feature`). After -cross-attention and MLP updates, features for different proposals can be very similar -even when decoded trajectories differ substantially. The scorer was ranking based on -features that didn't fully capture trajectory-level differences. - -**Fix:** Concatenate projected trajectory coordinates with pooled BEV features: -- `traj_proj`: MLP mapping (T*2) to 64-dim trajectory embedding per proposal -- `score_mlp` input: `[max_pooled_feat | traj_feat]` with LayerNorm before the MLP -- Both `bev_feature` and `proposals` are detached so score loss is fully isolated - -This gives the scorer explicit spatial information (velocity profiles, endpoint -locations, curvature) for more discriminative ranking. - -### Summary of Phase 5 changes - -| Parameter | Phase 4 | Phase 5 | Reason | -|-----------|---------|---------|--------| -| Score loss | BCE with `exp(-L1/tau)` target | Cross-entropy on `argmin(L1)` | Directly optimizes ranking for `argmax` selection | -| Scorer input gradient | Flows back to proposals | Detached | Eliminates gradient conflict with MoN | -| Scorer input features | Max-pooled BEV only | BEV + projected trajectory coords | Explicit spatial discrimination | -| Scorer MLP input | `d_model` | `d_model + traj_dim` with LayerNorm | Richer, normalized input | -| `score_temperature` | Used (tau=5.0) | Unused (kept as CLI arg) | CE loss doesn't need soft targets | - -### Phase 5 files changed - -| File | Change | -|------|--------| -| `models/scorer.py` | Rewritten: detached inputs, traj_proj, LayerNorm, updated forward signature | -| `models/proposal_planner.py` | Pass `proposals` to `scorer(bev_feature, proposals)` | -| `models/base_model.py` | BCE to cross-entropy with `argmin(L1)` label | -| `models/debug_callbacks.py` | Added `scorer_traj_proj` gradient norm logging | -| `CHANGES.md` | This section | - ---- - -## Phase 6: Multi-loss scorer experiments - -To enable parallel scorer-loss ablations from one codebase, scorer loss is now -configurable via CLI: - -- `bce`: iPad-faithful BCE with soft quality target `exp(-L1/tau)` -- `ce`: hard top-1 cross-entropy on `argmin(L1)` -- `bce_pairwise`: BCE + pairwise margin-ranking auxiliary -- `listnet`: listwise KL objective with `softmax(-L1/tau)` target distribution - -New args in `train.py`: -- `--score_loss_type` -- `--score_rank_weight` -- `--score_margin` -- `--score_topk` - -Additional scorer diagnostics are logged to metrics: -- `*_score_top1_acc`: `argmax(score) == argmin(L1)` accuracy -- `*_score_gap_best_second`: average logit gap between oracle-best and second-best mode - ---- - -## Phase 7: NAVSIM-style quality target - -### Motivation (iPad paper Section 3.3) - -The iPad scorer uses a multi-factor ground-truth score from log-replay simulation -(NAVSIM Eq. 5): `S = NC * DAC * (5*EP + 5*TTC + 2*Comf) / 12`, covering safety, -efficiency, and comfort — not just trajectory closeness. Our `exp(-L1/tau)` target -only captures geometric accuracy, which may explain why the scorer struggles to -discriminate between proposals that are close in L1 but differ in driving quality. - -### Fix 18: Geometric NAVSIM approximation (`models/base_model.py`) - -Without agent trajectories or road boundaries in the current data pipeline, we -approximate the NAVSIM sub-metrics from trajectory geometry alone: - -| Sub-metric | iPad weight | Our approximation | -|---|---|---| -| EP (Ego Progress) | 5/12 | `clamp(dot(prop_disp, gt_dir) / gt_dist, 0, 1)` | -| TTC (Time-to-Collision) | 5/12 | 1.0 (no agent data) | -| Comf (Comfort) | 2/12 | fraction of timesteps with jerk < threshold | -| NC (No at-fault Collision) | gate | 1.0 (no agent data) | -| DAC (Drivable Area Compliance) | gate | 1.0 (no map data loaded) | - -Combined: `quality = (5*EP + 5*1.0 + 2*Comf) / 12`, clamped to [0, 1]. - -New `_compute_quality_target()` method dispatches between `l1` (existing) and -`navsim` (new) based on `--score_target_type` CLI arg. - -### New CLI args - -- `--score_target_type`: `l1` (default, backward-compatible) or `navsim` -- `--comfort_jerk_threshold`: jerk threshold in m/s^3 for comfort metric (default 5.0) - -### Files changed - -| File | Change | -|------|--------| -| `models/base_model.py` | Added `_compute_navsim_score()`, `_compute_quality_target()`, new hparams | -| `train.py` | New `--score_target_type`, `--comfort_jerk_threshold` args | -| `scripts/run_train_proposal.sbatch` | Added new default args | -| `score_experiment_viz.py` | Tracks `score_target_type` in approach labels and summary | -| `CHANGES.md` | This section | - ---- - -## Phase 8: RFS-based scorer quality target - -### Motivation - -For datasets where full NAVSIM-style supervision is unavailable, we want the scorer -to learn proposal ranking from an RFS-style quality signal (longitudinal/lateral -deviation with speed-aware thresholds) instead of only `exp(-L1/tau)`. - -### Fix 19: Add per-proposal RFS quality target (`models/base_model.py`) - -Added a new scorer target path that computes `(B, K)` quality directly from RFS -logic for **all proposals** (not only top-1): - -- New method: `_compute_rfs_quality(proposals, reference, past)` - - Uses the same directional decomposition as existing `rfs_loss`: - - longitudinal and lateral error components - - time thresholds at 3s / 5s (indices 11, 19) - - speed scaling from the existing RFS formula - - Converts deviation to score with the same piecewise definition: - - `1.0` when within threshold - - `0.1 ** (deviation - 1)` beyond threshold - - Averages over selected timesteps to get per-proposal quality in `[0,1]` - - Optional comfort multiplier from jerk-based comfort term -- Extended `_compute_quality_target(...)`: - - Supports `score_target_type == "rfs"` - - Requires `past` for speed scaling -- Updated scorer call site in `_shared_step(...)` to pass `past` into - `_compute_quality_target(...)`. - -### New/updated hyperparameters and CLI - -- `train.py` - - `--score_target_type`: now supports `l1 | navsim | rfs` - - `--no_rfs_target_comfort`: disables jerk comfort multiplier for RFS targets -- `LitModel` hparams - - `rfs_target_use_comfort: bool = True` - -### Training script default - -- `scripts/run_train_proposal.sbatch` - - Default scorer target changed from `l1` to `rfs` - - Added inline note that this can be overridden via `EXTRA_ARGS` - -### Notes - -- Existing standalone trajectory-side `loss_rfs` path remains intact. -- If using `--score_target_type rfs`, keep `rfs_weight=0` unless you explicitly - want both scorer-target RFS and trajectory regularization RFS active together. diff --git a/src/camera-based-e2e/debug_viz.py b/src/camera-based-e2e/debug_viz.py deleted file mode 100644 index b3f9204..0000000 --- a/src/camera-based-e2e/debug_viz.py +++ /dev/null @@ -1,369 +0,0 @@ -""" -Generate debug diagnostic plots from training logs. - -Reads metrics.csv produced by PyTorch Lightning's CSVLogger and renders: - 1. Gradient norms per module group (scene_encoder, proposal_init, refinement, scorer) - 2. Gradient norms for refinement internals (cross_attn, mlp, traj_residual per block) - 3. Gradient dominance ratio over time - 4. Activation statistics (mean, std, dead fraction) per module - 5. Proposal diagnostics (spread, diversity, score distribution) - 6. Loss component breakdown (stacked area) - -Usage: - python debug_viz.py --log_dir /path/to/logs/camera_e2e_YYYYMMDD_HHMM/version_0 - python debug_viz.py --log_dir /path/to/logs # auto-picks newest run -""" -import argparse -from pathlib import Path - -import pandas as pd -import numpy as np -import matplotlib.pyplot as plt -import matplotlib.gridspec as gridspec - - -def find_metrics_csv(log_dir: Path) -> Path: - if (log_dir / "metrics.csv").exists(): - return log_dir / "metrics.csv" - candidates = sorted(log_dir.rglob("metrics.csv")) - if not candidates: - raise FileNotFoundError(f"No metrics.csv found under {log_dir}") - return candidates[-1] - - -def load_and_merge(csv_path: Path) -> pd.DataFrame: - """ - Lightning CSVLogger writes one row per log call, with NaN in columns - not logged in that call. Forward-fill within each step so we can - correlate gradient norms with losses at the same step. - """ - df = pd.read_csv(csv_path) - if "step" not in df.columns: - raise KeyError("metrics.csv must contain a 'step' column") - return df - - -def _get_cols(df: pd.DataFrame, prefix: str) -> list: - return sorted([c for c in df.columns if c.startswith(prefix)]) - - -def _plot_series(ax, df, columns, title, ylabel, legend_strip="", logy=False, ewm_span=20): - for col in columns: - series = df.set_index("step")[col].dropna() - if series.empty: - continue - label = col.replace(legend_strip, "").strip("/") - smoothed = series.ewm(span=ewm_span, min_periods=1).mean() - ax.plot(smoothed.index, smoothed.values, label=label, linewidth=1.2) - ax.fill_between(series.index, series.values, alpha=0.08) - ax.set_title(title, fontsize=11, fontweight="bold") - ax.set_ylabel(ylabel) - ax.set_xlabel("Step") - if logy: - ax.set_yscale("log") - ax.legend(fontsize=7, ncol=2) - ax.grid(True, alpha=0.3) - - -# ====================================================================== -# Individual plot functions -# ====================================================================== - -def plot_gradient_norms_by_module(df: pd.DataFrame, out_dir: Path): - """Bar-style time series of gradient norms for the four main modules.""" - cols = [c for c in _get_cols(df, "grad_norm/") if c in ( - "grad_norm/scene_encoder", "grad_norm/proposal_init", - "grad_norm/scorer", "grad_norm/total", - ) or c.startswith("grad_norm/refine_block_")] - if not cols: - return - - fig, ax = plt.subplots(figsize=(14, 5)) - _plot_series(ax, df, cols, "Gradient Norms by Module", "L2 Norm", "grad_norm/", logy=True) - fig.tight_layout() - fig.savefig(out_dir / "grad_norms_modules.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'grad_norms_modules.png'}") - - -def plot_gradient_norms_refinement(df: pd.DataFrame, out_dir: Path): - """Detailed gradient norms inside refinement blocks.""" - cols = [c for c in _get_cols(df, "grad_norm/refine_") if "block" not in c] - if not cols: - return - - fig, ax = plt.subplots(figsize=(14, 5)) - _plot_series(ax, df, cols, "Refinement Internals — Gradient Norms", "L2 Norm", "grad_norm/", logy=True) - fig.tight_layout() - fig.savefig(out_dir / "grad_norms_refinement_detail.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'grad_norms_refinement_detail.png'}") - - -def plot_gradient_norms_scorer_propinit(df: pd.DataFrame, out_dir: Path): - """Scorer and ProposalInit sub-module gradient norms.""" - cols = _get_cols(df, "grad_norm/scorer_") + _get_cols(df, "grad_norm/propinit_") - if not cols: - return - - fig, ax = plt.subplots(figsize=(14, 5)) - _plot_series(ax, df, cols, "Scorer & ProposalInit Internals — Gradient Norms", "L2 Norm", "grad_norm/", logy=True) - fig.tight_layout() - fig.savefig(out_dir / "grad_norms_scorer_propinit.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'grad_norms_scorer_propinit.png'}") - - -def plot_gradient_dominance(df: pd.DataFrame, out_dir: Path): - """Which module dominates the gradient signal?""" - col = "grad_norm/dominant_ratio" - if col not in df.columns: - return - - fig, ax = plt.subplots(figsize=(14, 4)) - series = df.set_index("step")[col].dropna() - smoothed = series.ewm(span=20, min_periods=1).mean() - ax.plot(smoothed.index, smoothed.values, color="crimson", linewidth=1.5) - ax.fill_between(series.index, series.values, alpha=0.15, color="crimson") - ax.axhline(0.5, color="gray", linestyle="--", alpha=0.5, label="50% dominance") - ax.set_title("Gradient Dominance Ratio (max module / total)", fontsize=11, fontweight="bold") - ax.set_ylabel("Ratio") - ax.set_xlabel("Step") - ax.set_ylim(0, 1.05) - ax.legend() - ax.grid(True, alpha=0.3) - fig.tight_layout() - fig.savefig(out_dir / "grad_dominance.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'grad_dominance.png'}") - - -def plot_activation_stats(df: pd.DataFrame, out_dir: Path): - """Activation mean/std/dead fraction for each hooked module output.""" - mean_cols = [c for c in _get_cols(df, "act/") if c.endswith("_mean")] - std_cols = [c for c in _get_cols(df, "act/") if c.endswith("_std")] - dead_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_dead")] - sat_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_saturated")] - - if not mean_cols: - return - - fig, axes = plt.subplots(2, 2, figsize=(16, 10)) - - _plot_series(axes[0, 0], df, mean_cols, "Activation Mean", "Mean", "act/") - _plot_series(axes[0, 1], df, std_cols, "Activation Std", "Std", "act/") - _plot_series(axes[1, 0], df, dead_cols, "Dead Neuron Fraction (|x| < 1e-6)", "Fraction", "act/") - _plot_series(axes[1, 1], df, sat_cols, "Saturated Fraction (|x| > 10)", "Fraction", "act/") - - fig.suptitle("Activation Statistics by Module Output", fontsize=13, fontweight="bold") - fig.tight_layout(rect=[0, 0, 1, 0.96]) - fig.savefig(out_dir / "activation_stats.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'activation_stats.png'}") - - -def plot_proposal_diagnostics(df: pd.DataFrame, out_dir: Path): - """Proposal spread, diversity, trajectory lengths, and score stats.""" - prop_cols = _get_cols(df, "proposals/") - if not prop_cols: - return - - spread_cols = [c for c in prop_cols if "std" in c or "dist" in c] - length_cols = [c for c in prop_cols if "length" in c] - score_cols = [c for c in prop_cols if "score" in c] - - n_panels = sum(1 for g in [spread_cols, length_cols, score_cols] if g) - if n_panels == 0: - return - - fig, axes = plt.subplots(1, n_panels, figsize=(6 * n_panels, 5)) - if n_panels == 1: - axes = [axes] - - idx = 0 - if spread_cols: - _plot_series(axes[idx], df, spread_cols, "Proposal Spread & Diversity", "Value", "proposals/") - idx += 1 - if length_cols: - _plot_series(axes[idx], df, length_cols, "Trajectory Lengths", "Length (m)", "proposals/") - idx += 1 - if score_cols: - _plot_series(axes[idx], df, score_cols, "Score Distribution", "Score", "proposals/") - idx += 1 - - fig.suptitle("Proposal Diagnostics", fontsize=13, fontweight="bold") - fig.tight_layout(rect=[0, 0, 1, 0.95]) - fig.savefig(out_dir / "proposal_diagnostics.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'proposal_diagnostics.png'}") - - -def plot_loss_breakdown(df: pd.DataFrame, out_dir: Path): - """Stacked area chart of all loss components.""" - loss_cols = [c for c in df.columns if c.startswith("train_loss_") and c != "train_loss"] - if not loss_cols: - return - - fig, axes = plt.subplots(1, 2, figsize=(16, 5)) - - # Individual loss curves - _plot_series(axes[0], df, loss_cols, "Training Loss Components", "Loss", "train_loss_", logy=True) - - # Stacked area (absolute contribution) - sub = df[["step"] + loss_cols].dropna(subset=loss_cols, how="all").copy() - if sub.empty: - plt.close(fig) - return - sub = sub.set_index("step").interpolate().fillna(0).clip(lower=0) - sub_smooth = sub.ewm(span=30, min_periods=1).mean() - axes[1].stackplot(sub_smooth.index, *[sub_smooth[c].values for c in loss_cols], - labels=[c.replace("train_loss_", "") for c in loss_cols], alpha=0.7) - axes[1].set_title("Loss Component Breakdown (stacked)", fontsize=11, fontweight="bold") - axes[1].set_ylabel("Loss") - axes[1].set_xlabel("Step") - axes[1].legend(fontsize=7, loc="upper right") - axes[1].grid(True, alpha=0.3) - - fig.tight_layout() - fig.savefig(out_dir / "loss_breakdown.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'loss_breakdown.png'}") - - -def plot_ade_regret(df: pd.DataFrame, out_dir: Path): - """Oracle ADE vs predicted-best ADE, and regret over time.""" - cols = ["train_ade_pred", "train_ade_oracle", "train_ade_regret"] - present = [c for c in cols if c in df.columns] - if not present: - return - - fig, ax = plt.subplots(figsize=(14, 5)) - _plot_series(ax, df, present, "ADE: Oracle vs Predicted-Best & Regret", "ADE (m)", "train_") - fig.tight_layout() - fig.savefig(out_dir / "ade_regret.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'ade_regret.png'}") - - -def plot_param_norms(df: pd.DataFrame, out_dir: Path): - """Parameter norms to detect weight explosion / vanishing.""" - cols = _get_cols(df, "param_norm/") - if not cols: - return - - fig, ax = plt.subplots(figsize=(14, 5)) - _plot_series(ax, df, cols, "Parameter Norms by Module", "L2 Norm", "param_norm/", logy=True) - fig.tight_layout() - fig.savefig(out_dir / "param_norms.png", dpi=180) - plt.close(fig) - print(f" -> {out_dir / 'param_norms.png'}") - - -# ====================================================================== -# Snapshot: single-image summary for quick glance -# ====================================================================== - -def plot_summary_dashboard(df: pd.DataFrame, out_dir: Path): - """Single 2x3 dashboard with the most important signals.""" - fig = plt.figure(figsize=(20, 12)) - gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.3) - - # 1. Module gradient norms - ax1 = fig.add_subplot(gs[0, 0]) - cols = [c for c in _get_cols(df, "grad_norm/") if c in ( - "grad_norm/scene_encoder", "grad_norm/proposal_init", - "grad_norm/scorer", "grad_norm/total", - ) or c.startswith("grad_norm/refine_block_")] - if cols: - _plot_series(ax1, df, cols, "Gradient Norms (modules)", "L2 Norm", "grad_norm/", logy=True, ewm_span=30) - else: - ax1.text(0.5, 0.5, "No grad_norm data", ha="center", va="center", transform=ax1.transAxes) - - # 2. Loss breakdown - ax2 = fig.add_subplot(gs[0, 1]) - loss_cols = [c for c in df.columns if c.startswith("train_loss_") and c != "train_loss"] - if loss_cols: - _plot_series(ax2, df, loss_cols, "Loss Components", "Loss", "train_loss_", logy=True, ewm_span=30) - else: - ax2.text(0.5, 0.5, "No loss data", ha="center", va="center", transform=ax2.transAxes) - - # 3. ADE regret - ax3 = fig.add_subplot(gs[0, 2]) - ade_cols = [c for c in ("train_ade_pred", "train_ade_oracle", "train_ade_regret") if c in df.columns] - if ade_cols: - _plot_series(ax3, df, ade_cols, "ADE: Oracle vs Predicted", "ADE (m)", "train_", ewm_span=30) - else: - ax3.text(0.5, 0.5, "No ADE data", ha="center", va="center", transform=ax3.transAxes) - - # 4. Activation dead fraction - ax4 = fig.add_subplot(gs[1, 0]) - dead_cols = [c for c in _get_cols(df, "act/") if c.endswith("_frac_dead")] - if dead_cols: - _plot_series(ax4, df, dead_cols, "Dead Neuron Fraction", "Fraction", "act/", ewm_span=30) - else: - ax4.text(0.5, 0.5, "No activation data", ha="center", va="center", transform=ax4.transAxes) - - # 5. Proposal diversity - ax5 = fig.add_subplot(gs[1, 1]) - prop_cols = [c for c in _get_cols(df, "proposals/") if "dist" in c or "std" in c] - if prop_cols: - _plot_series(ax5, df, prop_cols, "Proposal Diversity", "Value", "proposals/", ewm_span=30) - else: - ax5.text(0.5, 0.5, "No proposal data", ha="center", va="center", transform=ax5.transAxes) - - # 6. Gradient dominance - ax6 = fig.add_subplot(gs[1, 2]) - dom_col = "grad_norm/dominant_ratio" - if dom_col in df.columns: - series = df.set_index("step")[dom_col].dropna() - if not series.empty: - smoothed = series.ewm(span=30, min_periods=1).mean() - ax6.plot(smoothed.index, smoothed.values, color="crimson", linewidth=1.5) - ax6.fill_between(series.index, series.values, alpha=0.12, color="crimson") - ax6.axhline(0.5, color="gray", linestyle="--", alpha=0.5) - ax6.set_title("Gradient Dominance", fontsize=11, fontweight="bold") - ax6.set_ylabel("Ratio (max module / total)") - ax6.set_xlabel("Step") - ax6.set_ylim(0, 1.05) - ax6.grid(True, alpha=0.3) - - fig.suptitle("Training Debug Dashboard", fontsize=15, fontweight="bold", y=0.98) - fig.savefig(out_dir / "debug_dashboard.png", dpi=200) - plt.close(fig) - print(f" -> {out_dir / 'debug_dashboard.png'}") - - -def main(): - parser = argparse.ArgumentParser(description="Generate debug diagnostic plots from training logs") - parser.add_argument("--log_dir", type=str, required=True, - help="Path to log directory (version_0/) or parent (auto-picks newest)") - parser.add_argument("--out_dir", type=str, default=None, - help="Output directory for plots (default: /debug_plots/)") - args = parser.parse_args() - - log_dir = Path(args.log_dir) - csv_path = find_metrics_csv(log_dir) - print(f"Reading {csv_path}") - df = load_and_merge(csv_path) - print(f" {len(df)} rows, {len(df.columns)} columns") - - out_dir = Path(args.out_dir) if args.out_dir else csv_path.parent / "debug_plots" - out_dir.mkdir(parents=True, exist_ok=True) - print(f"Output: {out_dir}\n") - - plot_gradient_norms_by_module(df, out_dir) - plot_gradient_norms_refinement(df, out_dir) - plot_gradient_norms_scorer_propinit(df, out_dir) - plot_gradient_dominance(df, out_dir) - plot_activation_stats(df, out_dir) - plot_proposal_diagnostics(df, out_dir) - plot_loss_breakdown(df, out_dir) - plot_ade_regret(df, out_dir) - plot_param_norms(df, out_dir) - plot_summary_dashboard(df, out_dir) - - print(f"\nDone. {len(list(out_dir.glob('*.png')))} plots saved to {out_dir}") - - -if __name__ == "__main__": - main() diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py index 90be2b4..338e29c 100644 --- a/src/camera-based-e2e/models/base_model.py +++ b/src/camera-based-e2e/models/base_model.py @@ -367,13 +367,18 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: raw_output = self.forward(model_inputs) pred_depth = None - pred_scores = None + pred_scores: torch.Tensor = None + pred_traj_flat: torch.Tensor = None + query_for_score: torch.Tensor = None proposal_list = None if isinstance(raw_output, dict): - proposal_list = raw_output.get("proposal_list", None) - pred_scores = raw_output.get("scores", None) - pred_depth = raw_output.get("depth", None) - pred_future = raw_output["trajectory"] + outputs = raw_output + proposal_list = outputs.get("proposal_list", None) + pred_scores = outputs.get("scores", None) + pred_depth = outputs.get("depth", None) + pred_traj_flat = outputs.get("trajectory_flat", None) + query_for_score = outputs.get("query_for_score", None) + pred_future = outputs["trajectory"] else: pred_future = raw_output @@ -383,14 +388,14 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: k_modes = self.model.n_proposals if hasattr(self.model, "n_proposals") else 1 if pred.ndim != 2: - raise ValueError(f"Unexpected pred shape {pred.shape}") + raise ValueError(f"Unexpected pred shape {pred.shape}; expected (B, T*2) or (B, {k_modes}*T*2).") if pred.shape[1] == t2: pred = pred.view(pred.size(0), 1, t_steps, 2) elif pred.shape[1] == k_modes * t2: pred = pred.view(pred.size(0), k_modes, t_steps, 2) else: - raise ValueError(f"Unexpected pred shape {pred.shape}") + raise ValueError(f"Unexpected pred shape {pred.shape}; expected (B, T*2) or (B, {k_modes}*T*2).") if not torch.isfinite(pred).all(): pred = torch.nan_to_num(pred, nan=0.0, posinf=1e3, neginf=-1e3) diff --git a/src/camera-based-e2e/models/debug_callbacks.py b/src/camera-based-e2e/models/debug_callbacks.py deleted file mode 100644 index a9921a0..0000000 --- a/src/camera-based-e2e/models/debug_callbacks.py +++ /dev/null @@ -1,5 +0,0 @@ -import pytorch_lightning as pl - -class GradientDebugCallback(pl.Callback): - def __init__(self, log_every=10): - self.log_every = log_every diff --git a/src/camera-based-e2e/models/monocular.py b/src/camera-based-e2e/models/monocular.py index 289e21a..6c2035f 100644 --- a/src/camera-based-e2e/models/monocular.py +++ b/src/camera-based-e2e/models/monocular.py @@ -163,16 +163,21 @@ def forward(self, x): # Doesn't need no_grad b/c DINO/SAMFeatures will freeze if needed feats_vit = self.features(front_cam) # list or tensor - if isinstance(feats_vit, (list, tuple)): - if len(feats_vit) != 1: - raise ValueError(f"Expected single feature map, got {len(feats_vit)}") + if len(feats_vit) == 1 and isinstance(feats_vit, list): feats_vit = feats_vit[0] + feats = self.visual_adapter(feats_vit) # (B, C, H, W) # Depth Supervision output_depth = F.softplus(self.depth_gen(feats).squeeze(1)) # (B, 128, 128) - tokens = feats.flatten(2).permute(0, 2, 1) + self.positional_encoding # (B, N, C_total) + # tokens: handle list of features or single tensor + # TODO: is this made redundant by if statement above? + if isinstance(feats, (list, tuple)): + tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) + else: + tokens = feats.flatten(2) # (B, C, N) + tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) # copy procedure to build query_0 from MonocularModel intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() diff --git a/src/camera-based-e2e/models/transfuser b/src/camera-based-e2e/models/transfuser new file mode 160000 index 0000000..9d413b2 --- /dev/null +++ b/src/camera-based-e2e/models/transfuser @@ -0,0 +1 @@ +Subproject commit 9d413b2ad2d2d56c112b34a4a799be081800d77f diff --git a/src/camera-based-e2e/scripts/run_train_proposal.sbatch b/src/camera-based-e2e/scripts/run_train_proposal.sbatch deleted file mode 100644 index 2a7b429..0000000 --- a/src/camera-based-e2e/scripts/run_train_proposal.sbatch +++ /dev/null @@ -1,84 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=ipad -#SBATCH --output=logs/slurm-%j.out -#SBATCH --error=logs/slurm-%j.err -#SBATCH --time=56:00:00 -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --gpus-per-node=1 -#SBATCH --cpus-per-task=20 -#SBATCH --mem=64G -#SBATCH --account=csso -#SBATCH --partition=a10 - -module load conda -# Sourcing conda.sh makes activate work in non-interactive (sbatch) shells -CONDA_BASE=$(conda info --base 2>/dev/null) -if [ -n "$CONDA_BASE" ] && [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then - source "$CONDA_BASE/etc/profile.d/conda.sh" -else - eval "$(conda shell.bash hook)" -fi -conda activate robotvision -pip install pytorch-lightning -# Use the python that activation puts on PATH. -# Some clusters don't mount the original conda env prefix on compute nodes. -PYTHON="$(command -v python || true)" -echo "CONDA_PREFIX=${CONDA_PREFIX}" -echo "PYTHON=${PYTHON}" -if [ -z "$PYTHON" ] || [ ! -x "$PYTHON" ]; then - echo "ERROR: python not found after conda activate." - echo "If this is a shared-filesystem issue, recreate the env in a path mounted on compute nodes (e.g. under \$HOME)." - conda info --envs || true - exit 2 -fi - -# Ensure pytorch_lightning is available (required by train.py). -if ! "$PYTHON" -c "import pytorch_lightning as pl; print('pytorch_lightning', pl.__version__)" 2>/dev/null; then - echo "ERROR: pytorch_lightning not found in env: ${CONDA_PREFIX:-unknown}" - echo "Install it in the same env you use for this job, e.g.:" - echo " conda activate robotvision" - echo " pip install pytorch-lightning" - exit 3 -fi -# Waymo E2E data; index_train.pkl and index_val.pkl must be in PROJECT_DIR -DATA_DIR="/scratch/gilbreth/mathur91/waymo/waymo_open_dataset_end_to_end_camera_v_1_0_0" -PROJECT_DIR="/home/mathur91/robotvision/src/camera-based-e2e" - -mkdir -p "$PROJECT_DIR/logs" -cd "$PROJECT_DIR" - -# Scorer targets: rfs = per-proposal RFS quality at 3s/5s (× jerk comfort). Override with EXTRA_ARGS, e.g. --score_target_type l1 -"$PYTHON" train.py \ - --data_dir "$DATA_DIR" \ - --model_type proposal \ - --backbone resnet \ - --no_wandb \ - --batch_size 16 \ - --lr 1e-4 \ - --max_epochs 20 \ - --num_proposals 16 \ - --num_refinement_steps 4 \ - --smoothness_weight 0.0 \ - --comfort_weight 0.0 \ - --rfs_weight 0.2 \ - --diversity_weight 0.0 \ - --score_weight 1.0 \ - --score_warmup_epochs 2 \ - --score_temperature 5.0 \ - --score_loss_type bce \ - --score_target_type rfs \ - --score_rank_weight 0.2 \ - --score_margin 0.2 \ - --score_topk 0 \ - --comfort_jerk_threshold 5.0 \ - --prev_weight 0.1 \ - --grad_clip 1.0 \ - --log_every_n_steps 100 \ - ${EXTRA_ARGS:-} - -# Auto-refresh cross-run scorer experiment ranking plots/report. -LOGS_ROOT="$(dirname "$DATA_DIR")/logs" -if [ -d "$LOGS_ROOT" ]; then - "$PYTHON" score_experiment_viz.py --logs_root "$LOGS_ROOT" --latest_n 24 || true -fi diff --git a/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh b/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh deleted file mode 100755 index 7d7019d..0000000 --- a/src/camera-based-e2e/scripts/run_train_proposal_login_gpu.sh +++ /dev/null @@ -1,47 +0,0 @@ -#!/bin/bash -# Run proposal training on the current machine (e.g. login) with visible GPU. -# Same args as run_train_proposal.sbatch, without Slurm. Does not pip-install PL. - -set -euo pipefail - -module load conda 2>/dev/null || true -CONDA_BASE=$(conda info --base 2>/dev/null) -if [ -n "$CONDA_BASE" ] && [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then - source "$CONDA_BASE/etc/profile.d/conda.sh" -else - eval "$(conda shell.bash hook)" -fi -conda activate robotvision - -DATA_DIR="/scratch/gilbreth/mathur91/waymo/waymo_open_dataset_end_to_end_camera_v_1_0_0" -PROJECT_DIR="/home/mathur91/robotvision/src/camera-based-e2e" - -mkdir -p "$PROJECT_DIR/logs" -cd "$PROJECT_DIR" - -exec python train.py \ - --data_dir "$DATA_DIR" \ - --model_type proposal \ - --backbone resnet \ - --no_wandb \ - --batch_size 16 \ - --lr 1e-4 \ - --max_epochs 20 \ - --num_proposals 16 \ - --num_refinement_steps 4 \ - --smoothness_weight 0.0 \ - --comfort_weight 0.0 \ - --diversity_weight 0.0 \ - --score_weight 1.0 \ - --score_warmup_epochs 2 \ - --score_temperature 5.0 \ - --score_loss_type bce \ - --score_target_type l1 \ - --score_rank_weight 0.2 \ - --score_margin 0.2 \ - --score_topk 0 \ - --comfort_jerk_threshold 5.0 \ - --prev_weight 0.1 \ - --grad_clip 1.0 \ - --log_every_n_steps 100 \ - ${EXTRA_ARGS:-} diff --git a/src/camera-based-e2e/viz.py b/src/camera-based-e2e/viz.py new file mode 100644 index 0000000..78b8c2f --- /dev/null +++ b/src/camera-based-e2e/viz.py @@ -0,0 +1,337 @@ +import argparse, os, pathlib +import torch +import matplotlib.pyplot as plt +import matplotlib.animation as animation +import numpy as np +from typing import List, Tuple, Optional +from tqdm import tqdm + +from loader import WaymoE2E +from models.monocular import MonocularModel, SAMFeatures +from torch import nn + +def gen_viz_data(model: nn.Module, data_root: str, num_samples: int, device: torch.device): + """Generate the gt and pred trajectories along with past states""" + dataset = WaymoE2E( + batch_size=1, + indexFile="index_val.pkl", + data_dir=data_root, + images=True, + n_items=num_samples, + ) + loader = torch.utils.data.DataLoader( + dataset, + batch_size=1, + num_workers=4, + pin_memory=True, + persistent_workers=False, + ) + model.eval() + + pred_trajectories = [] + gt_trajectories = [] + past_trajectories = [] + + with torch.no_grad(): + for batch in tqdm(loader, desc="Generating trajectory data"): + past, future, images, intent = batch['PAST'], batch['FUTURE'], batch['IMAGES'], batch['INTENT'] + past = past.to(device, non_blocking=True) + future = future.to(device, non_blocking=True) + intent = intent.to(device, non_blocking=True) + images = [img.to(device, non_blocking=True) for img in images] + B, T, F = past.shape + + # Extract past positions (x, y) from the first 2 features + past_positions = past[:, :, :2].squeeze(0).cpu() # (T, 2) + past_trajectories.append(past_positions) + + pred_future = model({"PAST": past, "IMAGES": images, "INTENT": intent}) + pred_future = pred_future.view(B, -1, 2) + pred_trajectories.append(pred_future.squeeze(0).cpu()) + gt_trajectories.append(future.squeeze(0).cpu()) + + return past_trajectories, pred_trajectories, gt_trajectories + + +def create_animated_trajectory_plot( + past_trajectories: List[torch.Tensor], + pred_trajectories: List[torch.Tensor], + gt_trajectories: List[torch.Tensor], + save_path: Optional[str] = None, + interval: int = 200, +): + """ + Args: + past_trajectories: List of past position tensors [(T_past, 2), ...] + pred_trajectories: List of predicted future position tensors [(T_future, 2), ...] + gt_trajectories: List of ground truth future position tensors [(T_future, 2), ...] + save_path: path to save + interval: 1/framerate (ms) + """ + fig, ax = plt.subplots(figsize=(12, 8)) + + # Calculate global bounds once from all trajectories to prevent jumping + all_x_coords = [] + all_y_coords = [] + for past, pred, gt in zip(past_trajectories, pred_trajectories, gt_trajectories): + all_x_coords.extend([past[:, 0].numpy(), pred[:, 0].numpy(), gt[:, 0].numpy()]) + all_y_coords.extend([past[:, 1].numpy(), pred[:, 1].numpy(), gt[:, 1].numpy()]) + + global_x = np.concatenate(all_x_coords) + global_y = np.concatenate(all_y_coords) + margin = 5.0 + x_min, x_max = global_x.min() - margin, global_x.max() + margin + y_min, y_max = global_y.min() - margin, global_y.max() + margin + + past_color = "blue" + pred_color = "red" + gt_color = "green" + + (past_line,) = ax.plot( + [], + [], + "o-", + color=past_color, + linewidth=2, + markersize=4, + label="Past Trajectory", + alpha=0.8, + ) + (pred_line,) = ax.plot( + [], + [], + "o-", + color=pred_color, + linewidth=2, + markersize=4, + label="Predicted Future", + alpha=0.8, + ) + (gt_line,) = ax.plot( + [], + [], + "o-", + color=gt_color, + linewidth=2, + markersize=4, + label="Ground Truth Future", + alpha=0.8, + ) + + (current_pos,) = ax.plot([], [], "ko", markersize=8, label="Current Position") + + text = ax.text( + 0.02, + 0.98, + "", + transform=ax.transAxes, + fontsize=12, + verticalalignment="top", + bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8), + ) + + def init(): + ax.clear() + ax.set_xlabel("X Position (m)") + ax.set_ylabel("Y Position (m)") + ax.set_title("Vehicle Trajectory Animation: Past → Present → Future") + ax.grid(True, alpha=0.3) + ax.legend() + return past_line, pred_line, gt_line, current_pos, text + + def animate(frame): + # Update progress bar for frame generation + if not hasattr(animate, 'pbar'): + animate.pbar = tqdm(total=len(past_trajectories), desc="Generating animation frames", unit="frame", position=1, leave=False) + animate.pbar.update(1) + + ax.clear() + ax.set_xlabel("X Position (m)") + ax.set_ylabel("Y Position (m)") + ax.set_title("Vehicle Trajectory Animation: Past → Present → Future") + ax.grid(True, alpha=0.3) + + sample_idx = frame % len(past_trajectories) + + past = past_trajectories[sample_idx].numpy() + pred = pred_trajectories[sample_idx].numpy() + gt = gt_trajectories[sample_idx].numpy() + + ax.plot( + past[:, 0], + past[:, 1], + "o-", + color=past_color, + linewidth=2, + markersize=4, + label="Past Trajectory", + alpha=0.8, + ) + ax.plot( + pred[:, 0], + pred[:, 1], + "o-", + color=pred_color, + linewidth=2, + markersize=4, + label="Predicted Future", + alpha=0.8, + ) + ax.plot( + gt[:, 0], + gt[:, 1], + "o-", + color=gt_color, + linewidth=2, + markersize=4, + label="Ground Truth Future", + alpha=0.8, + ) + + current_x, current_y = past[-1, 0], past[-1, 1] + ax.plot(current_x, current_y, "ko", markersize=8, label="Current Position") + + # Use static bounds + ax.set_xlim(x_min, x_max) + ax.set_ylim(y_min, y_max) + + ax.text( + 0.02, + 0.98, + f"Sample {sample_idx + 1}/{len(past_trajectories)}", + transform=ax.transAxes, + fontsize=12, + verticalalignment="top", + bbox=dict(boxstyle="round", facecolor="lightblue", alpha=0.8), + ) + + ax.legend() + ax.set_aspect("equal", adjustable="box") + + return [] + + # Create animation + anim = animation.FuncAnimation( + fig, + animate, + init_func=init, + frames=len(past_trajectories), + interval=interval, + blit=False, + repeat=True, + ) + + if save_path: + anim.save(save_path, writer="pillow", fps=3) + + return anim + + +def calculate_metrics( + pred_trajectories: List[torch.Tensor], gt_trajectories: List[torch.Tensor] +) -> dict: + """Calculate trajectory prediction metrics""" + ade_scores = [] # Average Displacement Error + fde_scores = [] # Final Displacement Error + + for pred, gt in zip(pred_trajectories, gt_trajectories): + # ADE: Average Euclidean distance over all time steps + distances = torch.norm(pred - gt, dim=1) + ade = torch.mean(distances).item() + ade_scores.append(ade) + + # FDE: Euclidean distance at final time step + fde = torch.norm(pred[-1] - gt[-1]).item() + fde_scores.append(fde) + + return { + "ADE_mean": np.mean(ade_scores), + "ADE_std": np.std(ade_scores), + "ADE_p1": np.percentile(ade_scores, 1), + "ADE_p50": np.percentile(ade_scores, 50), + "ADE_p99": np.percentile(ade_scores, 99), + "FDE_mean": np.mean(fde_scores), + "FDE_std": np.std(fde_scores), + "FDE_p1": np.percentile(fde_scores, 1), + "FDE_p50": np.percentile(fde_scores, 50), + "FDE_p99": np.percentile(fde_scores, 99), + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--data_root", + type=str, + required=True, + help="Path to the root directory of the Waymo E2E dataset.", + ) + parser.add_argument( + "--model_path", + type=str, + required=True, + help="Path to the trained model checkpoint.", + ) + parser.add_argument( + "--num_samples", type=int, default=100, help="Number of samples to animate." + ) + parser.add_argument( + "--output_dir", + type=str, + default="./visualizations", + help="Directory to save visualizations.", + ) + parser.add_argument( + "--animation_interval", + type=int, + default=1000, + help="Animation interval in milliseconds.", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + in_dim = 16 * 6 # Past: (B, 16, 6) + out_dim = 20 * 2 # Future: (B, 20, 2) + model = MonocularModel(in_dim=in_dim, out_dim=out_dim, feature_extractor=SAMFeatures()) + + # HACK: fix for loading torch.compile()d models + ckpt = torch.load(args.model_path, map_location="cpu") + state = ckpt["state_dict"] if "state_dict" in ckpt else ckpt + mapped = {} + for k, v in state.items(): + k = k.replace("model._orig_mod.", "").replace("model.", "") + if k.startswith("features.sam_pos_embed"): + k = k.replace("features.sam_pos_embed", "features.sam_model.model.pos_embed") + elif k.startswith("features.sam_pos_embed_window"): + k = k.replace("features.sam_pos_embed_window", "features.sam_model.model.pos_embed_window") + elif k.startswith("features.sam_patch_embed"): + k = k.replace("features.sam_patch_embed", "features.sam_model.model.patch_embed") + elif k.startswith("features.sam_blocks"): + k = k.replace("features.sam_blocks", "features.sam_model.model.blocks") + mapped[k] = v + + model.load_state_dict(mapped, strict=True) + model.to(device) + model.eval() + + + past_trajectories, pred_trajectories, gt_trajectories = gen_viz_data( + model, args.data_root, args.num_samples, device + ) + + metrics = calculate_metrics(pred_trajectories, gt_trajectories) + print(f"ADE @ 5s: {metrics['ADE_mean']:.3f} ± {metrics['ADE_std']:.3f} meters") + print(f"ADE percentiles - 1%: {metrics['ADE_p1']:.3f}, 50%: {metrics['ADE_p50']:.3f}, 99%: {metrics['ADE_p99']:.3f}") + print(f"FDE: {metrics['FDE_mean']:.3f} ± {metrics['FDE_std']:.3f} m") + print(f"FDE percentiles - 1%: {metrics['FDE_p1']:.3f}, 50%: {metrics['FDE_p50']:.3f}, 99%: {metrics['FDE_p99']:.3f}") + + anim = create_animated_trajectory_plot( + past_trajectories, + pred_trajectories, + gt_trajectories, + save_path=os.path.join(args.output_dir, "trajectory_animation.gif"), + interval=args.animation_interval, + ) From 0da4658e72b8a492d18e8298f5fa8051d4e8701c Mon Sep 17 00:00:00 2001 From: Pratyush Date: Thu, 30 Apr 2026 17:01:48 -0400 Subject: [PATCH 6/6] Add IPadConfig dataclass for iPad-style hyperparams; restore IMAGES priority - Add IPadConfig dataclass in proposal_planner.py bundling all iPad-specific loss/scorer hyperparameters (matches reviewer request to keep config in same file as the iPad model class). - LitModel now accepts a single ipad_config argument instead of 17 individual kwargs. Defaults are preserved via the dataclass defaults; hparams are populated from the config so all internal getattr lookups still work. - train.py constructs an IPadConfig only for the proposal model. - Restore original IMAGES-first batch key priority in _shared_step so existing IMAGES batches keep their previous code path (per reviewer feedback). - Drop --debug / --debug_log_every flags and references to deleted debug_callbacks / debug_viz modules from train.py. Made-with: Cursor --- src/camera-based-e2e/models/base_model.py | 45 +++------- .../models/proposal_planner.py | 32 ++++++++ src/camera-based-e2e/train.py | 82 ++++++------------- 3 files changed, 67 insertions(+), 92 deletions(-) diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py index 338e29c..a6f570a 100644 --- a/src/camera-based-e2e/models/base_model.py +++ b/src/camera-based-e2e/models/base_model.py @@ -1,3 +1,4 @@ +from dataclasses import asdict from typing import List, Optional import torch import torch.nn as nn @@ -6,6 +7,7 @@ import torchvision from .losses.depth_loss import DepthLoss +from .proposal_planner import IPadConfig class BaseModel(nn.Module): def __init__(self, in_dim, out_dim): @@ -27,43 +29,16 @@ def __init__( model: nn.Module, lr: float, lr_vision: Optional[float] = None, - 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, - score_weight: float = 1.0, - score_warmup_epochs: int = 2, - score_temperature: float = 5.0, - score_loss_type: str = "bce", - score_target_type: str = "l1", - score_rank_weight: float = 0.0, - score_margin: float = 0.2, - score_topk: int = 0, - comfort_jerk_threshold: float = 5.0, - prev_weight: float = 0.1, - rfs_target_use_comfort: bool = True, + ipad_config: Optional[IPadConfig] = None, ): super(LitModel, self).__init__() self.model = model self.hparams.lr = lr self.hparams.lr_vision = lr_vision - self.hparams.rfs_weight = rfs_weight - self.hparams.smoothness_weight = smoothness_weight - self.hparams.collision_weight = collision_weight - self.hparams.comfort_weight = comfort_weight - self.hparams.diversity_weight = diversity_weight - self.hparams.score_weight = score_weight - self.hparams.score_warmup_epochs = score_warmup_epochs - self.hparams.score_temperature = score_temperature - self.hparams.score_loss_type = score_loss_type - self.hparams.score_target_type = score_target_type - self.hparams.score_rank_weight = score_rank_weight - self.hparams.score_margin = score_margin - self.hparams.score_topk = score_topk - self.hparams.comfort_jerk_threshold = comfort_jerk_threshold - self.hparams.prev_weight = prev_weight - self.hparams.rfs_target_use_comfort = rfs_target_use_comfort + + cfg = ipad_config if ipad_config is not None else IPadConfig() + for field, value in asdict(cfg).items(): + setattr(self.hparams, field, value) self.example_input_array = ({ 'PAST': torch.zeros((1, 16, 6)), # PAST @@ -355,11 +330,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: past, future, intent = batch['PAST'], batch['FUTURE'], batch['INTENT'] - if "IMAGES_JPEG" in batch: + if "IMAGES" in batch: + images = batch["IMAGES"] + elif "IMAGES_JPEG" in batch: images_jpeg = batch["IMAGES_JPEG"] images = self.decode_batch_jpeg(images_jpeg) - elif "IMAGES" in batch: - images = batch["IMAGES"] else: raise KeyError("Batch must contain either 'IMAGES_JPEG' or 'IMAGES' key.") diff --git a/src/camera-based-e2e/models/proposal_planner.py b/src/camera-based-e2e/models/proposal_planner.py index 2194dec..56e88d4 100644 --- a/src/camera-based-e2e/models/proposal_planner.py +++ b/src/camera-based-e2e/models/proposal_planner.py @@ -5,6 +5,7 @@ scene encoder → proposal init → iterative refinement → scorer (with intermediate proposal supervision) """ +from dataclasses import dataclass from typing import Dict, List import torch @@ -16,6 +17,37 @@ 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__( diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py index 7754d94..8c94263 100644 --- a/src/camera-based-e2e/train.py +++ b/src/camera-based-e2e/train.py @@ -18,9 +18,8 @@ # Replace with your model defined in models/ from models.base_model import LitModel, collate_with_images from models.monocular import DeepMonocularModel -from models.proposal_planner import ProposalPlanner +from models.proposal_planner import ProposalPlanner, IPadConfig from models.feature_extractors import SAMFeatures, DINOFeatures, ResNetFeatures -from models.debug_callbacks import GradientDebugCallback if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -71,8 +70,6 @@ parser.add_argument('--no_wandb', action='store_true', help='Disable wandb logging (use CSV only)') parser.add_argument('--compile', action='store_true', help='Whether to compile the model with torch.compile') parser.add_argument('--profile', action='store_true', help='Whether to run the profiler') - parser.add_argument('--debug', action='store_true', help='Enable debug visualizations (gradient norms, activation stats, proposal diagnostics)') - parser.add_argument('--debug_log_every', type=int, default=10, help='How often (steps) to log debug metrics') args = parser.parse_args() # Data @@ -106,30 +103,36 @@ else SAMFeatures(frozen=True) ) model = DeepMonocularModel(feature_extractor=backbone, out_dim=out_dim, n_blocks=8) - if args.debug: - model._debug = True if args.compile: model = torch.compile(model, mode="max-autotune") + + if args.model_type == 'proposal': + ipad_config = IPadConfig( + rfs_weight=args.rfs_weight, + smoothness_weight=args.smoothness_weight, + collision_weight=args.collision_weight, + comfort_weight=args.comfort_weight, + diversity_weight=args.diversity_weight, + score_weight=args.score_weight, + score_warmup_epochs=args.score_warmup_epochs, + score_temperature=args.score_temperature, + score_loss_type=args.score_loss_type, + score_target_type=args.score_target_type, + score_rank_weight=args.score_rank_weight, + score_margin=args.score_margin, + score_topk=args.score_topk, + comfort_jerk_threshold=args.comfort_jerk_threshold, + prev_weight=args.prev_weight, + rfs_target_use_comfort=not args.no_rfs_target_comfort, + ) + else: + ipad_config = None + lit_model = LitModel( model=model, lr=args.lr, - smoothness_weight=args.smoothness_weight if args.model_type == 'proposal' else 0.0, - collision_weight=args.collision_weight if args.model_type == 'proposal' else 0.0, - comfort_weight=args.comfort_weight if args.model_type == 'proposal' else 0.0, - rfs_weight=args.rfs_weight if args.model_type == 'proposal' else 0.0, - diversity_weight=args.diversity_weight if args.model_type == 'proposal' else 0.0, - score_weight=args.score_weight if args.model_type == 'proposal' else 0.0, - score_warmup_epochs=args.score_warmup_epochs if args.model_type == 'proposal' else 0, - score_temperature=args.score_temperature if args.model_type == 'proposal' else 5.0, - score_loss_type=args.score_loss_type if args.model_type == 'proposal' else 'bce', - score_target_type=args.score_target_type if args.model_type == 'proposal' else 'l1', - score_rank_weight=args.score_rank_weight if args.model_type == 'proposal' else 0.0, - score_margin=args.score_margin if args.model_type == 'proposal' else 0.2, - score_topk=args.score_topk if args.model_type == 'proposal' else 0, - comfort_jerk_threshold=args.comfort_jerk_threshold if args.model_type == 'proposal' else 5.0, - prev_weight=args.prev_weight if args.model_type == 'proposal' else 0.1, - rfs_target_use_comfort=not args.no_rfs_target_comfort if args.model_type == 'proposal' else True, + ipad_config=ipad_config, ) # We don't want to save logs or checkpoints in the home directory - it'll fill up fast @@ -150,8 +153,6 @@ filename='camera-e2e-{epoch:02d}-{val_loss:.2f}' ), ] - if args.debug: - callbacks.append(GradientDebugCallback(log_every=args.debug_log_every)) trainer = pl.Trainer( max_epochs=args.max_epochs, @@ -188,37 +189,4 @@ except Exception as e: print(f"Could not save loss plot: {e}") - if args.debug: - try: - from debug_viz import find_metrics_csv, load_and_merge - from debug_viz import ( - plot_gradient_norms_by_module, plot_gradient_norms_refinement, - plot_gradient_norms_scorer_propinit, plot_gradient_dominance, - plot_activation_stats, plot_proposal_diagnostics, - plot_loss_breakdown, plot_ade_regret, plot_param_norms, - plot_summary_dashboard, - ) - csv_path = find_metrics_csv(run_dir / "version_0") - df = load_and_merge(csv_path) - debug_out = base_path / "debug_plots" - debug_out.mkdir(parents=True, exist_ok=True) - print(f"\nGenerating debug plots to {debug_out} ...") - plot_gradient_norms_by_module(df, debug_out) - plot_gradient_norms_refinement(df, debug_out) - plot_gradient_norms_scorer_propinit(df, debug_out) - plot_gradient_dominance(df, debug_out) - plot_activation_stats(df, debug_out) - plot_proposal_diagnostics(df, debug_out) - plot_loss_breakdown(df, debug_out) - plot_ade_regret(df, debug_out) - plot_param_norms(df, debug_out) - plot_summary_dashboard(df, debug_out) - print(f"Debug plots saved to {debug_out}") - except Exception as e: - print(f"Could not generate debug plots: {e}") - - - - -