Post-Training the UniAD Planner with LoRA and GRPO

1. One question drives the whole experiment

UniAD unifies perception, prediction, occupancy, and planning in one BEV-centric model. After stage-2 end-to-end training, the planning head outputs six future ego waypoints — but the objective is primarily supervised L2 / ADE. That leaves little room to optimize comfort, collision avoidance, or multi-modal uncertainty without expensive full-model fine-tuning.

Following the spirit of Thinking Machines' LoRA work, we freeze the entire UniAD backbone and post-train only low-rank adapters plus a small Gaussian policy head on the planner. GRPO (group-relative policy optimization) lets us sample multiple waypoint trajectories per frame and reinforce the best ones relative to the group — without re-running the frozen perception stack.

The throughline is simple: preserve the scene representation, introduce the smallest trainable policy correction, branch it into candidate futures, learn from their relative quality, and finally judge the one deterministic path that deployment will use.

The experiment revealed a distinction that is easy to lose in a training dashboard. GRPO can move probability mass toward a good future without moving the center of the distribution by the same amount. A purple rollout can nearly touch the expert path while the blue trajectory actually shipped at inference remains several meters away. The curriculum after phase 3 is an attempt to close precisely that gap.

Thread continuesTo change only the planner, the codebase needs an explicit boundary between frozen scene understanding and trainable future selection.

2. Preserve the invariant: fork only the planner path

We kept the default UniAD training path intact. GRPO post-training is an opt-in fork activated by config (PlanningHeadGRPO, planner_posttrain_only=True). No changes to PlanningHeadSingleMode or stage-2 configs unless you explicitly switch. The integration follows the repository’s MMDetection3D registry and config conventions ^[mmdet3d].

File-by-file inventory

Integration checklist

1. Step 1 — Add LoRALinear. Wrap selected nn.Linear modules in the planning head. Freeze base weights; train only A and B matrices.
2. Step 2 — Subclass planning head. PlanningHeadGRPO extends PlanningHeadSingleMode. Reuse mlp_fuser → TransformerDecoder → reg_branch path.
3. Step 3 — Gaussian waypoint policy. reg_branch output (pre-cumsum) = mean μ of 6×2 increments. Learn log_std[T,2]. Trajectory = cumsum(Δ).
4. Step 4 — GRPO training forward. Sample G trajectories without re-running frozen backbone. Compute group-relative advantages and clipped surrogate.
5. Step 5 — Reference policy KL. Disable LoRA adapters; compute ref log-prob on same samples. Anchor to stage-2 SFT planner.
6. Step 6 — planner_posttrain_only on UniAD. Freeze all parameters; re-enable LoRA + log_std. Run track/motion under torch.no_grad(); optimize planning only.
7. Step 7 — Configs + hooks. Register head/loss/hooks in mmcv. Add phased configs and JSONL experiment logging.
8. Step 8 — Validation. pytest unit tests + mini campaign + waypoint export for qualitative BEV plots.

Design constraints we preserved

• Frozen backbone: track, map, motion, occupancy run under torch.no_grad().
• Reference policy = SFT: LoRA disabled → identical to uniad_base_e2e.pth planner.
• BEV adapter: stays in the forward path, frozen, no LoRA (matches reference).
• Inference default: deterministic mean waypoints (cumsum of μ), not stochastic samples.
• Trainable state: LoRA matrices plus log_std only; peak VRAM ~7.3 GiB (G=8, FP32, 4060 Ti).

Thread continuesThe boundary is safe. Now we need to understand the exact shape of the small correction allowed to cross it.

3. Adapt the map, not the world: LoRA on the planner

Post-training a full UniAD stack is expensive and risks catastrophic forgetting of perception. Following the PEFT intuition in ^[lora] and the “low-regret LoRA” regime emphasized by ^{[lora-noregret]}, we freeze stage-2 weights and train only low-rank adapters on the planning head — plus a tiny log_std for the Gaussian policy.

