PDL for parametric AC-OPF: paper-faithful implementation with analytical dual

[xkhainguyen]

Summary

Implementation of Primal-Dual Learning (PDL) (Park & Van Hentenryck, AAAI 2023) for parametric AC-OPF, with several fixes and extensions over a naive baseline.

Fixes to the base implementation

Fix	Effect
Paper-faithful augmented Lagrangian loss (raw g, h from BNK.constraints!)	Correct Lagrangian gradient; sign of μᵀg term pushes primal toward feasibility
Clamp inequality multipliers to μ ≥ 0	KKT sign correctness
BoundedOutput sigmoid head	max_bound_viol ≡ 0 by construction; no dead gradient zone (vs hardsigmoid)
FixRefBus layer	Reference-bus angle fixed architecturally; GPU-safe
Analytical primal loss	Eliminates dual tracking gap — see below

Analytical dual correction

The dual network lags the true multipliers by orders of magnitude at high ρ. Instead of using λ̂(θ) directly, we apply one ALM update analytically at every primal gradient step:

μ_eff = clamp(μ̂(θ) + ρ·g(ŷ),  0,  max_dual)
λ_eff = clamp(λ̂(θ) + ρ·h(ŷ), −max_dual, max_dual)

Gradient flows through g(ŷ) and h(ŷ) — not through the dual net — giving the correct ρg² and ρh² penalty terms without tracking lag. The dual network still trains in parallel, providing an improving warm-start base over time.

API

Both options live in ALMMethod (defaults shown):

method = ALMMethod(;
    batch_model         = bm,
    num_equal           = num_equal,
    use_analytical_dual = true,   # recommended; false = use dual net output directly
    use_dual_learning   = true,   # false = freeze dual network (ablation only)
    ρmax = 1e4, max_dual = 1e6, τ = 0.8, α = 2.0,
)
train!(method, trainer, data; K=100, L_primal=2500, L_dual=2500, warmup_epochs=25000)

Results on case57 (5000 held-out test samples)

Config	max_eq	max_ineq
use_analytical_dual=true, ρmax=1e4, max_dual=1e6	1.165	0.000
use_analytical_dual=false	~1.25	0.000
max_dual=1e4 (any ρ schedule)	~1.80	0.000

Variable bounds and inequalities are satisfied exactly in all runs.

Test plan

• julia --project=. -e 'using Pkg; Pkg.test()' passes on CPU
• BNK_TEST_CUDA=1 julia --project=. -e 'using Pkg; Pkg.test()' passes on GPU
• julia --project=. examples/case57_train.jl runs to completion

[andrewrosemberg]

Amazing!! Let us know when you are ready for a review

[xkhainguyen]

Hi @andrewrosemberg,

I wanted to mention that this is my personal attempt to reproduce the paper's approach. I tried to fix the branch's implementation and make it work on the 57-bus case. The results are reported above. To be honest, I was hoping for even better results, but it may be a limitation of the approach. Feel free to review and even improve or tune it further.

Many thanks!
Khai

[codecov]

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 ☂️