Skip to content

Add target-penalty annealing support (penalty_schedule) to the training loops#43

Merged
andrewrosemberg merged 7 commits into
LearningToOptimize:mainfrom
DJLee1208:issue/penalty-annealing
Jun 15, 2026
Merged

Add target-penalty annealing support (penalty_schedule) to the training loops#43
andrewrosemberg merged 7 commits into
LearningToOptimize:mainfrom
DJLee1208:issue/penalty-annealing

Conversation

@DJLee1208

Copy link
Copy Markdown
Contributor

Fixes #41

Summary

Training with a constant target-slack penalty weight either hides target violations (low/moderate weight:
the deterministic-equivalent view looks great while the stage-wise rollout violates its own targets) or
degrades the learned policy (high weight from epoch 1). This PR adds the penalty-schedule option proposed
in #41: a multiplier ladder ×0.1 → ×1 → ×10 → ×30 applied during training, relative to the penalty the
models were built with (so with penalty=:auto, relative to the auto-calibrated base
= max |objective coefficient|).

What changed

  • src/utils.jl: default_annealed_schedule(num_batches) (phase lengths proportional to 2/2/4/16, every
    phase ≥ 1 batch, < 4 batches falls back to the tail of the ladder so runs always end at ×30), schedule
    validation (contiguous from batch 1, finite positive multipliers, past-end hold), and internal helpers
    that capture the objective coefficients of the norm_deficit variables once at training start and
    rescale them at phase boundaries. Scaling that single coefficient covers all three create_deficit!
    modes (L1-only, L2-only, and the combined mode where the L1/L2 weights live in the linking constraint),
    as well as the renamed copies inside deterministic-equivalent and shooting-window models — no model
    rebuild is needed.
  • src/simulate_multistage.jl / src/multiple_shooting.jl: a new penalty_schedule keyword on both
    train_multistage overloads and on train_multiple_shooting. nothing (default) leaves behavior
    exactly as today; :default_annealed builds the ladder over num_batches; an explicit vector of
    (first_batch, last_batch, multiplier) phases is validated and used as given.
  • examples/HydroPowerModels: the three trainers default to the annealed schedule. They now make a single
    training call spanning num_epochs*num_batches batches so the schedule covers the whole run. Deliberate
    example-behavior notes: this keeps one optimizer state for the whole run (previously Flux.setup re-ran
    every epoch, silently resetting Adam moments), makes a true return from the record callback stop the
    whole run (previously it only ended the current epoch), and drops the iter % 2100 batch-doubling hook
    (which the collapsed iteration count would have newly activated). Constant-penalty training remains
    available via penalty_schedule = nothing.
  • test/runtests.jl: schedule-resolution exactness (including the 24-batch reference ladder
    1:2 ×0.1, 3:4 ×1, 5:8 ×10, 9:24 ×30), validation rejections, past-end hold, coefficient scaling in all
    three create_deficit! modes, untouched coefficients on the default path, and an end-to-end
    train_multistage run crossing a phase boundary on DiffOpt models.

Validation

  • Pkg.test() green on this branch (210/210, includes the new testsets).
  • Reduced-scale demo (Bolivia DCP, 12 stages, 240 batches per arm, vanilla Ipopt, identical seeds and a
    fixed 6-scenario held-out set): with the constant as-built penalty the trained policy ends at a
    stage-wise-rollout target-violation share of 0.220 — far outside the ≤ ~0.05 trust gate — while its
    deterministic-equivalent view looks unremarkable; with the annealed default at the same budget the share
    drops to 0.052 — near but still slightly above the ≤ ~0.05 gate (4.2× lower). This is partial,
    reduced-scale evidence for the core mechanism of the issue (constant penalties learn targets that are
    not followable stage by stage; annealing restores followability); it is not reduced-scale gap closure.
    The full gap-closure endpoint of the issue table (det-eq ≡ rollout to ~1e-7 relative) requires
    full-scale training and is documented in the issue itself.

Notes

…tiple_shooting (LearningToOptimize#41)

A new penalty_schedule keyword accepts nothing (default, unchanged behavior),
:default_annealed, or an explicit vector of (first_batch, last_batch, multiplier)
phases. Multipliers scale the objective coefficients of the norm_deficit
variables relative to the penalty the models were built with, covering the
L1-only, L2-only, and combined modes of create_deficit! as well as
deterministic-equivalent and shooting-window model copies.
…y schedule (LearningToOptimize#41)

The examples now make a single training call spanning num_epochs*num_batches
batches so the schedule covers the whole run. This also keeps one optimizer
state throughout (previously Flux.setup re-ran every epoch), makes a true
return from the record callback stop the whole run, and drops the
iter-2100 batch-doubling hook that the collapsed iteration count would have
activated. Constant penalties remain available via penalty_schedule = nothing.
Comment thread src/simulate_multistage.jl
…r) into penalty-annealing branch

# Conflicts:
#	examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl
#	src/DecisionRules.jl
#	src/simulate_multistage.jl
#	test/runtests.jl
@codecov

codecov Bot commented Jun 15, 2026

Copy link
Copy Markdown

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

ℹ️ You can also turn on project coverage checks and project coverage reporting on Pull Request comment

Thanks for integrating Codecov - We've got you covered ☂️

@andrewrosemberg andrewrosemberg merged commit 3cee1cc into LearningToOptimize:main Jun 15, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Target-penalty annealing: a constant slack-penalty weight either hides target violations or degrades the learned policy

2 participants