The useful mental shape is a narrow corridor. $A$ first compresses a planner feature from $d_{\mathrm{in}}$ coordinates into only $r$ coordinates; $B$ expands that compact correction back to $d_{\mathrm{out}}$. The base transformation $W$ remains the wide, frozen road. The adapter can steer the output, but only through an $r$-dimensional subspace.

For a frozen linear layer $W$:

$$W' = W + \frac{\alpha}{r} B A,\quad B\in\mathbb{R}^{d_{\mathrm{out}}\times r},\; A\in\mathbb{R}^{r\times d_{\mathrm{in}}}$$

Where we attach LoRA

• mlp_fuser.0 — fuses trajectory / track / command
• Decoder ×3: linear1/2, self_attn.out_proj, multihead_attn.out_proj
• reg_branch — predicts waypoint-increment means $\mu$

bev_adapter stays in the forward path but is frozen and LoRA-free, so the reference policy (LoRA off) matches deployed SFT ^[uniad].

Thread continuesLoRA gives the planner a narrow steering surface. GRPO turns that surface into a distribution over complete futures and asks which branch is better.

4. Explore futures, then turn preference into a gradient

Group Relative Policy Optimization ^[grpo][trl] lets us improve a planner with reinforcement learning without a value network. We treat each nuScenes frame ^[nuscenes] as a contextual bandit: one state, $G$ sampled trajectories, group-normalized advantages, PPO-style clipping ^[ppo], and a LoRA-off SFT reference for KL.

4.1 Why GRPO (not L2-only, not full PPO)?

• Supervised ADE fits the mean expert path — no comparison among futures.
• PPO needs $V(s)$; GRPO uses the group mean reward as baseline ^[grpo].
• Rewards can blend ADE / FDE / collision / comfort in one scalar.

4.2 State, action, policy

$$\Delta^{(g)}=\mu+\sigma\odot\varepsilon^{(g)},\quad \varepsilon^{(g)}\sim\mathcal{N}(0,I),\quad \tau^{(g)}=\mathrm{cumsum}(\Delta^{(g)})$$

4.3 Rollout pipeline (the interesting part)

Yes: run UniAD once → define $s$ → encode plan_query → sample $G$ noisy trajectories → reward → group advantages → losses. The backbone is not re-executed per sample.

Live examples from our checkpoint

4.4 Reward function

For sample $g$, GT path $\tau^{\star}$, mask $m_t$:

$$ \mathrm{ADE}_g=\frac{1}{\sum_t m_t}\sum_t m_t \left\|\tau^{(g)}_t-\tau^{\star}_t\right\|_2 $$

$$ \mathrm{FDE}_g= \left\|\tau^{(g)}_{t_{\mathrm{last}}} -\tau^{\star}_{t_{\mathrm{last}}}\right\|_2 $$

$$ \mathrm{Comfort}_g= \mathrm{mean}_{t\ge 3} \left\|\Delta_t-2\Delta_{t-1}+\Delta_{t-2}\right\|_2^{2} $$

$$ R_g=-w_{\mathrm{ade}}\mathrm{ADE}_g -w_{\mathrm{fde}}\mathrm{FDE}_g -w_{\mathrm{col}}\mathrm{Collision}_g -w_{\mathrm{com}}\mathrm{Comfort}_g $$

minADE in logs = $\mathbb{E}_b[\min_g\mathrm{ADE}_{b,g}]$ (best sample). Deployed ADE uses $\mu$ — a different statistic.

4.5 Group-relative advantages

Per batch row, over the $G$ samples ^[grpo]:

$$ \hat A_{b,g}= \frac{R_{b,g}-\mathrm{mean}_{g'}R_{b,g'}} {\mathrm{std}_{g'}R_{b,g'}+\varepsilon} $$

4.6 Clipped surrogate (PPO-style)

Importance ratio $r=\exp(\log\pi_\theta-\log\pi_{\mathrm{old}})$ with $\log\pi_{\mathrm{old}}$ detached ^[ppo]:

$$ \mathcal{L}_{\mathrm{GRPO}}=-\mathbb{E}\big[ \min\big(r\hat A,\; \mathrm{clip}(r,1-\varepsilon,1+\varepsilon)\hat A\big)\big] $$

