fix: Fix CosineAnnealingWarmUpRestarts

[yamsam]

Fix the post-warmup LR schedule to cosine-decay

Summary

The scheduler CosineAnnealingWarmUpRestarts did not cosine-anneal — despite the name, its post-warmup phase was MultiplicativeLR(lambda=1.0), i.e. it held the learning rate constant at its peak for the entire run. This let training collapse right after warmup: models peaked during the warmup epochs and then degraded, and best_model (by valid loss) tended to be a warmup-epoch fluke.

This PR makes the schedule a real linear-warmup → cosine-decay, while keeping the old behavior available under an accurate name for backward-compatible reproduction of pre-existing runs.

Based on dev.

Changes

diffusion_planner/utils/lr_schedule.py:

• WarmupCosineAnnealingLR (new): linear warmup for warm_up_epoch epochs, then CosineAnnealingLR decay to eta_min (1e-6) over the remaining epochs. Matches original planTF.
• WarmupConstantLR: the legacy behavior (warmup → constant), renamed to describe what it actually does. Kept for reproducing old runs / resumed checkpoints.
• CosineAnnealingWarmUpRestarts: retained as a backward-compatible alias bound to WarmupConstantLR (the exact legacy/constant behavior), so existing imports and resumed checkpoints reproduce unchanged. Marked deprecated in favor of WarmupCosineAnnealingLR.
• build_lr_scheduler(...): small dispatcher selecting "cosine" (default) or "constant".

train.py, train_config.py, train_predictor.py:

• train.py now builds the scheduler via build_lr_scheduler(..., schedule_type=args.lr_schedule_type).
• New config/CLI knob --lr_schedule_type {cosine,constant} (default cosine).

Behavior change & compatibility

• Default training changes: post-warmup LR now cosine-decays instead of staying at peak. This is the intended fix.
• Old behavior is fully reachable: --lr_schedule_type constant, or importing WarmupConstantLR / the CosineAnnealingWarmUpRestarts alias.
• No breaking API change: the old symbol still imports and behaves exactly as before.

Before vs. after (peak_lr=1e-3, warmup=5, 25 epochs), lr at the final epoch:

schedule	epoch 0	epoch 5 (peak)	epoch 24
constant (legacy)	1e-4	1e-3	1e-3
cosine (new default)	1e-4	1e-3	→ ~eta_min

Verification

• Unit tests (diffusion_planner/tests/test_lr_schedule.py, 7 cases): warmup ramp from start_factor·peak to peak; constant holds peak post-warmup; cosine strictly decreases post-warmup and lands near eta_min; cosine vs. constant identical through warmup then diverge; the alias reproduces the constant behavior exactly; build_lr_scheduler dispatch + invalid type raises; and the warm_up_epoch==1 / epoch==warm_up_epoch edge cases don't raise. All pass.
• TrainConfig builds with lr_schedule_type default cosine and --lr_schedule_type constant override; train.py imports and wires the dispatcher.

🤖 Generated with Claude Code