HUGSIM PostTrain

Post-training pipeline for HUGSIM adversarial closed-loop E2E autonomous driving.

Uses ~11,000 adversarial simulation episodes (UniAD/VAD/LTF x {idm, attack, standard, mild}, 372 scenes) to improve end-to-end driving planners via Replay-SFT, MOSAIC scaling-aware data selection, MPA-lite, and offline RL.

Quick Start

# 1. Activate environment (reuses HUGSIM pixi env + adds PYTHONPATH)
source setup_env.sh

# 2. Run Phase 0 audit (must pass before any training)
python scripts/phase0_a/audit.py --simulation_root links/outputs/simulation/ --output_dir audit/

# 3. Build cache and split
python scripts/phase0_bc/build_cache.py --audit_dir audit/ --output_dir cache/
python scripts/phase0_bc/build_split.py --cache_dir cache/

Project Structure

• links/ -- Symlinks to HUGSIM data, code, and dependencies (read-only)
• modified_src/ -- Copies of HUGSIM files that need structural changes
• posttrain/ -- Core Python package for all post-training methods
  • posttrain/mosaic/ -- MOSAIC scaling-aware data selection (domain partitioner, pilot runner, scaling law, mixture optimizer, sampler)
  • posttrain/optimizers/ -- Custom optimizers (AdamW-OLOR: weight rollback toward pretrained weights)
• scripts/ -- Executable scripts organized by phase
  • scripts/mosaic/ -- MOSAIC pipeline scripts (end-to-end pipeline, sub-metric evaluation)
• configs/ -- YAML configs per phase
• audit/ -- Phase 0 audit outputs
• cache/ -- Structured cache (metadata + aligned labels)
• teacher_pool/ -- Phase 2 replay teacher outputs
• results/ -- Evaluation results and checkpoints

Pipeline Phases

Phase	Description	Duration
0	Data audit + cache + scene-based split	~1 week
1	G-revised baseline (risk-weighted BC)	~1 week
2	Replay-SFT with teacher priority	~2 weeks
2.5	MOSAIC scaling-aware data selection	~2 days
3	MPA-lite (main contribution)	~2 weeks
4	Gate decision + E-lite/C-lite	~1 week
5	Evaluation + VAD transfer + ablation	~1 week

MOSAIC Data Selection (Phase 2.5)

MOSAIC optimizes training data mixture by fitting neural scaling laws per data domain and greedily allocating budget to maximize PDMS.

# Run full MOSAIC pipeline (after Phase 2 teacher pool is built)
sbatch sbatch/mosaic_pilot.sh

# Train with MOSAIC-optimized weights
sbatch sbatch/phase2_train_mosaic.sh

# Or manually:
python scripts/mosaic/run_mosaic_pipeline.py \
    --cache_dir cache/ --teacher_pool_dir teacher_pool/ \
    --output_dir results/mosaic/ --budget_ratio 0.7
torchrun --nproc_per_node=4 scripts/phase2/finetune_A.py \
    --use_mosaic --mosaic_weights results/mosaic/mosaic_weights.pt

OLOR Weight Rollback (Phase 1/2 Optimizer)

OLOR (One step Learning, One step Review) integrates weight rollback into the AdamW optimizer to mitigate knowledge forgetting during fine-tuning from UniAD pretrained weights. Based on Huang et al., AAAI 2024.

# Phase 1 with OLOR
sbatch sbatch/phase1_train_olor.sh

# Phase 2 with OLOR (after teacher pool is built)
sbatch sbatch/phase2_train_olor.sh

# Phase 2 with MOSAIC + OLOR (after mosaic_weights.pt exists)
sbatch sbatch/phase2_train_mosaic_olor.sh

# Or manually add OLOR to any training command:
torchrun --nproc_per_node=4 scripts/phase1/weighted_sft.py \
    --use_olor --olor_max_rollback 5e-3 --olor_min_rollback 0.0 --olor_power 2.0

Key parameters:

• --olor_max_rollback (iota_1): rollback strength for shallow layers (default: 5e-3)
• --olor_min_rollback (iota_2): rollback strength for deep layers (default: 0)
• --olor_power (gamma): layer-wise penalty decay rate (default: 2)

Dependencies

Reuses the HUGSIM pixi environment. Additional packages: wandb, faiss-cpu, pyarrow, decord