4.7 Reference KL to SFT

Disable LoRA → reference mean equals frozen stage-2 planner:

$$ \mathcal{L}_{\mathrm{KL}}=\beta\, \mathbb{E}[\log\pi_{\mathrm{ref}}-\log\pi_\theta] $$

Caveat: this MC estimate can go negative when $\pi_\theta$ is denser than $\pi_{\mathrm{ref}}$ on samples — consider a KL floor for trainval.

4.8 Total loss

$$ \mathcal{L}= \mathcal{L}_{\mathrm{GRPO}} +\mathcal{L}_{\mathrm{KL}} +\lambda\,\mathrm{PlanningLoss}(\mathrm{cumsum}(\mu),\tau^{\star}) $$

Gradients: LoRA $\{A,B\}$ + log_std only. Curriculum dials $(\lambda,\beta,G,\varepsilon)$ are documented in §6b.

Thread continuesThe equations now form one circuit. The next section follows that circuit in repository order before the experiment changes any of its dials.

5. Follow one batch through the actual implementation

UniAD (frozen) ──no_grad──► track / motion ──► bev_embed, sdc_traj_query, sdc_track_query, command
                                                      │
PlanningHeadGRPO ◄────────────────────────────────────┘
  mlp_fuser [LoRA] → TransformerDecoder [LoRA] → reg_branch [LoRA] → μ, log_std
  train: sample G increments → PlannerReward → GRPO + KL + aux ADE
  ref:   same forward, LoRA off → KL anchor
  test:  sdc_traj = cumsum(μ)   # deterministic waypoints

When planner_posttrain_only=True, every parameter is frozen except LoRA matrices and log_std. All non-planning loss weights are zeroed in config.

Thread continuesOne batch now has a complete learning circuit. The curriculum decides which part of that circuit dominates at each stage.

6. A curriculum organized around one metric gap

Multi-phase curriculum on nuScenes mini, then GT-fit iterations (phase4–6) targeting mean-path alignment.

Best mean-path checkpoint: projects/work_dirs/planner_posttrain/campaign/phase4_fit_gt/latest.pth

6b. Why the question changes phase by phase

The campaign is a deliberate curriculum, not a single long GRPO run. Early phases buy correctness and a usable LoRA basin; middle phases unlock group-relative RL at scale; late phases close the gap between sample-best training metrics and the deterministic mean path used at inference.

6c. Read the curriculum as signal allocation

Thinking Machines–style exploration of the knobs we actually swept across phases. Hover a point to read exact values. These are not a full factorial rank×G grid (that is future work on trainval), but they show how the curriculum moves $G$, rank, $\lambda$, $\beta$, and $\varepsilon$ as we go from smoke → full-mini sample search → mean-path GT fit.

How to read this. Capacity ($G$, rank) peaks at phase3 for sample-best hunting; GT-fit phases raise $\lambda$ and shrink $\varepsilon$/exploration. The right-hand ADE strip shows the remaining hard outliers (often idx 240 / 280) that dominate mean ADE even when the median is already near 1 m.

7. What the implementation learned from the evidence

Chronological changes and what we observed on mini:

8. The complete training trace (5666 iterations)

The rail marks curriculum boundaries; the curves are uniformly downsampled for display while epoch tables and report_data.json retain the full series.

9. Open the trace, phase by phase

10. The result depends on which path you measure

Columns are intentionally labeled by statistic: sample-best training minima are not directly comparable to deterministic mean-path averages. The table keeps both because the gap is the central campaign finding.

10b. Chase the metric that deployment actually uses

After phase3, we run additional GRPO phases optimizing the deterministic mean trajectory (inference path), not just sample-best minADE over G rollouts. Evaluation uses 12 held-out frame indices on nuScenes mini.

Stop target: mean ADE ≤ 0.5 m, max ADE ≤ 1.5 m, mean FDE ≤ 1.0 m.
Best so far: phase4_fit_gt — projects/work_dirs/planner_posttrain/campaign/phase4_fit_gt/latest.pth
Key finding: median ADE ≈ 1 m on most frames, but hard scenes (idx 240, 280) remain outliers.

11. Put the blue deployment path under a microscope

Each panel shows 6 future waypoints in the ego LiDAR frame (x forward, y left). The black dot is the ego origin. We compare three planners on identical frozen backbone features:

• Green solid — ground-truth expert trajectory from nuScenes.
• Gray dashed — frozen stage-2 SFT planner (uniad_base_e2e.pth, LoRA off).
• Blue solid — GRPO post-trained planner (deterministic mean μ, no sampling).
• Purple dashed — 8 stochastic GRPO rollouts (same μ, learned σ; used only during training for group advantages).

Reading the plots: The campaign headline sample-best metric comes from training JSONL. The purple minimum shown here is a separate, seeded export-time diagnostic over G rollouts. At inference we deploy the blue mean path only.

Quantitative comparison (open-loop, 6 steps)

Per-step waypoint tables

12. What the blue path lets us conclude

• Open-loop sample-best minADE improved from 9.80 m (phase0 smoke) to 0.00 m (phase3 best over G samples).
• Sample-best minFDE improved from 16.99 m → 0.052 m on nuScenes mini.
• Mean-path ADE on 12 eval frames: 3.60 m mean, 0.98 m median (best: phase4_fit_gt). Outliers idx 240/280 still ~16–20 m.
• Phase4 GT-fit: last-10% sample-best minADE 0.87 m.
• Phased curriculum matters: phase1 warmup on 50 frames drops minADE below 1 m before full GRPO.
• Collision reward stayed at zero on mini — collision signal is sparse; trainval needed for safety claims.
• Negative raw KL in late training indicates policy–reference density inversion; use KL floor or lower kl_coef.
• Only LoRA parameter tensors plus log_std train; peak VRAM was ~7.3 GiB with G=8 on the RTX 4060 Ti 16GB campaign host.

Roadmap

1. Make the policy update strict: cache a behavior-policy snapshot, detach actions / advantages, and replay multiple updates so $r$ can move and clipping can activate.
2. Add a KL floor or adaptive $\beta$ when raw KL goes negative (policy denser than reference).
3. Run nuScenes trainval and report L2 @ 3 s and collision rate with official eval scripts.
4. Hard-example mining for idx 240 / 280, then export LoRA-only deployment weights.

13. Reproduce

cd UniAD
# Full campaign + GT-fit iterations
/root/miniconda3/envs/uniad2.0/bin/python tools/run_planner_grpo_experiments.py \
  --work-root projects/work_dirs/planner_posttrain/campaign --gpu 0
/root/miniconda3/envs/uniad2.0/bin/python tools/run_grpo_until_fit.py --gpu 0

# Aggregate metrics + render blog
/root/miniconda3/envs/uniad2.0/bin/python tools/analyze_planner_grpo_campaign.py
/root/miniconda3/envs/uniad2.0/bin/python tools/export_planner_waypoints.py \
  --grpo-checkpoint projects/work_dirs/planner_posttrain/campaign/phase4_fit_gt/epoch_5.pth
/root/miniconda3/envs/uniad2.0/bin/python tools/render_planner_grpo_report.py

# Serve (WSL → Windows browser)
cd docs/planner_grpo_lora && python -m http.server 8765 --bind 0.0.0.0
# Open http://localhost:8765/index.html

References

1. ↑ [lora] Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022
2. ↑ [lora-noregret] Thinking Machines Lab, LoRA Without Regret, 2025
3. ↑ [grpo] Shao et al., DeepSeekMath: Pushing the Limits of Mathematical Reasoning via Reinforcement Learning of Process Supervision / GRPO, 2024
4. ↑ [ppo] Schulman et al., Proximal Policy Optimization Algorithms, 2017
5. ↑ [uniad] Hu et al., Planning-oriented Autonomous Driving (UniAD), CVPR 2023
6. ↑ [mmdet3d] OpenMMLab, MMDetection3D
7. ↑ [nuscenes] Caesar et al., nuScenes: A multimodal dataset for autonomous driving, CVPR 2020
8. ↑ [trl] Hugging Face TRL — Group Relative Policy Optimization docs