diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 2124afe..44cfebb 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -11,3 +11,9 @@ updates: directory: "/" schedule: interval: "weekly" + ignore: + # The v5→v7 bump silently broke coverage uploads ("Missing Head Commit" + # on PRs). Keep codecov-action pinned until a deliberate, verified + # migration — see the comment in .github/workflows/CI.yml. + - dependency-name: "codecov/codecov-action" + update-types: ["version-update:semver-major"] diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 852dc10..c3e707b 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,12 +19,13 @@ jobs: permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created actions: write contents: read + id-token: write # OIDC token for tokenless Codecov uploads (see codecov step) strategy: fail-fast: false matrix: version: - '1.10' - # - 'nightly' + - '1.12' os: - ubuntu-latest arch: @@ -39,8 +40,22 @@ jobs: - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 - - uses: codecov/codecov-action@v7 + # Pinned to v5: the dependabot bump to v7 (2026-06-25) silently broke + # uploads — Codecov has no commit newer than 2026-06-15, which is what + # produces "Missing Head Commit" on PRs. v5 is the last version verified + # to upload from this workflow. Before re-bumping, migrate deliberately + # (e.g. OIDC: `use_oidc: true` + `id-token: write` permission) and + # confirm a commit appears on Codecov. + # Authentication uses OIDC (`use_oidc` + the job's `id-token: write` + # permission) because the CODECOV_TOKEN secret is not set in this repo + # ("Token length: 0" in CI) and tokenless uploads are rejected on + # protected branches. OIDC requires the Codecov GitHub App to be + # installed for the organization; if uploads fail with an OIDC error, + # either install the app or set the CODECOV_TOKEN secret and replace + # `use_oidc` with `token: ${{ secrets.CODECOV_TOKEN }}`. + - uses: codecov/codecov-action@v5 with: files: lcov.info - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: false + use_oidc: true + # Fail loudly: a silent upload failure hid this breakage for weeks. + fail_ci_if_error: true diff --git a/.gitignore b/.gitignore index c019973..131bb87 100644 --- a/.gitignore +++ b/.gitignore @@ -55,3 +55,4 @@ plan/ *_cuts.json settings.json *.sh +*-backup diff --git a/README.md b/README.md index 710b710..941c349 100644 --- a/README.md +++ b/README.md @@ -54,22 +54,25 @@ using DecisionRules, JuMP, DiffOpt, Flux using SCS # 1) Build per-stage subproblems (DiffOpt-enabled) and collect: -# subproblems, state_params_in, state_params_out, uncertainty_sampler, uncertainties_structure +# subproblems, state_params_in, state_params_out, uncertainty_samples # 2) Build the deterministic equivalent over the full horizon det = DiffOpt.diff_model(() -> DiffOpt.diff_optimizer(SCS.Optimizer)) -det, uncertainties_structure_det = DecisionRules.deterministic_equivalent!( +det, uncertainty_samples_det = DecisionRules.deterministic_equivalent!( det, subproblems, state_params_in, state_params_out, Float64.(initial_state), - uncertainties_structure, + uncertainty_samples, ) +# deterministic_equivalent! remaps state_params_in/state_params_out in place. +# Copy those arrays first if you also need the original stage-wise refs later. + # 3) Train a TS-DDR policy end-to-end -num_uncertainties = length(uncertainty_sampler()[1]) # number of uncertainty components per stage +num_uncertainties = length(uncertainty_samples[1]) # number of uncertainty components per stage policy = Chain( Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu), Dense(64, length(initial_state)), @@ -79,9 +82,9 @@ DecisionRules.train_multistage( policy, initial_state, det, - state_in_det, - state_out_det, - uncertainty_sampler; + state_params_in, + state_params_out, + uncertainty_samples_det; num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(1e-3), @@ -97,7 +100,7 @@ Single shooting solves one optimization per stage and rolls forward using the re ```julia using DecisionRules, Flux -num_uncertainties = length(uncertainty_sampler()[1]) +num_uncertainties = length(uncertainty_samples[1]) policy = Chain( Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu), Dense(64, length(initial_state)), @@ -109,7 +112,7 @@ DecisionRules.train_multistage( subproblems, state_params_in, state_params_out, - uncertainty_sampler; + uncertainty_samples; num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(1e-3), @@ -126,7 +129,7 @@ Multiple shooting partitions the horizon into windows of length `window_size`. E using DecisionRules, Flux, DiffOpt using SCS -num_uncertainties = length(uncertainty_sampler()[1]) +num_uncertainties = length(uncertainty_samples[1]) policy = Chain( Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu), Dense(64, length(initial_state)), @@ -149,10 +152,7 @@ DecisionRules.train_multiple_shooting( policy, initial_state, windows, - state_params_in, - state_params_out, - uncertainty_sampler; - window_size=24, # e.g., 6, 24, ... + uncertainty_samples; num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(1e-3), @@ -168,7 +168,8 @@ The training loops record metrics through a per-sample `SampleLog` cache and a p ```julia using DecisionRules, Random -# Materialize a FIXED held-out evaluation set once, before training +# Materialize a FIXED held-out evaluation set once, before training. +# Use stage-wise subproblems and parameter refs, not DE-remapped refs. Random.seed!(1234) eval_scenarios = [DecisionRules.sample(uncertainty_samples) for _ in 1:8] @@ -178,7 +179,7 @@ rollout_eval = RolloutEvaluation( policy_state=:realized, ) -train_multistage(policy, initial_state, det, state_in_det, state_out_det, uncertainty_sampler; +train_multistage(policy, initial_state, det, state_params_in, state_params_out, uncertainty_samples_det; num_batches=100, record=(sample_log, iter, model) -> begin rollout_eval(iter, model) @@ -202,6 +203,48 @@ Each evaluation reports (a) the rollout objective **excluding** the target-slack Per-sample debugging hooks can be attached with `SampleLog(on_sample=(s, models, log) -> ...)`; the training loop calls the hook after each sample's solve with the live JuMP model(s). The previous `record_loss=(iter, model, loss, tag) -> ...` keyword keeps working as a deprecated adapter. +## Strict mode and reachable policies + +The standard TS-DDR target constraint uses slack: + +```math +x_t + \delta_t = \hat{x}_t, +\qquad +\text{objective} += C_\delta \|\delta_t\|. +``` + +Slack makes training robust to unreachable targets, but it also makes the dual +signal depend on the target-penalty calibration. Strict mode removes the slack: + +```math +x_t = \hat{x}_t. +``` + +The resulting dual is the clean shadow price of imposing the target. The price +of that cleaner signal is feasibility: every policy target must be reachable +from the state used to condition the policy. + +This is automatic in the hydro strict subproblem path because each stage is +solved sequentially and the policy receives the realized previous reservoir +state. It is also possible in regular deterministic equivalents when the target +trajectory is rolled out from the true initial state using a reachable policy: + +```math +\hat{x}_0 = x_0,\qquad +\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1}),\qquad +\hat{x}_t \in R(\hat{x}_{t-1}, w_t). +``` + +By induction, all targets are feasible, and the strict equalities force the +realized trajectory to match that reachable path. See the hydro example and the +DecisionRulesExa.jl companion for the GPU strict regular-DE implementation. + +The policy helpers separate two architectural choices: + +- recurrent `layers` / `DR_ENCODER_LAYERS` process uncertainty history only; +- `combiner_layers` / `DR_HEAD_LAYERS` add a nonlinear feed-forward + state-to-target head without recurrence over the state input. + ## GPU acceleration with DecisionRulesExa.jl For large-scale problems where the inner NLP solve is the bottleneck (e.g., AC-OPF with hundreds of buses), [DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl) provides a GPU-accelerated backend that replaces JuMP with [ExaModels.jl](https://github.com/exanauts/ExaModels.jl) and solves with [MadNLP.jl](https://github.com/MadNLP/MadNLP.jl) + CUDSS on GPU. @@ -222,6 +265,26 @@ Examples live in `examples/`. Run tests with: julia --project -e 'using Pkg; Pkg.test()' ``` +## Repository Map + +| Path | Purpose | +|---|---| +| `src/DecisionRules.jl` | Module entrypoint and exports | +| `src/dense_multilayer_nn.jl` | MLP helpers, state-conditioned recurrent policies, nonlinear target heads | +| `src/simulate_multistage.jl` | Stage-wise and deterministic-equivalent simulation logic | +| `src/multiple_shooting.jl` | Windowed multiple-shooting setup, simulation, and training | +| `src/utils.jl` | Target-parameter utilities, deficit construction, rollout evaluation | +| `src/integer_strategies.jl` | Strategies for extracting gradients from integer/mixed-integer models | +| `src/score_function.jl` | Score-function gradient correction for nonsmooth/integer problems | +| `src/parameter_duals.jl` | Dual/sensitivity helpers for parameterized JuMP models | +| `docs/src/` | Documenter.jl manual pages | +| `examples/HydroPowerModels/` | Bolivia hydrothermal scheduling, strict reachable policies, SDDP comparisons | +| `examples/inventory_control/` | Inventory-control example and dynamic-programming/SDDP comparisons | +| `examples/rocket_control/` | Rocket MPC/control example | +| `examples/RL/` | Reinforcement-learning style hydro scripts | +| `examples/Experimental/` | Research prototypes and robotics/control explorations | +| `test/runtests.jl` | Package test suite | + ## Citation If you use this package in academic work, please cite: diff --git a/docs/make.jl b/docs/make.jl index 490de9b..870fa26 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -20,6 +20,7 @@ makedocs(; format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true", canonical="https://LearningToOptimize.github.io/DecisionRules.jl", + size_threshold=300 * 1024, ), pages=[ "Home" => "index.md", diff --git a/docs/src/algorithm.md b/docs/src/algorithm.md index 36c3d24..c1b6b35 100644 --- a/docs/src/algorithm.md +++ b/docs/src/algorithm.md @@ -27,9 +27,13 @@ Instead of mapping observations directly to actions, the policy outputs **target states**: ```math -\hat{x}_{1:T} = \pi_\theta(w_{1:T}) +\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1}), \qquad \hat{x}_0 = x_0, ``` +evaluated stage-wise with state feedback: each target is conditioned on the +previous target state during deterministic-equivalent training (or on the +realized state ``x_{t-1}`` in closed-loop rollouts). + A projection subproblem enforces feasibility by solving: ```math @@ -105,8 +109,11 @@ for k = 1, ..., ⌈T/W⌉: pass realized end-state to window k+1 ``` -**Pros**: balances coupling (within windows) with tractability; parallelizable windows. -**Cons**: continuity gaps between windows require penalty tuning. +**Pros**: balances coupling (within windows) with tractability; cheaper inner +solves than a full-horizon deterministic equivalent. +**Cons**: windows are chained sequentially during rollout/training because each +window needs the previous realized end-state; cross-window coupling is weaker +than in the full deterministic equivalent. ## Mixed gradient: score-function (REINFORCE) correction @@ -129,8 +136,11 @@ rollouts under perturbed targets. rollouts solve the models exactly as built (MIPs stay MIPs), so the costs reflect true integer-feasible decisions. -3. **Advantage**: center the costs ``A_m = R_m - \bar{R}`` (mean baseline - reduces variance without changing the expected gradient). +3. **Advantage**: center the costs ``A_m = R_m - \bar{R}``. Because the mean + baseline ``\bar{R}`` is computed from the same ``M`` rollouts, mean-centering + reduces variance but introduces a small ``O(1/M)`` bias (effectively scaling + the estimator by ``(M-1)/M``) that vanishes as `num_rollouts` grows; a + leave-one-out baseline would be exactly unbiased. 4. **Surrogate loss**: the differentiable scalar whose gradient recovers the REINFORCE estimate: @@ -203,6 +213,159 @@ Phase 4 (lock): λ × 30.0 — final precision This is the `default_annealed` schedule, activated with `penalty_schedule=:default_annealed`. +## Strict mode: penalty-free gradient signal + +The standard TS-DDR formulation uses a penalty ``C_\delta \|\delta_t\|`` to +penalize deviations from the policy's targets. While effective, the penalty +introduces a trade-off: the dual ``\lambda_t`` conflates the **economic shadow +price** with a **penalty-correction term**. At high penalty, the gradient +signal tells the policy "reduce ``\delta``" rather than "be economically +optimal." + +**Strict mode** eliminates this coupling entirely by replacing the slack +constraint ``x_t + \delta_t = \hat{x}_t`` with a **hard equality**: + +```math +x_t = \hat{x}_t \quad :\lambda_t +``` + +There are no deficit variables, no penalty term, and no penalty to tune. The +dual ``\lambda_t`` is the **pure shadow price** ``\partial Q_t / \partial +\hat{x}_t`` — the marginal value of changing the target, uncontaminated by +any regularization. + +### Feasibility-guaranteeing policies + +Strict mode requires that the policy always produce feasible targets — if the +target is outside the feasible set, the subproblem has no slack to absorb the +violation and becomes infeasible. This is guaranteed by construction through a +**reachable-set policy** that bounds its output to the one-stage reachable set. + +For hydro scheduling with reservoir ``r``, the one-stage reachable set from +current volume ``v_r`` under inflow ``w_r`` is: + +```math +\hat{v}_r \in +\Bigl[ + \max\bigl(\underline{v}_r,\; v_r + K w_r - K \overline{q}_r - \overline{s}_r + + \sum_{u \in U_r^{\text{turn}}} K \underline{q}_u\bigr),\; + \min\bigl(\overline{v}_r,\; v_r + K w_r - K \underline{q}_r + + \sum_{u \in U_r^{\text{turn}}} K \overline{q}_u\bigr) +\Bigr], +``` + +where ``K`` is the water-balance conversion factor, ``\underline{q}_r, +\overline{q}_r`` are turbine bounds, ``\overline{s}_r`` is the spill bound, +and ``U_r^{\text{turn}}`` is the set of upstream units feeding into ``r``. + +#### Cascade-aware clamping + +The fixed upstream term ``\sum_{u} K \overline{q}_u`` in the displayed upper +bound is an **overestimate** whenever an upstream unit stores water: its actual +release is then smaller than ``K \overline{q}_u``, so the fixed bound can exceed +the true reachable set and make strict subproblems infeasible. The policy +therefore applies a clamping step after computing all raw targets. For each +cascade link ``u \to d``, the release implied by the upstream target is + +```math +R_u = K w_u + v_u - \hat{v}_u , +``` + +and the maximum contribution to the downstream unit is ``\max(0, R_u)`` for +turbine-plus-spill links and ``\min(K \overline{q}_u, \max(0, R_u))`` for +turbine-only links. The downstream target is then clamped to + +```math +\hat{v}_d \le \min\bigl(\overline{v}_d,\; + v_d + K w_d - K \underline{q}_d + \text{max\_contrib}\bigr). +``` + +Two assumptions are documented for this scheme: + +- **Single-level cascades**: the release formula ``R_u`` omits the upstream + unit's own incoming cascade contribution, which is conservative + (underestimates the release) for multi-level chains. +- **No gradient through binding clamps**: the bounds and the clamping step carry + no gradient (`@non_differentiable`), so when a clamp binds, the dependence of + the downstream target on the upstream target is not differentiated. + +The [`HydroReachablePolicy`] implements this by passing the LSTM encoder output +through a sigmoid activation and scaling the result to ``[\ell_r, u_r]``: + +```math +\hat{v}_r = \ell_r + (u_r - \ell_r) \cdot \sigma(z_r). +``` + +The bounds ``\ell_r, u_r`` are computed from physics (no gradient flows through +them); the gradient path is solely through ``\sigma(z_r)``, exactly as in the +standard TS-DDR pipeline. + +### Why strict mode is usually closed-loop + +Strict target equality is naturally compatible with formulations where the +policy sees the state from which the target must be reached. + +- In **stage-wise subproblems**, stage ``t`` is solved after stage ``t-1`` has + produced a realized state. The next policy call receives that realized state, + so a reachable policy can produce a target that is feasible for the next + strict stage solve. +- In **embedded deterministic equivalents**, the policy is evaluated inside the + NLP against realized state decision variables. The strict equality then + couples policy output and realized next state directly. + +A plain deterministic equivalent is different. The training loop normally +computes all targets before solving the coupled multi-stage NLP. Except for +``x_0``, the policy receives its own previously predicted targets rather than +the realized states that the optimizer will eventually choose. With an arbitrary +policy this is open-loop target generation, so a strict equality can make the +DE infeasible. + +### Strict regular DE by induction + +The hydro reachable policy creates an important exception. Suppose +``x_0`` is feasible and the target rollout for a sampled inflow path is + +```math +\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1}), +\qquad \hat{x}_0 = x_0, +``` + +with + +```math +\hat{x}_t \in R(\hat{x}_{t-1}, w_t) +``` + +for every stage ``t``. Then the full strict deterministic-equivalent target +trajectory is feasible. The proof is induction: + +1. Stage 1 is feasible because ``\hat{x}_1`` is reachable from the known + feasible initial state ``x_0``. +2. If stages ``1,\ldots,t`` are feasible and strict equalities force + ``x_t = \hat{x}_t``, then stage ``t+1`` starts from a feasible state equal to + the state used by the policy rollout. Since + ``\hat{x}_{t+1} \in R(\hat{x}_t, w_{t+1})``, stage ``t+1`` is feasible. + +This is why strict mode can be used in the regular hydro DE when targets are +generated from a reachability-preserving policy starting at the true initial +state. The policy is still not reacting to optimizer deviations, but strict +feasibility removes those deviations: the realized state path must equal the +reachable target path. + +### When to use strict mode + +Strict mode is the preferred approach when: + +1. The initial state is feasible and every policy target is **one-stage + reachable** from the state used as policy input. +2. The reachable set is **easy to compute** — e.g., reservoir water balance with + known turbine/spill bounds. +3. You want to avoid **penalty tuning** — strict mode has no penalty hyperparameter. + +In the [Hydropower Scheduling](@ref) example, strict mode with a +`HydroReachablePolicy` achieves competitive simulation costs out of the box, +with no penalty schedule, no annealing, and no hyperparameter search. + ## Evaluation semantics A policy trained on the deterministic equivalent generates targets using **target-state diff --git a/docs/src/assets/hydro_paired_cost_distributions.png b/docs/src/assets/hydro_paired_cost_distributions.png new file mode 100644 index 0000000..45f4211 Binary files /dev/null and b/docs/src/assets/hydro_paired_cost_distributions.png differ diff --git a/docs/src/assets/hydro_training_convergence.png b/docs/src/assets/hydro_training_convergence.png index 3fd5455..f03f277 100644 Binary files a/docs/src/assets/hydro_training_convergence.png and b/docs/src/assets/hydro_training_convergence.png differ diff --git a/docs/src/assets/hydro_training_convergence_by_step.png b/docs/src/assets/hydro_training_convergence_by_step.png new file mode 100644 index 0000000..571e7e6 Binary files /dev/null and b/docs/src/assets/hydro_training_convergence_by_step.png differ diff --git a/docs/src/assets/hydro_training_convergence_by_time.png b/docs/src/assets/hydro_training_convergence_by_time.png new file mode 100644 index 0000000..84acdee Binary files /dev/null and b/docs/src/assets/hydro_training_convergence_by_time.png differ diff --git a/docs/src/assets/hydro_violation_share.png b/docs/src/assets/hydro_violation_share.png deleted file mode 100644 index 5f1d863..0000000 Binary files a/docs/src/assets/hydro_violation_share.png and /dev/null differ diff --git a/docs/src/assets/inventory_integer_results.png b/docs/src/assets/inventory_integer_results.png index 0f26dd2..9b80599 100644 Binary files a/docs/src/assets/inventory_integer_results.png and b/docs/src/assets/inventory_integer_results.png differ diff --git a/docs/src/assets/inventory_relaxed_results.png b/docs/src/assets/inventory_relaxed_results.png index 8a16f85..0556bfa 100644 Binary files a/docs/src/assets/inventory_relaxed_results.png and b/docs/src/assets/inventory_relaxed_results.png differ diff --git a/docs/src/examples/hydro.jl b/docs/src/examples/hydro.jl index e7bf3fd..57cf4f1 100644 --- a/docs/src/examples/hydro.jl +++ b/docs/src/examples/hydro.jl @@ -1,13 +1,25 @@ # # Hydropower Scheduling # # This example trains target-setting decision rules for the Bolivia -# long-term hydrothermal dispatch (LTHD) problem — both **TS-DDR** (deep, -# LSTM-based) and **TS-LDR** (linear) — and compares them against an SDDP -# baseline with inconsistent formulations. -# -# The Bolivia system has **10 hydro plants**, **96 monthly stages**, and -# **AC power flow** constraints. Inflow uncertainty is sampled from 47 -# historical scenarios. +# long-term hydrothermal dispatch (LTHD) problem — **TS-DDR** (deep, +# LSTM-based) and its linear counterpart **TS-LDR** — and compares them +# against an SDDP baseline with inconsistent formulations. The strict +# (penalty-free) TS-DDR variant is trained in two independent +# implementations that we verify are numerically equivalent: +# +# 1. **Stage-wise subproblems on CPU** (this package: JuMP + DiffOpt + Ipopt), +# solving one AC-OPF per stage in closed loop; and +# 2. **Full-horizon deterministic equivalent on GPU** +# ([DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl): +# ExaModels + MadNLP/cuDSS), solving one coupled 126-stage NLP per +# gradient sample. +# +# The Bolivia system has **28 buses**, **34 generators**, **11 hydro +# plants** (three of them in river cascades), and **AC power flow** +# constraints. Policies are trained on a 126-stage horizon and evaluated on +# **96 monthly stages**; inflow uncertainty is sampled from **47 historical +# joint scenarios** (spatially correlated across plants, tiled cyclically +# for horizons beyond the record). # # ## Overview of the TS-DDR approach # @@ -57,15 +69,21 @@ # ## Gradient computation: the envelope theorem # # By the envelope theorem, the sensitivity of the optimal value with respect -# to the target parameter is simply the dual: +# to the target parameter is available in closed form from the multiplier of +# the target constraint. Throughout this page we define ``\lambda_t`` as +# that sensitivity, # # ```math -# \frac{\partial q_t}{\partial \hat{x}_t} -# \;=\; -\lambda_t. +# \lambda_t +# \;:=\; +# \frac{\partial q_t}{\partial \hat{x}_t}, # ``` # -# Combined with backpropagation through the policy network, the full gradient -# of the expected cost is: +# i.e. the constraint multiplier reported by the solver, mapped through the +# sign convention of the modeling layer (both implementations extract it this +# way, and the two extractions are verified to agree numerically — see the +# equivalence audit in the Results section). Combined with backpropagation +# through the policy network, the full gradient of the expected cost is: # # ```math # \nabla_\theta \mathbb{E}[Q] @@ -110,7 +128,8 @@ using Statistics, Random # optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0, "linear_solver" => "mumps") # ) # -# subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = +# subproblems, state_params_in, state_params_out, uncertainty_samples, +# initial_state, max_volume, hydro_meta = # build_hydropowermodels( # "bolivia", "ACPPowerModel.mof.json"; # num_stages=96, @@ -125,8 +144,8 @@ using Statistics, Random # # 1. **Encoder** — a stack of LSTM cells that processes only the uncertainty # (inflow) sequence, capturing temporal dependencies across stages. -# 2. **Combiner** — a Dense layer that merges the encoded uncertainty with the -# previous state to produce the next target. +# 2. **Combiner** — a feed-forward target head that merges the encoded +# uncertainty with the previous state to produce the next target. # # At each stage the policy receives ``[w_t;\; x_{t-1}]`` and outputs # target reservoir volumes ``\hat{x}_t``: @@ -134,8 +153,8 @@ using Statistics, Random # ``` # ┌─────────┐ ┌────────────────┐ ┌──────────────┐ # │ w_t │─────▶│ LSTM encoder │─────▶│ │ -# └─────────┘ └────────────────┘ │ Dense │──▶ x̂_t -# ┌─────────┐ │ combiner │ +# └─────────┘ └────────────────┘ │ feed-forward │──▶ x̂_t +# ┌─────────┐ │ head │ # │ x_{t-1} │─────────────────────────────▶│ │ # └─────────┘ └──────────────┘ # ``` @@ -143,11 +162,32 @@ using Statistics, Random # The LSTM carries hidden state across stages, giving the policy memory of # past inflows. The activation is `sigmoid` (bounding outputs to ``[0,1]``, # which is then scaled by the feasibility mapping). +# +# The `layers` argument controls recurrence over inflows only. The optional +# `combiner_layers` argument adds hidden layers to the nonrecurrent +# state-conditioned target head. This is the preferred way to make TS-DDR +# nonlinear in the current reservoir state without adding recurrence over state. +# +# !!! warning "Recurrent state must be threaded explicitly (Flux ≥ 0.16)" +# Since Flux 0.16, calling an `LSTM` layer restarts it from +# `initialstates` on **every call**, and `Flux.reset!` is a deprecated +# no-op. A per-stage policy loop written naively against this API is +# silently *memoryless* — the "LSTM" degenerates to a per-stage MLP over +# the current inflow. Both packages therefore thread the hidden state +# explicitly across stages inside their policy types (and implement a +# real `Flux.reset!` for scenario boundaries). The effect is large: on +# this example, evaluating the same trained weights with and without +# state threading changes the 100-scenario mean cost by **19%** +# (361,056 vs 303,631). If you build custom recurrent policies on this +# framework, verify statefulness with the pattern used in the package +# test suites: the same input fed twice without `reset!` must produce +# different outputs. # ```julia # models = state_conditioned_policy( # num_uncertainties, num_hydro, num_hydro, [128, 128]; # activation=sigmoid, encoder_type=Flux.LSTM, +# combiner_layers=[128, 128], # ) # ``` @@ -240,7 +280,7 @@ using Statistics, Random # # train_multistage( # models, initial_state, det_equivalent, -# state_params_in, state_params_out, uncertainty_samples; +# state_params_in, state_params_out, uncertainty_samples_det; # num_batches=4000, optimizer=Flux.Adam(), # penalty_schedule=[(1,100,0.1), (101,210,1.0), (211,300,10.0), (301,4000,30.0)], # ) @@ -272,25 +312,36 @@ using Statistics, Random # # ### Gradient chain # -# The gradient must account for how the realized state at stage ``t`` -# depends on the targets at all earlier stages. By the chain rule: +# The gradient must account for two coupled feedback paths: the realized +# state at stage ``t`` depends on the targets at all earlier stages, and the +# *policy input* at stage ``t`` includes that realized state. Writing +# ``x_t = X_t(x_{t-1}, \hat{x}_t)`` for the realized-state map of the stage +# solve, the exact total-derivative recursion is # # ```math -# \frac{\partial Q}{\partial \hat{x}_t} -# \;=\; -# \lambda_t -# + \sum_{k>t} -# \frac{\partial q_k}{\partial x_{k-1}} -# \cdot \prod_{j=t+1}^{k-1} -# \frac{\partial x_j}{\partial x_{j-1}} -# \cdot \frac{\partial x_t}{\partial \hat{x}_t}. +# \frac{dQ}{d\theta} +# = \sum_{t=1}^{T} \left[ +# \frac{\partial q_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d\theta} +# + \frac{\partial q_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta} +# \right], +# \qquad +# \frac{d x_t}{d\theta} +# = \frac{\partial X_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta} +# + \frac{\partial X_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d\theta}, +# ``` +# +# ```math +# \frac{d \hat{x}_t}{d\theta} +# = \nabla_\theta \pi_\theta +# + \frac{\partial \pi_\theta}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta}. # ``` # -# In practice, automatic differentiation (Zygote + ChainRules `rrule`s -# defined on each stage solve) handles this chain automatically. +# Reverse-mode automatic differentiation (Zygote + ChainRules `rrule`s +# defined on each stage solve) computes exactly this chain, including the +# policy-feedback term ``\partial \pi_\theta / \partial x_{t-1}``. # The `rrule` for each stage solve reads the dual ``\lambda_t`` for the # target constraint and uses DiffOpt's implicit differentiation for the -# state-transition sensitivities. +# state-transition sensitivities ``\partial X_t / \partial \cdot``. # # **Advantages**: closed-loop — the policy sees realized states, matching # deployment semantics. Each solve is small (single-stage AC-OPF). @@ -380,9 +431,231 @@ using Statistics, Random # ) # ``` -# ## Penalty annealing +# ## Training pipeline 4: Strict subproblems with reachable policy +# +# The three formulations above use a **slack penalty** ``C_\delta \|\delta_t\|`` +# to handle the gap between the policy's targets and the feasible set. While +# effective, the penalty introduces a hyperparameter and can corrupt the gradient +# signal: at high ``C_\delta``, the dual ``\lambda_t`` reflects "reduce the +# slack" rather than "improve economic dispatch." +# +# **Strict mode** eliminates the penalty entirely by enforcing a **hard equality** +# between the target and the realized state: +# +# ```math +# x_t = \hat{x}_t \quad :\lambda_t \qquad \text{(no slack, no } \delta_t \text{)} +# ``` +# +# The dual ``\lambda_t`` is then the **pure shadow price** +# ``\partial q_t / \partial \hat{x}_t``: the economic value of a marginal change +# in the target, free of any penalty noise. +# +# ### Feasibility guarantee: HydroReachablePolicy +# +# Removing the slack requires that every target produced by the policy be +# **physically achievable**. For hydro scheduling, this means the target volume +# must lie within the one-stage reachable set — the range of volumes achievable +# from the current state ``v_{r,t-1}`` by choosing turbine flow ``q_r`` and +# spillage ``s_r`` within their physical bounds. +# +# #### Per-unit reachable bounds +# +# The water balance for reservoir ``r`` at stage ``t`` is +# +# ```math +# v_{r,t} = v_{r,t-1} + K\, w_{r,t} - K\, q_{r,t} - K\, s_{r,t} +# + \sum_{u \in \mathcal{U}_r} K\, q_{u,t} +# + \sum_{u \in \mathcal{S}_r} K\, s_{u,t}, +# ``` +# +# where ``K`` is the water-balance conversion factor extracted from the model, +# ``w_{r,t}`` is the inflow, +# ``q_{r,t}`` is the turbined flow, ``s_{r,t}`` is the spillage, +# ``\mathcal{U}_r`` is the set of upstream units connected by turbine flow, +# and ``\mathcal{S}_r`` is the set connected by spillage. +# +# The reachable bounds for unit ``r`` (ignoring cascade interactions) are: +# +# ```math +# \ell_{r,t} = \max\bigl(\underline{v}_r,\; +# v_{r,t-1} + K\, w_{r,t} - K\,\bar{q}_r - K\,\bar{s}_r +# + K \sum_{u \in \mathcal{U}_r} \underline{q}_u\bigr), +# ``` +# +# ```math +# u_{r,t} = \min\bigl(\bar{v}_r,\; +# v_{r,t-1} + K\, w_{r,t} - K\,\underline{q}_r +# + K \sum_{u \in \mathcal{U}_r} \bar{q}_u +# + K \sum_{u \in \mathcal{S}_r} \bar{s}_u\bigr). +# ``` +# +# These bounds assume worst-case upstream contributions (maximum turbine/spill +# capacity). The [`HydroReachablePolicy`] wraps the same LSTM uncertainty +# encoder plus feed-forward state-conditioned target head as +# [`StateConditionedPolicy`](@ref) but uses a **sigmoid** activation to bound +# the output to this reachable interval: +# +# ```math +# \hat{v}_{r,t} = \ell_{r,t} + (u_{r,t} - \ell_{r,t}) \cdot \sigma(z_{r,t}). +# ``` +# +# #### Cascade-aware clamping +# +# The per-unit upper bound ``u_{r,t}`` uses worst-case upstream contributions +# (``K \bar{q}_u``, ``K \bar{s}_u``). When an upstream unit ``u`` stores water +# (its target ``\hat{v}_{u,t}`` is high), the actual upstream release +# +# ```math +# R_u = K\, w_{u,t} + v_{u,t-1} - \hat{v}_{u,t} +# ``` +# +# can be much less than the assumed maximum. For cascaded systems, this means +# the downstream target may exceed the true reachable set, causing infeasibility +# in strict mode (no slack to absorb the gap). +# +# After computing the initial sigmoid targets for all units, the policy applies +# a **cascade clamping** step. For each upstream→downstream connection: +# +# - **Turn + spill** connection: the full release reaches downstream, +# so ``\text{max\_contrib} = \max(0,\, R_u)``. +# - **Turn-only** connection: only turbined flow reaches downstream, +# so ``\text{max\_contrib} = \min(K\,\bar{q}_u,\, \max(0,\, R_u))``. +# +# The downstream target is then clamped: +# +# ```math +# \hat{v}_{d,t} \;\le\; v_{d,t-1} + K\, w_{d,t} +# - K\,\underline{q}_d + \text{max\_contrib}. +# ``` +# +# This clamping is `@non_differentiable` — gradient flows through ``\sigma`` +# for unclamped targets, and is zero for clamped ones (correct projected-gradient +# signal). # -# The target penalty ``C_\delta`` controls the trade-off between following +# ### Setup +# +# Building strict subproblems requires only the `strict=true` flag: + +# ```julia +# subproblems, state_params_in, state_params_out, uncertainty_samples, +# initial_state, max_volume, hydro_meta = build_hydropowermodels( +# case_dir, "ACPPowerModel.mof.json"; +# num_stages=126, optimizer=diff_optimizer, +# strict=true, # ← no deficit, hard equality targets +# ) +# ``` + +# The reachable policy is constructed from the hydro metadata returned by +# `build_hydropowermodels`: + +# ```julia +# models = hydro_reachable_policy(hydro_meta, [128, 128]) +# models_with_deep_state_head = hydro_reachable_policy( +# hydro_meta, +# [128, 128]; +# combiner_layers=[256, 256], +# ) +# ``` + +# Training uses the same `train_multistage` with no penalty schedule: + +# ```julia +# train_multistage( +# models, initial_state, subproblems, +# state_params_in, state_params_out, uncertainty_samples; +# num_batches=8000, optimizer=Flux.Adam(), +# penalty_schedule=nothing, # ← no penalty to tune +# ) +# ``` + +# !!! tip "Out-of-the-box convergence" +# Strict mode with `HydroReachablePolicy` requires **no penalty tuning**, +# no annealing schedule, and no hyperparameter search. The clean gradient +# signal allows the optimizer to directly minimize operational cost. + +# ## Training pipeline 5: Strict deterministic equivalent on GPU +# +# The strict formulation unlocks a second, much faster training route, +# implemented in the companion package +# [DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl): +# the **full-horizon deterministic equivalent solved on GPU** with +# [ExaModels.jl](https://github.com/exanauts/ExaModels.jl) (SIMD-friendly +# algebraic modeling) and [MadNLP.jl](https://github.com/MadNLP/MadNLP.jl) +# with the cuDSS sparse linear solver. +# +# ### Why strict mode makes the regular DE safe +# +# A regular (non-embedded) DE is normally *open-loop*: the policy produces +# all targets ``\hat{x}_{1:T}`` before the coupled solve, seeing its own +# previous target instead of a realized state. With an arbitrary policy this +# can render a strict DE infeasible. The reachable policy restores safety +# **by induction**: roll targets out as ``\hat{x}_0 = x_0`` and +# ``\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1})`` with every target inside the +# one-stage reachable set of its input state. Stage 1 is then feasible from +# the true initial state; and if stages ``1..t`` are feasible, the strict +# equalities force ``x_t = \hat{x}_t``, so stage ``t{+}1`` starts exactly at +# the state the policy planned from — making ``\hat{x}_{t+1}`` feasible too. +# Strict feasibility *removes* the open-loop/closed-loop gap: the realized +# state path must equal the reachable target path, so training-time DE +# solves and deployment-time stage-wise rollouts traverse identical +# trajectories (we verify this numerically below). +# +# ### Reservoir volumes become parameters +# +# Because strict equalities pin every reservoir volume to its target, the +# volume trajectory is **data, not a decision**, for the inner solver. The +# GPU formulation exploits this: the reservoir trajectory enters the NLP as +# a parameter vector, eliminating ``(T{+}1) \cdot n_{\text{hydro}}`` +# variables and all slack variables from the KKT system. The targets then +# appear only on the right-hand side of the ``T \cdot n_{\text{hydro}}`` +# water-balance rows. With ``\mu_t`` the multiplier of the stage-``t`` +# water-balance row, the envelope gradient follows from the two rows each +# target touches (``+1`` on ``v_{t+1}`` in row ``t``, ``-1`` on ``v_t`` in +# row ``t{+}1``): +# +# ```math +# \frac{\partial Q}{\partial \hat{x}_t} = \mu_t - \mu_{t+1} +# \quad (t < T), +# \qquad +# \frac{\partial Q}{\partial \hat{x}_T} = \mu_T . +# ``` +# +# ### What the GPU buys +# +# Each gradient sample requires one coupled 126-stage AC NLP solve. On an +# NVIDIA H200, MadNLP + cuDSS solves it fast enough that a full 8,000-solve +# training run completes in roughly a day — and, more importantly, the +# training loss enters the SDDP-forward-cost envelope within the **first +# hours** (see the wall-clock convergence figure in the Results). Solver +# state is reused across solves with a dual-snapshot warm-start scheme that +# prevents one failed solve from corrupting subsequent ones. +# +# ```julia +# ## In DecisionRulesExa.jl (see its examples/HydroPowerModels): +# de = build_hydro_de(power_data, hydro_data, 126; +# formulation=:ac_polar, strict_targets=true, +# backend=CUDABackend()) +# policy = hydro_reachable_policy(hydro_data, [128, 128]; +# combiner_layers=[128, 128]) +# train_tsddr(policy, x0, de, de.p_x0, de.p_target, de.p_inflow, sampler; +# num_batches=8000, madnlp_kwargs=(tol=1e-6,)) +# ``` +# +# !!! note "Cross-package equivalence audit" +# The two strict implementations are verified against each other on this +# exact case: with identical weights, data, and scenarios, (i) the two +# policy implementations produce bit-identical targets, (ii) the +# stage-wise CPU rollout (Ipopt) and the strict full-horizon DE (MadNLP) +# agree per scenario to ``10^{-9}`` relative cost, and (iii) the two +# packages' 100-scenario evaluations of the same checkpoint agree to +# 0.001% (303,631 vs 303,635). The audit scripts +# (`dump_paired_policy_reference.jl`, `eval_paired_exa_strict.jl`, +# `compare_paired_evals.jl`) ship with the packages and re-run on demand. + +# ## Penalty annealing (non-strict formulations) +# +# For the non-strict formulations (DE, stage-wise, multiple shooting), the +# target penalty ``C_\delta`` controls the trade-off between following # the policy's targets and minimizing operational cost. DecisionRules # supports a **penalty annealing schedule** that ramps the penalty multiplier # during training: @@ -396,6 +669,10 @@ using Statistics, Random # # This is activated with `penalty_schedule=:default_annealed` or by passing # an explicit list of `(start_iter, end_iter, multiplier)` tuples. +# +# The penalty schedule must be carefully tuned per problem. In contrast, +# strict mode bypasses this entirely when the problem admits an always-feasible +# policy (see above). # ## Evaluation # @@ -408,7 +685,30 @@ using Statistics, Random # # The **target-violation share** measures how much cost comes from the slack # penalty rather than actual operations — it should be small (``\le 5\%``) for -# a well-trained policy. +# a well-trained policy. In strict mode, the violation share is always **zero** +# by construction. +# +# ### Paired evaluation protocol +# +# All headline numbers in the Results section come from a **paired** +# protocol: a fixed 126×500 index matrix, generated deterministically from a +# documented seed (`paired_scenario_indices` in the example's +# `load_hydropowermodels.jl`, using `StableRNGs` so the stream is identical +# on every Julia version), selects the *same* joint inflow realization at +# every stage for every method. The protocol is therefore reproducible from +# code alone — there is no scenario data artifact to distribute. The SDDP +# policy is simulated +# with `SDDP.Historical` on those indices (`sddp/eval_paired_sddp.jl`); every +# decision-rule checkpoint is rolled out stage-wise on the identical +# trajectories (`eval_paired_tsddr.jl` here; `eval_paired_exa_strict.jl` in +# the GPU package). Pairing removes the between-scenario variance +# (per-scenario cost std ≈ 5,600) from the *comparison*: the standard error +# of the paired mean difference is ≈ 50, roughly two orders of magnitude +# tighter than comparing unpaired means. Data-file identity across +# implementations is enforced (byte-identical `inflows.csv`, `hydro.json`, +# `PowerModels.json`), and the stage-index-to-inflow-row mapping (cyclic +# tiling of the 47-row record) is asserted programmatically inside the +# cross-package evaluation. # ```julia # rollout_eval = RolloutEvaluation( @@ -431,51 +731,276 @@ using Statistics, Random # convex relaxation approximates the value function while the forward pass # evaluates under the true physics. # -# The SDDP policy is trained for up to 2000 iterations and the learned -# cuts are saved to a JSON file, which can be loaded to simulate the -# policy under the ACP formulation. +# The learned cuts are saved to a JSON file, which can be loaded to +# simulate the policy under the ACP formulation. On this case the SDDP +# training ran **441 iterations in ≈ 12 hours** (CPU, MadNLP subproblem +# solver), converging its lower bound to **378,207**; the forward-pass +# (ACP) simulation cost stabilizes around **380 K** on the 126-stage +# horizon. These two numbers frame everything below: no policy can have +# expected 126-stage cost below the bound, and SDDP's own policy sits +# roughly 0.5% above it. # ## Results # -# The plots below compare the TS-DDR and TS-LDR training formulations and -# the SDDP baseline on the Bolivia case. Training curves, out-of-sample -# cost distributions, reservoir volume trajectories, and thermal generation -# profiles are shown. -# -# ### Training convergence (TS-DDR methods) -# -# ![Training convergence](../assets/hydro_training_convergence.png) -# -# ### Out-of-sample cost (TS-DDR methods) -# -# ![Out-of-sample cost comparison](../assets/hydro_cost_comparison.png) -# -# ### Target-violation share (TS-DDR methods) -# -# ![Violation share](../assets/hydro_violation_share.png) -# -# ### Reservoir volume comparison (all methods) -# -# ![Volume comparison](../assets/hydro_volume_comparison.png) -# -# ### Thermal generation comparison (all methods) -# -# ![Generation comparison](../assets/hydro_generation_comparison.png) -# -# ### Summary -# -# | Method | Policy | Mean Cost | Std | N | -# |:-------|:------:|----------:|----:|--:| -# | TS-DDR (DE) | LSTM | 325 540 | 6 266 | 100 | -# | TS-DDR (DE, anneal) | LSTM | 324 445 | 6 134 | 100 | -# | TS-DDR (shooting w=12) | LSTM | 323 289 | 5 593 | 100 | -# | TS-DDR (shooting w=12, anneal) | LSTM | 322 812 | 6 081 | 100 | -# | TS-DDR (stage-wise, anneal) | LSTM | 321 543 | 6 214 | 100 | -# | SDDP (SOC-WR / ACP) | cuts | 303 684 | — | 100 | -# -# All three TS-DDR methods with penalty annealing converge to similar -# costs (321K–325K). SDDP trains on 126 stages (96 + 30 margin). -# -# !!! note "Preliminary results" -# These numbers reflect the current default training scripts. -# They will be updated as the package evolves. +# We evaluate the two **strict** TS-DDR implementations — stage-wise +# subproblems (CPU) and the full-horizon GPU deterministic equivalent — +# against the SDDP baseline on the Bolivia case with AC power flow (11 +# hydro plants, 47 inflow scenarios). The non-strict formulations (DE, +# stage-wise, multiple shooting) are available through the same API but +# require penalty scheduling and careful tuning; strict mode eliminates +# this entirely, so it is the configuration we benchmark. +# +# Three comparison disciplines keep the results honest: +# +# 1. **Same scenarios** — every number below uses the paired 100-scenario +# protocol described above. +# 2. **Same physics at evaluation** — all rollouts solve the identical +# ACP stage problem (hard reactive balance, mof.json objective); the +# cross-package audit pins the implementations to each other at +# solver-tolerance level. +# 3. **Documented training recipes** — every trained artifact corresponds +# to a from-scratch-reproducible configuration listed in the appendix +# and in the example READMEs (no ad-hoc checkpoints). +# +# ### Understanding the metrics +# +# - **SDDP lower bound** (126 stages): the expected-cost lower bound from +# the convex SOC-WR relaxation. This is a *relaxation bound* — it cannot +# be beaten by any feasible policy. For Bolivia, it converges to +# approximately **378 207**. +# +# - **SDDP forward-pass cost** (126 stages): the simulation cost of the +# SDDP policy evaluated under the true AC formulation during the forward +# pass. This is the 126-stage operational cost of the SDDP policy. +# +# - **Simulation cost** (96 stages): the operational cost obtained by +# rolling out a policy under AC power flow on the 500 seeded paired inflow +# scenarios. This is the primary metric for policy quality. SDDP's +# 96-stage simulation cost is **303 665** (mean over the 500 paired +# scenarios, std 5 921). +# +# During TS-DDR training, the logged loss is a *training-batch average* +# over a small number of sampled scenarios, and periodic rollout +# evaluations use small fixed held-out sets. Neither is comparable to the +# 500-scenario paired protocol: small evaluation sets carry offsets of +# several hundred cost units, so cross-method claims are made only on the +# paired protocol. +# +# ### Training convergence against wall-clock time (126 stages) +# +# The figure below shows all 126-stage training metrics against **wall-clock +# time** (log scale), which is the axis that exposes the computational +# trade-off between the methods: +# +# ![Training convergence vs wall-clock time](../assets/hydro_training_convergence_by_time.png) +# +# - **SDDP lower bound** (SOC-WR relaxation): converges to ~378.2 K over +# ≈ 12 hours (441 iterations). Dashed line — no policy can beat it. +# - **SDDP forward-pass cost**: the 126-stage simulation cost of the SDDP +# policy under the true AC formulation, ~380 K after convergence. +# - **TS-DDR strict subproblems (CPU)**: the final two-phase stage-wise +# schedule. Phase 1 uses single-sample gradients and training-loss +# checkpointing; phase 2 switches to larger batches, decayed learning +# rate, and held-out rollout checkpointing. The wall-clock axis counts +# both phases. +# - **TS-DDR strict DE (GPU)**: the analogous two-phase full-horizon +# deterministic-equivalent schedule on an H200. One coupled 126-stage +# solve provides each gradient sample; the second phase uses the same +# rollout-selected, lower-variance training discipline and hard reactive +# balance used at evaluation. +# +# In cumulative wall-clock time, the plotted two-phase schedules end at +# approximately **21.3 h** for strict subproblems and **24.2 h** for strict +# DE. The GPU DE curve reaches the lowest TS-DDR training objective, but this +# run is not a wall-clock win over SDDP. +# +# The plot is regenerated by +# `examples/HydroPowerModels/plot_hydro_strict_convergence.jl`, which parses +# the SDDP training log and pulls the policy-training histories from the +# experiment tracker. +# +# ### Two-phase TS-DDR training protocol +# +# The reported TS-DDR policies use a single, cumulative two-phase schedule +# (exact commands in the example README): +# +# 1. **Coarse training**: train from scratch with small gradient batches and +# checkpoint on training loss. This phase rapidly enters the SDDP cost +# envelope but the per-batch objective remains too noisy to rank policies +# that differ by a few hundred cost units. +# 2. **Selection phase**: continue from the selected checkpoint with larger +# gradient batches, a warmed-up and cosine-decayed learning rate, and +# checkpoint selection on a fixed held-out rollout objective. This phase +# optimizes the same quantity used for the paired evaluation below. +# +# The plotted wall-clock time is cumulative across both phases. +# +# ### 96-stage out-of-sample rollout cost (paired, 500 seeded scenarios) +# +# The primary evaluation metric is the **96-stage simulation cost** — +# total dispatch cost under AC power flow on the 500 paired inflow +# trajectories of the seeded protocol. +# +# | Method | Policy | Mean Cost | Std | Target violations | Training | +# |:-------|:------:|----------:|----:|:-----------------:|:---------| +# | SDDP (SOC-WR / ACP) | cuts | 303 665 | 5 921 | — | CPU, ≈ 12 h | +# | **TS-DDR strict DE (warm-continued)** | LSTM + reachable | 303 936 | 6 119 | 0.0% | H200 GPU, 22 h + 11 h | +# | **TS-DDR strict subproblems (two-stage)** | LSTM + reachable | 304 027 | 6 114 | 0.0% | CPU (Ipopt), ≈ 2 d + 21 h | +# | **TS-DDR strict DE (from scratch, SDDP-matched budget)** | LSTM + reachable | 304 460 | 6 052 | 0.0% | H200 GPU, 12.3 h | +# +# Because the protocol is paired, differences are measured per scenario and +# their standard errors are two orders of magnitude below the cost std: +# +# - Best TS-DDR (warm-continued GPU DE) − SDDP: **+270 ± 19** +# (``t \\approx 14.5``); TS-DDR dispatches cheaper than SDDP on **17.4%** +# of scenarios. +# - From-scratch GPU DE at SDDP's own ≈ 12 h training budget − SDDP: +# **+794 ± 19** (win rate 4.2%). +# +# The verdict on this benchmark is symmetric and honest: SDDP retains a +# statistically significant but operationally tiny advantage — **0.09%** +# against the best TS-DDR policy — while TS-DDR guarantees zero target +# violations by construction, requires no penalty tuning, and reaches +# within 0.26% of SDDP from scratch in the same wall-clock budget on one +# GPU. (On the inventory-control example, the same strict construction +# beats its SDDP baseline outright; see that example's page.) +# +# ### Cost distributions on the paired scenario set +# +# Means compress the comparison; the full per-scenario distributions show +# it. The top panel overlays each method's cost density over the shared +# paired scenarios (kernel density estimates of the per-scenario rollout +# costs; dashed lines mark means). The bottom panel is the statistically +# decisive view for paired data: the density of the **per-scenario paired +# difference** ``\\text{method} - \\text{SDDP}``. Pairing removes the +# common between-scenario variance (per-scenario cost std ≈ 5.6 K vs +# paired-difference std ≈ 0.5 K), so this panel resolves differences an +# order of magnitude smaller than the raw distributions can — probability +# mass left of the zero line is exactly the fraction of scenarios where the +# method dispatches cheaper than SDDP (the win rate annotated per series). +# +# ![Paired cost distributions](../assets/hydro_paired_cost_distributions.png) +# +# The figure is regenerated by +# `examples/HydroPowerModels/plot_hydro_paired_distributions.jl` from the +# tagged outputs of the paired evaluation scripts. +# +# The key advantage of strict mode is that it requires **no penalty tuning**: +# the gradient signal from the hard-equality duals avoids soft-deficit or +# reactive-balance penalty schedules. The two-phase protocol above uses only +# ordinary optimizer scheduling and held-out rollout checkpointing; the +# non-strict formulations require careful penalty schedules to achieve +# competitive results. +# +# ### Comparing the two strict implementations fairly +# +# The GPU DE trainer optimizes the same strict target-setting policy class, +# and the audit above shows its inner solve is numerically interchangeable +# with the CPU stage-wise solve. The table therefore compares policies under +# the same 96-stage paired rollout, hard reactive balance, and mof.json +# operating cost. The practical distinction is computational: the CPU path +# solves 126 small Ipopt subproblems sequentially per gradient sample, while +# the GPU path solves one coupled 126-stage MadNLP/cuDSS problem per sample. +# +# The current result is encouraging but not final for the "fastest method" +# ambition: the GPU DE schedule now reaches the best TS-DDR cost, but it has +# not yet delivered a policy that both beats SDDP and does so in less +# wall-clock time. Closing that remaining gap requires improving the DE +# training schedule or gradient estimator, not changing the evaluation +# protocol. + +# ## Appendix: experimental details +# +# Everything below is reproducible from the scripts in +# `examples/HydroPowerModels/` (this package) and the same-named folder of +# DecisionRulesExa.jl; the exact commands are in the two READMEs. +# +# ### A.1 Problem instance +# +# | Quantity | Value | +# |:---------|:------| +# | Network | Bolivia, 28 buses, 34 generators, AC polar (ACPPowerModel) | +# | Hydro plants | 11 (3 cascade links: 1 turbine-only, 2 turbine+spill) | +# | Load scaling | 0.6 × PowerModels.json loads (baked into the mof.json) | +# | Deficit cost | 6,000 per pu (= 60 $/MWh × baseMVA 100) | +# | Inflow record | 47 monthly joint scenarios, tiled cyclically beyond month 47 | +# | Training horizon | 126 stages | +# | Evaluation horizon | 96 stages, 500 seeded paired scenarios | +# | Water-balance factor | K = 0.0036 (flow → volume) | +# +# The 126/96 split is deliberate for **both** SDDP and TS-DDR: training on a +# longer horizon buffers end-of-horizon effects out of the reported window +# (SDDP additionally trains with 30 extra stages for the same reason). +# +# ### A.2 SDDP baseline +# +# | Setting | Value | +# |:--------|:------| +# | Cut generation | SOC-WR relaxation (convex), SDDP.jl | +# | Forward simulation | ACP (nonconvex, true physics) | +# | Subproblem solver | MadNLP (CPU) | +# | Iterations / wall time | 441 / ≈ 12 h | +# | Final lower bound | 378,207 (126 stages) | +# | Paired simulation | `SDDP.Historical` on the shared index matrix | +# +# ### A.3 TS-DDR strict subproblems (CPU) — phase 1 +# +# | Setting | Value | +# |:--------|:------| +# | Policy | `HydroReachablePolicy`: LSTM encoder [128, 128] over inflows; linear sigmoid head over `[encoding; state]` | +# | Stage solver | Ipopt (MUMPS), wrapped in `DiffOpt.diff_optimizer` | +# | Optimizer | Adam, constant ``10^{-3}``, no gradient clipping | +# | Samples per gradient step | 1 | +# | Iteration budget | 8,000 (80 epochs × 100 batches), stalling criterion | +# | Checkpoint selection | training-batch loss | +# | Rollout evaluation | every 25 iterations, 4 fixed held-out scenarios | +# | Penalty schedule | none (strict) | +# | Seeds | 8788 (initial evaluation), 8789 (held-out scenarios) | +# +# ### A.4 TS-DDR strict subproblems (CPU) — phase 2 +# +# | Setting | Value | +# |:--------|:------| +# | Initialization | selected phase-1 checkpoint (`DR_PRETRAINED_MODEL`) | +# | Samples per gradient step | 16 (`DR_NUM_TRAIN_PER_BATCH`) | +# | Learning rate | ``10^{-4} \to 10^{-5}`` cosine decay, 50-iteration linear warmup | +# | Checkpoint selection | held-out rollout objective (`DR_SAVE_METRIC=rollout`) | +# | Held-out set | 24 fixed scenarios (`DR_NUM_EVAL_SCENARIOS`) | +# | Iteration budget | 1,500 | +# +# ### A.5 TS-DDR strict DE (GPU) +# +# | Setting | Value | +# |:--------|:------| +# | Package | DecisionRulesExa.jl (ExaModels + MadNLP + cuDSS) | +# | Hardware | 1 × NVIDIA H200 | +# | Formulation | AC polar, strict targets, reservoir trajectory as parameter | +# | Policy | same `HydroReachablePolicy` family; encoder [128, 128]; head [256, 256] | +# | Phase 1 optimizer | Adam ``10^{-3}``, one sampled trajectory per gradient step | +# | Phase 2 optimizer | Adam ``10^{-4} \to 10^{-5}``, 50-iteration linear warmup, 4 sampled trajectories per gradient step | +# | Solver settings | tol ``10^{-6}``, max_iter 9,000, dual-snapshot warm starts | +# | Checkpoint selection | phase 1: training-batch loss; phase 2: held-out rollout objective | +# | Iteration budget | phase 1: 8,000; phase 2: 800 | +# | Deficit cost (training) | ``10^{5}`` (inert: deficit never activates; evaluation uses the mof.json 6,000) | +# | Reactive balance | phase 2 and evaluation use hard reactive balance (`reactive_deficit_cost = Inf`) | +# | Recurrent state | threaded explicitly across stages (see warning above) | +# +# ### A.6 Software +# +# Julia 1.10/1.12 CI matrix; JuMP + DiffOpt + Ipopt (CPU path); ExaModels + +# MadNLP + cuDSS (GPU path); Flux 0.16 (explicit recurrent-state threading); +# SDDP.jl for the baseline. Package versions are pinned in the respective +# `Manifest.toml` files. +# +# ### A.7 Artifact-to-script map +# +# | Artifact | Produced by | +# |:---------|:------------| +# | Strict CPU training runs | `train_dr_hydropowermodels_strict.jl` (stage 1 & 2 via env recipes) | +# | Strict GPU training runs | `train_hydro_exa_strict.jl` (DecisionRulesExa.jl) | +# | SDDP cuts + log | `sddp/run_sddp.jl` | +# | Paired SDDP simulation | `sddp/eval_paired_sddp.jl` | +# | Paired CPU-policy evaluation | `eval_paired_tsddr.jl` | +# | Paired GPU-policy evaluation | `eval_paired_exa_strict.jl` (DecisionRulesExa.jl) | +# | Cross-package equivalence audit | `dump_paired_policy_reference.jl` + `compare_paired_evals.jl` | +# | Convergence figures | `plot_hydro_strict_convergence.jl` | diff --git a/docs/src/examples/hydro.md b/docs/src/examples/hydro.md index 6c23019..0427427 100644 --- a/docs/src/examples/hydro.md +++ b/docs/src/examples/hydro.md @@ -5,13 +5,25 @@ EditURL = "hydro.jl" # Hydropower Scheduling This example trains target-setting decision rules for the Bolivia -long-term hydrothermal dispatch (LTHD) problem — both **TS-DDR** (deep, -LSTM-based) and **TS-LDR** (linear) — and compares them against an SDDP -baseline with inconsistent formulations. - -The Bolivia system has **10 hydro plants**, **96 monthly stages**, and -**AC power flow** constraints. Inflow uncertainty is sampled from 47 -historical scenarios. +long-term hydrothermal dispatch (LTHD) problem — **TS-DDR** (deep, +LSTM-based) and its linear counterpart **TS-LDR** — and compares them +against an SDDP baseline with inconsistent formulations. The strict +(penalty-free) TS-DDR variant is trained in two independent +implementations that we verify are numerically equivalent: + +1. **Stage-wise subproblems on CPU** (this package: JuMP + DiffOpt + Ipopt), + solving one AC-OPF per stage in closed loop; and +2. **Full-horizon deterministic equivalent on GPU** + ([DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl): + ExaModels + MadNLP/cuDSS), solving one coupled 126-stage NLP per + gradient sample. + +The Bolivia system has **28 buses**, **34 generators**, **11 hydro +plants** (three of them in river cascades), and **AC power flow** +constraints. Policies are trained on a 126-stage horizon and evaluated on +**96 monthly stages**; inflow uncertainty is sampled from **47 historical +joint scenarios** (spatially correlated across plants, tiled cyclically +for horizons beyond the record). ## Overview of the TS-DDR approach @@ -61,15 +73,21 @@ the dual multiplier that provides the gradient signal. ## Gradient computation: the envelope theorem By the envelope theorem, the sensitivity of the optimal value with respect -to the target parameter is simply the dual: +to the target parameter is available in closed form from the multiplier of +the target constraint. Throughout this page we define ``\lambda_t`` as +that sensitivity, ```math -\frac{\partial q_t}{\partial \hat{x}_t} -\;=\; -\lambda_t. +\lambda_t +\;:=\; +\frac{\partial q_t}{\partial \hat{x}_t}, ``` -Combined with backpropagation through the policy network, the full gradient -of the expected cost is: +i.e. the constraint multiplier reported by the solver, mapped through the +sign convention of the modeling layer (both implementations extract it this +way, and the two extractions are verified to agree numerically — see the +equivalence audit in the Results section). Combined with backpropagation +through the policy network, the full gradient of the expected cost is: ```math \nabla_\theta \mathbb{E}[Q] @@ -116,7 +134,8 @@ diff_optimizer = () -> DiffOpt.diff_optimizer( optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0, "linear_solver" => "mumps") ) -subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( "bolivia", "ACPPowerModel.mof.json"; num_stages=96, @@ -131,8 +150,8 @@ The policy is a [`StateConditionedPolicy`](@ref) with two components: 1. **Encoder** — a stack of LSTM cells that processes only the uncertainty (inflow) sequence, capturing temporal dependencies across stages. -2. **Combiner** — a Dense layer that merges the encoded uncertainty with the - previous state to produce the next target. +2. **Combiner** — a feed-forward target head that merges the encoded + uncertainty with the previous state to produce the next target. At each stage the policy receives ``[w_t;\; x_{t-1}]`` and outputs target reservoir volumes ``\hat{x}_t``: @@ -140,8 +159,8 @@ target reservoir volumes ``\hat{x}_t``: ``` ┌─────────┐ ┌────────────────┐ ┌──────────────┐ │ w_t │─────▶│ LSTM encoder │─────▶│ │ - └─────────┘ └────────────────┘ │ Dense │──▶ x̂_t - ┌─────────┐ │ combiner │ + └─────────┘ └────────────────┘ │ feed-forward │──▶ x̂_t + ┌─────────┐ │ head │ │ x_{t-1} │─────────────────────────────▶│ │ └─────────┘ └──────────────┘ ``` @@ -150,10 +169,31 @@ The LSTM carries hidden state across stages, giving the policy memory of past inflows. The activation is `sigmoid` (bounding outputs to ``[0,1]``, which is then scaled by the feasibility mapping). +The `layers` argument controls recurrence over inflows only. The optional +`combiner_layers` argument adds hidden layers to the nonrecurrent +state-conditioned target head. This is the preferred way to make TS-DDR +nonlinear in the current reservoir state without adding recurrence over state. + +!!! warning "Recurrent state must be threaded explicitly (Flux ≥ 0.16)" + Since Flux 0.16, calling an `LSTM` layer restarts it from + `initialstates` on **every call**, and `Flux.reset!` is a deprecated + no-op. A per-stage policy loop written naively against this API is + silently *memoryless* — the "LSTM" degenerates to a per-stage MLP over + the current inflow. Both packages therefore thread the hidden state + explicitly across stages inside their policy types (and implement a + real `Flux.reset!` for scenario boundaries). The effect is large: on + this example, evaluating the same trained weights with and without + state threading changes the 100-scenario mean cost by **19%** + (361,056 vs 303,631). If you build custom recurrent policies on this + framework, verify statefulness with the pattern used in the package + test suites: the same input fed twice without `reset!` must produce + different outputs. + ```julia models = state_conditioned_policy( num_uncertainties, num_hydro, num_hydro, [128, 128]; activation=sigmoid, encoder_type=Flux.LSTM, + combiner_layers=[128, 128], ) ``` @@ -246,7 +286,7 @@ det_equivalent, uncertainty_samples_det = DecisionRules.deterministic_equivalent train_multistage( models, initial_state, det_equivalent, - state_params_in, state_params_out, uncertainty_samples; + state_params_in, state_params_out, uncertainty_samples_det; num_batches=4000, optimizer=Flux.Adam(), penalty_schedule=[(1,100,0.1), (101,210,1.0), (211,300,10.0), (301,4000,30.0)], ) @@ -278,25 +318,36 @@ as input to the next stage. ### Gradient chain -The gradient must account for how the realized state at stage ``t`` -depends on the targets at all earlier stages. By the chain rule: +The gradient must account for two coupled feedback paths: the realized +state at stage ``t`` depends on the targets at all earlier stages, and the +*policy input* at stage ``t`` includes that realized state. Writing +``x_t = X_t(x_{t-1}, \hat{x}_t)`` for the realized-state map of the stage +solve, the exact total-derivative recursion is ```math -\frac{\partial Q}{\partial \hat{x}_t} -\;=\; -\lambda_t -+ \sum_{k>t} - \frac{\partial q_k}{\partial x_{k-1}} - \cdot \prod_{j=t+1}^{k-1} - \frac{\partial x_j}{\partial x_{j-1}} - \cdot \frac{\partial x_t}{\partial \hat{x}_t}. +\frac{dQ}{d\theta} += \sum_{t=1}^{T} \left[ + \frac{\partial q_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d\theta} + + \frac{\partial q_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta} + \right], +\qquad +\frac{d x_t}{d\theta} += \frac{\partial X_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta} ++ \frac{\partial X_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d\theta}, +``` + +```math +\frac{d \hat{x}_t}{d\theta} += \nabla_\theta \pi_\theta ++ \frac{\partial \pi_\theta}{\partial x_{t-1}} \frac{d x_{t-1}}{d\theta}. ``` -In practice, automatic differentiation (Zygote + ChainRules `rrule`s -defined on each stage solve) handles this chain automatically. +Reverse-mode automatic differentiation (Zygote + ChainRules `rrule`s +defined on each stage solve) computes exactly this chain, including the +policy-feedback term ``\partial \pi_\theta / \partial x_{t-1}``. The `rrule` for each stage solve reads the dual ``\lambda_t`` for the target constraint and uses DiffOpt's implicit differentiation for the -state-transition sensitivities. +state-transition sensitivities ``\partial X_t / \partial \cdot``. **Advantages**: closed-loop — the policy sees realized states, matching deployment semantics. Each solve is small (single-stage AC-OPF). @@ -386,9 +437,231 @@ train_multiple_shooting( ) ``` -## Penalty annealing +## Training pipeline 4: Strict subproblems with reachable policy + +The three formulations above use a **slack penalty** ``C_\delta \|\delta_t\|`` +to handle the gap between the policy's targets and the feasible set. While +effective, the penalty introduces a hyperparameter and can corrupt the gradient +signal: at high ``C_\delta``, the dual ``\lambda_t`` reflects "reduce the +slack" rather than "improve economic dispatch." + +**Strict mode** eliminates the penalty entirely by enforcing a **hard equality** +between the target and the realized state: + +```math +x_t = \hat{x}_t \quad :\lambda_t \qquad \text{(no slack, no } \delta_t \text{)} +``` + +The dual ``\lambda_t`` is then the **pure shadow price** +``\partial q_t / \partial \hat{x}_t``: the economic value of a marginal change +in the target, free of any penalty noise. + +### Feasibility guarantee: HydroReachablePolicy + +Removing the slack requires that every target produced by the policy be +**physically achievable**. For hydro scheduling, this means the target volume +must lie within the one-stage reachable set — the range of volumes achievable +from the current state ``v_{r,t-1}`` by choosing turbine flow ``q_r`` and +spillage ``s_r`` within their physical bounds. + +#### Per-unit reachable bounds + +The water balance for reservoir ``r`` at stage ``t`` is + +```math +v_{r,t} = v_{r,t-1} + K\, w_{r,t} - K\, q_{r,t} - K\, s_{r,t} + + \sum_{u \in \mathcal{U}_r} K\, q_{u,t} + + \sum_{u \in \mathcal{S}_r} K\, s_{u,t}, +``` + +where ``K`` is the water-balance conversion factor extracted from the model, +``w_{r,t}`` is the inflow, +``q_{r,t}`` is the turbined flow, ``s_{r,t}`` is the spillage, +``\mathcal{U}_r`` is the set of upstream units connected by turbine flow, +and ``\mathcal{S}_r`` is the set connected by spillage. + +The reachable bounds for unit ``r`` (ignoring cascade interactions) are: + +```math +\ell_{r,t} = \max\bigl(\underline{v}_r,\; + v_{r,t-1} + K\, w_{r,t} - K\,\bar{q}_r - K\,\bar{s}_r + + K \sum_{u \in \mathcal{U}_r} \underline{q}_u\bigr), +``` + +```math +u_{r,t} = \min\bigl(\bar{v}_r,\; + v_{r,t-1} + K\, w_{r,t} - K\,\underline{q}_r + + K \sum_{u \in \mathcal{U}_r} \bar{q}_u + + K \sum_{u \in \mathcal{S}_r} \bar{s}_u\bigr). +``` + +These bounds assume worst-case upstream contributions (maximum turbine/spill +capacity). The [`HydroReachablePolicy`] wraps the same LSTM uncertainty +encoder plus feed-forward state-conditioned target head as +[`StateConditionedPolicy`](@ref) but uses a **sigmoid** activation to bound +the output to this reachable interval: + +```math +\hat{v}_{r,t} = \ell_{r,t} + (u_{r,t} - \ell_{r,t}) \cdot \sigma(z_{r,t}). +``` -The target penalty ``C_\delta`` controls the trade-off between following +#### Cascade-aware clamping + +The per-unit upper bound ``u_{r,t}`` uses worst-case upstream contributions +(``K \bar{q}_u``, ``K \bar{s}_u``). When an upstream unit ``u`` stores water +(its target ``\hat{v}_{u,t}`` is high), the actual upstream release + +```math +R_u = K\, w_{u,t} + v_{u,t-1} - \hat{v}_{u,t} +``` + +can be much less than the assumed maximum. For cascaded systems, this means +the downstream target may exceed the true reachable set, causing infeasibility +in strict mode (no slack to absorb the gap). + +After computing the initial sigmoid targets for all units, the policy applies +a **cascade clamping** step. For each upstream→downstream connection: + +- **Turn + spill** connection: the full release reaches downstream, + so ``\text{max\_contrib} = \max(0,\, R_u)``. +- **Turn-only** connection: only turbined flow reaches downstream, + so ``\text{max\_contrib} = \min(K\,\bar{q}_u,\, \max(0,\, R_u))``. + +The downstream target is then clamped: + +```math +\hat{v}_{d,t} \;\le\; v_{d,t-1} + K\, w_{d,t} + - K\,\underline{q}_d + \text{max\_contrib}. +``` + +This clamping is `@non_differentiable` — gradient flows through ``\sigma`` +for unclamped targets, and is zero for clamped ones (correct projected-gradient +signal). + +### Setup + +Building strict subproblems requires only the `strict=true` flag: + +```julia +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( + case_dir, "ACPPowerModel.mof.json"; + num_stages=126, optimizer=diff_optimizer, + strict=true, # ← no deficit, hard equality targets +) +``` + +The reachable policy is constructed from the hydro metadata returned by +`build_hydropowermodels`: + +```julia +models = hydro_reachable_policy(hydro_meta, [128, 128]) +models_with_deep_state_head = hydro_reachable_policy( + hydro_meta, + [128, 128]; + combiner_layers=[256, 256], +) +``` + +Training uses the same `train_multistage` with no penalty schedule: + +```julia +train_multistage( + models, initial_state, subproblems, + state_params_in, state_params_out, uncertainty_samples; + num_batches=8000, optimizer=Flux.Adam(), + penalty_schedule=nothing, # ← no penalty to tune +) +``` + +!!! tip "Out-of-the-box convergence" + Strict mode with `HydroReachablePolicy` requires **no penalty tuning**, + no annealing schedule, and no hyperparameter search. The clean gradient + signal allows the optimizer to directly minimize operational cost. + +## Training pipeline 5: Strict deterministic equivalent on GPU + +The strict formulation unlocks a second, much faster training route, +implemented in the companion package +[DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl): +the **full-horizon deterministic equivalent solved on GPU** with +[ExaModels.jl](https://github.com/exanauts/ExaModels.jl) (SIMD-friendly +algebraic modeling) and [MadNLP.jl](https://github.com/MadNLP/MadNLP.jl) +with the cuDSS sparse linear solver. + +### Why strict mode makes the regular DE safe + +A regular (non-embedded) DE is normally *open-loop*: the policy produces +all targets ``\hat{x}_{1:T}`` before the coupled solve, seeing its own +previous target instead of a realized state. With an arbitrary policy this +can render a strict DE infeasible. The reachable policy restores safety +**by induction**: roll targets out as ``\hat{x}_0 = x_0`` and +``\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1})`` with every target inside the +one-stage reachable set of its input state. Stage 1 is then feasible from +the true initial state; and if stages ``1..t`` are feasible, the strict +equalities force ``x_t = \hat{x}_t``, so stage ``t{+}1`` starts exactly at +the state the policy planned from — making ``\hat{x}_{t+1}`` feasible too. +Strict feasibility *removes* the open-loop/closed-loop gap: the realized +state path must equal the reachable target path, so training-time DE +solves and deployment-time stage-wise rollouts traverse identical +trajectories (we verify this numerically below). + +### Reservoir volumes become parameters + +Because strict equalities pin every reservoir volume to its target, the +volume trajectory is **data, not a decision**, for the inner solver. The +GPU formulation exploits this: the reservoir trajectory enters the NLP as +a parameter vector, eliminating ``(T{+}1) \cdot n_{\text{hydro}}`` +variables and all slack variables from the KKT system. The targets then +appear only on the right-hand side of the ``T \cdot n_{\text{hydro}}`` +water-balance rows. With ``\mu_t`` the multiplier of the stage-``t`` +water-balance row, the envelope gradient follows from the two rows each +target touches (``+1`` on ``v_{t+1}`` in row ``t``, ``-1`` on ``v_t`` in +row ``t{+}1``): + +```math +\frac{\partial Q}{\partial \hat{x}_t} = \mu_t - \mu_{t+1} +\quad (t < T), +\qquad +\frac{\partial Q}{\partial \hat{x}_T} = \mu_T . +``` + +### What the GPU buys + +Each gradient sample requires one coupled 126-stage AC NLP solve. On an +NVIDIA H200, MadNLP + cuDSS solves it fast enough that a full 8,000-solve +training run completes in roughly a day — and, more importantly, the +training loss enters the SDDP-forward-cost envelope within the **first +hours** (see the wall-clock convergence figure in the Results). Solver +state is reused across solves with a dual-snapshot warm-start scheme that +prevents one failed solve from corrupting subsequent ones. + +```julia +## In DecisionRulesExa.jl (see its examples/HydroPowerModels): +de = build_hydro_de(power_data, hydro_data, 126; + formulation=:ac_polar, strict_targets=true, + backend=CUDABackend()) +policy = hydro_reachable_policy(hydro_data, [128, 128]; + combiner_layers=[128, 128]) +train_tsddr(policy, x0, de, de.p_x0, de.p_target, de.p_inflow, sampler; + num_batches=8000, madnlp_kwargs=(tol=1e-6,)) +``` + +!!! note "Cross-package equivalence audit" + The two strict implementations are verified against each other on this + exact case: with identical weights, data, and scenarios, (i) the two + policy implementations produce bit-identical targets, (ii) the + stage-wise CPU rollout (Ipopt) and the strict full-horizon DE (MadNLP) + agree per scenario to ``10^{-9}`` relative cost, and (iii) the two + packages' 100-scenario evaluations of the same checkpoint agree to + 0.001% (303,631 vs 303,635). The audit scripts + (`dump_paired_policy_reference.jl`, `eval_paired_exa_strict.jl`, + `compare_paired_evals.jl`) ship with the packages and re-run on demand. + +## Penalty annealing (non-strict formulations) + +For the non-strict formulations (DE, stage-wise, multiple shooting), the +target penalty ``C_\delta`` controls the trade-off between following the policy's targets and minimizing operational cost. DecisionRules supports a **penalty annealing schedule** that ramps the penalty multiplier during training: @@ -403,6 +676,10 @@ during training: This is activated with `penalty_schedule=:default_annealed` or by passing an explicit list of `(start_iter, end_iter, multiplier)` tuples. +The penalty schedule must be carefully tuned per problem. In contrast, +strict mode bypasses this entirely when the problem admits an always-feasible +policy (see above). + ## Evaluation After training, we evaluate the policy using stage-wise rollout on held-out @@ -414,7 +691,26 @@ scenarios. Two modes: The **target-violation share** measures how much cost comes from the slack penalty rather than actual operations — it should be small (``\le 5\%``) for -a well-trained policy. +a well-trained policy. In strict mode, the violation share is always **zero** +by construction. + +### Paired evaluation protocol + +All headline numbers in the Results section come from a **paired** +protocol: a fixed 126×100 matrix of scenario indices +(`bolivia/paired_scenario_indices.csv`) selects the *same* joint inflow +realization at every stage for every method. The SDDP policy is simulated +with `SDDP.Historical` on those indices (`sddp/eval_paired_sddp.jl`); every +decision-rule checkpoint is rolled out stage-wise on the identical +trajectories (`eval_paired_tsddr.jl` here; `eval_paired_exa_strict.jl` in +the GPU package). Pairing removes the between-scenario variance +(per-scenario cost std ≈ 5,600) from the *comparison*: the standard error +of the paired mean difference is ≈ 50, roughly two orders of magnitude +tighter than comparing unpaired means. Data-file identity across +implementations is enforced (byte-identical `inflows.csv`, `hydro.json`, +`PowerModels.json`), and the stage-index-to-inflow-row mapping (cyclic +tiling of the 47-row record) is asserted programmatically inside the +cross-package evaluation. ```julia rollout_eval = RolloutEvaluation( @@ -437,52 +733,246 @@ pass (simulation). This is a pragmatic approach when the true problem convex relaxation approximates the value function while the forward pass evaluates under the true physics. -The SDDP policy is trained for up to 2000 iterations and the learned -cuts are saved to a JSON file, which can be loaded to simulate the -policy under the ACP formulation. +The learned cuts are saved to a JSON file, which can be loaded to +simulate the policy under the ACP formulation. On this case the SDDP +training ran **441 iterations in ≈ 12 hours** (CPU, MadNLP subproblem +solver), converging its lower bound to **378,207**; the forward-pass +(ACP) simulation cost stabilizes around **380 K** on the 126-stage +horizon. These two numbers frame everything below: no policy can have +expected 126-stage cost below the bound, and SDDP's own policy sits +roughly 0.5% above it. ## Results -The plots below compare the TS-DDR and TS-LDR training formulations and -the SDDP baseline on the Bolivia case. Training curves, out-of-sample -cost distributions, reservoir volume trajectories, and thermal generation -profiles are shown. - -### Training convergence (TS-DDR methods) - -![Training convergence](../assets/hydro_training_convergence.png) - -### Out-of-sample cost (TS-DDR methods) - -![Out-of-sample cost comparison](../assets/hydro_cost_comparison.png) - -### Target-violation share (TS-DDR methods) - -![Violation share](../assets/hydro_violation_share.png) - -### Reservoir volume comparison (all methods) - -![Volume comparison](../assets/hydro_volume_comparison.png) - -### Thermal generation comparison (all methods) - -![Generation comparison](../assets/hydro_generation_comparison.png) - -### Summary - -| Method | Policy | Mean Cost | Std | N | -|:-------|:------:|----------:|----:|--:| -| TS-DDR (DE) | LSTM | 325 540 | 6 266 | 100 | -| TS-DDR (DE, anneal) | LSTM | 324 445 | 6 134 | 100 | -| TS-DDR (shooting w=12) | LSTM | 323 289 | 5 593 | 100 | -| TS-DDR (shooting w=12, anneal) | LSTM | 322 812 | 6 081 | 100 | -| TS-DDR (stage-wise, anneal) | LSTM | 321 543 | 6 214 | 100 | -| SDDP (SOC-WR / ACP) | cuts | 303 684 | — | 100 | - -All three TS-DDR methods with penalty annealing converge to similar -costs (321K–325K). SDDP trains on 126 stages (96 + 30 margin). - -!!! note "Preliminary results" - These numbers reflect the current default training scripts. - They will be updated as the package evolves. +We evaluate the two **strict** TS-DDR implementations — stage-wise +subproblems (CPU) and the full-horizon GPU deterministic equivalent — +against the SDDP baseline on the Bolivia case with AC power flow (11 +hydro plants, 47 inflow scenarios). The non-strict formulations (DE, +stage-wise, multiple shooting) are available through the same API but +require penalty scheduling and careful tuning; strict mode eliminates +this entirely, so it is the configuration we benchmark. + +Three comparison disciplines keep the results honest: + +1. **Same scenarios** — every number below uses the paired 100-scenario + protocol described above. +2. **Same physics at evaluation** — all rollouts solve the identical + ACP stage problem (hard reactive balance, mof.json objective); the + cross-package audit pins the implementations to each other at + solver-tolerance level. +3. **Documented training recipes** — every trained artifact corresponds + to a from-scratch-reproducible configuration listed in the appendix + and in the example READMEs (no ad-hoc checkpoints). + +### Understanding the metrics + +- **SDDP lower bound** (126 stages): the expected-cost lower bound from + the convex SOC-WR relaxation. This is a *relaxation bound* — it cannot + be beaten by any feasible policy. For Bolivia, it converges to + approximately **378 207**. + +- **SDDP forward-pass cost** (126 stages): the simulation cost of the + SDDP policy evaluated under the true AC formulation during the forward + pass. This is the 126-stage operational cost of the SDDP policy. + +- **Simulation cost** (96 stages): the operational cost obtained by + rolling out a policy under AC power flow on 100 held-out inflow + scenarios. This is the primary metric for policy quality. SDDP's + 96-stage simulation cost is **302 674** (mean over 100 paired scenarios, + std 5 572; re-evaluated 2026-07-03). + +During TS-DDR training, the logged loss is a *training-batch average* +over a small number of sampled scenarios (typically 1–4 per batch). +This number is noisy and not directly comparable to the 100-scenario +simulation cost. + +### Training convergence against wall-clock time (126 stages) + +The figure below shows all 126-stage training metrics against **wall-clock +time** (log scale), which is the axis that exposes the computational +trade-off between the methods: + +![Training convergence vs wall-clock time](../assets/hydro_training_convergence_by_time.png) + +- **SDDP lower bound** (SOC-WR relaxation): converges to ~378.2 K over + ≈ 12 hours (441 iterations). Dashed line — no policy can beat it. +- **SDDP forward-pass cost**: the 126-stage simulation cost of the SDDP + policy under the true AC formulation, ~380 K after convergence. +- **TS-DDR strict subproblems (CPU)**: the final two-phase stage-wise + schedule. Phase 1 uses single-sample gradients and training-loss + checkpointing; phase 2 switches to larger batches, decayed learning + rate, and held-out rollout checkpointing. The wall-clock axis counts + both phases. +- **TS-DDR strict DE (GPU)**: the analogous two-phase full-horizon + deterministic-equivalent schedule on an H200. One coupled 126-stage + solve provides each gradient sample; the second phase uses the same + rollout-selected, lower-variance training discipline and hard reactive + balance used at evaluation. + +In cumulative wall-clock time, the plotted two-phase schedules end at +approximately **21.3 h** for strict subproblems and **24.2 h** for strict +DE. The GPU DE curve reaches the lowest TS-DDR training objective, but this +run is not a wall-clock win over SDDP. + +The plot is regenerated by +`examples/HydroPowerModels/plot_hydro_strict_convergence.jl`, which parses +the SDDP training log and pulls the policy-training histories from the +experiment tracker. + +### Two-phase TS-DDR training protocol + +The reported TS-DDR policies use a single, cumulative two-phase schedule +(exact commands in the example README): + +1. **Coarse training**: train from scratch with small gradient batches and + checkpoint on training loss. This phase rapidly enters the SDDP cost + envelope but the per-batch objective remains too noisy to rank policies + that differ by a few hundred cost units. +2. **Selection phase**: continue from the selected checkpoint with larger + gradient batches, a warmed-up and cosine-decayed learning rate, and + checkpoint selection on a fixed held-out rollout objective. This phase + optimizes the same quantity used for the paired evaluation below. + +The plotted wall-clock time is cumulative across both phases. + +### 96-stage out-of-sample rollout cost (paired, 100 scenarios) + +The primary evaluation metric is the **96-stage simulation cost** — +total dispatch cost under AC power flow on the 100 paired inflow +trajectories. + +| Method | Policy | Mean Cost | Std | Target violations | Training hardware | +|:-------|:------:|----------:|----:|:-----------------:|:------------------| +| SDDP (SOC-WR / ACP) | cuts | 302 674 | 5 572 | — | CPU, ≈ 12 h | +| **TS-DDR strict DE** | LSTM + reachable | 302 910 | 5 771 | 0.0% | H200 GPU, ≈ 24 h | +| **TS-DDR strict subproblems** | LSTM + reachable | 303 010 | 5 738 | 0.0% | CPU (Ipopt), ≈ 21 h | + +The strict DE policy is the best TS-DDR policy in this run: it is within +**237 cost units** of SDDP on the paired mean (0.08%) and slightly improves +on the stage-wise subproblem policy. The strict subproblem policy is within +**337 cost units** of SDDP (0.11%); its paired SDDP − TS-DDR difference is +**−337 ± 396** across scenarios (standard error ≈ 40). These are small +operational gaps, but they are still gaps: on this benchmark the current +TS-DDR schedules match the SDDP cost envelope closely but do not yet beat +the SDDP policy. + +The key advantage of strict mode is that it requires **no penalty tuning**: +the gradient signal from the hard-equality duals avoids soft-deficit or +reactive-balance penalty schedules. The two-phase protocol above uses only +ordinary optimizer scheduling and held-out rollout checkpointing; the +non-strict formulations require careful penalty schedules to achieve +competitive results. + +### Comparing the two strict implementations fairly + +The GPU DE trainer optimizes the same strict target-setting policy class, +and the audit above shows its inner solve is numerically interchangeable +with the CPU stage-wise solve. The table therefore compares policies under +the same 96-stage paired rollout, hard reactive balance, and mof.json +operating cost. The practical distinction is computational: the CPU path +solves 126 small Ipopt subproblems sequentially per gradient sample, while +the GPU path solves one coupled 126-stage MadNLP/cuDSS problem per sample. + +The current result is encouraging but not final for the "fastest method" +ambition: the GPU DE schedule now reaches the best TS-DDR cost, but it has +not yet delivered a policy that both beats SDDP and does so in less +wall-clock time. Closing that remaining gap requires improving the DE +training schedule or gradient estimator, not changing the evaluation +protocol. + +## Appendix: experimental details + +Everything below is reproducible from the scripts in +`examples/HydroPowerModels/` (this package) and the same-named folder of +DecisionRulesExa.jl; the exact commands are in the two READMEs. + +### A.1 Problem instance + +| Quantity | Value | +|:---------|:------| +| Network | Bolivia, 28 buses, 34 generators, AC polar (ACPPowerModel) | +| Hydro plants | 11 (3 cascade links: 1 turbine-only, 2 turbine+spill) | +| Load scaling | 0.6 × PowerModels.json loads (baked into the mof.json) | +| Deficit cost | 6,000 per pu (= 60 $/MWh × baseMVA 100) | +| Inflow record | 47 monthly joint scenarios, tiled cyclically beyond month 47 | +| Training horizon | 126 stages | +| Evaluation horizon | 96 stages, 100 paired scenarios | +| Water-balance factor | K = 0.0036 (flow → volume) | + +The 126/96 split is deliberate for **both** SDDP and TS-DDR: training on a +longer horizon buffers end-of-horizon effects out of the reported window +(SDDP additionally trains with 30 extra stages for the same reason). + +### A.2 SDDP baseline + +| Setting | Value | +|:--------|:------| +| Cut generation | SOC-WR relaxation (convex), SDDP.jl | +| Forward simulation | ACP (nonconvex, true physics) | +| Subproblem solver | MadNLP (CPU) | +| Iterations / wall time | 441 / ≈ 12 h | +| Final lower bound | 378,207 (126 stages) | +| Paired simulation | `SDDP.Historical` on the shared index matrix | + +### A.3 TS-DDR strict subproblems (CPU) — phase 1 + +| Setting | Value | +|:--------|:------| +| Policy | `HydroReachablePolicy`: LSTM encoder [128, 128] over inflows; linear sigmoid head over `[encoding; state]` | +| Stage solver | Ipopt (MUMPS), wrapped in `DiffOpt.diff_optimizer` | +| Optimizer | Adam, constant ``10^{-3}``, no gradient clipping | +| Samples per gradient step | 1 | +| Iteration budget | 8,000 (80 epochs × 100 batches), stalling criterion | +| Checkpoint selection | training-batch loss | +| Rollout evaluation | every 25 iterations, 4 fixed held-out scenarios | +| Penalty schedule | none (strict) | +| Seeds | 8788 (initial evaluation), 8789 (held-out scenarios) | + +### A.4 TS-DDR strict subproblems (CPU) — phase 2 + +| Setting | Value | +|:--------|:------| +| Initialization | selected phase-1 checkpoint (`DR_PRETRAINED_MODEL`) | +| Samples per gradient step | 16 (`DR_NUM_TRAIN_PER_BATCH`) | +| Learning rate | ``10^{-4} \to 10^{-5}`` cosine decay, 50-iteration linear warmup | +| Checkpoint selection | held-out rollout objective (`DR_SAVE_METRIC=rollout`) | +| Held-out set | 24 fixed scenarios (`DR_NUM_EVAL_SCENARIOS`) | +| Iteration budget | 1,500 | + +### A.5 TS-DDR strict DE (GPU) + +| Setting | Value | +|:--------|:------| +| Package | DecisionRulesExa.jl (ExaModels + MadNLP + cuDSS) | +| Hardware | 1 × NVIDIA H200 | +| Formulation | AC polar, strict targets, reservoir trajectory as parameter | +| Policy | same `HydroReachablePolicy` family; encoder [128, 128]; head [256, 256] | +| Phase 1 optimizer | Adam ``10^{-3}``, one sampled trajectory per gradient step | +| Phase 2 optimizer | Adam ``10^{-4} \to 10^{-5}``, 50-iteration linear warmup, 4 sampled trajectories per gradient step | +| Solver settings | tol ``10^{-6}``, max_iter 9,000, dual-snapshot warm starts | +| Checkpoint selection | phase 1: training-batch loss; phase 2: held-out rollout objective | +| Iteration budget | phase 1: 8,000; phase 2: 800 | +| Deficit cost (training) | ``10^{5}`` (inert: deficit never activates; evaluation uses the mof.json 6,000) | +| Reactive balance | phase 2 and evaluation use hard reactive balance (`reactive_deficit_cost = Inf`) | +| Recurrent state | threaded explicitly across stages (see warning above) | + +### A.6 Software + +Julia 1.10/1.12 CI matrix; JuMP + DiffOpt + Ipopt (CPU path); ExaModels + +MadNLP + cuDSS (GPU path); Flux 0.16 (explicit recurrent-state threading); +SDDP.jl for the baseline. Package versions are pinned in the respective +`Manifest.toml` files. + +### A.7 Artifact-to-script map + +| Artifact | Produced by | +|:---------|:------------| +| Strict CPU training runs | `train_dr_hydropowermodels_strict.jl` (stage 1 & 2 via env recipes) | +| Strict GPU training runs | `train_hydro_exa_strict.jl` (DecisionRulesExa.jl) | +| SDDP cuts + log | `sddp/run_sddp.jl` | +| Paired SDDP simulation | `sddp/eval_paired_sddp.jl` | +| Paired CPU-policy evaluation | `eval_paired_tsddr.jl` | +| Paired GPU-policy evaluation | `eval_paired_exa_strict.jl` (DecisionRulesExa.jl) | +| Cross-package equivalence audit | `dump_paired_policy_reference.jl` + `compare_paired_evals.jl` | +| Convergence figures | `plot_hydro_strict_convergence.jl` | diff --git a/docs/src/gpu_acceleration.md b/docs/src/gpu_acceleration.md index 5cbf498..3e41b35 100644 --- a/docs/src/gpu_acceleration.md +++ b/docs/src/gpu_acceleration.md @@ -144,8 +144,9 @@ pattern for a full AC-OPF problem with reservoir dynamics: ```julia # In examples/HydroPowerModels/hydro_power_exa.jl prob = build_hydro_de( - data; - num_stages = 96, + power_data, + hydro_data, + 96; backend = CUDABackend(), formulation = :ac_polar, deficit_cost = 1e5, @@ -236,6 +237,183 @@ target-deficit penalty, and target-violation share. | Stage-wise decomposition | — | JuMP only | | Multiple shooting | — | JuMP only | +## Embedded deterministic equivalent + +The standard `DeterministicEquivalentProblem` treats the policy's target +trajectory as an external parameter: the training loop generates +``\hat{x}_{1:T}`` outside the NLP and passes it in via `set_targets!`. +This is **open-loop** — the policy does not see the realized states +from the coupled solve. + +`EmbeddedDeterministicEquivalentProblem` embeds the policy *inside* +the NLP via a `VectorNonlinearOracle`. The NLP constraint becomes: + +```math +\pi_\theta(w_t,\, x_{t-1}^*) - x_t - \delta_t = 0 \quad \forall t +``` + +where ``x_{t-1}^*`` is the solver's realized state. This is +**closed-loop**: the policy sees realized states from the coupled solve, +and the duals ``\lambda_t`` reflect the joint (policy + physics) system. + +```julia +prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + nu = nu, + nw = nw, + dynamics_eq = my_dynamics, + stage_cost = my_cost, + backend = CUDABackend(), +) + +train_tsddr_embedded( + policy, x0, prob, sampler; + num_batches = 500, + num_train_per_batch = 4, + optimizer = Flux.Adam(1f-3), + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), +) +``` + +The oracle closures capture the policy **by reference** — updating Flux +parameters between solves automatically changes the NLP without +rebuilding it. Use `invalidate_policy_cache!` if your oracle caches +policy-dependent intermediates. + +### Strict reachable targets + +When the policy is guaranteed to produce feasible targets (e.g., via a +reachable-set mapping), the slack variables ``\delta_t`` can be removed +entirely. This is strict mode: target constraints are hard equalities, the duals +are pure shadow prices, and there is no target penalty to tune. + +There are two safe strict paths in the hydro example. + +**Embedded strict DE** evaluates the policy inside the NLP against realized +reservoir decision variables: + +```julia +prob = build_embedded_hydro_de(policy, power_data, hydro_data, T; + formulation = :ac_polar, + strict_targets = true, +) +``` + +Its constraint is simply ``x_t = \pi_\theta(w_t, x_{t-1}^*)``. Because the +policy receives the realized previous state, a reachable-set map can guarantee +that the next strict equality is feasible. The current embedded hydro oracle +hand-codes the reachable policy Jacobian for the default single Dense head; use +`combiner_layers = Int[]` unless that oracle is extended. + +**Regular strict DE** keeps the policy outside the NLP but rolls out targets +from the known initial state: + +```math +\hat{x}_0 = x_0,\qquad +\hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1}). +``` + +If the policy returns ``\hat{x}_t \in R(\hat{x}_{t-1}, w_t)`` at every stage, +then the entire strict DE target trajectory is feasible by induction. The solve +then enforces ``x_t = \hat{x}_t`` for every stage, so the realized state path is +exactly the reachable target path. + +For cascaded hydro systems (upstream→downstream water flows), the per-unit +reachable set depends on the upstream unit's target at the same stage. The +`HydroReachablePolicy` handles this by computing initial targets for all units +using worst-case upstream bounds (``K \times \bar{q}_u``), then applying a +**cascade clamping** step that adjusts downstream targets based on the actual +upstream release implied by the upstream target. For turn+spill connections +the full release is available; for turn-only connections only +``\min(K \bar{q}_u, R_u)`` reaches the downstream unit. This clamping is +`@non_differentiable` — gradient flows through for non-clamped targets, and +is zero for clamped ones (correct projected-gradient signal). + +The ExaModels strict DE also applies cascade clamping in `prepare_solve!` as a +safety net: before each NLP solve, the reservoir parameter values are checked +and clamped stage-by-stage to ensure cascade feasibility. + +The Exa hydro script exposing this path is: + +```bash +DR_ENCODER_LAYERS=128,128 \ +DR_HEAD_LAYERS=128,128 \ +julia --project -t auto train_hydro_exa_strict.jl +``` + +`DR_ENCODER_LAYERS` controls recurrence over inflows. `DR_HEAD_LAYERS` controls +the nonrecurrent state-conditioned target head; it can be used to make the +reachable policy nonlinear in the current reservoir state without adding state +recurrence. + +## Sequential rollout evaluation + +`train_tsddr` solves the full deterministic equivalent in one shot. For +deployment diagnostics, DecisionRulesExa.jl provides `RolloutEvaluation`, which +solves a one-stage ExaModels problem sequentially over a materialized scenario: + +```julia +eval = RolloutEvaluation( + stage_problem, + x0, + eval_scenarios; + horizon = T, + n_uncertainty = nw, + set_stage_parameters! = my_setter!, + realized_state = my_state_reader, + policy_state = :realized, +) +``` + +This mirrors deployment semantics: the policy can be evaluated with the +realized previous state (`policy_state = :realized`) or with its previous target +(`policy_state = :target`) to match regular-DE target-generation semantics. + +## Critic control variate + +`train_tsddr` optionally trains a scalar critic ``C(w, \hat{x})`` that +provides a learned control variate for the dual gradient signal. The +critic does not replace the NLP solve — dual multipliers remain the +primary actor gradient. The critic reduces gradient variance by +subtracting a correlated baseline. + +```julia +critic = Chain(Dense(input_dim => 128, tanh), Dense(128 => 128, tanh), Dense(128 => 1)) + +cv = ScalarCriticControlVariate(critic; + featurizer = default_critic_featurizer, + value_loss_weight = 1.0, + gradient_loss_weight = 0.0, +) + +critic_target = RolloutCriticTarget(stage_problem; + horizon = T, + n_uncertainty = nw, + set_stage_parameters! = my_setter!, + realized_state = my_state_reader, + policy_state = :target, +) + +train_tsddr(policy, x0, prob, prob.p_x0, prob.p_target, prob.p_w, sampler; + control_variate = cv, + critic_training_target = critic_target, + actor_gradient_mode = :control_variate, + critic_cv_weight = 1.0, + critic_optimizer = Flux.Adam(1f-3), +) +``` + +Two actor modes are supported: + +- `:control_variate` — subtracts ``\nabla_{\hat{x}} C`` from the dual + signal and adds it back as a differentiable surrogate. Unbiased when + the critic is exact; reduces variance otherwise. +- `:surrogate` — blends dual and critic actor gradients via explicit + weights (`dual_actor_weight`, `critic_actor_weight`). Useful when raw + duals are noisy, but no longer strictly unbiased. + ## Full example: HydroPowerModels The `examples/HydroPowerModels/` directory in DecisionRulesExa.jl contains @@ -250,4 +428,6 @@ a complete AC-OPF hydrothermal scheduling example for the Bolivia test case - GPU training with parallel MadNLP solves - Warm-start caching to prevent cascade solver failures - Penalty and sample-count annealing schedules +- Embedded closed-loop training with strict reachable targets +- Optional critic control variate with rollout-based value targets - W&B metric logging diff --git a/docs/src/index.md b/docs/src/index.md index 7051de0..15c8c19 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -18,13 +18,20 @@ dynamics and constraints. Lagrange duals and implicit differentiation (via [DiffOpt.jl](https://github.com/jump-dev/DiffOpt.jl)) provide the gradient signal to update the policy end-to-end. -Three training formulations are supported: +Four training formulations are supported: | Formulation | Horizon coupling | Gradient source | |:---|:---|:---| | **Deterministic Equivalent** | Full horizon, one large NLP | Duals on the coupled problem | | **Stage-wise (single shooting)** | Sequential rollout | Duals + DiffOpt per stage | | **Multiple Shooting** | Windowed sub-horizons | DiffOpt per window, continuity penalties | +| **Strict subproblems** | Sequential rollout, no slack | Pure shadow-price duals | + +The **strict subproblems** formulation eliminates the target-slack penalty entirely by +enforcing hard equality constraints between the policy's targets and the realized state. +Combined with a feasibility-guaranteeing policy (e.g., `HydroReachablePolicy` for hydro +scheduling), this produces clean gradient signals with no penalty tuning — the dual +``\lambda_t`` is the pure shadow price, uncontaminated by any regularization term. ## Installation diff --git a/examples/HydroPowerModels/Project.toml b/examples/HydroPowerModels/Project.toml index 135292b..978d7d0 100644 --- a/examples/HydroPowerModels/Project.toml +++ b/examples/HydroPowerModels/Project.toml @@ -1,12 +1,15 @@ [deps] +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DecisionRules = "47937410-f832-486f-8300-12c95b225dfc" DiffOpt = "930fe3bc-9c6b-11ea-2d94-6184641e85e7" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" diff --git a/examples/HydroPowerModels/README.md b/examples/HydroPowerModels/README.md index 7a6d62b..94dfb50 100644 --- a/examples/HydroPowerModels/README.md +++ b/examples/HydroPowerModels/README.md @@ -2,8 +2,8 @@ This directory contains the primary application from the paper: training Two-Stage Deep Decision Rules (TS-DDR) for the Bolivia Long-Term -Hydrothermal Dispatching problem with 10 hydro units, 96 monthly stages, -and AC/SOC/DC power-flow formulations. +Hydrothermal Dispatching problem with 11 hydro units, 96 monthly stages, +and AC power-flow constraints. ## Problem overview @@ -13,136 +13,297 @@ drive reservoir levels; the decision rule maps (inflow history, current state) to reservoir-level targets, and an NLP optimizer dispatches generation to meet those targets at minimum cost. -## Training scripts +## Quick start: strict mode (recommended) + +Strict mode eliminates the target-penalty hyperparameter by enforcing hard +equality constraints between targets and realized states. A +`HydroReachablePolicy` guarantees that every target is physically achievable +from the state used as policy input. + +```bash +# Train with the dedicated strict stage-wise entrypoint. +# Configure horizon with DR_NUM_STAGES, for example: +DR_NUM_STAGES=126 julia --project -t auto train_dr_hydropowermodels_strict.jl + +# Evaluate on 100 held-out scenarios, 96 stages +julia --project -t auto eval_strict_rollout.jl \ + bolivia/ACPPowerModel/models/.jld2 +``` + +The training script logs to Weights & Biases and saves model checkpoints. +The evaluation script writes per-scenario costs, mean reservoir volumes, +and mean thermal generation to CSV files for comparison against SDDP. + +## Reproducing the paper results (strict, Bolivia ACP) + +Both stages below are ordinary configurations of the same entrypoint — no +side scripts. All knobs are environment variables documented in the header of +`train_dr_hydropowermodels_strict.jl`; defaults reproduce stage 1 exactly. + +**Stage 1 — train from scratch** (h126 training horizon, r96 rollout +evaluation, LSTM(128,128) encoder, linear sigmoid head): + +```bash +DR_NUM_STAGES=126 DR_NUM_ROLLOUT_STAGES=96 \ +julia --project -t auto train_dr_hydropowermodels_strict.jl +``` + +**Stage 2 — fine-tune the best stage-1 checkpoint** (variance-reduced +gradients, warmup + cosine-decayed learning rate, checkpoints selected by the +held-out 96-stage rollout objective rather than the noisy training loss): + +```bash +DR_PRETRAINED_MODEL=bolivia/ACPPowerModel/models/.jld2 \ +DR_NUM_STAGES=126 DR_NUM_ROLLOUT_STAGES=96 DR_NUM_EPOCHS=15 \ +DR_NUM_TRAIN_PER_BATCH=16 DR_LR=1e-4 DR_LR_FINAL=1e-5 DR_LR_WARMUP=50 \ +DR_SAVE_METRIC=rollout DR_NUM_EVAL_SCENARIOS=24 \ +julia --project -t auto train_dr_hydropowermodels_strict.jl +``` + +Why stage 2 is configured this way: near convergence the single-sample +training loss (per-scenario std ≈ 5–7k) cannot rank checkpoints that differ by +a few hundred cost units, so selection switches to the deployment metric +(`DR_SAVE_METRIC=rollout`) over 24 fixed scenarios; a fresh Adam takes +full-size steps while its second-moment estimates are still zero, so the +warmup protects the loaded optimum; and the decayed learning rate lets the +policy re-anneal instead of random-walking at a fixed step size. + +**Paired evaluation against SDDP** — both methods realize the same 500 +inflow trajectories, generated from a documented seed +(`paired_scenario_indices` in `load_hydropowermodels.jl`, StableRNG, so the +protocol is reproducible from code alone): + +```bash +julia --project -t auto eval_paired_tsddr.jl \ + bolivia/ACPPowerModel/models/.jld2 +cd sddp && julia --project -t auto eval_paired_sddp.jl +``` + +## Files in this folder + +Every script, what it does, and when to use it. All commands run from this +folder with `julia --project`; configuration is via the env vars documented +in each script's header. + +### Problem construction (used by everything else) + +| File | Purpose | +|------|---------| +| `load_hydropowermodels.jl` | Case loader: builds the stage subproblems from the MOF file + hydro data, reads inflows (tiled cyclically beyond the 47-month record), returns hydro metadata; `strict=true` for hard target equalities. Also defines `paired_scenario_indices`, the seeded paired-evaluation protocol. | +| `hydro_reachable_policy.jl` | `HydroReachablePolicy`: reachable-set targets (sigmoid scaled to physics bounds), cascade-aware clamping, explicit LSTM state threading, optional stage-context inputs; plus checkpoint loaders. | + +### Training + +| File | Purpose | +|------|---------| +| `train_dr_hydropowermodels_strict.jl` | **The strict trainer** (stage-wise, Ipopt). Header documents the two-stage recipe that produces the paper policies (stage-1 from scratch, stage-2 fine-tune) and all env knobs incl. `DR_CONTEXT`. | +| `train_dr_hydropowermodels.jl` | Non-strict deterministic-equivalent training (penalty formulation). | +| `train_dr_hydropowermodels_subproblems.jl` | Non-strict stage-wise (single-shooting) training. | +| `train_dr_hydropowermodels_multipleshooting.jl` | Non-strict multiple-shooting training. | +| `train_ldr_hydropowermodels.jl` | TS-LDR (linear decision rule) baseline training. | +| `train_dr_l2O_supervised.jl` | Supervised learning-to-optimize baseline (needs `gen_inputs_l2O_hydropowermodels.jl` outputs). | + +### Paired evaluation (produces the results tables) + +| File | Purpose | +|------|---------| +| `eval_paired_tsddr.jl` | Paired rollout evaluation of a checkpoint (`ARGS[1]`) on the seeded 500-scenario protocol; writes `paired_costs[_].csv` and trajectory CSVs. | +| `sddp/eval_paired_sddp.jl` | Paired SDDP simulation (`SDDP.Historical`) on the identical seeded scenarios; merges its column into the costs CSV. | +| `dump_paired_policy_reference.jl` | Dumps policy outputs + inflow values (JLD2) consumed by DecisionRulesExa.jl's cross-package equivalence evaluation. | + +### Other evaluation & validation + +| File | Purpose | +|------|---------| +| `eval_strict_rollout.jl` | Standalone 96-stage rollout evaluation of a strict checkpoint (unpaired scenario set). | +| `evaluate_hydro_policies.jl` | Batch evaluation of pre-trained TS-DDR/TS-LDR checkpoints on a fixed scenario set. | +| `eval_jump_de.jl` | Evaluates a policy through the full-horizon JuMP deterministic equivalent. | +| `check_consistent_state_paths.jl` | Consistency check: state paths agree across training formulations. | +| `test_strict_mode.jl` | Strict-mode construction checks (hard equalities, reachable targets). | +| `test_sampling_consistency.jl` | Verifies scenario sampling consistency across code paths. | +| `validate_sddp_vs_jump.jl` | Validates that the JuMP subproblems match SDDP's formulation (the audit that keeps the SDDP comparison fair). | + +### Figures (regenerate the docs assets) + +| File | Purpose | +|------|---------| +| `plot_hydro_strict_convergence.jl` | Wall-clock training-convergence figure (parses `sddp/SDDP.log`, pulls W&B histories). | +| `plot_hydro_paired_distributions.jl` | Paired cost-distribution figure (absolute densities + paired differences vs SDDP). | +| `compare_hydro_results.jl` | W&B comparison plots across training formulations. | -### TS-DDR (Deep Decision Rules — LSTM policy) +### SDDP baseline -| Script | Decomposition | Reference | -|--------|--------------|-----------| -| `train_dr_hydropowermodels.jl` | Deterministic equivalent (GPU-enabled) | Extension §1 | -| `train_dr_hydropowermodels_subproblems.jl` | Stage-wise (single shooting) | Extension §2 | -| `train_dr_hydropowermodels_multipleshooting.jl` | Windowed (multiple shooting) | Extension §3 | +| File | Purpose | +|------|---------| +| `sddp/run_sddp.jl` | Trains the SDDP baseline (SOC-WR cuts, ACP forward pass); writes cuts JSON + `SDDP.log`. | +| `sddp/run_sddp_inconsistent.jl` | SDDP with inconsistent backward/forward formulations (the published baseline configuration). | +| `sddp/simulate_sddp_policy.jl` | Simulates a trained SDDP policy from its cuts file. | +| `sddp/extract_sddp_trajectories.jl` | Extracts state/generation trajectories from SDDP simulations. | -These use a `StateConditionedPolicy` (LSTM encoder + state-conditioned dense -layers, `[128, 128]`, sigmoid activation). +### Utilities -### TS-LDR (Linear Decision Rules — linear policy) +| File | Purpose | +|------|---------| +| `export_subproblem_mof.jl` | Exports the single-stage OPF as `.mof.json` (how `bolivia/*.mof.json` was produced). | +| `gen_inputs_l2O_hydropowermodels.jl` | Generates supervised-learning inputs for the L2O baseline. | -| Script | Decomposition | Reference | -|--------|--------------|-----------| -| `train_ldr_hydropowermodels.jl` | Deterministic equivalent (GPU-enabled) | §3 | +### Data (`bolivia/`) -TS-LDR uses `dense_multilayer_nn` with identity activation — a composition -of linear layers that is equivalent to a single linear map from -(uncertainties, state) to targets. Same training pipeline as TS-DDR; the -only difference is the policy architecture. +`PowerModels.json` + `hydro.json` (network and hydro data), `inflows.csv` +(47 monthly joint inflow scenarios), `ACPPowerModel.mof.json` etc. (exported +stage subproblems), `ACPPowerModel/models/` (trained checkpoints) and +`ACPPowerModel/results/` (evaluation outputs — data, not tracked). + +## Strict Reachability Logic + +Strict mode is normally safe only when the policy can see the current state: + +- Stage-wise subproblems are closed-loop because the next policy call receives + the realized state from the previous solve. +- Embedded deterministic equivalents are closed-loop because the policy is + evaluated inside the NLP against realized state variables. +- A generic regular deterministic equivalent is open-loop after the initial + state, because all targets are computed before the multi-stage solve. + +The hydro reachable policy gives a special regular-DE exception used in the Exa +companion package. If targets are rolled out as +`target[0] = initial_state` and `target[t] = policy(inflow[t], target[t-1])`, +and the policy maps into the one-stage reachable set from its input state, then +all targets are feasible by induction. Strict equalities then force realized +states to equal the reachable target trajectory. + +## Training scripts + +### Strict subproblems (no penalty tuning) + +| Script | Description | +|--------|-------------| +| `train_dr_hydropowermodels_strict.jl` | Dedicated strict stage-wise training entrypoint with `DR_NUM_STAGES`, `DR_ENCODER_LAYERS`, and `DR_HEAD_LAYERS` environment knobs | + +### Non-strict formulations (require penalty tuning) + +| Script | Decomposition | +|--------|--------------| +| `train_dr_hydropowermodels.jl` | Deterministic equivalent (full-horizon coupled NLP) | +| `train_dr_hydropowermodels_subproblems.jl` | Stage-wise (single shooting) | +| `train_dr_hydropowermodels_multipleshooting.jl` | Windowed (multiple shooting, `W=12`) | +| `train_ldr_hydropowermodels.jl` | Linear decision rules (identity activation) | All training scripts share the data loader (`load_hydropowermodels.jl`), log to Weights & Biases, and save the best model to JLD2. +### Key files + +| File | Description | +|------|-------------| +| `load_hydropowermodels.jl` | Builds JuMP stage subproblems from MOF + hydro JSON + inflow CSV; supports `strict=true` for penalty-free targets | +| `hydro_reachable_policy.jl` | `HydroReachablePolicy` — bounds LSTM output to the one-stage reachable set via sigmoid; `load_policy_weights!` for checkpoint loading | + +### Policy architecture knobs + +`HydroReachablePolicy` and `state_conditioned_policy` separate temporal memory +from state conditioning: + +- `layers` / `DR_ENCODER_LAYERS`: recurrent LSTM layers over inflows only. +- `combiner_layers` / `DR_HEAD_LAYERS`: optional nonrecurrent feed-forward + hidden layers over `[encoded_inflow; reservoir_state]`. + +This lets TS-DDR depart from linear decision rules in the state-to-target map +without adding recurrence over the reservoir state. + ### GPU training `train_dr_hydropowermodels.jl` auto-detects CUDA and switches to -MadNLP+CUDSS on GPU when available. Submit via: +MadNLP+CUDSS on GPU when available. From this example environment, run: ```bash cd examples/HydroPowerModels -mkdir -p logs -sbatch run_train_deteq_gpu.sbatch +julia --project train_dr_hydropowermodels.jl ``` -### Penalty schedule - -All training scripts support `:default_annealed` penalty schedules that -gradually increase target-violation penalties during training, improving -convergence on the nonconvex AC formulation. - -### Rollout metrics - -For deterministic-equivalent training, `metrics/loss` is computed on the same -target-state history produced by the policy. The matching held-out metric is -`metrics/rollout_objective_no_deficit`, which now uses `RolloutEvaluation(...; -policy_state=:target)` in `train_dr_hydropowermodels.jl`. - -The same script also logs -`metrics/rollout_realized_objective_no_deficit` with `policy_state=:realized`. -That is the closed-loop deployment diagnostic: each stage passes the optimizer's -realized reservoir state back to the policy. It can be harder than the target-state -metric, especially while the policy is trained through the deterministic equivalent. +If you want to submit through Slurm, use a site-local batch script that loads +Julia and requests the GPU resources required by MadNLP+CUDSS. -All rollout objective metrics exclude the target-slack/deficit penalty term. Track -the paired target-violation share and `metrics/target_penalty_multiplier` to see -whether a low operational objective is coming from feasible targets or from the -policy relying on slack. +For GPU-accelerated training using ExaModels (recommended for large NLPs), +see the companion package +[DecisionRulesExa.jl](https://github.com/LearningToOptimize/DecisionRulesExa.jl). -## Evaluation and baselines +## Evaluation | Script | Purpose | |--------|---------| -| `evaluate_hydro_policies.jl` | Load all trained TS-DDR and TS-LDR models and evaluate on a common out-of-sample scenario set using stage-wise ACP rollout; writes `eval_costs.csv` | +| `eval_strict_rollout.jl` | 100-scenario stage-wise rollout of a strict-mode policy; writes costs, mean volumes, and mean thermal generation to CSVs | +| `evaluate_hydro_policies.jl` | Auto-discovers all saved checkpoints and evaluates them on a common scenario set; writes `eval_costs.csv` | | `eval_jump_de.jl` | Solve the DE with a constant policy and save a reference solution (JLD2) for cross-validation with ExaModels | -| `check_consistent_state_paths.jl` | Verify that stage-wise, deterministic equivalent, and multiple-shooting decompositions produce identical state trajectories under the same policy and inflows | +| `check_consistent_state_paths.jl` | Verify that stage-wise, DE, and multiple-shooting decompositions produce identical state trajectories | + +## File Inventory + +| Path | Purpose | +|---|---| +| `Project.toml` | Hydro example environment | +| `LocalPreferences.toml` | Local solver/GPU preferences when present | +| `README.md` | This guide | +| `load_hydropowermodels.jl` | JuMP/DiffOpt stage-problem builder with optional strict targets | +| `hydro_reachable_policy.jl` | Reachability-guaranteed hydro policy and checkpoint loading helpers | +| `train_dr_hydropowermodels.jl` | Regular deterministic-equivalent TS-DDR training | +| `train_dr_hydropowermodels_subproblems.jl` | Stage-wise TS-DDR training; supports strict mode | +| `train_dr_hydropowermodels_strict.jl` | Dedicated strict stage-wise training script | +| `train_dr_hydropowermodels_multipleshooting.jl` | Multiple-shooting TS-DDR training | +| `train_ldr_hydropowermodels.jl` | Linear decision-rule baseline | +| `train_dr_l2O_supervised.jl` | Supervised learning-to-optimize training utility | +| `eval_strict_rollout.jl` | Rollout evaluation for strict reachable policies | +| `evaluate_hydro_policies.jl` | Batch checkpoint evaluator | +| `eval_paired_tsddr.jl` | Paired TS-DDR scenario evaluation | +| `eval_jump_de.jl` | JuMP deterministic-equivalent reference evaluation | +| `compare_hydro_results.jl` | Result aggregation/comparison helper | +| `check_consistent_state_paths.jl` | State-path consistency diagnostic | +| `validate_sddp_vs_jump.jl` | Validation against SDDP/JuMP references | +| `test_strict_mode.jl` | Strict-mode and reachable-policy checks | +| `test_sampling_consistency.jl` | Scenario sampling consistency tests | +| `export_subproblem_mof.jl` | MOF template export utility | +| `gen_inputs_l2O_hydropowermodels.jl` | Dataset/input generation for supervised L2O experiments | +| `sddp/run_sddp.jl` | Consistent SDDP baseline training | +| `sddp/run_sddp_inconsistent.jl` | SOC backward / AC forward SDDP training | +| `sddp/simulate_sddp_policy.jl` | AC simulation of a trained SDDP policy | +| `sddp/eval_paired_sddp.jl` | Paired SDDP scenario evaluation | +| `sddp/extract_sddp_trajectories.jl` | Extract SDDP trajectory data for diagnostics | +| `bolivia/` | Bolivia case data, MOF templates, figures, and saved models/results | +| `case3/` | Small development case | ## SDDP baselines -These scripts use a dedicated Julia environment in `sddp/`. The inconsistent -SOC-backward/AC-forward baseline uses -[HydroPowerModels.jl](https://github.com/LAMPSPUC/HydroPowerModels.jl), SDDP.jl, -Clarabel for the SOC backward pass, and MadNLP for the AC forward pass. Training -runs log iteration and final simulation metrics to Weights & Biases using the -same keys as the DR runs: `metrics/loss` is the SDDP bound, and -`metrics/rollout_realized_objective_no_deficit` is the SDDP forward-pass -objective. SDDP iterations are logged as `batch` so W&B plots can share the same -x-axis as the DR training runs. Because SDDP solves the forward policy -stage-wise, that forward-pass objective is already the no-target-penalty -objective. +The SDDP baseline uses an **inconsistent formulation**: SOC-WR relaxation +for the backward pass (cut generation) and ACP for the forward pass +(simulation). Scripts are in `sddp/` with a dedicated Julia environment. | Script | Description | |--------|-------------| | `sddp/run_sddp.jl` | Train SDDP with a consistent convex (SOCWRConic) formulation | -| `sddp/run_sddp_inconsistent.jl` | Train SDDP with SOCWRConic backward pass and ACP forward pass | -| `sddp/run_sddp_inconsistent.sbatch` | Submit the SOC-backward/AC-forward run with a 12-hour wall time | -| `sddp/simulate_sddp_policy.jl` | Simulate a pre-trained SDDP policy under ACP and produce comparison plots | +| `sddp/run_sddp_inconsistent.jl` | Train SDDP with SOCWRConic backward / ACP forward | +| `sddp/simulate_sddp_policy.jl` | Simulate a pre-trained SDDP policy under ACP (100 scenarios, 96 stages); writes costs, mean volumes, mean thermal generation to CSVs | -## Learning-to-Optimize (L2O) pipeline +**SDDP 96-stage simulation cost**: 303 684 (mean, 100 scenarios, std 5 453). +**SDDP 126-stage lower bound**: 378 207 (SOC-WR relaxation — not beatable). -| Script | Description | -|--------|-------------| -| `gen_inputs_l2O_hydropowermodels.jl` | Generate input datasets for the L2O supervised pipeline (requires [L2O.jl](https://github.com/andrewrosemberg/L2O.jl)) | -| `train_dr_l2O_supervised.jl` | Supervised pre-training of a decision rule from L2O-generated optimal solutions | - -## Subproblem export (generating `.mof.json` files) +## Data -The training pipeline (`load_hydropowermodels.jl`) reads pre-exported `.mof.json` -subproblem templates rather than depending on HydroPowerModels.jl at training time. -These files already ship with the repository: +- `bolivia/` — Bolivia case: `hydro.json` (11 hydro units), `inflows.csv` + (historical scenarios), `ACPPowerModel.mof.json` / `SOCWRConicPowerModel.mof.json` / + `DCPPowerModel.mof.json` (subproblem templates), `PowerModels.json` (39 buses, + 55 branches, 34 generators) +- `case3/` — Small 3-bus test case for development -``` -bolivia/ACPPowerModel.mof.json -bolivia/SOCWRConicPowerModel.mof.json -bolivia/DCPPowerModel.mof.json -case3/ACPPowerModel.mof.json -``` +## Subproblem export -To regenerate them (e.g. after updating HydroPowerModels data or adding a new -formulation), use `export_subproblem_mof.jl`: +The training pipeline reads pre-exported `.mof.json` subproblem templates. +To regenerate (e.g., after updating data or adding a formulation): ```bash julia export_subproblem_mof.jl bolivia ACPPowerModel -julia export_subproblem_mof.jl bolivia SOCWRConicPowerModel ``` -This builds the full SDDP model via HydroPowerModels.jl, extracts one stage's -subproblem from the policy graph, removes the unnamed slack variable that -HydroPowerModels adds, and writes a clean JuMP `.mof.json` to disk. Requires -HydroPowerModels.jl and a solver (Mosek by default). - -## Data - -- `bolivia/` — Bolivia case: `hydro.json` (10 hydro units), `inflows.csv` (historical scenarios), `ACPPowerModel.mof.json` / `SOCWRConicPowerModel.mof.json` / `DCPPowerModel.mof.json` (subproblem templates) -- `case3/` — Small 3-bus test case for development - ## Dependencies See `Project.toml` in this directory. Key packages: DecisionRules, DiffOpt, -Ipopt+HSL, MadNLP+MadNLPGPU+CUDA (GPU), Flux, JuMP, Wandb. +Ipopt, Flux, JuMP, Wandb. diff --git a/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanGeneration.csv b/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanGeneration.csv index d0b560d..f3c97b9 100644 --- a/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanGeneration.csv +++ b/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanGeneration.csv @@ -1,97 +1,97 @@ -TS-DDR,TS-LDR,SDDP-DCLL,SDDP-SOC -147.0247133419527,299.7450624492452,100.8484585807442,209.86380882088514 -157.58814081813028,225.6675464897668,104.86242775104952,209.53839385906838 -185.43371546547627,180.7943849292062,107.48502989328503,209.52017504171553 -199.3134646164604,199.83572444735833,107.55475020324891,209.4783394659705 -206.61858480532382,231.35910431382658,106.18650990754062,209.52518234738324 -214.3697084291256,189.64814901428522,113.67317158934263,209.6034477390574 -187.9473304345247,241.6103392896273,116.58714528380276,209.65384702850434 -171.46565827999535,187.4542766440176,129.25366035880148,209.51238800566185 -170.2002528391322,233.150743363177,114.76063596215229,209.45259355985954 -182.98142893312487,185.52122552406792,109.93768423864336,209.34525985187597 -202.19901785621693,215.06714163113537,116.98303295186679,209.44039837517906 -193.96716199312775,199.9910790500363,121.32762963754153,209.3672289663951 -193.52811950138334,205.21889018439646,118.81044678071419,209.32163376554064 -195.13814426290378,184.97486998527484,140.68919880246779,209.6614054735699 -226.57039256458856,167.1372529242544,157.98870219622455,210.19631956661485 -248.08925112087996,177.07636222131532,182.13344732161957,210.70698165532528 -233.5593017377245,206.24124265365077,213.25680127887568,211.4401169112003 -212.98635880030474,203.58562602361448,219.65332808686793,211.48105832602775 -208.1182840666282,225.62114837627547,215.53856096654266,211.94334659983778 -205.87501153760303,282.5815380137001,218.69325537706618,211.8784711569352 -218.77529830243512,261.2652681819222,229.70401987822862,211.85037808628837 -219.5621182240177,228.86267450841828,233.80086503295337,212.00653608776338 -215.15355785846822,278.8379771241615,219.9540629779893,211.71653085812682 -225.44578342326264,208.6334543644992,227.25750181666496,211.72866526590903 -216.46699729157837,233.62842636523723,228.98664748977347,211.79095480066385 -210.99855030225916,236.77260170176334,245.24308942520355,211.9396216916039 -200.82120574954234,257.832718296366,266.63858091994916,212.2412014072736 -217.00901852303974,206.80316405135653,287.580525250383,211.96614587275292 -205.60116177559695,247.22715853816413,288.43119570669523,212.0593948656155 -204.48717949362427,242.7626603477798,295.40865231198757,212.13189528353013 -215.01700896149435,264.8798596812393,316.4539931784796,212.48730840572054 -194.39594990932568,253.05048627557306,329.7940889425171,212.170191104127 -212.44246770637898,281.106369677497,327.59228806131625,212.68099082676446 -218.88119237987075,264.11335838123404,330.87683426769047,212.9854607332948 -209.54103152313093,265.672452188459,320.54817802711284,212.90364434114954 -232.50511220115487,242.47928625172216,336.16459151820277,212.98400917882333 -246.55559517109822,250.15790720527644,328.4530145309637,212.90706046302626 -259.1804770347488,241.4580174607301,306.6623006700101,212.78491180246007 -243.65212407401336,216.2392032573588,311.05102021329515,213.50440427222605 -268.3327047244221,244.87238185642838,313.43043807340126,214.71903130733853 -259.25135590920064,212.7607093726518,289.6446657672623,215.71309550713656 -239.7596346992519,211.60973323507133,236.34688350288008,214.21514981807115 -225.1627151901455,224.31226414469103,227.41520899884554,214.14602835221518 -226.300220128706,190.17002437882283,237.60749837482643,213.05570166141047 -246.5897657016228,197.34726823968018,242.14935260524663,214.3397929974062 -200.45447214778832,159.42600322266816,217.0164054870767,212.87589299860522 -177.65036023699594,199.6529727342857,200.46699991530508,210.512610777988 -152.7449270008933,184.7732720253377,215.7682105980341,208.87860648289214 -116.54920043933797,185.46879418205066,111.6333307935832,205.95870346385465 -118.11835868140619,166.0896692911,101.80408752144571,205.70698904522408 -135.40331893488192,159.49003675324482,106.15347696391866,205.4160873909239 -165.39822644830028,182.5511233035246,108.92328402254266,205.59700933991252 -164.63498605950215,157.20350355656583,103.94596690664551,205.58479725277505 -166.92019582010437,188.70842110946302,103.82258613108722,205.77526373955115 -153.73417427187314,161.11055875216002,112.13281054476954,205.88135777815637 -164.778321612267,181.93462654464733,161.52068981849445,206.02649485451556 -160.9396855926067,146.59967315111336,169.88300890301332,205.8871588917278 -161.6813291137583,157.86943128110318,177.59131665764,206.70433459244137 -161.9520731591163,129.58260488127797,172.45550740547597,206.366999884319 -161.5166767553574,192.37516226972488,177.4638905803286,206.38870455127955 -171.23315428685606,169.81265452067095,176.73862367347002,206.249716219982 -203.84334411204006,189.01908062873738,182.19685255887867,206.63455688813036 -243.88748126609443,185.3643782473636,192.64656185901407,205.9797537356252 -255.36863069633904,189.10691135198536,210.14038242475175,206.03555271847185 -275.72231730419776,177.47926018459844,214.11405852350882,206.51953984726723 -231.14746381351682,203.55835450764843,216.8881876993481,206.12018860291096 -239.70981393223852,185.15015653078225,206.20963274723462,206.250132372938 -236.55408488066604,209.9345245251102,228.06281826263384,206.1138429495153 -232.50309001204022,247.09182945796942,232.53663833367236,206.4146828063853 -211.09249114427467,224.13559766576765,229.0267023271387,206.5605843384091 -222.7816613997703,251.3620122443193,232.90850460804563,206.40037604289813 -219.16351862667358,241.65598188514326,228.66399277324248,206.5295113131882 -207.22133488700123,260.0058643392084,266.0049119099325,206.8315496621826 -217.06224817213152,249.65113999794553,228.7078025295844,207.13270789037426 -210.07761177326952,201.43928958603075,257.88491102881073,207.77068168915991 -216.6438147791045,209.70322997573356,256.3000900779569,208.16087070287398 -201.41487156994177,254.6147418912871,260.0864123226794,208.40830707759466 -207.97135419475285,233.44163279290632,245.1593291381967,207.38456438434162 -211.26689971022438,206.86923643700797,266.2262958784643,206.9039133561306 -240.1582296197083,249.82887487767167,286.04078934122384,207.77917011114187 -233.02195021621122,218.9756433172769,302.4507102096619,207.558738022174 -237.29598744217932,254.3292021981702,310.8026658382449,208.03806284333123 -262.7454036444481,254.22526995036333,312.90905574443457,209.270249411708 -277.54946543623424,228.08210258025755,326.2879316740818,208.91105191830513 -259.14118496809317,237.00317892462166,325.13527205939647,208.38620445160527 -258.0603903695542,268.0249997632043,305.61306417662183,207.98642330663995 -263.1354853907414,238.83057850845063,298.7323501222888,209.43087634497854 -280.35492909981747,216.14279418347928,318.1451091369037,211.27160719784078 -271.56374567280704,229.39964558658272,286.6845872174376,210.0236414480106 -229.35820796650296,212.93348980868046,245.77044877320677,208.8917459728048 -230.49394051322878,203.67236378550288,224.55512987384654,205.8564923994707 -225.7601668320821,219.35714063162678,238.83881638414672,208.55931411570532 -230.75407630714457,198.06757965648347,238.36899164573157,208.37478598750155 -221.31628509636886,149.20055233004982,218.55183620363255,205.669374135596 -200.12300075458802,167.6799930782707,209.6087072026431,200.7031326575932 -138.65216945726675,143.81558253281227,219.38832428287833,201.49222415721096 +TS-DDR,TS-LDR,SDDP-DCLL,SDDP-SOC,TS-DDR (strict) +147.0247133419527,299.7450624492452,100.8484585807442,209.86380882088514,234.75510980001854 +157.58814081813028,225.6675464897668,104.86242775104952,209.53839385906838,205.21989985849396 +185.43371546547627,180.7943849292062,107.48502989328503,209.52017504171553,208.67816647888662 +199.3134646164604,199.83572444735833,107.55475020324891,209.4783394659705,212.22482561872755 +206.61858480532382,231.35910431382658,106.18650990754062,209.52518234738324,212.55023309260963 +214.3697084291256,189.64814901428522,113.67317158934263,209.6034477390574,214.5326077382685 +187.9473304345247,241.6103392896273,116.58714528380276,209.65384702850434,212.99020182480024 +171.46565827999535,187.4542766440176,129.25366035880148,209.51238800566185,212.91458967596776 +170.2002528391322,233.150743363177,114.76063596215229,209.45259355985954,208.08679618187475 +182.98142893312487,185.52122552406792,109.93768423864336,209.34525985187597,205.24871965873155 +202.19901785621693,215.06714163113537,116.98303295186679,209.44039837517906,203.72363105705048 +193.96716199312775,199.9910790500363,121.32762963754153,209.3672289663951,200.20688302172906 +193.52811950138334,205.21889018439646,118.81044678071419,209.32163376554064,196.2013393001419 +195.13814426290378,184.97486998527484,140.68919880246779,209.6614054735699,197.97739148855652 +226.57039256458856,167.1372529242544,157.98870219622455,210.19631956661485,204.08186126947413 +248.08925112087996,177.07636222131532,182.13344732161957,210.70698165532528,211.34080332566114 +233.5593017377245,206.24124265365077,213.25680127887568,211.4401169112003,214.69061359629825 +212.98635880030474,203.58562602361448,219.65332808686793,211.48105832602775,216.9537815779875 +208.1182840666282,225.62114837627547,215.53856096654266,211.94334659983778,216.71582598859962 +205.87501153760303,282.5815380137001,218.69325537706618,211.8784711569352,215.89358174848533 +218.77529830243512,261.2652681819222,229.70401987822862,211.85037808628837,216.02104060586097 +219.5621182240177,228.86267450841828,233.80086503295337,212.00653608776338,214.14628642224545 +215.15355785846822,278.8379771241615,219.9540629779893,211.71653085812682,212.6755011855709 +225.44578342326264,208.6334543644992,227.25750181666496,211.72866526590903,211.23043099837014 +216.46699729157837,233.62842636523723,228.98664748977347,211.79095480066385,210.62464082200245 +210.99855030225916,236.77260170176334,245.24308942520355,211.9396216916039,207.20024451574864 +200.82120574954234,257.832718296366,266.63858091994916,212.2412014072736,206.69466050470263 +217.00901852303974,206.80316405135653,287.580525250383,211.96614587275292,208.17775695655894 +205.60116177559695,247.22715853816413,288.43119570669523,212.0593948656155,204.3539449275232 +204.48717949362427,242.7626603477798,295.40865231198757,212.13189528353013,202.55758447677093 +215.01700896149435,264.8798596812393,316.4539931784796,212.48730840572054,199.11725065355355 +194.39594990932568,253.05048627557306,329.7940889425171,212.170191104127,195.65711788490407 +212.44246770637898,281.106369677497,327.59228806131625,212.68099082676446,194.0820072522097 +218.88119237987075,264.11335838123404,330.87683426769047,212.9854607332948,192.80649624294819 +209.54103152313093,265.672452188459,320.54817802711284,212.90364434114954,193.0831439581086 +232.50511220115487,242.47928625172216,336.16459151820277,212.98400917882333,196.4114335422718 +246.55559517109822,250.15790720527644,328.4530145309637,212.90706046302626,195.7154775743919 +259.1804770347488,241.4580174607301,306.6623006700101,212.78491180246007,204.36404255970675 +243.65212407401336,216.2392032573588,311.05102021329515,213.50440427222605,212.99514061661503 +268.3327047244221,244.87238185642838,313.43043807340126,214.71903130733853,236.7242644845188 +259.25135590920064,212.7607093726518,289.6446657672623,215.71309550713656,241.09270312524015 +239.7596346992519,211.60973323507133,236.34688350288008,214.21514981807115,220.01759022231735 +225.1627151901455,224.31226414469103,227.41520899884554,214.14602835221518,209.14147269754804 +226.300220128706,190.17002437882283,237.60749837482643,213.05570166141047,210.69197950461563 +246.5897657016228,197.34726823968018,242.14935260524663,214.3397929974062,228.55661937313394 +200.45447214778832,159.42600322266816,217.0164054870767,212.87589299860522,212.37433911366224 +177.65036023699594,199.6529727342857,200.46699991530508,210.512610777988,202.5465164778937 +152.7449270008933,184.7732720253377,215.7682105980341,208.87860648289214,219.84564553814513 +116.54920043933797,185.46879418205066,111.6333307935832,205.95870346385465,216.92116471391626 +118.11835868140619,166.0896692911,101.80408752144571,205.70698904522408,215.35769556857792 +135.40331893488192,159.49003675324482,106.15347696391866,205.4160873909239,210.9321344651103 +165.39822644830028,182.5511233035246,108.92328402254266,205.59700933991252,215.10732360430717 +164.63498605950215,157.20350355656583,103.94596690664551,205.58479725277505,213.53337186715123 +166.92019582010437,188.70842110946302,103.82258613108722,205.77526373955115,212.85090553779278 +153.73417427187314,161.11055875216002,112.13281054476954,205.88135777815637,209.902570025936 +164.778321612267,181.93462654464733,161.52068981849445,206.02649485451556,210.89012895639436 +160.9396855926067,146.59967315111336,169.88300890301332,205.8871588917278,205.50333692724715 +161.6813291137583,157.86943128110318,177.59131665764,206.70433459244137,204.32269653063892 +161.9520731591163,129.58260488127797,172.45550740547597,206.366999884319,201.43948777354547 +161.5166767553574,192.37516226972488,177.4638905803286,206.38870455127955,201.08952940134793 +171.23315428685606,169.81265452067095,176.73862367347002,206.249716219982,197.19467385465967 +203.84334411204006,189.01908062873738,182.19685255887867,206.63455688813036,196.69745986831143 +243.88748126609443,185.3643782473636,192.64656185901407,205.9797537356252,203.70645768573323 +255.36863069633904,189.10691135198536,210.14038242475175,206.03555271847185,210.0315143766425 +275.72231730419776,177.47926018459844,214.11405852350882,206.51953984726723,213.6098442211835 +231.14746381351682,203.55835450764843,216.8881876993481,206.12018860291096,215.6095194115856 +239.70981393223852,185.15015653078225,206.20963274723462,206.250132372938,213.42327196978673 +236.55408488066604,209.9345245251102,228.06281826263384,206.1138429495153,213.92786115893045 +232.50309001204022,247.09182945796942,232.53663833367236,206.4146828063853,215.32875670807255 +211.09249114427467,224.13559766576765,229.0267023271387,206.5605843384091,213.56203072051485 +222.7816613997703,251.3620122443193,232.90850460804563,206.40037604289813,211.33883588842454 +219.16351862667358,241.65598188514326,228.66399277324248,206.5295113131882,211.21103171947453 +207.22133488700123,260.0058643392084,266.0049119099325,206.8315496621826,209.84379993269118 +217.06224817213152,249.65113999794553,228.7078025295844,207.13270789037426,207.37243503995663 +210.07761177326952,201.43928958603075,257.88491102881073,207.77068168915991,206.3273067775563 +216.6438147791045,209.70322997573356,256.3000900779569,208.16087070287398,208.4308976434162 +201.41487156994177,254.6147418912871,260.0864123226794,208.40830707759466,204.48310839464855 +207.97135419475285,233.44163279290632,245.1593291381967,207.38456438434162,202.30900516056604 +211.26689971022438,206.86923643700797,266.2262958784643,206.9039133561306,199.0957215686977 +240.1582296197083,249.82887487767167,286.04078934122384,207.77917011114187,195.65265290918552 +233.02195021621122,218.9756433172769,302.4507102096619,207.558738022174,193.82165670656616 +237.29598744217932,254.3292021981702,310.8026658382449,208.03806284333123,194.10120155140936 +262.7454036444481,254.22526995036333,312.90905574443457,209.270249411708,193.48457104609795 +277.54946543623424,228.08210258025755,326.2879316740818,208.91105191830513,197.68279482028152 +259.14118496809317,237.00317892462166,325.13527205939647,208.38620445160527,198.27172714513736 +258.0603903695542,268.0249997632043,305.61306417662183,207.98642330663995,204.66198840076032 +263.1354853907414,238.83057850845063,298.7323501222888,209.43087634497854,213.63776171148848 +280.35492909981747,216.14279418347928,318.1451091369037,211.27160719784078,239.5404812525324 +271.56374567280704,229.39964558658272,286.6845872174376,210.0236414480106,240.58356433154577 +229.35820796650296,212.93348980868046,245.77044877320677,208.8917459728048,221.74490244061576 +230.49394051322878,203.67236378550288,224.55512987384654,205.8564923994707,207.85492498180685 +225.7601668320821,219.35714063162678,238.83881638414672,208.55931411570532,208.23135036415238 +230.75407630714457,198.06757965648347,238.36899164573157,208.37478598750155,227.9515285095849 +221.31628509636886,149.20055233004982,218.55183620363255,205.669374135596,214.2964276752777 +200.12300075458802,167.6799930782707,209.6087072026431,200.7031326575932,201.64848038288213 +138.65216945726675,143.81558253281227,219.38832428287833,201.49222415721096,219.54922660407325 diff --git a/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanVolume.csv b/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanVolume.csv index 6b0f78f..2074544 100644 --- a/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanVolume.csv +++ b/examples/HydroPowerModels/bolivia/ACPPowerModel/MeanVolume.csv @@ -1,97 +1,97 @@ -TS-DDR,TS-LDR,SDDP-DCLL,SDDP-SOC -22.51732496848532,37.919375441643616,15.010226053947687,29.88884403814745 -47.20258811999907,71.04409366252557,27.799866146567854,58.85398607133367 -72.83176865479658,87.01935425652064,41.809771632008996,82.61277155868845 -91.6036862853376,110.9546703800459,54.779619691636846,105.02407149811434 -111.75802365829193,124.62427072370988,67.41758956472614,127.58313808354393 -132.06225985386516,140.7114941122255,76.61311783876522,147.63328200911607 -151.88292311509423,163.09183590390472,85.64543943322955,168.65289410472099 -167.62574263315506,176.9949496837815,90.09866940534309,189.75454495424944 -184.1432711637443,192.12358945548075,97.21306040484191,218.11088982453344 -199.10712684414443,211.72514898332275,100.754649277593,246.93492472290026 -213.03560361830984,234.5026497996314,107.46682144224346,272.21531353974717 -224.76897734579637,261.6236901648443,111.30922724569007,292.90711864275784 -238.39447203378649,282.426645017808,114.62593628727466,311.83209342410777 -245.37549219043413,285.94648083840264,113.43065917802855,322.5848167767246 -252.4494932363904,296.90814187996295,110.20680607797237,322.2425933406325 -250.45907641085287,296.9831633048792,106.20779976718435,320.00444979375743 -248.8524629192563,300.7632599043702,101.76735381105871,312.8087106999844 -245.3285907233545,290.30183231928174,95.88402955956306,303.1966369847269 -237.90416834064467,290.15844458648814,89.15470564095466,296.46398240261857 -230.9413064277374,281.77833258453296,82.17436482998376,286.69479612549327 -221.1118971810357,281.1300748339776,74.6379921116223,275.7387873489527 -209.82043242344366,273.15550833073513,66.46062403663097,264.90582250183024 -197.17838472323572,255.4079970540009,58.1801555890566,252.08676360007652 -184.5082417295772,243.31108770222295,49.643407346969596,237.85231993230704 -171.364986731152,236.8863043514778,40.88630847779404,224.51757587451243 -158.67507679180278,228.75659130962944,32.42172641872411,211.62777515165263 -145.70893099014296,228.60197560222855,25.915677587474477,197.57381524080466 -132.80165423375078,219.4870486506905,17.949546658738377,182.8952051783611 -118.76446158949864,212.842623503242,10.91236331395401,169.0595060088052 -105.47673733968136,197.82735475771423,6.342665426831365,156.70187722582688 -92.28262527930282,183.60880119824964,2.9637113646432924,143.70426904824114 -80.08178289428561,178.21417462426308,1.4904346490510063,131.34345908273133 -68.05075167883604,174.38024725627528,0.6542767511259123,119.6560155151396 -57.33892828225943,163.75151223638323,-2.1629451753485584e-7,107.89499970347059 -47.91122996837452,158.30124789510745,0.01699339492037168,95.99205089892341 -39.71203072777365,146.6433762840449,0.013683141695575645,81.98936536386094 -32.08903884626882,138.74956149885196,-2.776369796429933e-5,70.22278074065235 -25.567879137873327,117.64643250247984,0.02029723087689403,57.957500191547524 -20.226109971544595,105.87824541767543,0.02942785780510852,48.218945817433884 -15.822029319315941,95.26193888634329,0.0015619407570932227,36.867704663165426 -11.86482403506727,82.88471848733366,0.0139202251526318,29.26424780181313 -11.41444238123061,89.84998297693535,0.0666095021126761,25.92639157504742 -11.760709387675428,80.56903002626132,0.02911519239806684,26.098422131593825 -15.608653812842075,84.4812008303581,0.1451610159592241,22.394312574066713 -18.46652374296125,75.10873866189522,0.257529864394238,19.330147292287013 -19.154852416425946,79.66424881896181,0.34310578392406726,19.064203280951464 -17.76382346268219,82.33214023398799,0.04020249114668608,20.09172473445721 -11.770417835270539,77.16028138611318,0.1601258590007935,19.35396878018543 -28.575902113160843,100.30120799881679,8.290616244150943,46.258638018871736 -44.6500656371219,133.7662232684492,18.826473984628194,73.04283772380003 -58.93359681990834,144.39896871097727,26.967830366959436,95.94410009405843 -78.81921579909906,157.05027238202555,36.00161441777166,117.71853421648952 -96.31823829542331,172.46207271242636,47.67488678462672,140.18906733662052 -113.68465589999252,199.6980903165839,56.815618070614086,161.20146693953137 -131.44733746695687,218.31345641963154,65.0298903498588,181.69307761993895 -143.57022598996326,248.0669589387606,78.12103160217615,199.8999392076057 -160.2545160514829,276.4730547448535,95.53937029048053,220.31192988484966 -173.71546658765067,294.32632069405474,107.5932705161822,239.09635346381702 -188.87717865623495,289.0248763377732,118.23444331568548,255.85211397472423 -200.72440725109396,310.47111738812987,125.58297208702875,272.7401007551301 -211.54578497851003,321.62467412938156,135.04096202155375,292.5534588326327 -219.67787470209927,322.9377982306577,138.23259862143857,305.2552176250861 -225.5911712053431,313.0818205198798,137.52220454328514,306.9497613201085 -225.39958719867292,303.0418832885243,134.56161843454288,303.58072451918565 -223.0003148921021,295.41511330313085,130.39480512697176,295.56194440948934 -219.3358500411617,292.2241103571665,124.0594015578791,285.94793100031006 -215.7865012983455,282.52479046838965,117.20362201093069,278.35828561341475 -208.9688074571684,270.8475139202375,109.91183630093335,268.7469595331369 -200.25820568279497,268.14577165979705,102.18413675040794,256.9288932038997 -189.3589474950667,274.8028403967962,94.29202902742942,246.23379738459545 -179.11347981019935,269.33461930444184,86.51972305194262,235.28235566184617 -166.46173064664583,261.0404132262943,77.8672879103141,222.9405685439762 -154.36450314445943,258.2200123763158,69.45071513510558,209.73529860734234 -142.61948735542381,250.57594722312155,61.04749888339048,197.4097594424145 -130.39403275688142,237.56983972063503,53.57905190860542,183.16067786345798 -118.09172135231981,232.9065992805976,45.31949038232733,167.08432725260056 -107.2702477985387,215.22999411819978,37.242727577570335,153.20481806438926 -97.0065973471623,205.80014155651824,28.961340379753675,140.27220558777037 -84.6333338632984,192.9275305544331,20.748060187265494,128.88089859059698 -74.44115642907373,175.65153529061962,13.951363389962466,115.32978413235291 -65.89676047851572,159.63588608030236,9.971531743391353,102.90650396427401 -59.101026640990746,157.6032563807769,5.649949808269948,91.77112364358631 -50.4628042074048,151.14613237451394,2.8422167357923183,81.38224893770428 -43.380434756627224,142.1297446607122,0.8089371675179539,69.40743842003496 -37.2402039626919,131.15093810046022,-2.7757075579902497e-5,58.51697840372556 -32.6377410503266,127.61272326799181,0.014576637404480192,49.72583740775145 -29.407339371204802,122.7377666202245,0.020931438505678426,41.99241630483234 -24.06393930942189,85.84155793018172,0.0062571966267834424,29.914239653894793 -23.127901917898463,65.85303059430143,0.03619878202007641,21.935606087997716 -22.924839043559988,62.40507538369134,0.23070127422069103,18.650205536148956 -22.007801580404273,59.154416378847564,0.49709652305249996,16.31324545284173 -21.136620549198433,54.197233720503064,0.12869349623927082,12.457190721911122 -19.031142363960612,38.14249299045677,0.2798932995958727,8.11780690376344 -17.40204173868524,30.7482859887398,0.6275685251939801,6.793617876065368 -12.246953654446704,19.1380483634613,0.7553403609073082,6.373971458330337 -0.40074148030873114,1.9740202514084515,-2.7761152226093412e-5,4.45582104449006 +TS-DDR,TS-LDR,SDDP-DCLL,SDDP-SOC,TS-DDR (strict) +22.51732496848532,37.919375441643616,15.010226053947687,29.88884403814745,30.458765766600905 +47.20258811999907,71.04409366252557,27.799866146567854,58.85398607133367,57.76343794462871 +72.83176865479658,87.01935425652064,41.809771632008996,82.61277155868845,79.55942523544871 +91.6036862853376,110.9546703800459,54.779619691636846,105.02407149811434,102.3640045482149 +111.75802365829193,124.62427072370988,67.41758956472614,127.58313808354393,127.23314750703618 +132.06225985386516,140.7114941122255,76.61311783876522,147.63328200911607,149.28395529505624 +151.88292311509423,163.09183590390472,85.64543943322955,168.65289410472099,169.75042070770465 +167.62574263315506,176.9949496837815,90.09866940534309,189.75454495424944,187.6893129283803 +184.1432711637443,192.12358945548075,97.21306040484191,218.11088982453344,206.98022476486403 +199.10712684414443,211.72514898332275,100.754649277593,246.93492472290026,225.59237229075973 +213.03560361830984,234.5026497996314,107.46682144224346,272.21531353974717,240.12435421501593 +224.76897734579637,261.6236901648443,111.30922724569007,292.90711864275784,252.04369533669177 +238.39447203378649,282.426645017808,114.62593628727466,311.83209342410777,265.09242625889686 +245.37549219043413,285.94648083840264,113.43065917802855,322.5848167767246,271.33120381927586 +252.4494932363904,296.90814187996295,110.20680607797237,322.2425933406325,273.33920353732105 +250.45907641085287,296.9831633048792,106.20779976718435,320.00444979375743,270.04043722632156 +248.8524629192563,300.7632599043702,101.76735381105871,312.8087106999844,265.38504262887227 +245.3285907233545,290.30183231928174,95.88402955956306,303.1966369847269,259.0130759486473 +237.90416834064467,290.15844458648814,89.15470564095466,296.46398240261857,252.12555542602587 +230.9413064277374,281.77833258453296,82.17436482998376,286.69479612549327,243.77847669015543 +221.1118971810357,281.1300748339776,74.6379921116223,275.7387873489527,234.0300699231589 +209.82043242344366,273.15550833073513,66.46062403663097,264.90582250183024,224.2307970115964 +197.17838472323572,255.4079970540009,58.1801555890566,252.08676360007652,213.74441729427633 +184.5082417295772,243.31108770222295,49.643407346969596,237.85231993230704,202.78706726330324 +171.364986731152,236.8863043514778,40.88630847779404,224.51757587451243,190.4835295366333 +158.67507679180278,228.75659130962944,32.42172641872411,211.62777515165263,178.93558488118396 +145.70893099014296,228.60197560222855,25.915677587474477,197.57381524080466,165.6894736001225 +132.80165423375078,219.4870486506905,17.949546658738377,182.8952051783611,152.04457076413573 +118.76446158949864,212.842623503242,10.91236331395401,169.0595060088052,139.79789850230674 +105.47673733968136,197.82735475771423,6.342665426831365,156.70187722582688,127.65358321143336 +92.28262527930282,183.60880119824964,2.9637113646432924,143.70426904824114,114.79802785727885 +80.08178289428561,178.21417462426308,1.4904346490510063,131.34345908273133,100.93687977706415 +68.05075167883604,174.38024725627528,0.6542767511259123,119.6560155151396,87.76382284286755 +57.33892828225943,163.75151223638323,-2.1629451753485584e-7,107.89499970347059,74.90399069532566 +47.91122996837452,158.30124789510745,0.01699339492037168,95.99205089892341,61.85270078016027 +39.71203072777365,146.6433762840449,0.013683141695575645,81.98936536386094,48.82119456836235 +32.08903884626882,138.74956149885196,-2.776369796429933e-5,70.22278074065235,37.94921398550819 +25.567879137873327,117.64643250247984,0.02029723087689403,57.957500191547524,28.26309568696599 +20.226109971544595,105.87824541767543,0.02942785780510852,48.218945817433884,20.65473538878002 +15.822029319315941,95.26193888634329,0.0015619407570932227,36.867704663165426,13.177382267227502 +11.86482403506727,82.88471848733366,0.0139202251526318,29.26424780181313,8.90534370422652 +11.41444238123061,89.84998297693535,0.0666095021126761,25.92639157504742,7.522263709728443 +11.760709387675428,80.56903002626132,0.02911519239806684,26.098422131593825,5.850846433137121 +15.608653812842075,84.4812008303581,0.1451610159592241,22.394312574066713,3.5447955720063016 +18.46652374296125,75.10873866189522,0.257529864394238,19.330147292287013,2.0999732627326524 +19.154852416425946,79.66424881896181,0.34310578392406726,19.064203280951464,1.8250334808778776 +17.76382346268219,82.33214023398799,0.04020249114668608,20.09172473445721,1.211979687666863 +11.770417835270539,77.16028138611318,0.1601258590007935,19.35396878018543,0.8408710530490285 +28.575902113160843,100.30120799881679,8.290616244150943,46.258638018871736,25.68369049047177 +44.6500656371219,133.7662232684492,18.826473984628194,73.04283772380003,53.77022334956752 +58.93359681990834,144.39896871097727,26.967830366959436,95.94410009405843,77.38194434940218 +78.81921579909906,157.05027238202555,36.00161441777166,117.71853421648952,100.76000446648004 +96.31823829542331,172.46207271242636,47.67488678462672,140.18906733662052,125.91256006951619 +113.68465589999252,199.6980903165839,56.815618070614086,161.20146693953137,147.66561942053426 +131.44733746695687,218.31345641963154,65.0298903498588,181.69307761993895,168.44823138836287 +143.57022598996326,248.0669589387606,78.12103160217615,199.8999392076057,186.97956897371287 +160.2545160514829,276.4730547448535,95.53937029048053,220.31192988484966,207.86295400946065 +173.71546658765067,294.32632069405474,107.5932705161822,239.09635346381702,225.47126269672492 +188.87717865623495,289.0248763377732,118.23444331568548,255.85211397472423,241.13236595003633 +200.72440725109396,310.47111738812987,125.58297208702875,272.7401007551301,252.9735009904081 +211.54578497851003,321.62467412938156,135.04096202155375,292.5534588326327,265.7166212598034 +219.67787470209927,322.9377982306577,138.23259862143857,305.2552176250861,272.87372052066604 +225.5911712053431,313.0818205198798,137.52220454328514,306.9497613201085,274.7942798565622 +225.39958719867292,303.0418832885243,134.56161843454288,303.58072451918565,272.15925221784215 +223.0003148921021,295.41511330313085,130.39480512697176,295.56194440948934,266.9798090916902 +219.3358500411617,292.2241103571665,124.0594015578791,285.94793100031006,260.39700073573346 +215.7865012983455,282.52479046838965,117.20362201093069,278.35828561341475,253.634969687372 +208.9688074571684,270.8475139202375,109.91183630093335,268.7469595331369,245.5294693796638 +200.25820568279497,268.14577165979705,102.18413675040794,256.9288932038997,235.13414146568158 +189.3589474950667,274.8028403967962,94.29202902742942,246.23379738459545,225.55640475182324 +179.11347981019935,269.33461930444184,86.51972305194262,235.28235566184617,215.37355080618886 +166.46173064664583,261.0404132262943,77.8672879103141,222.9405685439762,203.81449685097138 +154.36450314445943,258.2200123763158,69.45071513510558,209.73529860734234,191.70728842690482 +142.61948735542381,250.57594722312155,61.04749888339048,197.4097594424145,179.63985165695263 +130.39403275688142,237.56983972063503,53.57905190860542,183.16067786345798,166.84081909849618 +118.09172135231981,232.9065992805976,45.31949038232733,167.08432725260056,153.4635393785389 +107.2702477985387,215.22999411819978,37.242727577570335,153.20481806438926,140.99387306129591 +97.0065973471623,205.80014155651824,28.961340379753675,140.27220558777037,129.00590608741652 +84.6333338632984,192.9275305544331,20.748060187265494,128.88089859059698,115.79991834248376 +74.44115642907373,175.65153529061962,13.951363389962466,115.32978413235291,102.11766760572887 +65.89676047851572,159.63588608030236,9.971531743391353,102.90650396427401,88.89344582756674 +59.101026640990746,157.6032563807769,5.649949808269948,91.77112364358631,76.09868001212048 +50.4628042074048,151.14613237451394,2.8422167357923183,81.38224893770428,63.06719571114836 +43.380434756627224,142.1297446607122,0.8089371675179539,69.40743842003496,50.21450504411685 +37.2402039626919,131.15093810046022,-2.7757075579902497e-5,58.51697840372556,38.97364417660644 +32.6377410503266,127.61272326799181,0.014576637404480192,49.72583740775145,28.915285549722427 +29.407339371204802,122.7377666202245,0.020931438505678426,41.99241630483234,21.03021397587421 +24.06393930942189,85.84155793018172,0.0062571966267834424,29.914239653894793,13.40218700173315 +23.127901917898463,65.85303059430143,0.03619878202007641,21.935606087997716,9.268440492946311 +22.924839043559988,62.40507538369134,0.23070127422069103,18.650205536148956,7.830520814578429 +22.007801580404273,59.154416378847564,0.49709652305249996,16.31324545284173,6.1674695208565185 +21.136620549198433,54.197233720503064,0.12869349623927082,12.457190721911122,3.809203547797979 +19.031142363960612,38.14249299045677,0.2798932995958727,8.11780690376344,1.9133917353552115 +17.40204173868524,30.7482859887398,0.6275685251939801,6.793617876065368,1.583261510930154 +12.246953654446704,19.1380483634613,0.7553403609073082,6.373971458330337,1.2736466082355737 +0.40074148030873114,1.9740202514084515,-2.7761152226093412e-5,4.45582104449006,0.8413817689226626 diff --git a/examples/HydroPowerModels/check_consistent_state_paths.jl b/examples/HydroPowerModels/check_consistent_state_paths.jl index 349573a..c45c426 100644 --- a/examples/HydroPowerModels/check_consistent_state_paths.jl +++ b/examples/HydroPowerModels/check_consistent_state_paths.jl @@ -34,7 +34,7 @@ function build_problem() end # ---- Stage-wise simulation ---- -sub_s, state_in_s, state_out_s, uncert_s, initial_state, max_volume = build_problem() +sub_s, state_in_s, state_out_s, uncert_s, initial_state, max_volume, _ = build_problem() num_uncertainties = length(uncert_s[1]) # Constant policy so targets do not depend on state or uncertainty. @@ -60,7 +60,7 @@ for t in 1:num_stages end # ---- Deterministic equivalent ---- -sub_d, state_in_d, state_out_d, uncert_d, initial_state_d, _ = build_problem() +sub_d, state_in_d, state_out_d, uncert_d, initial_state_d, _, _ = build_problem() Det_model = JuMP.Model(det_optimizer) Det_model, uncert_d = DecisionRules.deterministic_equivalent!( @@ -83,7 +83,7 @@ for t in 1:num_stages end # ---- Multiple shooting ---- -sub_w, state_in_w, state_out_w, uncert_w, initial_state_w, _ = build_problem() +sub_w, state_in_w, state_out_w, uncert_w, initial_state_w, _, _ = build_problem() windows = DecisionRules.setup_shooting_windows( sub_w, diff --git a/examples/HydroPowerModels/dump_paired_policy_reference.jl b/examples/HydroPowerModels/dump_paired_policy_reference.jl new file mode 100644 index 0000000..cb4ae4d --- /dev/null +++ b/examples/HydroPowerModels/dump_paired_policy_reference.jl @@ -0,0 +1,260 @@ +# Dump the paired-evaluation policy reference for cross-package equivalence checks. +# +# This script produces the "ground truth" artifacts that the ExaModels-based +# companion package (DecisionRulesExa.jl) uses to verify that its policy, +# scenario indexing, and units reproduce this repository's behavior exactly: +# +# 1. The paired scenarios' inflow VALUE arrays, reconstructed from the +# seeded protocol (`paired_scenario_indices` in load_hydropowermodels.jl) +# and `build_hydropowermodels`'s `uncertainty_samples` — the same +# construction as `eval_paired_tsddr.jl`. +# 2. The OPEN-LOOP policy target trajectories +# +# x̂_0 = x_0, x̂_t = π_θ(w_t, x̂_{t-1}), +# +# computed WITHOUT any solver. In strict mode the stage subproblem enforces +# the hard equality `reservoir_out == x̂_t`, so the closed-loop realized +# state equals the target up to solver tolerance (~1e-8): the open-loop +# recursion doubles as the closed-loop state-path reference AND the +# policy-parity reference. +# 3. Three deterministic probe evaluations of the policy (single call after +# `Flux.reset!`), which test weight loading independently of any +# recurrent-state threading across stages. +# 4. The policy's frozen hydro metadata (bounds, upstream contribution, K), +# to validate the companion package's bounds construction. +# +# No subproblem is ever solved — only built (needed for `uncertainty_samples` +# and `hydro_meta`, whose `K` is extracted from the MOF hydro_balance row). +# +# Usage: +# julia --project -t auto dump_paired_policy_reference.jl [MODEL_PATH] +# +# Environment overrides: +# DR_NUM_EVAL_STAGES=96 +# DR_NUM_SCENARIOS=500 +# DR_ENCODER_LAYERS=128,128 +# DR_HEAD_LAYERS= +# DR_CONTEXT= ""/"none", "phase", or "phase+progress" +# DR_CONTEXT_HORIZON=126 denominator/horizon used for progress context +# DR_OUTPUT_TAG= suffix for paired_policy_reference_.jld2 +# +# Output: +# bolivia/ACPPowerModel/results/paired_policy_reference.jld2 +using DecisionRules +using Flux +using Statistics +using Random +using JuMP +using JLD2 +using DelimitedFiles + +HydroPowerModels_dir = dirname(@__FILE__) +include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) +include(joinpath(HydroPowerModels_dir, "hydro_reachable_policy.jl")) + +# ── Configuration ────────────────────────────────────────────────────────────── +# Default to the stage-wise strict TS-DDR checkpoint evaluated by +# eval_paired_tsddr.jl (the run that nearly matched SDDP on paired scenarios). +default_model_path = joinpath( + HydroPowerModels_dir, + "bolivia", + "ACPPowerModel", + "models", + "bolivia-ACPPowerModel-h126-r96-subproblems-strict-2026-07-01T09:41:53.026.jld2", +) +model_path = length(ARGS) >= 1 ? ARGS[1] : default_model_path +num_eval_stages = parse(Int, get(ENV, "DR_NUM_EVAL_STAGES", "96")) +num_scenarios = parse(Int, get(ENV, "DR_NUM_SCENARIOS", "500")) +output_tag = strip(get(ENV, "DR_OUTPUT_TAG", "")) + +parse_layers(s::AbstractString) = + isempty(strip(s)) ? Int64[] : [parse(Int64, strip(x)) for x in split(s, ",") if !isempty(strip(x))] + +function canonical_context_mode(raw_mode::AbstractString) + mode = lowercase(strip(raw_mode)) + mode in ("", "none", "off", "false") && return "" + mode in ("phase", "phase+progress") && return mode + error("DR_CONTEXT must be \"\", \"phase\", or \"phase+progress\"; got \"$raw_mode\"") +end + +function build_stage_context(mode::AbstractString, horizon::Int, period::Int) + isempty(mode) && return nothing + include_progress = mode == "phase+progress" + return DecisionRules.stage_phase_context( + horizon; + period=period, + include_progress=include_progress, + ) +end + +layers = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +head_layers = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +context_mode = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +context_horizon = parse(Int, get(ENV, "DR_CONTEXT_HORIZON", "126")) +context_period = countlines(joinpath(HydroPowerModels_dir, "bolivia", "inflows.csv")) +context_horizon >= num_eval_stages || + error("DR_CONTEXT_HORIZON=$context_horizon must cover DR_NUM_EVAL_STAGES=$num_eval_stages") +stage_context = build_stage_context(context_mode, context_horizon, context_period) +n_context = isnothing(stage_context) ? 0 : size(stage_context, 1) + +println("=" ^ 60) +println("Paired Policy Reference Dump (no solves)") +println(" Model: $model_path") +println(" Stages: $num_eval_stages") +println(" Scenarios: $num_scenarios") +println(" Layers: $layers") +println(" Head: $head_layers") +println(" Context: $(isempty(context_mode) ? "none" : context_mode)") +isempty(output_tag) || println(" Output tag: $output_tag") +println("=" ^ 60) + +# ── Build strict subproblems (structure only — no optimizer, no solves) ──────── +# Built exactly like eval_paired_tsddr.jl (num_stages=96, strict=true) so that +# `uncertainty_samples` and `hydro_meta` (including K parsed from the MOF +# hydro_balance coefficient) are bit-identical to the ground-truth evaluation. +# `optimizer=nothing` skips `set_optimizer` — safe because nothing is solved. +case_name = "bolivia" +formulation = "ACPPowerModel" +formulation_file = formulation * ".mof.json" + +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( + joinpath(HydroPowerModels_dir, case_name), + formulation_file; + num_stages=num_eval_stages, + optimizer=nothing, + strict=true, +) + +num_hydro = length(initial_state) +nCen = length(uncertainty_samples[1]) +println("nHyd=$num_hydro, nCen=$nCen, K=$(hydro_meta.K)") + +# Paired scenario indices: pure function of the protocol seed (entry [t, s] +# indexes `uncertainty_samples[t][ω]`; see load_hydropowermodels.jl). Fixed +# 126-row shape; this dump uses rows 1:num_eval_stages. +@assert num_eval_stages <= PAIRED_NUM_STAGES +all_indices = paired_scenario_indices(num_scenarios, nCen) +println("Paired protocol: seed=$(PAIRED_SCENARIO_SEED), rows 1:$(num_eval_stages) of $(PAIRED_NUM_STAGES)×$(num_scenarios), nCen=$nCen") + +# ── Build policy and load checkpoint weights ─────────────────────────────────── +# Same construction as eval_paired_tsddr.jl lines 86-91: encoder+combiner +# weights are restored; hydro bounds come from hydro_meta. +base_model = hydro_reachable_policy( + hydro_meta, + layers; + combiner_layers=head_layers, + n_context=n_context, +) +models = isnothing(stage_context) ? base_model : ContextualPolicy(base_model, stage_context) +model_save = JLD2.load(model_path) +model_state = model_save["model_state"] +load_policy_weights!(models, model_state) +println("Loaded model weights from $model_path") + +# ── Reconstruct inflow VALUE arrays for the paired scenarios ─────────────────── +# inflow_values[t, r, s] is the inflow of hydro unit r at stage t under paired +# scenario s, i.e. the second element of `uncertainty_samples[t][idx[t,s]][r]`. +# This is exactly the value fed to both the policy and the subproblem parameter +# in eval_paired_tsddr.jl. +inflow_values = zeros(Float64, num_eval_stages, num_hydro, num_scenarios) +for s in 1:num_scenarios + for t in 1:num_eval_stages + w_t = uncertainty_samples[t][all_indices[t, s]] + for (r, (_, val)) in enumerate(w_t) + inflow_values[t, r, s] = val + end + end +end + +# ── Open-loop policy target trajectories ─────────────────────────────────────── +# Recursion (no solver): +# +# x̂_0 = x_0, x̂_t = π_θ(Float32.(w_t), Float32.(x̂_{t-1})). +# +# In strict mode the closed-loop rollout of eval_paired_tsddr.jl realizes +# x_t = x̂_t exactly (hard equality), so this trajectory is the reference for +# both the policy outputs and the realized reservoir path (up to the Ipopt +# solver tolerance with which the closed loop reads back the realized state). +xhat_trajectories = zeros(Float64, num_eval_stages, num_hydro, num_scenarios) +for s in 1:num_scenarios + # Fresh recurrent state per scenario, matching the closed-loop evaluation. + Flux.reset!(models) + state = Float64.(initial_state) + for t in 1:num_eval_stages + w_vals = Float32.(inflow_values[t, :, s]) + # π_θ(w_t, x̂_{t-1}) — identical input construction to eval_paired_tsddr.jl + x_hat = models(vcat(w_vals, Float32.(state))) + xhat_trajectories[t, :, s] = Float64.(x_hat) + # Open-loop: feed the target back as the next previous state. + state = Float64.(x_hat) + end + if s % 20 == 0 || s == num_scenarios + println(" open-loop trajectories: [$s/$num_scenarios]") + end +end + +# ── Deterministic probe evaluations ──────────────────────────────────────────── +# Each probe is a SINGLE policy call after Flux.reset!, i.e. from the zero +# initial recurrent state. These test pure weight parity: any cross-package +# difference in recurrent-state threading across stages cannot affect them. +# +# probe 1: w = stage-1 inflows of paired scenario 1, x = initial_state +# probe 2: w = zeros, x = initial_state +# probe 3: w = mean over ω of stage-1 inflows, x = initial_state +w_probe_1 = Float32.(inflow_values[1, :, 1]) +w_probe_2 = zeros(Float32, num_hydro) +w_probe_3 = Float32.([ + mean(uncertainty_samples[1][ω][r][2] for ω in 1:nCen) for r in 1:num_hydro +]) +x_probe = Float32.(initial_state) + +probe_inputs = zeros(Float32, 2 * num_hydro, 3) +probe_inputs[:, 1] = vcat(w_probe_1, x_probe) +probe_inputs[:, 2] = vcat(w_probe_2, x_probe) +probe_inputs[:, 3] = vcat(w_probe_3, x_probe) + +probe_outputs = zeros(Float64, num_hydro, 3) +for p in 1:3 + # Reset before every probe: single call from the zero recurrent state. + Flux.reset!(models) + probe_outputs[:, p] = Float64.(models(probe_inputs[:, p])) +end +println("Probe outputs computed.") + +# ── Save reference JLD2 ──────────────────────────────────────────────────────── +out_dir = joinpath(HydroPowerModels_dir, case_name, formulation, "results") +mkpath(out_dir) +out_suffix = isempty(output_tag) ? "" : "_" * output_tag +out_file = joinpath(out_dir, "paired_policy_reference$(out_suffix).jld2") +metadata_policy = models isa ContextualPolicy ? models.policy : models +jldsave(out_file; + # Scenario data (authoritative values for the cross-package rollout) + inflow_values=inflow_values, # [T × nHyd × S] + scenario_indices=all_indices[1:num_eval_stages, 1:num_scenarios], + # Policy references + xhat_trajectories=xhat_trajectories, # [T × nHyd × S] open-loop targets + probe_inputs=probe_inputs, # [2 nHyd × 3] Float32 + probe_outputs=probe_outputs, # [nHyd × 3] + # System data for metadata cross-checks + initial_state=Float64.(initial_state), + max_volume=Float64.(max_volume), + # Frozen hydro metadata exactly as used by the policy forward pass + policy_K=metadata_policy.K, + policy_min_vol=Float64.(metadata_policy.min_vol), + policy_max_vol=Float64.(metadata_policy.max_vol), + policy_min_turn=Float64.(metadata_policy.min_turn), + policy_max_turn=Float64.(metadata_policy.max_turn), + policy_upstream_max=Float64.(metadata_policy.upstream_max), + # Provenance + model_path=model_path, + num_eval_stages=num_eval_stages, + num_scenarios=num_scenarios, + encoder_layers=layers, + head_layers=head_layers, + context_mode=isempty(context_mode) ? "none" : context_mode, + context_period=context_period, + context_horizon=context_horizon, + n_context=n_context, +) +println("Saved: $out_file") diff --git a/examples/HydroPowerModels/eval_jump_de.jl b/examples/HydroPowerModels/eval_jump_de.jl index 1a72412..299b935 100644 --- a/examples/HydroPowerModels/eval_jump_de.jl +++ b/examples/HydroPowerModels/eval_jump_de.jl @@ -34,7 +34,7 @@ const TARGET_FRAC = 0.6 # constant target = TARGET_FRAC × max_volume @info "Building HydroPowerModels ($FORMULATION, T=$NUM_STAGES)..." -sub, state_in, state_out, uncert, initial_state, max_volume = build_hydropowermodels( +sub, state_in, state_out, uncert, initial_state, max_volume, _ = build_hydropowermodels( CASE_DIR, FORMULATION * ".mof.json"; num_stages=NUM_STAGES, penalty_l2=:auto ) nHyd = length(initial_state) diff --git a/examples/HydroPowerModels/eval_paired_tsddr.jl b/examples/HydroPowerModels/eval_paired_tsddr.jl new file mode 100644 index 0000000..15a4c50 --- /dev/null +++ b/examples/HydroPowerModels/eval_paired_tsddr.jl @@ -0,0 +1,248 @@ +# Paired TS-DDR strict rollout evaluation on the seeded paired protocol. +# +# Scenario indices are a pure function of PAIRED_SCENARIO_SEED (see +# paired_scenario_indices in load_hydropowermodels.jl), so this script and the +# SDDP Historical simulation realize the exact same inflows with no shared +# data file. +# +# Usage: +# julia --project -t auto eval_paired_tsddr.jl MODEL_PATH +# +# Environment overrides: +# DR_NUM_EVAL_STAGES=96 +# DR_NUM_SCENARIOS=500 +# DR_ENCODER_LAYERS=128,128 +# DR_HEAD_LAYERS= +# DR_CONTEXT= ""/"none", "phase", or "phase+progress" +# DR_CONTEXT_HORIZON=126 denominator/horizon used for progress context +# DR_OUTPUT_TAG="" (when set, ALL output filenames are suffixed with +# "_" — paired_costs_.csv, etc. — so evaluating +# a new checkpoint never clobbers the untagged +# ground-truth results; unset → historical filenames) +using DecisionRules +using Flux +using Statistics +using Random +using JuMP, DiffOpt, Ipopt +using JLD2 +using CSV, DataFrames +using JSON +using DelimitedFiles + +HydroPowerModels_dir = dirname(@__FILE__) +include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) +include(joinpath(HydroPowerModels_dir, "hydro_reachable_policy.jl")) + +model_path = ARGS[1] +num_eval_stages = parse(Int, get(ENV, "DR_NUM_EVAL_STAGES", "96")) +num_scenarios = parse(Int, get(ENV, "DR_NUM_SCENARIOS", "500")) +# Optional output tag: suffixes every output filename with "_" so a new +# checkpoint's evaluation cannot overwrite the untagged ground-truth files. +output_tag = strip(get(ENV, "DR_OUTPUT_TAG", "")) +tag_suffix = isempty(output_tag) ? "" : "_" * output_tag + +parse_layers(s::AbstractString) = + isempty(strip(s)) ? Int64[] : [parse(Int64, strip(x)) for x in split(s, ",") if !isempty(strip(x))] + +function canonical_context_mode(raw_mode::AbstractString) + mode = lowercase(strip(raw_mode)) + mode in ("", "none", "off", "false") && return "" + mode in ("phase", "phase+progress") && return mode + error("DR_CONTEXT must be \"\", \"phase\", or \"phase+progress\"; got \"$raw_mode\"") +end + +function build_stage_context(mode::AbstractString, horizon::Int, period::Int) + isempty(mode) && return nothing + include_progress = mode == "phase+progress" + return DecisionRules.stage_phase_context( + horizon; + period=period, + include_progress=include_progress, + ) +end + +layers = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +head_layers = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +context_mode = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +context_horizon = parse(Int, get(ENV, "DR_CONTEXT_HORIZON", "126")) +context_period = countlines(joinpath(HydroPowerModels_dir, "bolivia", "inflows.csv")) +context_horizon >= num_eval_stages || + error("DR_CONTEXT_HORIZON=$context_horizon must cover DR_NUM_EVAL_STAGES=$num_eval_stages") +stage_context = build_stage_context(context_mode, context_horizon, context_period) +n_context = isnothing(stage_context) ? 0 : size(stage_context, 1) + +println("=" ^ 60) +println("Paired TS-DDR Strict Rollout Evaluation") +println(" Model: $model_path") +println(" Stages: $num_eval_stages") +println(" Scenarios: $num_scenarios") +println(" Layers: $layers") +println(" Head: $head_layers") +println(" Context: $(isempty(context_mode) ? "none" : context_mode)") +isempty(output_tag) || println(" Output tag: $output_tag") +println("=" ^ 60) + +# ── Build strict subproblems ─────────────────────────────────────────────── +case_name = "bolivia" +formulation = "ACPPowerModel" +formulation_file = formulation * ".mof.json" + +diff_optimizer = + () -> DiffOpt.diff_optimizer( + optimizer_with_attributes( + Ipopt.Optimizer, + "print_level" => 0, + "linear_solver" => "mumps", + ), + ) + +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( + joinpath(HydroPowerModels_dir, case_name), + formulation_file; + num_stages=num_eval_stages, + optimizer=diff_optimizer, + strict=true, +) + +num_hydro = length(initial_state) +nCen = length(uncertainty_samples[1]) +println("nHyd=$num_hydro, nCen=$nCen") + +# Paired scenario indices: a pure function of the protocol seed (see +# paired_scenario_indices in load_hydropowermodels.jl). Fixed 126-row shape; +# this evaluation uses rows 1:num_eval_stages. +@assert num_eval_stages <= PAIRED_NUM_STAGES +all_indices = paired_scenario_indices(num_scenarios, nCen) +println("Paired protocol: seed=$(PAIRED_SCENARIO_SEED), rows 1:$(num_eval_stages) of $(PAIRED_NUM_STAGES)×$(num_scenarios), nCen=$nCen") + +# ── Identify thermal generators ──────────────────────────────────────────── +hydro_data = JSON.parsefile(joinpath(HydroPowerModels_dir, case_name, "hydro.json")) +power_data = JSON.parsefile(joinpath(HydroPowerModels_dir, case_name, "PowerModels.json")) +baseMVA = power_data["baseMVA"] +hydro_grid_idx = Set(hg["index_grid"] for hg in hydro_data["Hydrogenerators"]) +num_gen = length(power_data["gen"]) +thermal_idx = [i for i in 1:num_gen if !(i in hydro_grid_idx)] + +volume_to_mw(volume; k=0.0036) = volume / k + +pg_vars_per_stage = [DecisionRules.find_variables(subproblems[t], ["pg"]) for t in 1:num_eval_stages] + +# ── Build policy and load weights ────────────────────────────────────────── +base_model = hydro_reachable_policy( + hydro_meta, + layers; + combiner_layers=head_layers, + n_context=n_context, +) +models = isnothing(stage_context) ? base_model : ContextualPolicy(base_model, stage_context) +model_save = JLD2.load(model_path) +model_state = model_save["model_state"] +load_policy_weights!(models, model_state) +println("Loaded model weights from $model_path") + +# ── Construct scenarios from pre-sampled indices ─────────────────────────── +eval_scenarios = Vector{Vector{Vector{Tuple{eltype(uncertainty_samples[1][1][1][1]), eltype(uncertainty_samples[1][1][1][2])}}}}(undef, num_scenarios) +for s in 1:num_scenarios + eval_scenarios[s] = [uncertainty_samples[t][all_indices[t, s]] for t in 1:num_eval_stages] +end + +# Verify: print first scenario's first stage inflows +println("\nFirst scenario, stage 1 inflows:") +for (param, val) in eval_scenarios[1][1] + println(" $(JuMP.name(param)) = $val") +end + +# ── Run rollout evaluation ───────────────────────────────────────────────── +println("\nEvaluating $num_scenarios scenarios on $num_eval_stages stages...") + +costs = Float64[] +vol_trajectories = zeros(num_eval_stages, num_scenarios) +gen_trajectories = zeros(num_eval_stages, num_scenarios) + +for (i, scenario) in enumerate(eval_scenarios) + Flux.reset!(models) + + state = Float64.(initial_state) + scenario_cost = 0.0 + + for t in 1:num_eval_stages + for (j, param) in enumerate(state_params_in[t]) + set_parameter_value(param, state[j]) + end + + w_t = scenario[t] + for (param, val) in w_t + set_parameter_value(param, val) + end + + w_vals = Float32.([val for (_, val) in w_t]) + x_hat = models(vcat(w_vals, Float32.(state))) + + for j in 1:num_hydro + target_param = state_params_out[t][j][1] + set_parameter_value(target_param, Float64(x_hat[j])) + end + + optimize!(subproblems[t]) + scenario_cost += objective_value(subproblems[t]) + + for j in 1:num_hydro + state[j] = value(state_params_out[t][j][2]) + end + + vol_trajectories[t, i] = sum(volume_to_mw(state[j]) for j in 1:num_hydro) + gen_trajectories[t, i] = sum( + value(pg_vars_per_stage[t][j]) * baseMVA for j in thermal_idx + ) + end + + push!(costs, scenario_cost) + if i % 10 == 0 || i == num_scenarios + println(" [$i/$num_scenarios] cost = $(round(scenario_cost; digits=1)), running mean = $(round(mean(costs); digits=1))") + end +end + +# ── Report results ───────────────────────────────────────────────────────── +println("\n" * "=" ^ 60) +println("Results: Paired TS-DDR Strict ($num_eval_stages stages, $num_scenarios scenarios)") +println("=" ^ 60) +println(" Mean cost: $(round(mean(costs); digits=1))") +println(" Std: $(round(std(costs); digits=1))") +println(" Min: $(round(minimum(costs); digits=1))") +println(" Max: $(round(maximum(costs); digits=1))") +println(" Median: $(round(median(costs); digits=1))") +println(" Violation: 0.0% (strict mode)") +println("=" ^ 60) + +# ── Save results ─────────────────────────────────────────────────────────── +out_dir = joinpath(HydroPowerModels_dir, case_name, formulation) + +const COL_NAME = "TS-DDR (strict, paired)" +costs_file = joinpath(out_dir, "paired_costs$(tag_suffix).csv") +df = DataFrame(Symbol(COL_NAME) => costs) +CSV.write(costs_file, df) +println("Saved: $costs_file") + +mean_vol = vec(mean(vol_trajectories; dims=2)) +vol_file = joinpath(out_dir, "paired_MeanVolume$(tag_suffix).csv") +df_vol = DataFrame(Symbol(COL_NAME) => mean_vol) +CSV.write(vol_file, df_vol) +println("Saved: $vol_file") + +mean_gen = vec(mean(gen_trajectories; dims=2)) +gen_file = joinpath(out_dir, "paired_MeanGeneration$(tag_suffix).csv") +df_gen = DataFrame(Symbol(COL_NAME) => mean_gen) +CSV.write(gen_file, df_gen) +println("Saved: $gen_file") + +results_dir = joinpath(out_dir, "results") +mkpath(results_dir) +results_file = joinpath(results_dir, "paired_strict_rollout$(tag_suffix).jld2") +jldsave(results_file; + costs=costs, + vol_trajectories=vol_trajectories, + gen_trajectories=gen_trajectories, + scenario_indices=all_indices[1:num_eval_stages, 1:num_scenarios], +) +println("Saved: $results_file") diff --git a/examples/HydroPowerModels/eval_strict_rollout.jl b/examples/HydroPowerModels/eval_strict_rollout.jl new file mode 100644 index 0000000..ff1cefd --- /dev/null +++ b/examples/HydroPowerModels/eval_strict_rollout.jl @@ -0,0 +1,269 @@ +# Evaluate a trained strict-mode HydroReachablePolicy on 96-stage rollouts. +# +# Loads a saved model checkpoint, builds 96-stage strict subproblems, and +# runs stage-wise rollout evaluation on 100 held-out inflow scenarios. +# Reports per-scenario operational costs (no deficit/penalty — strict mode +# has none) and summary statistics. Also records per-stage reservoir volumes +# and thermal generation for comparison plots (MeanVolume.csv, +# MeanGeneration.csv, costs.csv). +# +# Usage: +# julia --project -t auto eval_strict_rollout.jl MODEL_PATH +# +# Environment overrides: +# DR_NUM_EVAL_STAGES=96 number of rollout stages +# DR_NUM_SCENARIOS=100 number of evaluation scenarios +# DR_EVAL_SEED=1221 random seed (matches SDDP evaluation) +# DR_ENCODER_LAYERS=128,128 +# DR_HEAD_LAYERS= +# DR_CONTEXT= ""/"none", "phase", or "phase+progress" +# DR_CONTEXT_HORIZON=126 denominator/horizon used for progress context +using DecisionRules +using Flux +using Statistics +using Random +using JuMP, DiffOpt, Ipopt +using JLD2 +using CSV, DataFrames +using JSON + +HydroPowerModels_dir = dirname(@__FILE__) +include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) +include(joinpath(HydroPowerModels_dir, "hydro_reachable_policy.jl")) + +# ── Parse arguments ───────────────────────────────────────────────────────── + +model_path = ARGS[1] +num_eval_stages = parse(Int, get(ENV, "DR_NUM_EVAL_STAGES", "96")) +num_scenarios = parse(Int, get(ENV, "DR_NUM_SCENARIOS", "100")) +seed = parse(Int, get(ENV, "DR_EVAL_SEED", "1221")) + +parse_layers(s::AbstractString) = + isempty(strip(s)) ? Int64[] : [parse(Int64, strip(x)) for x in split(s, ",") if !isempty(strip(x))] + +function canonical_context_mode(raw_mode::AbstractString) + mode = lowercase(strip(raw_mode)) + mode in ("", "none", "off", "false") && return "" + mode in ("phase", "phase+progress") && return mode + error("DR_CONTEXT must be \"\", \"phase\", or \"phase+progress\"; got \"$raw_mode\"") +end + +function build_stage_context(mode::AbstractString, horizon::Int, period::Int) + isempty(mode) && return nothing + include_progress = mode == "phase+progress" + return DecisionRules.stage_phase_context( + horizon; + period=period, + include_progress=include_progress, + ) +end + +layers = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +head_layers = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +context_mode = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +context_horizon = parse(Int, get(ENV, "DR_CONTEXT_HORIZON", "126")) +context_period = countlines(joinpath(HydroPowerModels_dir, "bolivia", "inflows.csv")) +context_horizon >= num_eval_stages || + error("DR_CONTEXT_HORIZON=$context_horizon must cover DR_NUM_EVAL_STAGES=$num_eval_stages") +stage_context = build_stage_context(context_mode, context_horizon, context_period) +n_context = isnothing(stage_context) ? 0 : size(stage_context, 1) + +println("=" ^ 60) +println("Strict Rollout Evaluation") +println(" Model: $model_path") +println(" Stages: $num_eval_stages") +println(" Scenarios: $num_scenarios") +println(" Seed: $seed") +println(" Layers: $layers") +println(" Head: $head_layers") +println(" Context: $(isempty(context_mode) ? "none" : context_mode)") +println("=" ^ 60) + +# ── Build strict subproblems for evaluation ───────────────────────────────── + +case_name = "bolivia" +formulation = "ACPPowerModel" +formulation_file = formulation * ".mof.json" + +diff_optimizer = + () -> DiffOpt.diff_optimizer( + optimizer_with_attributes( + Ipopt.Optimizer, + "print_level" => 0, + "linear_solver" => "mumps", + ), + ) + +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( + joinpath(HydroPowerModels_dir, case_name), + formulation_file; + num_stages=num_eval_stages, + optimizer=diff_optimizer, + strict=true, +) + +num_hydro = length(initial_state) + +# ── Identify thermal generators ───────────────────────────────────────────── + +hydro_data = JSON.parsefile(joinpath(HydroPowerModels_dir, case_name, "hydro.json")) +power_data = JSON.parsefile(joinpath(HydroPowerModels_dir, case_name, "PowerModels.json")) +baseMVA = power_data["baseMVA"] +hydro_grid_idx = Set(hg["index_grid"] for hg in hydro_data["Hydrogenerators"]) +num_gen = length(power_data["gen"]) +thermal_idx = [i for i in 1:num_gen if !(i in hydro_grid_idx)] + +volume_to_mw(volume; k=0.0036) = volume / k + +# Precompute pg variable references for each stage +pg_vars_per_stage = [DecisionRules.find_variables(subproblems[t], ["pg"]) for t in 1:num_eval_stages] + +# ── Build policy and load trained weights ─────────────────────────────────── + +base_model = hydro_reachable_policy( + hydro_meta, + layers; + combiner_layers=head_layers, + n_context=n_context, +) +models = isnothing(stage_context) ? base_model : ContextualPolicy(base_model, stage_context) + +model_save = JLD2.load(model_path) +model_state = model_save["model_state"] +load_policy_weights!(models, model_state) +println("Loaded model weights from $model_path") + +# ── Run rollout evaluation ────────────────────────────────────────────────── + +Random.seed!(seed) +eval_scenarios = [DecisionRules.sample(uncertainty_samples) for _ in 1:num_scenarios] + +println("\nEvaluating $num_scenarios scenarios on $num_eval_stages stages...") + +costs = Float64[] +vol_trajectories = zeros(num_eval_stages, num_scenarios) +gen_trajectories = zeros(num_eval_stages, num_scenarios) + +for (i, scenario) in enumerate(eval_scenarios) + Flux.reset!(models) + + state = Float64.(initial_state) + scenario_cost = 0.0 + + for t in 1:num_eval_stages + # Set initial state parameters + for (i, param) in enumerate(state_params_in[t]) + set_parameter_value(param, state[i]) + end + + # Set uncertainty parameters + w_t = scenario[t] + for (param, val) in w_t + set_parameter_value(param, val) + end + + # Policy forward pass + w_vals = Float32.([val for (_, val) in w_t]) + x_hat = models(vcat(w_vals, Float32.(state))) + + # Set targets + for j in 1:num_hydro + target_param = state_params_out[t][j][1] + set_parameter_value(target_param, Float64(x_hat[j])) + end + + # Solve + optimize!(subproblems[t]) + scenario_cost += objective_value(subproblems[t]) + + # Read realized reservoir volumes + for j in 1:num_hydro + state[j] = value(state_params_out[t][j][2]) + end + + # Record per-stage data + vol_trajectories[t, i] = sum(volume_to_mw(state[j]) for j in 1:num_hydro) + + # Thermal generation: sum pg * baseMVA for thermal generators + gen_trajectories[t, i] = sum( + value(pg_vars_per_stage[t][j]) * baseMVA for j in thermal_idx + ) + end + + push!(costs, scenario_cost) + if i % 10 == 0 || i == num_scenarios + println(" [$i/$num_scenarios] cost = $(round(scenario_cost; digits=1)), running mean = $(round(mean(costs); digits=1))") + end +end + +# ── Report results ────────────────────────────────────────────────────────── + +println("\n" * "=" ^ 60) +println("Results: Strict TS-DDR Rollout ($num_eval_stages stages, $num_scenarios scenarios)") +println("=" ^ 60) +println(" Mean cost: $(round(mean(costs); digits=1))") +println(" Std: $(round(std(costs); digits=1))") +println(" Min: $(round(minimum(costs); digits=1))") +println(" Max: $(round(maximum(costs); digits=1))") +println(" Median: $(round(median(costs); digits=1))") +println(" Violation: 0.0% (strict mode)") +println("=" ^ 60) + +# ── Save results ──────────────────────────────────────────────────────────── + +out_dir = joinpath(HydroPowerModels_dir, case_name, formulation) + +# Per-scenario costs → costs.csv +const STRICT_COL = "TS-DDR (strict)" +costs_file = joinpath(out_dir, "costs.csv") +if isfile(costs_file) + df = CSV.read(costs_file, DataFrame) + df[!, STRICT_COL] = costs +else + df = DataFrame(Symbol(STRICT_COL) => costs) +end +CSV.write(costs_file, df) +println("Updated: $costs_file") + +# Mean volume trajectory → MeanVolume.csv +mean_vol = vec(mean(vol_trajectories; dims=2)) +vol_file = joinpath(out_dir, "MeanVolume.csv") +if isfile(vol_file) + df_vol = CSV.read(vol_file, DataFrame; header=true) + df_vol[!, STRICT_COL] = mean_vol +else + df_vol = DataFrame(Symbol(STRICT_COL) => mean_vol) +end +CSV.write(vol_file, df_vol) +println("Updated: $vol_file") + +# Mean thermal generation → MeanGeneration.csv +mean_gen = vec(mean(gen_trajectories; dims=2)) +gen_file = joinpath(out_dir, "MeanGeneration.csv") +if isfile(gen_file) + df_gen = CSV.read(gen_file, DataFrame; header=true) + df_gen[!, STRICT_COL] = mean_gen +else + df_gen = DataFrame(Symbol(STRICT_COL) => mean_gen) +end +CSV.write(gen_file, df_gen) +println("Updated: $gen_file") + +# Full results → JLD2 +results_dir = joinpath(out_dir, "results") +mkpath(results_dir) +results_file = joinpath(results_dir, "strict_rollout_$(num_eval_stages)stages_$(num_scenarios)scenarios.jld2") +jldsave(results_file; + costs=costs, + mean_cost=mean(costs), + std_cost=std(costs), + vol_trajectories=vol_trajectories, + gen_trajectories=gen_trajectories, + mean_vol=mean_vol, + mean_gen=mean_gen, + num_stages=num_eval_stages, + num_scenarios=num_scenarios, + model_path=model_path, +) +println("Saved: $results_file") diff --git a/examples/HydroPowerModels/evaluate_hydro_policies.jl b/examples/HydroPowerModels/evaluate_hydro_policies.jl index cf59032..01eae0a 100644 --- a/examples/HydroPowerModels/evaluate_hydro_policies.jl +++ b/examples/HydroPowerModels/evaluate_hydro_policies.jl @@ -32,6 +32,7 @@ using DataFrames const HYDRO_DIR = dirname(@__FILE__) include(joinpath(HYDRO_DIR, "load_hydropowermodels.jl")) +include(joinpath(HYDRO_DIR, "hydro_reachable_policy.jl")) const CASE_NAME = "bolivia" const FORMULATION = "ACPPowerModel" @@ -63,7 +64,7 @@ diff_optimizer = () -> DiffOpt.diff_optimizer( ), ) -subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = +subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume, hydro_meta = build_hydropowermodels( CASE_DIR, FORMULATION_FILE; num_stages=NUM_STAGES, @@ -97,6 +98,7 @@ struct PolicySpec label::String model_file::String is_ldr::Bool + is_strict::Bool end function _method_variant(base) @@ -104,6 +106,8 @@ function _method_variant(base) "ldr" elseif contains(base, "shooting") "shooting" + elseif contains(base, "subproblems-strict") || contains(base, "strict") + "subproblems-strict" elseif contains(base, "subproblems") "subproblems" elseif contains(base, "deteq") @@ -124,6 +128,8 @@ function _variant_label(variant) "subproblems-const" => "Subproblems (const)", "subproblems-clip-const" => "Subproblems (clip, const)", "subproblems" => "Subproblems", + "subproblems-strict" => "Subproblems (strict)", + "subproblems-strict-clip" => "Subproblems (strict, clip)", "shooting-anneal" => "Shooting w=12 (anneal)", "shooting-clip-anneal" => "Shooting w=12 (clip, anneal)", "shooting" => "Shooting w=12", @@ -137,24 +143,27 @@ end function discover_policies(model_dir) files = sort(filter(f -> endswith(f, ".jld2"), readdir(model_dir; join=true))) - best = Dict{String,Tuple{String,Bool}}() + best = Dict{String,Tuple{String,Bool,Bool}}() for f in files base = basename(f) variant = _method_variant(base) isnothing(variant) && continue is_ldr = contains(base, "ldr") - best[variant] = (f, is_ldr) + is_strict = contains(base, "strict") + best[variant] = (f, is_ldr, is_strict) end specs = PolicySpec[] - for (variant, (path, is_ldr)) in sort(collect(best); by=first) - push!(specs, PolicySpec(_variant_label(variant), path, is_ldr)) + for (variant, (path, is_ldr, is_strict)) in sort(collect(best); by=first) + push!(specs, PolicySpec(_variant_label(variant), path, is_ldr, is_strict)) end return specs end -function build_policy(spec::PolicySpec, num_inputs, num_hydro, num_uncertainties) +function build_policy(spec::PolicySpec, num_inputs, num_hydro, num_uncertainties, hydro_meta) if spec.is_ldr return dense_multilayer_nn(num_inputs, num_hydro, Int64[64, 64]; activation=identity) + elseif spec.is_strict + return hydro_reachable_policy(hydro_meta, Int64[128, 128]) else return state_conditioned_policy( num_uncertainties, num_hydro, num_hydro, Int64[128, 128]; @@ -166,7 +175,7 @@ end policies = discover_policies(MODEL_DIR) println("\nDiscovered policies:") for p in policies - tag = p.is_ldr ? " (LDR)" : " (DDR)" + tag = p.is_ldr ? " (LDR)" : p.is_strict ? " (DDR-strict)" : " (DDR)" println(" ", p.label, tag, " → ", basename(p.model_file)) end @@ -177,9 +186,13 @@ results = DataFrame() for spec in policies println("\nEvaluating: ", spec.label) - models = build_policy(spec, num_inputs, num_hydro, num_uncertainties) + models = build_policy(spec, num_inputs, num_hydro, num_uncertainties, hydro_meta) model_state = JLD2.load(spec.model_file, "model_state") - Flux.loadmodel!(models, model_state) + if spec.is_strict + load_policy_weights!(models, model_state) + else + Flux.loadmodel!(models, model_state) + end objectives_no_deficit = Vector{Float64}(undef, NUM_SIMULATIONS) objectives_total = Vector{Float64}(undef, NUM_SIMULATIONS) diff --git a/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl b/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl index 775fffd..60566cd 100644 --- a/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl +++ b/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl @@ -17,7 +17,7 @@ HydroPowerModels_dir = dirname(@__FILE__) include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) case_dir = joinpath(HydroPowerModels_dir, case_name) -subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels( +subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages ) diff --git a/examples/HydroPowerModels/hydro_reachable_policy.jl b/examples/HydroPowerModels/hydro_reachable_policy.jl new file mode 100644 index 0000000..f69acde --- /dev/null +++ b/examples/HydroPowerModels/hydro_reachable_policy.jl @@ -0,0 +1,522 @@ +# Hydro Reachable Policy — feasibility-guaranteeing target policy for strict subproblems +# +# This file defines HydroReachablePolicy, a policy architecture that guarantees +# one-stage reachability for hydro reservoir targets. It wraps an LSTM encoder + +# feed-forward combiner (same architecture family as StateConditionedPolicy) but +# bounds the output to the one-stage reachable set via sigmoid activation. +# +# Depends on: DecisionRules (for _step_encoder, _init_recurrent_state, _state_eltype) +# Must be included AFTER `using DecisionRules`. + +using Functors +using ChainRulesCore + +# ── Cascade link: upstream→downstream water-balance coupling ────────────────── + +struct CascadeLink + downstream::Int # array position of downstream unit + upstream::Int # array position of upstream unit + turn_only::Bool # true if only turbine outflow (not spill) reaches downstream + K_max_turn::Float32 # K × max_turn of the upstream unit +end + +""" + HydroReachablePolicy{E,C,S,V,SM} + +A policy that guarantees one-stage reachability for hydro reservoir targets. + +The policy architecture mirrors [`StateConditionedPolicy`](@ref): an LSTM +encoder processes only the uncertainty (inflow) sequence, then a feed-forward +combiner maps `[encoded_inflow; current_reservoir_state]` to normalized targets +in `[0, 1]` via sigmoid activation. These normalized targets are scaled to the +one-stage **reachable set** `[lower, upper]` for each reservoir: + +```math +target_r = lower_r + (upper_r - lower_r) \\cdot \\sigma(z_r) +``` + +where `lower_r` and `upper_r` are the minimum and maximum reservoir volumes achievable +in one stage from the current state `x_r` given inflow `w_r`, turbine bounds +`[min\\_turn_r, max\\_turn_r]`, and upstream cascade inflows. + +# Reachable bounds (per reservoir r) + +The upper reachable bound assumes minimum outflow (no turbine, no spill) plus maximum +upstream inflow from cascade connections: + +```math +upper_r = \\min(max\\_vol_r,\\; x_r + K \\cdot w_r - K \\cdot min\\_turn_r + upstream\\_max_r) +``` + +The lower reachable bound assumes maximum outflow (full turbine + max spill): + +```math +lower_r = \\max(min\\_vol_r,\\; x_r + K \\cdot w_r - K \\cdot max\\_turn_r - spill\\_max_r) +``` + +When `spill_max === nothing` (unlimited spillage), `lower_r = min_vol_r` since the +reservoir can always be emptied to its physical minimum. + +The bounds are marked `@non_differentiable` — gradients flow only through the sigmoid +path `σ(z_r)`, not through the bounds themselves. + +# Cascade-aware target clamping + +For downstream units receiving water from upstream cascade connections, the initial +upper bound uses `upstream_max_r = Σ K × max_turn_u`, which can overestimate the +actual upstream contribution when the upstream unit stores water (target increases). + +After computing initial targets for all units, a **cascade clamping** step adjusts +downstream targets using the actual upstream release implied by the upstream target: + +```math +R_u = K \\cdot w_u + x_u - \\hat{x}_u +``` + +- **Turn + spill connection**: max upstream contribution = ``\\max(0, R_u)`` +- **Turn-only connection**: max contribution = ``\\min(K \\cdot max\\_turn_u, \\max(0, R_u))`` + +The downstream target is clamped: ``\\hat{x}_r ← \\min(\\hat{x}_r, true\\_upper_r)``. +This clamping is `@non_differentiable` — gradient flows through when not active, +zero when clamped (correct projected-gradient signal). + +# Strict-mode guarantee + +If `x₀` is a feasible initial reservoir state and each policy call returns +``\\hat{x}_t ∈ R(x_{t-1}, w_t)`` (including cascade-consistent bounds), then the +strict equality ``x_t = \\hat{x}_t`` is feasible for every stage solved in sequence. +The proof is by induction: stage 1 is feasible because ``\\hat{x}_1`` is reachable from +``x_0``; if stage ``t`` is feasible and realizes ``x_t = \\hat{x}_t``, then the +policy computes ``\\hat{x}_{t+1}`` from a feasible previous state with cascade-clamped +bounds, so stage ``t+1`` is feasible. + +# Fields +- `encoder::E`: Recurrent cell or Chain of cells (processes inflow only) +- `combiner::C`: Feed-forward head combining encoder output with previous state +- `state::S`: Recurrent state (threaded across stages) +- `n_context::Int`: Number of context dimensions prepended before inflow +- `n_uncertainty::Int`: Number of inflow dimensions (= nHyd) +- `n_state::Int`: Number of state dimensions (= nHyd) +- `min_vol::V`: Per-unit minimum reservoir volume +- `max_vol::V`: Per-unit maximum reservoir volume +- `min_turn::V`: Per-unit minimum turbine outflow +- `max_turn::V`: Per-unit maximum turbine outflow +- `upstream_max::V`: Pre-computed maximum upstream inflow contribution per unit +- `spill_max::SM`: Per-unit max spillage (`nothing` = unlimited) +- `K::Float64`: Water-balance conversion factor from flow to volume + +See also: [`hydro_reachable_policy`](@ref), [`StateConditionedPolicy`](@ref) +""" +mutable struct HydroReachablePolicy{E,C,S,V,SM} + encoder::E # Recurrent encoder (LSTM/GRU chain) processing inflow + combiner::C # Feed-forward [encoder_out; state] => normalized target + state::S # Carried recurrent state, threaded across stages + n_context::Int # Number of context dimensions prepended before inflow + n_uncertainty::Int # Number of uncertainty (inflow) dimensions + n_state::Int # Number of state (reservoir) dimensions + min_vol::V # Per-unit minimum reservoir volume [nHyd] + max_vol::V # Per-unit maximum reservoir volume [nHyd] + min_turn::V # Per-unit minimum turbine outflow [nHyd] + max_turn::V # Per-unit maximum turbine outflow [nHyd] + upstream_max::V # Pre-computed K × Σ(upstream max_turn) per unit [nHyd] + spill_max::SM # Per-unit max spill, or nothing for unlimited + K::Float64 # Water-balance conversion factor from flow to volume + cascade::Vector{CascadeLink} # Upstream→downstream connections for target clamping +end + +# Only encoder and combiner are trainable. Bounds, state, dimensions are frozen. +Functors.@functor HydroReachablePolicy (encoder, combiner) + +""" + _hydro_reachable_bounds(policy::HydroReachablePolicy, inflow, x_prev) + +Compute the one-stage reachable reservoir bounds given current state `x_prev` +and per-unit inflow `inflow`. Returns `(lower, upper)` vectors. + +The upper bound is the maximum volume achievable in one stage: current volume +plus inflow minus minimum turbine outflow plus maximum upstream cascade inflow, +clamped to `max_vol`. The lower bound is the minimum volume achievable: current +volume plus inflow minus maximum turbine outflow minus maximum spill, clamped +to `min_vol`. + +Marked `@non_differentiable` — gradients do not flow through the bounds. The gradient +path is solely through the sigmoid-scaled output `σ(z_r)`. + +# Arguments +- `policy::HydroReachablePolicy`: policy containing hydro bounds and parameters +- `inflow`: per-unit inflow vector for this stage +- `x_prev`: current reservoir volumes (state from previous stage) + +# Returns +- `(lower, upper)`: tuple of vectors, each of length `n_state` +""" +function _hydro_reachable_bounds(policy::HydroReachablePolicy, inflow, x_prev) + # Cast all bound vectors to match the element type of x_prev for type stability + T = eltype(x_prev) + K = T(policy.K) + min_vol = T.(policy.min_vol) + max_vol = T.(policy.max_vol) + min_turn = T.(policy.min_turn) + max_turn = T.(policy.max_turn) + upstream = T.(policy.upstream_max) + + # Upper reachable: minimum outflow (min turbine, no spill) + max upstream + upper_raw = x_prev .+ K .* inflow .- K .* min_turn .+ upstream + # Clamp to physical maximum volume + upper = min.(max_vol, upper_raw) + + # Lower reachable: maximum outflow (max turbine + max spill) + lower = if policy.spill_max === nothing + # Unlimited spill → can always dump down to min_vol + min_vol + else + spill_max = T.(policy.spill_max) + # Volume after maximum discharge and maximum spill + lower_raw = x_prev .+ K .* inflow .- K .* max_turn .- spill_max + # Clamp to physical minimum volume + max.(min_vol, lower_raw) + end + + # Ensure lower ≤ upper (numerical safety for edge cases like CHJ with max_vol=0) + upper = max.(upper, lower) + + return lower, upper +end +# Gradients do not flow through the bounds — only through σ(z) +ChainRulesCore.@non_differentiable _hydro_reachable_bounds(::Any, ::Any, ::Any) + +""" + _cascade_upper_bounds(policy, target, inflow, x_prev) + +Compute the true reachable upper bound for downstream units given the actual +upstream targets. Returns a vector of upper bounds (Inf for units with no +upstream connections). Marked `@non_differentiable`. + +# Documented assumptions +- **Single-level cascades**: the implied upstream release + ``R_u = K w_u + x_u - \\hat{x}_u`` omits the upstream unit's own incoming + cascade contribution, which is conservative (underestimates the release) + for multi-level chains. +- **No gradient through binding clamps**: this function is + `@non_differentiable`, so when the resulting clamp binds, the dependence of + the downstream target on the upstream target is not differentiated. +""" +function _cascade_upper_bounds(policy::HydroReachablePolicy, target, inflow, x_prev) + cascade = policy.cascade + T = eltype(target) + K = T(policy.K) + n = length(target) + upper = fill(T(Inf), n) + for conn in cascade + u = conn.upstream + d = conn.downstream + R_u = K * inflow[u] + x_prev[u] - target[u] + if conn.turn_only + max_contrib = min(T(conn.K_max_turn), max(zero(T), R_u)) + else + max_contrib = max(zero(T), R_u) + end + true_upper = x_prev[d] + K * inflow[d] - K * T(policy.min_turn[d]) + max_contrib + true_upper = min(T(policy.max_vol[d]), true_upper) + upper[d] = min(upper[d], true_upper) + end + return upper +end +ChainRulesCore.@non_differentiable _cascade_upper_bounds(::Any, ::Any, ::Any, ::Any) + +""" + (m::HydroReachablePolicy)(x) + +Forward pass: given input `x = [context; inflow₁..nHyd; x_prev₁..nHyd]`, produce +one-stage reachable reservoir targets. + +1. Split input into context, inflow, and previous state +2. Encode `[context; inflow]` through recurrent encoder, carrying state across stages +3. Combine encoder output with previous state via a sigmoid head → y_norm ∈ [0,1] +4. Compute reachable bounds [lower, upper] from physics (no gradient) +5. Scale: target = lower + (upper - lower) × y_norm +6. Clamp downstream targets to cascade-aware upper bounds (no gradient through bounds) + +# Documented assumptions +- **Single-level cascades**: the cascade clamp uses the implied upstream release + ``R_u = K w_u + x_u - \\hat{x}_u``, which omits the upstream unit's own + incoming cascade contribution — conservative for multi-level chains. +- **No gradient through binding clamps**: reachable bounds and cascade clamps + are `@non_differentiable`; when a clamp binds, the downstream target's + dependence on the upstream target is not differentiated. +- **Physically-infeasible edge case**: if the cascade upper bound falls below + the reachable lower bound, the clamped target may fall below `lower`. No + policy-level remedy exists in that case — the underlying problem is + infeasible. + +# Arguments +- `x`: concatenated input vector `[context..., inflow..., previous_state...]` + +# Returns +- `Vector`: target reservoir volumes, guaranteed within one-stage reachable set +""" +function (m::HydroReachablePolicy)(x) + # Split input: optional context first, then true inflow, then previous state. + # Physics bounds must use only the true inflow slice. + c_end = m.n_context + w_start = c_end + 1 + w_end = c_end + m.n_uncertainty + context = c_end == 0 ? x[1:0] : x[1:c_end] + inflow = x[w_start:w_end] + x_prev = x[w_end+1:end] + encoder_input = c_end == 0 ? inflow : vcat(context, inflow) + + # Encode inflow through the recurrent encoder, carrying state across calls. + # Cast to encoder precision for type stability (avoids Zygote codegen bugs). + T = DecisionRules._state_eltype(m.state) + encoded, new_state = DecisionRules._step_encoder(m.encoder, T.(encoder_input), m.state) + # Thread recurrent state to the next call + m.state = new_state + + # Raw output from combiner (sigmoid activation → values in [0, 1]) + y_norm = m.combiner(vcat(encoded, x_prev)) + + # Compute reachable bounds from current state and inflow (no gradient) + lower, upper = _hydro_reachable_bounds(m, inflow, x_prev) + + # Scale normalized output to the reachable interval [lower, upper] + raw_target = lower .+ (upper .- lower) .* y_norm + + # Clamp downstream targets to cascade-aware reachable bounds + if !isempty(m.cascade) + cascade_upper = _cascade_upper_bounds(m, raw_target, inflow, x_prev) + return min.(raw_target, cascade_upper) + end + return raw_target +end + +""" + Flux.reset!(m::HydroReachablePolicy) + +Reset the encoder's recurrent state to `Flux.initialstates`, e.g. before starting +a new rollout. The hydro bounds (min_vol, max_vol, etc.) are unchanged. +""" +function Flux.reset!(m::HydroReachablePolicy) + # Reinitialize recurrent state from the encoder's initial states + m.state = DecisionRules._init_recurrent_state(m.encoder) + return nothing +end + +""" + hydro_reachable_policy(hydro_meta, layers; encoder_type=Flux.LSTM, + spill_max=nothing, combiner_layers=Int[]) + +Create a [`HydroReachablePolicy`](@ref) from hydro metadata (as returned by +the 7th return value of `build_hydropowermodels`). + +The architecture mirrors [`state_conditioned_policy`](@ref): an LSTM encoder +processes inflows, then a feed-forward combiner produces normalized targets in +`[0, 1]`. These are scaled to the one-stage reachable interval. The combiner +uses sigmoid activation — this is mandatory and cannot be overridden. Set +`combiner_layers` to add hidden layers to the nonrecurrent state-conditioned +target map. This is the preferred way to depart from linear decision rules +without adding recurrence over the reservoir-state input. + +# Arguments +- `hydro_meta::NamedTuple`: hydro system metadata with fields `nHyd`, `min_vol`, + `max_vol`, `min_turn`, `max_turn`, `K`, `upstream_turn` +- `layers::Vector{Int}`: hidden layer sizes for the LSTM encoder + (e.g. `[128, 128]` for a 2-layer LSTM) +- `encoder_type`: recurrent layer/cell type (default: `Flux.LSTM`). Must support + `Flux.initialstates` and the stateful `(x, state) -> (output, new_state)` call +- `spill_max`: per-unit maximum spillage vector (`nothing` = unlimited spillage, + meaning `lower = min_vol` always) +- `combiner_layers::Vector{Int}`: hidden widths for the nonrecurrent target head + +# Returns +- `HydroReachablePolicy`: ready-to-train policy with sigmoid-bounded outputs + +# Examples +```julia +subproblems, _, _, _, _, _, hydro_meta = build_hydropowermodels( + case_dir, formulation_file; strict=true, optimizer=diff_opt +) +policy = hydro_reachable_policy(hydro_meta, [128, 128]) +policy_with_deep_state_head = hydro_reachable_policy( + hydro_meta, + [128, 128]; + combiner_layers=[256, 256], +) +``` + +See also: [`HydroReachablePolicy`](@ref), [`state_conditioned_policy`](@ref), +[`load_policy_weights!`](@ref) +""" +function hydro_reachable_policy( + hydro_meta::NamedTuple, + layers::Vector{Int}; + encoder_type=Flux.LSTM, + spill_max=nothing, + combiner_layers=Int[], + n_context::Int=0, +) + nHyd = hydro_meta.nHyd + # Validate layer sizes + isempty(layers) && throw(ArgumentError("layers must be non-empty")) + n_context >= 0 || throw(ArgumentError("n_context must be nonnegative")) + + # Build encoder: stack of recurrent cells processing [context; inflow]. + encoder_input_dim = nHyd + n_context + if length(layers) == 1 + # Single-layer encoder + encoder = DecisionRules._as_cell(encoder_type(encoder_input_dim => layers[1])) + else + # Multi-layer encoder: chain of recurrent cells + encoder_layers = [DecisionRules._as_cell(encoder_type(encoder_input_dim => layers[1]))] + for i in 1:(length(layers) - 1) + push!( + encoder_layers, + DecisionRules._as_cell(encoder_type(layers[i] => layers[i + 1])), + ) + end + encoder = Chain(encoder_layers...) + end + + # Build combiner — sigmoid activation is MANDATORY for reachable bounds guarantee + # The sigmoid output ∈ [0, 1] is scaled to [lower, upper] in the forward pass + combiner = DecisionRules.dense_policy_head( + layers[end] + nHyd, + nHyd, + collect(Int, combiner_layers); + activation=sigmoid, + ) + + # Pre-compute maximum upstream inflow contribution per unit: + # upstream_max[r] = Σ_{u ∈ upstream(r)} K × max_turn_u + K = hydro_meta.K + upstream_max = zeros(Float32, nHyd) + for (r, upstream_list) in enumerate(hydro_meta.upstream_turn) + for (u_pos, u_max_turn) in upstream_list + upstream_max[r] += Float32(K * u_max_turn) + end + end + + # Build cascade connections for target clamping + spill_dests = Dict{Int,Set{Int}}() + for (r, upstream_list) in enumerate(hydro_meta.upstream_spill) + for (u_pos, _) in upstream_list + push!(get!(spill_dests, u_pos, Set{Int}()), r) + end + end + cascade = CascadeLink[] + for (r, upstream_list) in enumerate(hydro_meta.upstream_turn) + for (u_pos, u_max_turn) in upstream_list + has_spill = haskey(spill_dests, u_pos) && r in spill_dests[u_pos] + push!(cascade, CascadeLink(r, u_pos, !has_spill, Float32(K * u_max_turn))) + end + end + for (r, upstream_list) in enumerate(hydro_meta.upstream_spill) + for (u_pos, _) in upstream_list + already = any(c -> c.downstream == r && c.upstream == u_pos, cascade) + already || push!(cascade, CascadeLink(r, u_pos, false, Float32(K * hydro_meta.max_turn[u_pos]))) + end + end + + # Validate spill_max dimensions if provided + if spill_max !== nothing && length(spill_max) != nHyd + throw(ArgumentError("spill_max length must be nHyd=$nHyd; got $(length(spill_max))")) + end + + return HydroReachablePolicy( + encoder, + combiner, + DecisionRules._init_recurrent_state(encoder), # initial recurrent state + n_context, # context dimensions prepended before inflow + nHyd, # n_uncertainty = nHyd (one inflow per unit) + nHyd, # n_state = nHyd (one reservoir per unit) + Float32.(hydro_meta.min_vol), # per-unit min volume + Float32.(hydro_meta.max_vol), # per-unit max volume + Float32.(hydro_meta.min_turn), # per-unit min turbine outflow + Float32.(hydro_meta.max_turn), # per-unit max turbine outflow + upstream_max, # pre-computed upstream contribution + spill_max === nothing ? nothing : Float32.(collect(spill_max)), # spill bounds + K, # water-balance conversion factor + cascade, # cascade connections for target clamping + ) +end + +""" + load_policy_weights!(policy::HydroReachablePolicy, state) + +Load encoder/combiner weights from a saved model state (e.g., from a +[`StateConditionedPolicy`](@ref) checkpoint). Hydro bounds are preserved. + +This enables warmstarting: train a `StateConditionedPolicy` with non-strict +subproblems, then load its encoder/combiner weights into a `HydroReachablePolicy` +for strict fine-tuning. + +# Arguments +- `policy::HydroReachablePolicy`: target policy (bounds are preserved) +- `state`: saved model state (from `Flux.state(model)` or JLD2 checkpoint) + +# Returns +- `policy`: the modified policy (mutated in place) + +See also: [`hydro_reachable_policy`](@ref) +""" +function load_policy_weights!(policy::HydroReachablePolicy, state) + # Load only the encoder and combiner weights, keeping hydro bounds unchanged + try + Flux.loadmodel!(policy.encoder, state.encoder) + Flux.loadmodel!(policy.combiner, state.combiner) + catch err + throw(ArgumentError( + "Could not load HydroReachablePolicy weights. The checkpoint architecture " * + "must match encoder/head widths and n_context=$(policy.n_context); " * + "contextual policies require newly trained checkpoints. Original error: $err", + )) + end + return policy +end + +function load_policy_weights!(policy::DecisionRules.ContextualPolicy, state) + inner_state = hasproperty(state, :policy) ? getproperty(state, :policy) : state + load_policy_weights!(policy.policy, inner_state) + return policy +end + +""" + load_hydro_reachable_policy(checkpoint_path, hydro_meta, layers; + encoder_type=Flux.LSTM, spill_max=nothing, + combiner_layers=Int[]) + +Load a [`HydroReachablePolicy`](@ref) from a JLD2 checkpoint, reconstructing +the hydro bounds from `hydro_meta` (since JLD2 may not preserve exact types). + +# Arguments +- `checkpoint_path::String`: path to JLD2 file with `"model_state"` key +- `hydro_meta::NamedTuple`: hydro metadata from `build_hydropowermodels` +- `layers::Vector{Int}`: encoder hidden layer sizes (must match checkpoint) +- `encoder_type`: recurrent layer type (default: `Flux.LSTM`) +- `spill_max`: per-unit max spillage, or `nothing` for unlimited +- `combiner_layers::Vector{Int}`: hidden widths for the nonrecurrent target head + +# Returns +- `HydroReachablePolicy`: policy with loaded weights and fresh hydro bounds + +See also: [`hydro_reachable_policy`](@ref), [`load_policy_weights!`](@ref) +""" +function load_hydro_reachable_policy( + checkpoint_path::String, + hydro_meta::NamedTuple, + layers::Vector{Int}; + encoder_type=Flux.LSTM, + spill_max=nothing, + combiner_layers=Int[], + n_context::Int=0, +) + # Build fresh policy with correct bounds from hydro_meta + policy = hydro_reachable_policy(hydro_meta, layers; encoder_type=encoder_type, + spill_max=spill_max, + combiner_layers=combiner_layers, + n_context=n_context) + # Load saved weights into the fresh policy + model_state = JLD2.load(checkpoint_path, "model_state") + load_policy_weights!(policy, model_state) + return policy +end diff --git a/examples/HydroPowerModels/load_hydropowermodels.jl b/examples/HydroPowerModels/load_hydropowermodels.jl index 613b425..5279261 100644 --- a/examples/HydroPowerModels/load_hydropowermodels.jl +++ b/examples/HydroPowerModels/load_hydropowermodels.jl @@ -2,6 +2,36 @@ using JuMP using CSV using Tables using JSON +using StableRNGs + +# Paired evaluation protocol: at stage t of paired scenario s, EVERY method +# (SDDP.Historical, TS-DDR CPU/GPU rollouts) realizes joint inflow scenario +# `paired_scenario_indices(N, nCen)[t, s]`. StableRNG streams are stable +# across Julia versions, so the protocol is fully defined by this seed — no +# data file to distribute or track. +const PAIRED_SCENARIO_SEED = 20260706 +# The generated matrix ALWAYS has this many stage rows (the full SDDP horizon: +# 96 reported + 30 end-of-horizon buffer). Consumers that need fewer stages +# slice rows — they must never generate a smaller matrix, because arrays of +# different shapes consume the RNG stream differently and pairing across +# methods would silently break. +const PAIRED_NUM_STAGES = 126 + +""" + paired_scenario_indices(num_scenarios, nCen; + seed = PAIRED_SCENARIO_SEED) -> Matrix{Int} + +Deterministic paired-evaluation index matrix of fixed shape +`PAIRED_NUM_STAGES × num_scenarios`: entry `[t, s]` is uniform on `1:nCen` +(the per-stage joint inflow support, `nCen = ncol(inflows.csv) ÷ nHyd`). +Slice rows for shorter horizons; never regenerate at a different shape. +""" +function paired_scenario_indices( + num_scenarios::Integer, nCen::Integer; + seed::Integer = PAIRED_SCENARIO_SEED, +) + return rand(StableRNG(seed), 1:Int(nCen), PAIRED_NUM_STAGES, Int(num_scenarios)) +end function find_reservoirs_and_inflow(model::JuMP.Model) reservoir_in = find_variables(model, ["reservoir", "_in"]) @@ -27,6 +57,64 @@ function read_inflow(file::String, nHyd::Int; num_stages=nothing) return vector_inflows, nCen, num_stages end +""" + build_hydropowermodels(case_folder, subproblem_file; num_stages, penalty, penalty_l1, + penalty_l2, optimizer, strict) -> (subproblems, state_params_in, + state_params_out, uncertainty_samples, initial_state, max_volume, + hydro_meta) + +Build multi-stage hydro power subproblems from a case folder containing `hydro.json`, +`inflows.csv`, and a MOF subproblem file. Each stage gets its own JuMP model with +parameterized incoming state, outgoing target (with or without deficit slack), and +uncertainty (inflow) samples. + +When `strict=true`, the outgoing state is bound to the target via a hard equality +constraint (`reservoir_out == target`) with **no deficit variables** and **no penalty +term**. The dual of this equality is the clean shadow price ∂Q/∂target — pure economic +signal without penalty noise. This requires a feasibility-guaranteeing policy (e.g. +[`HydroReachablePolicy`]) to avoid infeasible subproblems. + +When `strict=false` (default), deficit variables are created via [`create_deficit!`](@ref) +and penalized in the objective, allowing the solver to deviate from the target. + +# Arguments +- `case_folder::AbstractString`: path to the case directory (must contain `hydro.json` + and `inflows.csv`) +- `subproblem_file::AbstractString`: MOF filename for the stage subproblem (e.g. + `"ACPPowerModel.mof.json"`) +- `num_stages`: number of stages (default: number of rows in `inflows.csv`) +- `penalty`: legacy L1 penalty coefficient (use `penalty_l1`/`penalty_l2` instead) +- `penalty_l1`: L1 norm penalty coefficient, or `:auto` +- `penalty_l2`: L2 squared norm penalty coefficient, or `:auto` +- `optimizer`: optimizer factory for DiffOpt, e.g. `() -> DiffOpt.diff_optimizer(...)` +- `strict::Bool=false`: if `true`, use hard equality target constraints (no deficit) + +# Returns +A 7-tuple `(subproblems, state_params_in, state_params_out, uncertainty_samples, +initial_state, max_volume, hydro_meta)` where: +- `subproblems::Vector{JuMP.Model}`: one JuMP model per stage +- `state_params_in::Vector{Vector{Any}}`: incoming state parameters per stage +- `state_params_out::Vector{Vector{Tuple{Any,VariableRef}}}`: `(parameter, variable)` + tuples for outgoing state per stage +- `uncertainty_samples`: joint inflow scenarios per stage +- `initial_state::Vector{Float64}`: initial reservoir volumes +- `max_volume::Vector{Float64}`: maximum reservoir volumes +- `hydro_meta::NamedTuple`: hydro system metadata for policy construction (see below) + +## `hydro_meta` fields +- `nHyd::Int`: number of hydro units +- `min_vol`, `max_vol`: per-unit volume bounds +- `min_turn`, `max_turn`: per-unit turbine outflow bounds +- `initial_volume`: initial reservoir volumes +- `downstream_turn`, `downstream_spill`: downstream connectivity (by hydro index) +- `upstream_turn`: `Vector{Vector{Tuple{Int,Float64}}}` — for each unit, list of + `(upstream_array_pos, upstream_max_turn)` pairs feeding into it +- `upstream_spill`: same structure for spill connections +- `K::Float64`: water-balance conversion factor from flow units to volume units +- `production_factor`: per-unit production factors + +See also: [`create_deficit!`](@ref), [`variable_to_parameter`](@ref) +""" function build_hydropowermodels( case_folder::AbstractString, subproblem_file::AbstractString; @@ -35,15 +123,71 @@ function build_hydropowermodels( penalty_l1=nothing, penalty_l2=nothing, optimizer=nothing, + strict::Bool=false, ) - hydro_file = JSON.parsefile(joinpath(case_folder, "hydro.json"))["Hydrogenerators"] + # Parse the hydro system data file + hydro_json = JSON.parsefile(joinpath(case_folder, "hydro.json")) + hydro_file = hydro_json["Hydrogenerators"] nHyd = length(hydro_file) + # Extract water-balance conversion factor K from the MOF model's hydro_balance + # constraint. K converts flow units (m³/s) to volume units (hm³) per stage. + # hydro.json["stage_hours"] is the stage duration, NOT the water-balance K. + _tmp_model = JuMP.read_from_file( + joinpath(case_folder, subproblem_file); use_nlp_block=false + ) + _hb_con = JuMP.constraint_by_name(_tmp_model, "hydro_balance[1]") + _hb_func = JuMP.constraint_object(_hb_con).func + _inflow_var = first(filter( + v -> occursin("inflow", JuMP.name(v)), JuMP.all_variables(_tmp_model) + )) + K = abs(JuMP.coefficient(_hb_func, _inflow_var)) + # Read historical inflow scenarios from CSV vector_inflows, nCen, num_stages = read_inflow( joinpath(case_folder, "inflows.csv"), nHyd; num_stages=num_stages ) + # Extract initial volumes and volume bounds initial_state = [hydro["initial_volume"] for hydro in hydro_file] max_volume = [hydro["max_volume"] for hydro in hydro_file] + min_volume = [hydro["min_volume"] for hydro in hydro_file] + # Build upstream connectivity: for each unit, who feeds into it? + # The hydro.json stores downstream references; we invert them here. + # index_to_pos maps the hydro "index" field to the array position (1-based) + index_to_pos = Dict(hydro["index"] => i for (i, hydro) in enumerate(hydro_file)) + # upstream_turn[r] = [(upstream_array_pos, upstream_max_turn), ...] + upstream_turn = [Tuple{Int,Float64}[] for _ in 1:nHyd] + # upstream_spill[r] = [(upstream_array_pos, Inf), ...] — spill is unbounded + upstream_spill = [Tuple{Int,Float64}[] for _ in 1:nHyd] + for (i, hydro) in enumerate(hydro_file) + # Turbine outflow from unit i feeds into each downstream unit + for ds_idx in hydro["downstream_turn"] + ds_pos = index_to_pos[ds_idx] + push!(upstream_turn[ds_pos], (i, hydro["max_turn"])) + end + # Spillage from unit i feeds into each downstream unit + for ds_idx in hydro["downstream_spill"] + ds_pos = index_to_pos[ds_idx] + push!(upstream_spill[ds_pos], (i, Inf)) + end + end + + # Assemble hydro metadata for policy construction (e.g. HydroReachablePolicy) + hydro_meta = ( + nHyd = nHyd, + min_vol = min_volume, + max_vol = max_volume, + min_turn = [hydro["min_turn"] for hydro in hydro_file], + max_turn = [hydro["max_turn"] for hydro in hydro_file], + initial_volume = initial_state, + downstream_turn = [hydro["downstream_turn"] for hydro in hydro_file], + downstream_spill = [hydro["downstream_spill"] for hydro in hydro_file], + upstream_turn = upstream_turn, + upstream_spill = upstream_spill, + K = K, + production_factor = [hydro["production_factor"] for hydro in hydro_file], + ) + + # Allocate per-stage containers subproblems = Vector{JuMP.Model}(undef, num_stages) state_params_in = Vector{Vector{Any}}(undef, num_stages) state_params_out = Vector{Vector{Tuple{Any,VariableRef}}}(undef, num_stages) @@ -52,6 +196,7 @@ function build_hydropowermodels( ) for t in 1:num_stages + # Read the stage subproblem from MOF file subproblems[t] = JuMP.read_from_file( joinpath(case_folder, subproblem_file); use_nlp_block=false ) @@ -59,25 +204,52 @@ function build_hydropowermodels( if !isnothing(optimizer) set_optimizer(subproblems[t], optimizer) end - norm_deficit, _deficit = create_deficit!( - subproblems[t], - nHyd; - penalty=penalty, - penalty_l1=penalty_l1, - penalty_l2=penalty_l2, - ) - # delete fix constraints + + if strict + # Strict mode: no deficit variables, no penalty — hard equality + # reservoir_out[i] == target[i] enforced directly + else + # Default mode: create deficit variables with penalty + norm_deficit, _deficit = create_deficit!( + subproblems[t], + nHyd; + penalty=penalty, + penalty_l1=penalty_l1, + penalty_l2=penalty_l2, + ) + end + + # Delete fix constraints (fixed-value equality constraints on variables) for con in JuMP.all_constraints(subproblems[t], VariableRef, MOI.EqualTo{Float64}) delete(subproblems[t], con) end + # Identify reservoir and inflow variables by name pattern state_params_in[t], state_param_out, inflow = find_reservoirs_and_inflow( subproblems[t] ) + # Convert incoming state variables to parameters state_params_in[t] = variable_to_parameter.(subproblems[t], state_params_in[t]) - state_params_out[t] = [ - variable_to_parameter(subproblems[t], state_param_out[i]; deficit=_deficit[i]) - for i in 1:nHyd - ] + + if strict + # Strict mode: hard equality constraint (no deficit slack) + # variable_to_parameter without deficit creates: reservoir_out[i] == parameter + # Returns just the parameter; we manually pair it with the variable + state_params_out[t] = [ + let param = variable_to_parameter(subproblems[t], state_param_out[i]) + (param, state_param_out[i]) + end + for i in 1:nHyd + ] + else + # Default mode: variable_to_parameter with deficit returns (parameter, variable) + state_params_out[t] = [ + variable_to_parameter( + subproblems[t], state_param_out[i]; deficit=_deficit[i] + ) + for i in 1:nHyd + ] + end + # Joint scenarios: all hydro units share the same scenario index ω, # preserving the spatial correlation in the historical inflow data. inflow_params = [variable_to_parameter(subproblems[t], inflow[i]) for i in 1:nHyd] @@ -90,7 +262,7 @@ function build_hydropowermodels( return subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, - max_volume + max_volume, hydro_meta end function ensure_feasibility_cap(state_out, state_in, uncertainty, max_volume) diff --git a/examples/HydroPowerModels/plot_hydro_paired_distributions.jl b/examples/HydroPowerModels/plot_hydro_paired_distributions.jl new file mode 100644 index 0000000..31adb03 --- /dev/null +++ b/examples/HydroPowerModels/plot_hydro_paired_distributions.jl @@ -0,0 +1,209 @@ +# plot_hydro_paired_distributions.jl +# +# Distributional comparison of per-scenario 96-stage rollout costs on the +# PAIRED scenario set (all methods solve the exact same inflow trajectories; +# see paired_scenario_indices in load_hydropowermodels.jl and the +# eval_paired_* scripts). +# +# Produces docs/src/assets/hydro_paired_cost_distributions.png with two panels: +# +# (a) Absolute cost densities — one transparent filled density per method, +# overlaid. Shows location and spread of each method's cost +# distribution over the shared scenarios. +# (b) Paired-difference densities — density of (method − SDDP) per +# scenario, with a reference line at 0. Because the evaluation is +# paired, this panel removes the common between-scenario variance +# (per-scenario cost std ≈ 5.6k vs paired-difference std ≈ 0.5k) and +# makes the comparison decidable at a glance: probability mass left of +# zero is scenarios where the method is cheaper than SDDP. +# +# Densities are Gaussian kernel density estimates computed directly from the +# per-scenario samples (Silverman's bandwidth), no extra packages required. +# +# Data sources (outputs of the paired evaluation scripts; tags identify +# policies; any missing source is skipped so the plot degrades gracefully): +# MAIN bolivia/ACPPowerModel/paired_costs.csv (SDDP column) +# MAIN bolivia/ACPPowerModel/paired_costs_ft16.csv (CPU ft16 column) +# EXA .../results/paired_exa_strict_warm_*.jld2 (costs_stagewise) +# +# Usage (inside a SLURM job; examples/HydroPowerModels project): +# julia --project plot_hydro_paired_distributions.jl +# +# Colors: fixed Okabe–Ito assignments, consistent with +# plot_hydro_strict_convergence.jl (SDDP green, CPU subproblems vermillion, +# GPU DE blue, second GPU variant orange). Identity is encoded by color + +# legend + direct mean annotations, never color alone. + +using CSV, DataFrames +using JLD2 +using Plots +using Statistics + +const HPM_DIR = dirname(@__FILE__) +const OUT_DIR = joinpath(HPM_DIR, "bolivia", "ACPPowerModel") +const EXA_RESULTS = "/storage/scratch1/9/arosemberg3/DecisionRulesExa.jl/" * + "examples/HydroPowerModels/bolivia/ACPPowerModel/results" +const DOCS_ASSETS = joinpath(HPM_DIR, "..", "..", "docs", "src", "assets") +mkpath(DOCS_ASSETS) + +# Tags identify POLICIES (the paired protocol itself is fixed by seed; +# see paired_scenario_indices in load_hydropowermodels.jl). + +# Fixed categorical assignments (Okabe–Ito; validated CVD-safe set). The +# mapping follows the entity, matching the convergence plot, and is never +# reassigned when a series is missing. +const COLORS = Dict( + "SDDP (SOC-WR / ACP)" => colorant"#009E73", # green + "TS-DDR strict subproblems (CPU)" => colorant"#D55E00", # vermillion + "TS-DDR strict DE (GPU, H256x2)" => colorant"#0072B2", # blue + "TS-DDR strict DE (GPU, H128x3)" => colorant"#E69F00", # orange +) + +""" + gaussian_kde(samples; npoints=400) -> (xs, density) + +Gaussian kernel density estimate with Silverman's rule-of-thumb bandwidth + +```math +h = 0.9 \\, \\min(\\hat\\sigma, \\mathrm{IQR}/1.34) \\; n^{-1/5}. +``` + +Evaluated on an even grid spanning the sample range padded by three +bandwidths, so the tails fall smoothly to zero inside the plot. +""" +function gaussian_kde(samples::AbstractVector{<:Real}; npoints::Int=400) + n = length(samples) + σ = std(samples) + iqr = quantile(samples, 0.75) - quantile(samples, 0.25) + # Guard degenerate spread (all-equal samples) with a tiny bandwidth. + h = max(0.9 * min(σ, iqr / 1.34) * n^(-1/5), eps(Float64) + 1e-9 * max(abs(mean(samples)), 1.0)) + lo, hi = minimum(samples) - 3h, maximum(samples) + 3h + xs = range(lo, hi; length=npoints) + dens = [sum(exp(-0.5 * ((x - s) / h)^2) for s in samples) / (n * h * sqrt(2π)) for x in xs] + return collect(xs), dens +end + +# ── Load per-scenario cost vectors (skip missing sources gracefully) ───────── + +""" + load_series() -> Vector{Pair{String,Vector{Float64}}} + +Collect every available method's per-scenario paired cost vector, in the +fixed display order. SDDP first (it is the comparison baseline for panel b). +""" +function load_series() + series = Pair{String,Vector{Float64}}[] + + # MAIN tagged CSV: SDDP column + TS-DDR (.026/ft16) columns. + main_csv = joinpath(OUT_DIR, "paired_costs.csv") + if isfile(main_csv) + df = CSV.read(main_csv, DataFrame) + for name in names(df) + if occursin("SDDP", name) + push!(series, "SDDP (SOC-WR / ACP)" => Float64.(df[!, name])) + end + end + else + @warn "missing $main_csv — SDDP/anchor series skipped" + end + + ft_csv = joinpath(OUT_DIR, "paired_costs_ft16.csv") + if isfile(ft_csv) + df = CSV.read(ft_csv, DataFrame) + col = first(names(df)) + push!(series, "TS-DDR strict subproblems (CPU)" => Float64.(df[!, col])) + else + @warn "missing $ft_csv — CPU subproblems series skipped" + end + + for (label, exa_tag) in ( + "TS-DDR strict DE (GPU, H256x2)" => "warm_H256x2", + "TS-DDR strict DE (GPU, H128x3)" => "warm_H128x3", + ) + f = joinpath(EXA_RESULTS, "paired_exa_strict_$(exa_tag).jld2") + if isfile(f) + d = JLD2.load(f) + key = haskey(d, "costs_stagewise") ? "costs_stagewise" : "costs" + push!(series, label => Float64.(vec(d[key]))) + else + @warn "missing $f — $label skipped" + end + end + + return series +end + +series = load_series() +isempty(series) && error("No paired cost sources found") + +# All series must share the paired scenario count for panel (b) to be valid. +n_scen = length(last(first(series))) +@assert all(length(v) == n_scen for (_, v) in series) "Series lengths differ — not the same paired protocol" +println("Loaded $(length(series)) methods × $n_scen paired scenarios") + +# ── Panel (a): absolute cost densities ──────────────────────────────────────── + +plt_abs = plot(; + xlabel="96-stage rollout cost", + ylabel="Density", + title="Per-scenario cost distributions ($n_scen paired scenarios)", + legend=:topright, + grid=:y, gridalpha=0.15, + size=(900, 400), +) +for (label, costs) in series + xs, dens = gaussian_kde(costs) + c = COLORS[label] + plot!(plt_abs, xs, dens; label=label, color=c, linewidth=2, + fill=(0, 0.30, c)) + # Direct mean annotation as a short tick, text in muted ink (not series + # color), placed above the curve peak region. + vline!(plt_abs, [mean(costs)]; color=c, linestyle=:dash, linewidth=1, + alpha=0.7, label="") +end + +# ── Panel (b): paired differences vs SDDP ───────────────────────────────────── + +sddp = last(first(series)) +plt_diff = plot(; + xlabel="Paired cost difference vs SDDP (method − SDDP, same scenario)", + ylabel="Density", + title="Paired differences (mass left of 0 = scenarios beating SDDP)", + legend=:topright, + grid=:y, gridalpha=0.15, + size=(900, 400), +) +vline!(plt_diff, [0.0]; color=:gray30, linestyle=:solid, linewidth=1, label="") +for (label, costs) in series[2:end] # skip SDDP itself + diffs = costs .- sddp + xs, dens = gaussian_kde(diffs) + c = COLORS[label] + winrate = round(100 * count(<(0), diffs) / n_scen; digits=1) + plot!(plt_diff, xs, dens; + label="$label (win $winrate%)", + color=c, linewidth=2, fill=(0, 0.30, c)) + vline!(plt_diff, [mean(diffs)]; color=c, linestyle=:dash, linewidth=1, + alpha=0.7, label="") +end + +final = plot(plt_abs, plt_diff; layout=(2, 1), size=(900, 800), + left_margin=8Plots.mm, bottom_margin=6Plots.mm) +outfile = joinpath(DOCS_ASSETS, "hydro_paired_cost_distributions.png") +savefig(final, outfile) +println("Saved: $outfile") + +# ── Summary table ───────────────────────────────────────────────────────────── + +println("\n", "="^78) +println(rpad("Method", 36), rpad("Mean", 12), rpad("Std", 10), + rpad("Δ vs SDDP", 12), "SEM(Δ)") +println("-"^78) +for (label, costs) in series + d = costs .- sddp + println(rpad(label, 36), + rpad(string(round(mean(costs); digits=1)), 12), + rpad(string(round(std(costs); digits=1)), 10), + rpad(label == first(first(series)) ? "—" : string(round(mean(d); digits=1)), 12), + label == first(first(series)) ? "—" : string(round(std(d) / sqrt(n_scen); digits=1))) +end +println("="^78) diff --git a/examples/HydroPowerModels/plot_hydro_strict_convergence.jl b/examples/HydroPowerModels/plot_hydro_strict_convergence.jl new file mode 100644 index 0000000..67d3707 --- /dev/null +++ b/examples/HydroPowerModels/plot_hydro_strict_convergence.jl @@ -0,0 +1,462 @@ +# plot_hydro_strict_convergence.jl +# +# Regenerate the docs convergence figure for the 126-stage Bolivia hydro case, +# comparing the final training schedules against WALL-CLOCK TIME: +# +# 1. SDDP (SOC-WR relaxation of the AC OPF), CPU: +# - forward-pass simulation cost per iteration (statistical upper estimate) +# - lower bound per iteration, plus a horizontal dashed line at its final +# value ("SDDP lower bound (SOC-WR)") — no policy can go below it. +# Source: OFFLINE parse of sddp/SDDP.log (SDDP.jl iteration table). +# 2. TS-DDR strict, stage-wise subproblems (CPU, Ipopt): +# phase 1 uses single-sample gradients and training-loss checkpointing; +# phase 2 switches to 16-sample gradients, decayed LR, and held-out +# rollout checkpointing. Phase-2 runtime is offset by phase-1 runtime. +# 3. TS-DDR strict, full-horizon deterministic equivalent (GPU, ExaModels+MadNLP): +# phase 1 uses the strict GPU DE trainer; phase 2 switches to +# hard-reactive physics, larger mini-batches, decayed LR, and held-out +# rollout checkpointing. Phase-2 runtime is offset by phase-1 runtime. +# +# Outputs (both saved into docs/src/assets/, path relative to this script): +# - hydro_training_convergence_by_time.png (x = wall-clock hours, log10) +# - hydro_training_convergence_by_step.png (x = training iteration, log10) +# +# Usage: +# cd examples/HydroPowerModels +# julia --project plot_hydro_strict_convergence.jl +# +# The script must run under the examples/HydroPowerModels environment +# (examples/HydroPowerModels/Project.toml), which carries both `Plots` and +# `Wandb` — the root DecisionRules Project.toml does NOT have them. +# Network access + W&B credentials (~/.netrc or WANDB_API_KEY) are required +# for series 2 and 3; series 1 (SDDP) is parsed fully offline. +# +# Environment knobs: +# DR_PLOT_RUNS Override the W&B run-name substrings, format: +# "main=,;exa=,,..." +# Either part may be omitted to keep its default. +# DR_PLOT_SMOOTH Rolling-mean window for the noisy single-sample losses +# (default 25). +# +# The script is idempotent: it only reads the log/W&B histories and overwrites +# the two PNGs; re-running it produces the same figures for the same data. + +using Plots # plotting backend (GR) — writes the PNGs +using Statistics # mean/median/quantile for smoothing and axis limits +using Printf # formatted summary table + +# W&B access follows the exact pattern of compare_hydro_results.jl: +# Wandb.jl wraps the Python `wandb` module via PythonCall; `PC` is PythonCall +# (recovered as the parent module of the wrapped Python object type) and gives +# us `pyconvert`/`pylist` for Julia<->Python conversion. +import Wandb +const wb = Wandb.wandb +const PC = parentmodule(typeof(wb)) + +# ── Paths ───────────────────────────────────────────────────────────────────── + +# Docs asset directory, relative to this script (same convention as +# compare_hydro_results.jl) so the script works from any checkout location. +const DOCS_ASSETS = joinpath(@__DIR__, "..", "..", "docs", "src", "assets") +mkpath(DOCS_ASSETS) + +# SDDP.jl training log (offline source for series 1). +const SDDP_LOG = joinpath(@__DIR__, "sddp", "SDDP.log") + +# ── Run-name configuration (overridable via DR_PLOT_RUNS) ───────────────────── + +# Series 2: strict stage-wise subproblems (CPU, Ipopt), represented as a +# two-phase schedule. Runtime is cumulative across phases. +DEFAULT_MAIN_RUNS = [ + "bolivia-ACPPowerModel-h126-r96-subproblems-strict-2026-07-01T09:41:53.026", + "bolivia-ACPPowerModel-h126-r96-subproblems-strict-Hlinear-nt16-warm-2026-07-04T21:13:01.712", +] + +# Series 3: strict full-horizon DE (GPU), represented as a two-phase schedule. +# Runtime is cumulative across phases. +DEFAULT_EXA_RUNS = [ + "bolivia-ACPPowerModel-h126-r96-deteq-strict-gpu-H256_256-20260704-131003", + "bolivia-ACPPowerModel-h126-r96-deteq-strict-gpu-H256_256-nt4-rollout-warm-rqhard-20260705-121353", +] + +""" + parse_run_overrides(spec) -> (main::Vector{String}, exa::Vector{String}) + +Parse the `DR_PLOT_RUNS` override string of the form +`"main=,;exa=,,..."`. Each `key=value` +segment is optional; omitted keys keep the defaults above. Unknown keys raise +an error so typos fail loudly instead of silently plotting the default runs. +""" +function parse_run_overrides(spec::AbstractString) + main = copy(DEFAULT_MAIN_RUNS) # start from defaults … + exa = copy(DEFAULT_EXA_RUNS) + isempty(strip(spec)) && return main, exa # empty knob → defaults + for part in split(spec, ';'; keepempty=false) # each "key=value" segment + kv = split(part, '='; limit=2) + length(kv) == 2 || error("DR_PLOT_RUNS segment '$part' is not key=value") + key, val = strip(kv[1]), strip(kv[2]) + if key == "main" + main = String.(strip.(split(val, ','; keepempty=false))) + elseif key == "exa" + exa = String.(strip.(split(val, ','; keepempty=false))) # comma list + else + error("DR_PLOT_RUNS: unknown key '$key' (expected 'main' or 'exa')") + end + end + return main, exa +end + +const MAIN_RUN_SUBSTRS, EXA_RUN_SUBSTRS = parse_run_overrides(get(ENV, "DR_PLOT_RUNS", "")) + +# Rolling-mean window for the single-sample (batch size 1) training losses. +const SMOOTH_WINDOW = parse(Int, get(ENV, "DR_PLOT_SMOOTH", "25")) + +# ── SDDP.log parsing (offline) ──────────────────────────────────────────────── + +""" + parse_sddp_log(path) -> (iteration, simulation, bound, time_s) + +Parse the standard SDDP.jl iteration table from a training log. The table +header is `iteration simulation bound time (s) solves pid` and each data +row is `iter simulation bound cumulative_time_s total_solves pid` +(times are CUMULATIVE seconds since training start; values are in the model's +objective units, here \$). + +The log may contain several training blocks (restarted jobs each print a fresh +header). This file (sddp/SDDP.log) contains THREE blocks: the first two are +short aborted trial runs (2 and 1 iterations); the LAST block is the real +441-iteration training run whose final bound is ≈ 378,207. We therefore locate +the LAST table header and parse the data rows that follow it. Note the last +block has no closing "status/total time" footer (the job hit its wall limit), +which is fine: parsing simply stops at the first non-data line or EOF. +""" +function parse_sddp_log(path::AbstractString) + lines = readlines(path) # whole log in memory (small) + # Locate every iteration-table header; we will use the LAST one. + header_re = r"^\s*iteration\s+simulation\s+bound\s+time" + header_idxs = findall(l -> occursin(header_re, l), lines) + isempty(header_idxs) && error("No SDDP iteration table found in $path") + start = header_idxs[end] + 2 # +1 header, +1 dashed rule + # Data row: iter (int), simulation (float), bound (float), time (float), solves (int), pid. + row_re = r"^\s*(\d+)\s+([0-9.eE+-]+)\s+([0-9.eE+-]+)\s+([0-9.eE+-]+)\s+(\d+)\s+(\d+)\s*$" + iters, sims, bounds, times = Int[], Float64[], Float64[], Float64[] + for l in lines[start:end] # scan until table ends + m = match(row_re, l) + m === nothing && break # dashed rule / footer / EOF + push!(iters, parse(Int, m.captures[1])) + push!(sims, parse(Float64, m.captures[2])) + push!(bounds, parse(Float64, m.captures[3])) + push!(times, parse(Float64, m.captures[4])) + end + isempty(iters) && error("SDDP iteration table at line $(start-2) had no data rows") + return iters, sims, bounds, times +end + +# ── Smoothing ───────────────────────────────────────────────────────────────── + +""" + rolling_mean(v, w) -> Vector{Float64} + +Trailing (causal) rolling mean: `out[i] = mean(v[max(1, i-w+1) : i])`, i.e. the +average of the last `w` samples up to and including `i`. A trailing window is +used (rather than centered) so the smoothed curve never uses future samples — +appropriate for a convergence trace. The first `w-1` points average over the +shorter available prefix. +""" +function rolling_mean(v::AbstractVector{<:Real}, w::Int) + n = length(v) + out = Vector{Float64}(undef, n) + for i in 1:n # O(n·w); n ≲ 10⁴ so fine + lo = max(1, i - w + 1) + out[i] = mean(@view v[lo:i]) + end + return out +end + +# ── W&B helpers (patterned on compare_hydro_results.jl) ────────────────────── + +""" + find_runs_by_substring(api, substrs; project="RL", max_scan=400) -> Dict + +Scan the project's runs in `-created_at` order (most recent first) and return a +`Dict(substring => run)` mapping each requested name-substring to the MOST +RECENT run whose name contains it. Substrings with no match within `max_scan` +runs are absent from the result (caller warns + skips). +""" +function find_runs_by_substring(api, substrs::Vector{String}; project::String="RL", max_scan::Int=400) + all_runs = api.runs(project, order="-created_at") # lazy paginated iterator + found = Dict{String,Any}() + for i in 0:(max_scan - 1) # python 0-based indexing + r = try + all_runs[i] + catch + break # ran off the end of the project + end + name = try + PC.pyconvert(String, r.name) + catch + continue # unreadable run entry — skip + end + for s in substrs # first (= newest) match wins + if !haskey(found, s) && occursin(s, name) + found[s] = r + end + end + length(found) == length(substrs) && break # all found — stop scanning + end + return found +end + +""" + get_history_with_runtime(r, metric) -> (t_seconds, values) + +Fetch the full logged history of `metric` from W&B run `r` via `scan_history`, +together with the built-in `_runtime` field (wall-clock seconds since run +start, attached by W&B to every logged row). Rows missing either field, or +with non-finite metric values, are dropped. Returns two aligned vectors. +""" +function get_history_with_runtime(r, metric::String) + keys_list = PC.pylist([metric, "_runtime"]) # request both columns + hist = r.scan_history(keys=keys_list) + ts, vals = Float64[], Float64[] + for row in hist # rows stream from the API + v = try PC.pyconvert(Float64, get(row, metric, nothing)) catch; nothing end + t = try PC.pyconvert(Float64, get(row, "_runtime", nothing)) catch; nothing end + (v === nothing || t === nothing || !isfinite(v)) && continue + push!(ts, t); push!(vals, v) + end + return ts, vals +end + +""" + prefix_to_best_training_loss(t, v) -> (t, v) + +Return the prefix ending at the best raw training-loss row. The phase-1 strict +CPU checkpoint is selected by this metric. +""" +function prefix_to_best_training_loss(t::Vector{Float64}, v::Vector{Float64}) + isempty(v) && return t, v + stop = argmin(v) + return t[1:stop], v[1:stop] +end + +""" + prefix_to_best_rollout(r, t, v) -> (t, v) + +Second phases use `DR_SAVE_METRIC=rollout`, so the selected checkpoint is the +best held-out rollout objective rather than the noisy training loss. If the +rollout metric is available, cut the trace at that runtime. +""" +function prefix_to_best_rollout(r, t::Vector{Float64}, v::Vector{Float64}) + rt, rv = get_history_with_runtime(r, "metrics/rollout_objective_no_deficit") + isempty(rv) && return t, v + cutoff = rt[argmin(rv)] + keep = findall(x -> x <= cutoff, t) + isempty(keep) && return t[1:1], v[1:1] + return t[keep], v[keep] +end + +""" + append_continuation!(all_t, all_v, t, v) + +Append a later phase after the current last wall-clock time. Each W&B run +records runtime from its own start; this converts phases to total runtime. +""" +function append_continuation!(all_t::Vector{Float64}, all_v::Vector{Float64}, + t::Vector{Float64}, v::Vector{Float64}) + isempty(v) && return all_t, all_v + offset = isempty(all_t) ? 0.0 : all_t[end] + append!(all_t, offset .+ t) + append!(all_v, v) + return all_t, all_v +end + +# ── Load all series ─────────────────────────────────────────────────────────── + +@info "Parsing SDDP log (offline): $SDDP_LOG" +sddp_iter, sddp_sim, sddp_bound, sddp_time = parse_sddp_log(SDDP_LOG) +sddp_hours = sddp_time ./ 3600 # cumulative s → h +bound_final = sddp_bound[end] # ≈ 3.782e5 — the SOC-WR lower bound +@info " $(length(sddp_iter)) SDDP iterations, final bound = $(round(bound_final; digits=1)), " * + "total time = $(round(sddp_hours[end]; digits=2)) h" + +@info "Connecting to W&B (project RL)..." +api = wb.Api() +wanted = vcat(MAIN_RUN_SUBSTRS, EXA_RUN_SUBSTRS) # all substrings in one scan +found = find_runs_by_substring(api, wanted) + +# Series 2: MAIN strict stage-wise subproblems (CPU). `metrics/training_loss` +# is the per-batch mean 126-stage objective. The default is a continuation +# curve: stage 1 up to its selected training-loss checkpoint, then phase 2 up +# to its best held-out rollout checkpoint, with runtimes offset. +main_hours = Float64[]; main_loss = Float64[] +for (i, s) in enumerate(MAIN_RUN_SUBSTRS) + if !haskey(found, s) + @warn "MAIN subproblems-strict run not found (substring '$s') — continuation segment skipped" + continue + end + r = found[s] + name = PC.pyconvert(String, r.name) + @info " MAIN run segment $i: $name ($(PC.pyconvert(String, r.state)))" + t, v = get_history_with_runtime(r, "metrics/training_loss") + if isempty(v) + @warn " no metrics/training_loss rows yet for $name — segment skipped" + continue + end + if i == 1 + t, v = prefix_to_best_training_loss(t, v) + else + t, v = prefix_to_best_rollout(r, t, v) + end + append_continuation!(main_hours, main_loss, t ./ 3600, v) +end + +# Series 3: strict-DE GPU schedule. `metrics/training_loss` is the +# highest-resolution loss logged by DecisionRulesExa. +exa_hours = Float64[]; exa_loss = Float64[] +for (i, s) in enumerate(EXA_RUN_SUBSTRS) + if !haskey(found, s) + @warn "Exa GPU run not found (substring '$s') — continuation segment skipped" + continue + end + r = found[s] + name = PC.pyconvert(String, r.name) + @info " Exa run segment $i: $name ($(PC.pyconvert(String, r.state)))" + t, v = get_history_with_runtime(r, "metrics/training_loss") + if isempty(v) + @warn " no metrics/training_loss rows yet for $name — segment skipped" + continue + end + if i > 1 + t, v = prefix_to_best_rollout(r, t, v) + end + append_continuation!(exa_hours, exa_loss, t ./ 3600, v) +end + +main_smooth = rolling_mean(main_loss, SMOOTH_WINDOW) +exa_smooth = rolling_mean(exa_loss, SMOOTH_WINDOW) + +# ── Axis limits & styling ───────────────────────────────────────────────────── + +# Colors: Okabe–Ito colorblind-safe hues, kept in the same blue/red/green +# family order as the existing docs plots (compare_hydro_results.jl uses +# :blue/:red/:green for DE/subproblems/other): +# SDDP simulation → blue, SDDP bound → neutral dark gray, +# MAIN subproblems (CPU) → vermillion, Exa DE (GPU) family → bluish green. +const C_SDDP = "#0072B2" # blue +const C_BOUND = "#4D4D4D" # neutral dark gray (bound is a reference, not a competitor) +const C_MAIN = "#D55E00" # vermillion +const C_EXA = "#009E73" # bluish green + +# Y-limits: the interesting band is near the SDDP lower bound. Early-training +# TS-DDR losses can be orders of magnitude larger (deficit-penalty transients); +# cropping them keeps the convergence region readable. Lower limit sits just +# below the bound (nothing can be below it); upper limit covers the SDDP +# forward-pass scatter and every method's converged level (robust tail median), +# hard-capped at 3× the bound so a diverged run cannot blow up the scale. +""" + tail_level(v; frac=0.1) -> Float64 + +Robust "final level" of a series: the median of its last `max(1, frac·n)` +samples. Used only for choosing the y-axis upper limit. +""" +tail_level(v::AbstractVector{<:Real}; frac::Float64=0.1) = + isempty(v) ? NaN : median(@view v[max(1, end - max(1, round(Int, frac*length(v))) + 1):end]) + +finals = Float64[tail_level(sddp_sim)] # SDDP forward-pass level +isempty(main_smooth) || push!(finals, tail_level(main_smooth)) +isempty(exa_smooth) || push!(finals, tail_level(exa_smooth)) +y_lo = 0.985 * bound_final +y_hi = min(3.0 * bound_final, + 1.10 * max(maximum(sddp_sim), maximum(filter(isfinite, finals)))) + +# X-axis choice: log10. Justification: the series span ~3 orders of magnitude +# of wall time (GPU DE runs log their first losses within minutes, while the +# SDDP and CPU subproblem schedules run much longer). On a +# linear axis the fast-GPU story is squashed into an invisible sliver at the +# left edge; log10 lets both the early GPU descent and the long CPU tails read. +# Same reasoning for the by-iteration plot (441 SDDP iterations vs ~8000 +# TS-DDR batches). Zero/near-zero times are clamped to 10 s to stay on-axis. +clamp_hours(h) = max.(h, 10 / 3600) + +# ── Plotting ────────────────────────────────────────────────────────────────── + +""" + convergence_plot(xsel, xlabel) -> Plots.Plot + +Build one convergence figure. `xsel(series_kind, i)` is inlined below instead — +this helper simply takes, per series, precomputed x-vectors, so both figures +(by-time and by-step) share identical styling and differ only in x data. +`xs` is a NamedTuple with fields `sddp`, `main`, `exa::Vector` holding the +x-vectors for each series. +""" +function convergence_plot(xs, xlabel::String) + plt = plot(; size = (900, 540), # docs figure size + xlabel = xlabel, + ylabel = "Operational Cost (126 stages)", # docs-style cost label + title = "Training Convergence (126-stage Bolivia AC): GPU vs CPU", + legend = :topright, + xscale = :log10, # see justification above + ylims = (y_lo, y_hi), + left_margin = 5Plots.mm, bottom_margin = 4Plots.mm) + + # SDDP forward-pass simulation cost: per-iteration Monte-Carlo estimate of + # the current policy's cost (an upper, noisy companion to the bound). + plot!(plt, xs.sddp, sddp_sim; color = C_SDDP, alpha = 0.55, lw = 1.2, + label = "SDDP forward simulation (SOC-WR, CPU)") + # SDDP lower bound: monotone deterministic series. + plot!(plt, xs.sddp, sddp_bound; color = C_BOUND, lw = 2.0, + label = "SDDP lower bound (SOC-WR)") + # TS-DDR curves use a trailing rolling mean (window $SMOOTH_WINDOW) to make + # small-sample training losses readable without using future samples. + if !isempty(main_loss) + plot!(plt, xs.main, main_smooth; color = C_MAIN, lw = 2.0, + label = "TS-DDR strict subproblems (CPU, Ipopt)") + end + + if !isempty(exa_loss) + plot!(plt, xs.exa, exa_smooth; color = C_EXA, lw = 2.3, + label = "TS-DDR strict DE (GPU, MadNLP/cuDSS)") + end + return plt +end + +# Figure 1 — by wall-clock time (hours). This is the headline GPU-vs-CPU plot. +xs_time = (sddp = clamp_hours(sddp_hours), + main = clamp_hours(main_hours), + exa = clamp_hours(exa_hours)) +plt_time = convergence_plot(xs_time, "Wall-clock time (hours, log scale)") +savefig(plt_time, joinpath(DOCS_ASSETS, "hydro_training_convergence_by_time.png")) +println("Saved hydro_training_convergence_by_time.png") + +# Figure 2 — by training iteration (SDDP iterations vs TS-DDR batches). +# Iteration costs differ wildly across methods (one SDDP iteration = 1890 +# solves; one TS-DDR batch = 1 DE solve or 126 stage solves) — this companion +# shows per-iteration progress, while Figure 1 shows the honest time story. +xs_step = (sddp = Float64.(sddp_iter), + main = collect(1.0:length(main_loss)), + exa = collect(1.0:length(exa_loss))) +plt_step = convergence_plot(xs_step, "Training iteration (log scale)") +savefig(plt_step, joinpath(DOCS_ASSETS, "hydro_training_convergence_by_step.png")) +println("Saved hydro_training_convergence_by_step.png") + +# ── Summary table ───────────────────────────────────────────────────────────── + +println("\n", "="^88) +println("Summary — final value & wall time per series") +println("="^88) +@printf("%-48s %14s %10s %8s\n", "Series", "Final value", "Hours", "Points") +println("-"^88) +@printf("%-48s %14.1f %10.2f %8d\n", "SDDP lower bound (SOC-WR)", + bound_final, sddp_hours[end], length(sddp_bound)) +@printf("%-48s %14.1f %10.2f %8d\n", "SDDP forward simulation", + tail_level(sddp_sim), sddp_hours[end], length(sddp_sim)) +if !isempty(main_loss) + @printf("%-48s %14.1f %10.2f %8d\n", "TS-DDR strict subproblems (CPU) [smoothed]", + main_smooth[end], main_hours[end], length(main_loss)) +end +if !isempty(exa_loss) + @printf("%-48s %14.1f %10.2f %8d\n", "TS-DDR strict DE GPU [smoothed]", + exa_smooth[end], exa_hours[end], length(exa_loss)) +end +println("="^88) diff --git a/examples/HydroPowerModels/sddp/Project.toml b/examples/HydroPowerModels/sddp/Project.toml index 2aac97c..f3c87c5 100644 --- a/examples/HydroPowerModels/sddp/Project.toml +++ b/examples/HydroPowerModels/sddp/Project.toml @@ -1,4 +1,5 @@ [deps] +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" Clarabel = "61c947e1-3e6d-4ee4-985a-eec8c727bd6e" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" diff --git a/examples/HydroPowerModels/sddp/eval_paired_sddp.jl b/examples/HydroPowerModels/sddp/eval_paired_sddp.jl new file mode 100644 index 0000000..fd66fd4 --- /dev/null +++ b/examples/HydroPowerModels/sddp/eval_paired_sddp.jl @@ -0,0 +1,226 @@ +# Paired SDDP simulation on the seeded paired protocol via SDDP.Historical. +# +# Scenario indices are generated from PAIRED_SCENARIO_SEED (identical to +# eval_paired_tsddr.jl's protocol), so the SDDP policy is simulated under the +# exact inflow realizations every TS-DDR evaluation uses. +# +# Usage: +# julia --project -t auto eval_paired_sddp.jl +using MadNLP +using StableRNGs +using HydroPowerModels +using JuMP +using PowerModels +using Statistics +using SDDP: SDDP +using DelimitedFiles +using CSV, DataFrames + +const CASE = "bolivia" +const SDDP_DIR = dirname(@__FILE__) +const HYDRO_DIR = dirname(SDDP_DIR) +const CASE_DIR = joinpath(HYDRO_DIR, CASE) +const RM_STAGES = 30 +const REPORT_STAGES = 96 +const NUM_STAGES = REPORT_STAGES + RM_STAGES +const FORMULATION = ACPPowerModel +const FORMULATION_B = SOCWRConicPowerModel + +# ── Paired scenario indices (seeded protocol) ────────────────────────────── +# Identical generation to load_hydropowermodels.jl's paired_scenario_indices: +# entry [t, s] is uniform on 1:nCen from StableRNG(PAIRED_SCENARIO_SEED), so +# this script and every TS-DDR evaluation realize the same inflow at the same +# stage of the same paired scenario — with no shared data file. nCen is +# derived from the inflow data (columns ÷ hydro units), never hardcoded. +const PAIRED_SCENARIO_SEED = 20260706 +# Fixed generated shape shared by ALL consumers (see load_hydropowermodels.jl: +# arrays of different shapes consume the RNG stream differently, so every +# script must generate exactly this shape and slice what it needs). +const PAIRED_NUM_STAGES = 126 +const N_HYDRO = 11 +@assert NUM_STAGES == PAIRED_NUM_STAGES "SDDP horizon must equal the paired protocol shape" +num_scenarios = parse(Int, get(ENV, "DR_NUM_SCENARIOS", "500")) +nCen = div(size(readdlm(joinpath(HYDRO_DIR, CASE, "inflows.csv"), ','), 2), N_HYDRO) +# ALWAYS generate the full protocol matrix (fixed shape — see the note above), +# then optionally simulate only a shard of its columns so the 500 scenarios +# can run in parallel across nodes (DR_SCENARIO_FIRST/LAST, 1-based inclusive). +all_indices = rand(StableRNG(PAIRED_SCENARIO_SEED), 1:nCen, PAIRED_NUM_STAGES, num_scenarios) +scen_first = parse(Int, get(ENV, "DR_SCENARIO_FIRST", "1")) +scen_last = parse(Int, get(ENV, "DR_SCENARIO_LAST", string(num_scenarios))) +@assert 1 <= scen_first <= scen_last <= num_scenarios +scen_range = scen_first:scen_last +println("Paired protocol: seed=$PAIRED_SCENARIO_SEED, $(PAIRED_NUM_STAGES)×$(num_scenarios), nCen=$nCen") +println("Evaluating scenarios $scen_first:$scen_last, $REPORT_STAGES reported stages (of $NUM_STAGES total)") + +# ── Build SDDP model and load cuts ──────────────────────────────────────── +alldata = HydroPowerModels.parse_folder(CASE_DIR) +for load in values(alldata[1]["powersystem"]["load"]) + load["qd"] = load["qd"] * 0.6 + load["pd"] = load["pd"] * 0.6 +end + +params = create_param(; + stages=NUM_STAGES, + model_constructor_grid=FORMULATION, + post_method=PowerModels.build_opf, + # Hardened solver settings for the 500-scenario simulation: with bare + # defaults one nonconvex ACP node failed to converge on one seeded draw + # (SDDP aborts the whole simulation on any node failure). Larger + # iteration budget + explicit tolerance make every node solvable; the + # cuts themselves are unaffected (they were trained separately). + optimizer=() -> MadNLP.Optimizer(; + print_level=0, + max_iter=9000, + tol=1e-6, + ), +) + +m = hydro_thermal_operation(alldata, params) + +cuts_file = joinpath( + CASE_DIR, + string(FORMULATION), + string(FORMULATION_B) * "-" * string(FORMULATION) * ".cuts.json", +) +SDDP.read_cuts_from_file(m.forward_graph, cuts_file) +println("Loaded cuts: $cuts_file") + +# ── Build SDDP.Historical sampling scheme ────────────────────────────────── +# Each scenario is a vector of (node, noise_term) pairs. +# node = stage index, noise_term = scenario column ω ∈ 1:nCen +historical_scenarios = [ + [(t, all_indices[t, s]) for t in 1:NUM_STAGES] + for s in scen_range +] +n_sim = length(scen_range) + +sampling_scheme = SDDP.Historical(historical_scenarios) + +# ── Simulate ─────────────────────────────────────────────────────────────── +println("\nSimulating $n_sim scenarios with SDDP.Historical...") +results = HydroPowerModels.simulate( + m, n_sim; + sampling_scheme=sampling_scheme, +) + +# Verify noise terms match our indices +for (si, s) in enumerate(first(scen_range, min(3, n_sim))), t in 1:min(5, REPORT_STAGES) + recorded_ω = results[:simulations][si][t][:noise_term] + expected_ω = all_indices[t, s] + if recorded_ω != expected_ω + error("Mismatch at scenario $s, stage $t: got ω=$recorded_ω, expected $expected_ω") + end +end +println("Noise term verification passed (spot-checked)") + +# ── Extract results ──────────────────────────────────────────────────────── +nhyd = alldata[1]["hydro"]["nHyd"] +volume_to_mw(volume; k=0.0036) = volume / k + +objective_values = [ + sum(results[:simulations][i][t][:stage_objective] for t in 1:REPORT_STAGES) + for i in 1:n_sim +] + +hydro_vol = [ + mean( + sum( + volume_to_mw(results[:simulations][i][t][:reservoirs][:reservoir][j].out) for + j in 1:nhyd + ) for i in 1:n_sim + ) for t in 1:REPORT_STAGES +] + +num_gen = length(results[:simulations][1][1][:powersystem]["solution"]["gen"]) +hydro_idx = HydroPowerModels.idx_hydro(results[:data][1]) +thermal_gen = [ + mean( + sum( + results[:simulations][i][t][:powersystem]["solution"]["gen"]["$j"]["pg"] * + results[:data][1]["powersystem"]["baseMVA"] for + j in 1:num_gen if !(j in hydro_idx) + ) for i in 1:n_sim + ) for t in 1:REPORT_STAGES +] + +# ── Report ───────────────────────────────────────────────────────────────── +println("\n" * "=" ^ 60) +println("Results: Paired SDDP ($REPORT_STAGES stages, $num_scenarios scenarios)") +println("=" ^ 60) +println(" Mean cost: $(round(mean(objective_values); digits=1))") +println(" Std: $(round(std(objective_values); digits=1))") +println(" Min: $(round(minimum(objective_values); digits=1))") +println(" Max: $(round(maximum(objective_values); digits=1))") +println(" Median: $(round(median(objective_values); digits=1))") +println("=" ^ 60) + +# ── Save results ─────────────────────────────────────────────────────────── +out_dir = joinpath(CASE_DIR, string(FORMULATION)) + +# Shard mode: emit only this shard's per-scenario costs (merged afterwards by +# merge_sddp_shards.jl); the full-run outputs below are skipped. +if n_sim != num_scenarios + shard_file = joinpath(out_dir, "sddp_shard_$(scen_first)_$(scen_last).csv") + CSV.write(shard_file, DataFrame(scenario = collect(scen_range), cost = objective_values)) + println("Shard written: $shard_file") + exit(0) +end + +# Optional output tag (mirrors eval_paired_tsddr.jl): when DR_OUTPUT_TAG is +# set, every output filename gets an _ suffix so re-evaluations at a +# different scenario count never overwrite existing result files. Note that a +# tagged run starts fresh tagged files; the merge-with-existing-columns logic +# only applies within the same tag. +tag_suffix = let tag = get(ENV, "DR_OUTPUT_TAG", "") + isempty(tag) ? "" : "_$(tag)" +end +isempty(tag_suffix) || println("Output tag suffix: $tag_suffix") + +const COL_NAME = "SDDP-SOC (paired)" +costs_file = joinpath(out_dir, "paired_costs$(tag_suffix).csv") +if isfile(costs_file) + df = CSV.read(costs_file, DataFrame) + df[!, COL_NAME] = objective_values +else + df = DataFrame(Symbol(COL_NAME) => objective_values) +end +CSV.write(costs_file, df) +println("Updated: $costs_file") + +vol_file = joinpath(out_dir, "paired_MeanVolume$(tag_suffix).csv") +if isfile(vol_file) + df_vol = CSV.read(vol_file, DataFrame; header=true) + df_vol[!, COL_NAME] = hydro_vol +else + df_vol = DataFrame(Symbol(COL_NAME) => hydro_vol) +end +CSV.write(vol_file, df_vol) +println("Updated: $vol_file") + +gen_file = joinpath(out_dir, "paired_MeanGeneration$(tag_suffix).csv") +if isfile(gen_file) + df_gen = CSV.read(gen_file, DataFrame; header=true) + df_gen[!, COL_NAME] = thermal_gen +else + df_gen = DataFrame(Symbol(COL_NAME) => thermal_gen) +end +CSV.write(gen_file, df_gen) +println("Updated: $gen_file") + +# Per-scenario paired costs for direct comparison +println("\nPer-scenario cost comparison (first 10):") +if isfile(costs_file) + df_all = CSV.read(costs_file, DataFrame) + if hasproperty(df_all, Symbol("TS-DDR (strict, paired)")) + tsddr_costs = df_all[!, "TS-DDR (strict, paired)"] + sddp_costs = df_all[!, COL_NAME] + diffs = sddp_costs .- tsddr_costs + println(" Scenario | SDDP | TS-DDR | Diff") + for s in 1:min(10, num_scenarios) + println(" $(lpad(s, 7)) | $(lpad(round(sddp_costs[s]; digits=1), 9)) | $(lpad(round(tsddr_costs[s]; digits=1), 9)) | $(round(diffs[s]; digits=1))") + end + println("\n Mean diff (SDDP - TS-DDR): $(round(mean(diffs); digits=1))") + println(" Std diff: $(round(std(diffs); digits=1))") + println(" Paired t-test p-value: (compute externally)") + end +end diff --git a/examples/HydroPowerModels/sddp/extract_sddp_trajectories.jl b/examples/HydroPowerModels/sddp/extract_sddp_trajectories.jl new file mode 100644 index 0000000..e13df0b --- /dev/null +++ b/examples/HydroPowerModels/sddp/extract_sddp_trajectories.jl @@ -0,0 +1,170 @@ +# Extract per-stage SDDP trajectory data for cross-validation with JuMP subproblems. +# +# Simulates a small number of scenarios under the trained SDDP policy, then saves +# per-stage reservoir states, inflows, thermal generation, and stage costs to CSV. +# The companion script `validate_sddp_vs_jump.jl` replays these trajectories through +# the DecisionRules JuMP subproblems and compares stage-by-stage costs. +# +# Runs in the sddp/ project environment. +using MadNLP +using HydroPowerModels +using JuMP +using PowerModels +using Statistics +using SDDP: SDDP +using CSV, DataFrames +using DelimitedFiles +using Random + +const CASE = "bolivia" +const SDDP_DIR = dirname(@__FILE__) +const HYDRO_DIR = dirname(SDDP_DIR) +const CASE_DIR = joinpath(HYDRO_DIR, CASE) +const RM_STAGES = 30 +const REPORT_STAGES = parse(Int, get(ENV, "DR_VALIDATE_STAGES", "96")) +const NUM_STAGES = REPORT_STAGES + RM_STAGES +const FORMULATION = ACPPowerModel +const FORMULATION_B = SOCWRConicPowerModel +const NUM_SCENARIOS = parse(Int, get(ENV, "DR_VALIDATE_SCENARIOS", "4")) +const SEED = parse(Int, get(ENV, "DR_VALIDATE_SEED", "42")) + +println("=" ^ 60) +println("Extract SDDP trajectories for cross-validation") +println(" Stages: $REPORT_STAGES (of $NUM_STAGES total)") +println(" Scenarios: $NUM_SCENARIOS") +println(" Seed: $SEED") +println("=" ^ 60) + +# ── Build SDDP model and load cuts ──────────────────────────────────────── +alldata = HydroPowerModels.parse_folder(CASE_DIR) +for load in values(alldata[1]["powersystem"]["load"]) + load["qd"] = load["qd"] * 0.6 + load["pd"] = load["pd"] * 0.6 +end + +params = create_param(; + stages=NUM_STAGES, + model_constructor_grid=FORMULATION, + post_method=PowerModels.build_opf, + optimizer=() -> MadNLP.Optimizer(; print_level=0), +) + +m = hydro_thermal_operation(alldata, params) + +cuts_file = joinpath( + CASE_DIR, string(FORMULATION), + string(FORMULATION_B) * "-" * string(FORMULATION) * ".cuts.json", +) +SDDP.read_cuts_from_file(m.forward_graph, cuts_file) +println("Loaded cuts: $cuts_file") + +# ── Simulate ─────────────────────────────────────────────────────────────── +Random.seed!(SEED) +results = HydroPowerModels.simulate(m, NUM_SCENARIOS) + +nhyd = alldata[1]["hydro"]["nHyd"] +hydro_data = alldata[1]["hydro"] +n_inflow_rows = hydro_data["size_inflow"][1] +num_gen = length(results[:simulations][1][1][:powersystem]["solution"]["gen"]) +hydro_idx = HydroPowerModels.idx_hydro(results[:data][1]) +baseMVA_val = alldata[1]["powersystem"]["baseMVA"] + +function cidx(i, n) + return mod(i, n) == 0 ? n : mod(i, n) +end + +# ── Extract per-stage data ───────────────────────────────────────────────── +rows = NamedTuple[] +for s in 1:NUM_SCENARIOS + for t in 1:REPORT_STAGES + stage = results[:simulations][s][t] + noise_term = stage[:noise_term] + res = stage[:reservoirs] + + # Reservoir states + res_in = [Float64(res[:reservoir][r].in) for r in 1:nhyd] + res_out = [Float64(res[:reservoir][r].out) for r in 1:nhyd] + outflow = [Float64(res[:outflow][r]) for r in 1:nhyd] + spill = [Float64(res[:spill][r]) for r in 1:nhyd] + + # Inflows: read from the same data that SDDP used + row_idx = cidx(t, n_inflow_rows) + inflows = [Float64(hydro_data["Hydrogenerators"][r]["inflow"][row_idx, noise_term]) for r in 1:nhyd] + + # Generator dispatch + sol = stage[:powersystem]["solution"] + num_gen = length(sol["gen"]) + pg_vals = [Float64(sol["gen"]["$g"]["pg"]) for g in 1:num_gen] + + row = ( + scenario = s, + stage = t, + noise_term = noise_term, + stage_objective = Float64(stage[:stage_objective]), + ) + + # Add per-hydro fields + for r in 1:nhyd + row = merge(row, NamedTuple{( + Symbol("res_in_$r"), Symbol("res_out_$r"), + Symbol("inflow_$r"), Symbol("outflow_$r"), Symbol("spill_$r"), + )}((res_in[r], res_out[r], inflows[r], outflow[r], spill[r]))) + end + + # Add per-gen pg fields + for g in 1:num_gen + row = merge(row, NamedTuple{(Symbol("pg_$g"),)}((pg_vals[g],))) + end + + push!(rows, row) + end +end + +# ── Save ─────────────────────────────────────────────────────────────────── +out_file = joinpath(CASE_DIR, string(FORMULATION), "sddp_trajectories_validate.csv") +CSV.write(out_file, DataFrame(rows)) +println("Saved: $out_file ($(length(rows)) rows)") + +# Summary +for s in 1:NUM_SCENARIOS + total = sum(r.stage_objective for r in rows if r.scenario == s) + println(" Scenario $s: total cost = $(round(total; digits=1))") +end + +# Also save initial state for reference +initial_volumes = [Float64(hydro_data["Hydrogenerators"][r]["initial_volume"]) for r in 1:nhyd] +println("\nInitial volumes: $initial_volumes") + +load_data = alldata[1]["powersystem"]["load"] +println("baseMVA: $baseMVA_val") +println("Number of loads: $(length(load_data))") +println("\nHydro grid indices: $hydro_idx") +println("nHyd: $nhyd, num_gen: $num_gen") +println("Inflow rows: $n_inflow_rows") + +# Save metadata +meta_file = joinpath(CASE_DIR, string(FORMULATION), "sddp_validate_meta.csv") +meta = DataFrame( + key = ["baseMVA", "nhyd", "n_inflow_rows", "num_gen", "load_scaler", + [string("initial_vol_", r) for r in 1:nhyd]..., + [string("hydro_grid_idx_", r) for r in 1:nhyd]...], + value = [Float64(baseMVA_val), Float64(nhyd), Float64(n_inflow_rows), Float64(num_gen), 0.6, + initial_volumes..., + [Float64(hydro_idx[r]) for r in 1:nhyd]...], +) +CSV.write(meta_file, meta) +println("Saved metadata: $meta_file") + +# Save per-bus demand (after 0.6 scaling) for comparison with MOF +demand_file = joinpath(CASE_DIR, string(FORMULATION), "sddp_validate_demands.csv") +demand_rows = NamedTuple[] +for (bus_key, ld) in load_data + push!(demand_rows, ( + bus = bus_key, + pd = Float64(ld["pd"]), + qd = Float64(ld["qd"]), + status = get(ld, "status", 1), + )) +end +CSV.write(demand_file, DataFrame(demand_rows)) +println("Saved demands (post-scaling): $demand_file") diff --git a/examples/HydroPowerModels/test_strict_mode.jl b/examples/HydroPowerModels/test_strict_mode.jl new file mode 100644 index 0000000..cf74dbb --- /dev/null +++ b/examples/HydroPowerModels/test_strict_mode.jl @@ -0,0 +1,370 @@ +# Test suite for strict-mode subproblems and HydroReachablePolicy +# +# Tests: +# 1. Build strict subproblems — no deficit variables +# 2. Solve strict subproblem with feasible target +# 3. Solve strict subproblem with infeasible target +# 4. HydroReachablePolicy output bounds +# 5. HydroReachablePolicy gradient +# 6. End-to-end strict rollout +# 7. Backward compatibility (non-strict unchanged) +# +# Usage: +# julia --project test_strict_mode.jl + +using DecisionRules +using Test +using JuMP +using Flux +using Random +using Ipopt +using DiffOpt +using Statistics + +const HYDRO_DIR = dirname(@__FILE__) +include(joinpath(HYDRO_DIR, "load_hydropowermodels.jl")) +include(joinpath(HYDRO_DIR, "hydro_reachable_policy.jl")) + +const CASE_NAME = "bolivia" +const FORMULATION_FILE = "ACPPowerModel.mof.json" +const NUM_STAGES = 4 # small for fast testing + +# DiffOpt optimizer factory +diff_optimizer = + () -> DiffOpt.diff_optimizer( + optimizer_with_attributes( + Ipopt.Optimizer, + "print_level" => 0, + "linear_solver" => "mumps", + ), + ) + +println("="^60) +println("Strict Mode Test Suite") +println("="^60) + +@testset "Strict Mode Test Suite" begin + +# ── Test 1: Build strict subproblems ──────────────────────────────────────── + +@testset "Test 1: Build strict subproblems — no deficit variables" begin + # Build with strict=true — should have no norm_deficit variables + sub, sp_in, sp_out, uncert, x0, maxvol, hmeta = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + nHyd = length(x0) + + for t in 1:NUM_STAGES + # No variable named "norm_deficit" should exist in strict subproblems + deficit_vars = filter( + v -> occursin("norm_deficit", JuMP.name(v)), + JuMP.all_variables(sub[t]), + ) + @test isempty(deficit_vars) + + # No variable named "_deficit" should exist either + deficit_inner = filter( + v -> occursin("_deficit", JuMP.name(v)), + JuMP.all_variables(sub[t]), + ) + @test isempty(deficit_inner) + + # state_params_out[t] should be a vector of (param, var) tuples + @test length(sp_out[t]) == nHyd + for i in 1:nHyd + @test sp_out[t][i] isa Tuple + @test length(sp_out[t][i]) == 2 + end + end + + # hydro_meta should have all required fields + @test hmeta.nHyd == nHyd + @test length(hmeta.min_vol) == nHyd + @test length(hmeta.max_vol) == nHyd + @test length(hmeta.min_turn) == nHyd + @test length(hmeta.max_turn) == nHyd + @test length(hmeta.upstream_turn) == nHyd + @test length(hmeta.upstream_spill) == nHyd + @test hmeta.K > 0 + + println(" Test 1 PASSED: strict subproblems have no deficit variables") +end + +# ── Test 2: Solve strict subproblem with feasible target ──────────────────── + +@testset "Test 2: Solve strict subproblem with feasible target" begin + # Build non-strict first to get a feasible solution as reference target + sub_ref, sp_in_ref, sp_out_ref, uncert_ref, x0, maxvol, _ = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, + penalty_l1=:auto, penalty_l2=:auto, + ) + + nHyd = length(x0) + + # Solve stage 1 of the non-strict model to get a feasible realized state + # Set incoming state = initial state + for i in 1:nHyd + JuMP.set_parameter_value(sp_in_ref[1][i], x0[i]) + end + # Set inflow to first scenario values + for (param, val) in uncert_ref[1][1] + JuMP.set_parameter_value(param, val) + end + # Set target = initial state (feasible since it's within bounds) + for i in 1:nHyd + # target is the parameter in the tuple + JuMP.set_parameter_value(sp_out_ref[1][i][1], x0[i]) + end + optimize!(sub_ref[1]) + ref_status = termination_status(sub_ref[1]) + @test ref_status in (MOI.OPTIMAL, MOI.LOCALLY_SOLVED, MOI.ALMOST_LOCALLY_SOLVED, MOI.ALMOST_OPTIMAL) + + # Get the realized reservoir volumes from the non-strict solve + feasible_target = [JuMP.value(sp_out_ref[1][i][2]) for i in 1:nHyd] + + # Now build strict subproblems and solve with the feasible target + sub_s, sp_in_s, sp_out_s, uncert_s, _, _, _ = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + # Set incoming state + for i in 1:nHyd + JuMP.set_parameter_value(sp_in_s[1][i], x0[i]) + end + # Set inflow + for (param, val) in uncert_s[1][1] + JuMP.set_parameter_value(param, val) + end + # Set target = feasible values from non-strict solve + for i in 1:nHyd + JuMP.set_parameter_value(sp_out_s[1][i][1], feasible_target[i]) + end + optimize!(sub_s[1]) + strict_status = termination_status(sub_s[1]) + @test strict_status in (MOI.OPTIMAL, MOI.LOCALLY_SOLVED, MOI.ALMOST_LOCALLY_SOLVED, MOI.ALMOST_OPTIMAL) + + # Verify reservoir_out ≈ target for all units + for i in 1:nHyd + realized = JuMP.value(sp_out_s[1][i][2]) + @test isapprox(realized, feasible_target[i]; atol=1e-4) + end + + println(" Test 2 PASSED: strict subproblem solves with feasible target") +end + +# ── Test 3: Solve strict subproblem with infeasible target ────────────────── + +@testset "Test 3: Solve strict subproblem with infeasible target" begin + sub_s, sp_in_s, sp_out_s, uncert_s, x0, maxvol, _ = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + nHyd = length(x0) + + # Set incoming state + for i in 1:nHyd + JuMP.set_parameter_value(sp_in_s[1][i], x0[i]) + end + # Set inflow + for (param, val) in uncert_s[1][1] + JuMP.set_parameter_value(param, val) + end + # Set target = max_volume * 10 (clearly infeasible — exceeds physical limits) + for i in 1:nHyd + infeasible_target = maxvol[i] * 10.0 + 1000.0 + JuMP.set_parameter_value(sp_out_s[1][i][1], infeasible_target) + end + optimize!(sub_s[1]) + strict_status = termination_status(sub_s[1]) + # Should NOT be optimal — the target is physically impossible + @test strict_status in (MOI.INFEASIBLE, MOI.LOCALLY_INFEASIBLE, + MOI.INFEASIBLE_OR_UNBOUNDED, MOI.NUMERICAL_ERROR, + MOI.OTHER_ERROR) + + println(" Test 3 PASSED: strict subproblem detects infeasible target") +end + +# ── Test 4: HydroReachablePolicy output bounds ───────────────────────────── + +@testset "Test 4: HydroReachablePolicy output bounds" begin + _, _, _, _, x0, maxvol, hmeta = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + nHyd = hmeta.nHyd + + # Create policy + policy = hydro_reachable_policy(hmeta, Int64[32, 32]) + Flux.reset!(policy) + + # Forward pass with known state (initial volumes) and arbitrary inflow + Random.seed!(42) + inflow = rand(Float32, nHyd) .* 5.0f0 # random inflows + x_prev = Float32.(x0) + input = vcat(inflow, x_prev) + + target = policy(input) + + # Verify every output is within [min_vol, max_vol] + for r in 1:nHyd + @test target[r] >= hmeta.min_vol[r] - 1e-5 + @test target[r] <= hmeta.max_vol[r] + 1e-5 + end + + # Verify output is within the computed reachable bounds + lower, upper = _hydro_reachable_bounds(policy, inflow, x_prev) + for r in 1:nHyd + @test target[r] >= lower[r] - 1e-5 + @test target[r] <= upper[r] + 1e-5 + end + + # Test CHJ (index 10, max_vol=0) — output should be exactly 0 + chj_idx = findfirst(v -> v == 0.0, hmeta.max_vol) + if chj_idx !== nothing + @test abs(target[chj_idx]) < 1e-6 + end + + println(" Test 4 PASSED: policy outputs within reachable bounds") +end + +# ── Test 5: HydroReachablePolicy gradient ────────────────────────────────── + +@testset "Test 5: HydroReachablePolicy gradient" begin + _, _, _, _, x0, _, hmeta = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + nHyd = hmeta.nHyd + + # Create policy + policy = hydro_reachable_policy(hmeta, Int64[32, 32]) + Flux.reset!(policy) + + # Forward + gradient + Random.seed!(42) + inflow = rand(Float32, nHyd) .* 5.0f0 + x_prev = Float32.(x0) + input = vcat(inflow, x_prev) + + grads = Flux.gradient(policy) do m + sum(m(input)) + end + + # Gradient should not be nothing + @test grads[1] !== nothing + + # Check encoder and combiner gradients via Flux.state + grad_state = grads[1] + + # Verify all gradient values are finite (no NaN or Inf) + for p in Flux.trainables(policy) + # Find corresponding gradient + g = Flux.gradient(m -> sum(m(input)), policy)[1] + break # just verify the gradient computation works + end + + # A more direct check: compute gradient and verify nonzero for active units + Flux.reset!(policy) + loss_val, grads2 = Flux.withgradient(policy) do m + sum(m(input)) + end + @test isfinite(loss_val) + @test grads2[1] !== nothing + + println(" Test 5 PASSED: policy gradient is finite and computable") +end + +# ── Test 6: End-to-end strict rollout ────────────────────────────────────── + +@testset "Test 6: End-to-end strict rollout" begin + sub, sp_in, sp_out, uncert, x0, maxvol, hmeta = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, strict=true, + ) + + nHyd = hmeta.nHyd + + # Build reachable policy + policy = hydro_reachable_policy(hmeta, Int64[32, 32]) + + # Sample a scenario + Random.seed!(42) + scenario = DecisionRules.sample(uncert) + + # Run a full rollout — all stages should be feasible with the reachable policy + Flux.reset!(policy) + obj = nothing + try + obj = simulate_multistage( + sub, sp_in, sp_out, x0, scenario, policy; + ) + catch e + # If it fails, print the error for debugging + @error "Rollout failed" exception=(e, catch_backtrace()) + end + + # The rollout should complete successfully + @test obj !== nothing + @test isfinite(obj) + + # In strict mode, get_objective_no_target_deficit should equal the raw objective + # (no deficit to subtract) + obj_no_deficit = DecisionRules.get_objective_no_target_deficit(sub) + # In strict mode these should be equal (no norm_deficit variables) + @test isapprox(obj, obj_no_deficit; rtol=1e-4) + + println(" Test 6 PASSED: end-to-end strict rollout completes successfully") +end + +# ── Test 7: Backward compatibility ───────────────────────────────────────── + +@testset "Test 7: Backward compatibility (non-strict unchanged)" begin + sub, sp_in, sp_out, uncert, x0, maxvol, hmeta = build_hydropowermodels( + joinpath(HYDRO_DIR, CASE_NAME), FORMULATION_FILE; + num_stages=NUM_STAGES, optimizer=diff_optimizer, + penalty_l1=:auto, penalty_l2=:auto, + ) + + nHyd = length(x0) + + for t in 1:NUM_STAGES + # Deficit variables SHOULD exist in non-strict mode + deficit_vars = filter( + v -> occursin("norm_deficit", JuMP.name(v)), + JuMP.all_variables(sub[t]), + ) + @test !isempty(deficit_vars) + end + + # hydro_meta should still be returned as 7th element + @test hmeta.nHyd == nHyd + @test hmeta.K > 0 + + # Standard policy should work with non-strict subproblems + num_uncertainties = length(uncert[1][1]) + models = state_conditioned_policy( + num_uncertainties, nHyd, nHyd, Int64[32, 32]; + activation=sigmoid, encoder_type=Flux.LSTM, + ) + + Random.seed!(42) + scenario = DecisionRules.sample(uncert) + Flux.reset!(models) + obj = simulate_multistage(sub, sp_in, sp_out, x0, scenario, models) + @test isfinite(obj) + + println(" Test 7 PASSED: non-strict mode unchanged (backward compatible)") +end + +end # top-level @testset + +println("\n", "="^60) +println("All tests completed!") +println("="^60) diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels.jl b/examples/HydroPowerModels/train_dr_hydropowermodels.jl index ad6876c..71bd557 100644 --- a/examples/HydroPowerModels/train_dr_hydropowermodels.jl +++ b/examples/HydroPowerModels/train_dr_hydropowermodels.jl @@ -74,7 +74,7 @@ diff_optimizer = "linear_solver" => "mumps", ), ) -subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, initial_state, max_volume = build_hydropowermodels( +subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, initial_state, max_volume, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, @@ -84,7 +84,7 @@ subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, ) # Build det-eq for training -subproblems_de, state_params_in, state_params_out, uncertainty_samples, _, _ = build_hydropowermodels( +subproblems_de, state_params_in, state_params_out, uncertainty_samples, _, _, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl b/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl index 906f8e4..f0d7505 100644 --- a/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl +++ b/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl @@ -71,7 +71,7 @@ diff_model = ), ) -subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels( +subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels_strict.jl b/examples/HydroPowerModels/train_dr_hydropowermodels_strict.jl new file mode 100644 index 0000000..8c80617 --- /dev/null +++ b/examples/HydroPowerModels/train_dr_hydropowermodels_strict.jl @@ -0,0 +1,372 @@ +# Train HydroPowerModels using strict subproblems with reachable policy +# +# Strict mode removes the deficit slack variables from the target constraint: +# reservoir_out[r] = target[r] (hard equality, no penalty) +# The dual λ_r is the clean shadow price ∂Q/∂target — pure economic signal. +# +# This requires HydroReachablePolicy, which guarantees every target is within +# the one-stage reachable set via sigmoid-bounded outputs scaled to physical +# reservoir limits. Stage-wise strict mode is closed-loop: each policy call sees +# the realized state from the previous strict stage solve. +# +# Usage: +# julia --project -t auto train_dr_hydropowermodels_strict.jl +# +# Environment overrides: +# DR_NUM_STAGES=126 number of training stages +# DR_NUM_ROLLOUT_STAGES=96 number of rollout evaluation stages (default: DR_NUM_STAGES) +# DR_NUM_EPOCHS=80 number of epochs +# DR_ENCODER_LAYERS=128,128 recurrent inflow encoder sizes +# DR_HEAD_LAYERS= nonrecurrent state-conditioned target head sizes +# DR_GRAD_CLIP=0 gradient clipping (0 = disabled) +# DR_NUM_TRAIN_PER_BATCH=1 sampled trajectories per gradient step (variance reduction) +# DR_PRETRAINED_MODEL=path warmstart from a StateConditionedPolicy checkpoint +# DR_CONTEXT= optional known context prepended to the policy input: +# ""/"none" = off, "phase" = seasonal sin/cos, +# "phase+progress" = seasonal sin/cos plus t/T +# DR_NUM_EVAL_SCENARIOS=4 fixed held-out scenarios for rollout evaluation +# DR_EVAL_EVERY=25 rollout-evaluate every this many batches +# DR_SAVE_METRIC=training checkpoint-selection metric: +# "training" — per-batch training loss (historical +# behavior; noisy for small DR_NUM_TRAIN_PER_BATCH) +# "rollout" — mean deficit-free objective of the +# fixed held-out rollout evaluation (the metric +# policies are ultimately judged on; evaluated +# every DR_EVAL_EVERY batches) +# DR_LR=0.001 Adam learning rate +# DR_LR_FINAL=DR_LR final learning rate of a cosine decay across the +# full run (equal to DR_LR → constant, historical) +# DR_LR_WARMUP=0 linear warmup iterations from DR_LR/100 to DR_LR. +# Protects a warmstarted policy from the initial +# full-size Adam steps taken while its second-moment +# estimates are still zero. +# +# Reproducible recipes (also listed in this folder's README): +# From scratch (paper configuration — all defaults): +# julia --project -t auto train_dr_hydropowermodels_strict.jl +# Fine-tune from a converged checkpoint (variance-reduced, decayed LR, +# rollout-selected checkpoints): +# DR_PRETRAINED_MODEL= DR_NUM_TRAIN_PER_BATCH=16 \ +# DR_LR=1e-4 DR_LR_FINAL=1e-5 DR_LR_WARMUP=50 \ +# DR_SAVE_METRIC=rollout DR_NUM_EVAL_SCENARIOS=24 DR_NUM_EPOCHS=15 \ +# julia --project -t auto train_dr_hydropowermodels_strict.jl +using DecisionRules +using Statistics +using Random +using Flux + +using Ipopt +using Wandb, Dates, Logging +using JLD2 +using DiffOpt + +HydroPowerModels_dir = dirname(@__FILE__) +include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) +include(joinpath(HydroPowerModels_dir, "hydro_reachable_policy.jl")) + +# ── Parameters ─────────────────────────────────────────────────────────────── + +case_name = "bolivia" +formulation = "ACPPowerModel" +num_stages = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +num_rollout_stages = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", string(num_stages))) +model_dir = joinpath(HydroPowerModels_dir, case_name, formulation, "models") +mkpath(model_dir) +formulation_file = formulation * ".mof.json" +num_epochs = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) +num_batches = 100 +# Trajectories sampled per gradient step; >1 averages the per-sample dual +# gradients, reducing estimator variance at proportionally higher solve cost. +_num_train_per_batch = parse(Int, get(ENV, "DR_NUM_TRAIN_PER_BATCH", "1")) +""" + parse_layers(s::AbstractString) -> Vector{Int64} + +Parse comma-separated policy architecture settings from environment variables. + +`DR_ENCODER_LAYERS` controls the recurrent inflow encoder. `DR_HEAD_LAYERS` +controls optional hidden layers in the nonrecurrent state-conditioned target +head. An empty string means no extra hidden head layers, preserving the +historical single sigmoid head. + +# Arguments +- `s::AbstractString`: comma-separated layer widths. + +# Returns +- `Vector{Int64}`: parsed hidden widths; `Int64[]` when `s` is empty. + +# Examples +```julia +parse_layers("256, 256") == Int64[256, 256] +parse_layers("") == Int64[] +``` +""" +parse_layers(s::AbstractString) = + isempty(strip(s)) ? Int64[] : [parse(Int64, strip(x)) for x in split(s, ",") if !isempty(strip(x))] +layers = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +head_layers = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +grad_clip = parse(Float32, get(ENV, "DR_GRAD_CLIP", "0")) + +function canonical_context_mode(raw_mode::AbstractString) + mode = lowercase(strip(raw_mode)) + mode in ("", "none", "off", "false") && return "" + mode in ("phase", "phase+progress") && return mode + error("DR_CONTEXT must be \"\", \"phase\", or \"phase+progress\"; got \"$raw_mode\"") +end + +function build_stage_context(mode::AbstractString, horizon::Int, period::Int) + isempty(mode) && return nothing + include_progress = mode == "phase+progress" + return DecisionRules.stage_phase_context( + horizon; + period=period, + include_progress=include_progress, + ) +end + +function context_run_tag(mode::AbstractString) + isempty(mode) && return "" + return "-ctx" * replace(mode, "+" => "p") +end + +context_mode = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +context_period = countlines(joinpath(HydroPowerModels_dir, case_name, "inflows.csv")) +stage_context = build_stage_context(context_mode, num_stages, context_period) +n_context = isnothing(stage_context) ? 0 : size(stage_context, 1) + +# ── Learning-rate schedule ─────────────────────────────────────────────────── +# lr(iter) = linear warmup from lr_init/100 over `lr_warmup` iterations, then +# cosine decay from lr_init to lr_final across the remaining budget. With the +# defaults (warmup = 0, lr_final = lr_init) this is a constant lr_init, +# reproducing the historical behavior exactly. +lr_init = parse(Float64, get(ENV, "DR_LR", "0.001")) +lr_final = parse(Float64, get(ENV, "DR_LR_FINAL", string(lr_init))) +lr_warmup = parse(Int, get(ENV, "DR_LR_WARMUP", "0")) + +""" + lr_schedule(iter, total_iters) -> Float64 + +Learning rate at one-based training iteration `iter`. + +Linear warmup from `lr_init / 100` to `lr_init` over the first `lr_warmup` +iterations, then cosine decay from `lr_init` to `lr_final`: + +```math +\\eta(k) = \\eta_f + \\tfrac{1}{2} (\\eta_0 - \\eta_f) + \\bigl(1 + \\cos(\\pi \\rho_k)\\bigr), +``` + +where ``\\rho_k`` is the post-warmup progress fraction. Constant when +`lr_warmup == 0` and `lr_final == lr_init` (the defaults). +""" +function lr_schedule(iter, total_iters) + if iter <= lr_warmup + # Warmup guards a warmstarted policy against full-size Adam steps + # taken while the optimizer's second-moment estimates are near zero. + return lr_init * (0.01 + 0.99 * iter / max(lr_warmup, 1)) + end + # Post-warmup progress in [0, 1] over the remaining iteration budget. + ρ = clamp((iter - lr_warmup) / max(total_iters - lr_warmup, 1), 0.0, 1.0) + return lr_final + 0.5 * (lr_init - lr_final) * (1 + cos(π * ρ)) +end + +optimizers = if grad_clip > 0 + [Flux.Optimisers.OptimiserChain(Flux.Optimisers.ClipGrad(grad_clip), Flux.Adam(lr_init))] +else + [Flux.Adam(lr_init)] +end +pre_trained_model = get(ENV, "DR_PRETRAINED_MODEL", nothing) +clip_tag = grad_clip > 0 ? "-clip$(Int(grad_clip))" : "" +head_tag = isempty(head_layers) ? "-Hlinear" : "-H$(join(head_layers, "_"))" +_rollout_tag = num_rollout_stages != num_stages ? "-r$(num_rollout_stages)" : "" +# Tag runs with a non-default batch size so checkpoints are distinguishable. +nt_tag = _num_train_per_batch > 1 ? "-nt$(_num_train_per_batch)" : "" +# Tag warmstarted runs: their result is a fine-tune of another checkpoint, not +# a from-scratch training (the parent checkpoint is recorded in wandb config). +warm_tag = (isnothing(pre_trained_model) || pre_trained_model == "nothing") ? "" : "-warm" +save_file = "$(case_name)-$(formulation)-h$(num_stages)$(_rollout_tag)-subproblems-strict$(clip_tag)$(head_tag)$(nt_tag)$(context_run_tag(context_mode))$(warm_tag)-$(now())" +num_eval_scenarios = parse(Int, get(ENV, "DR_NUM_EVAL_SCENARIOS", "4")) +eval_every = parse(Int, get(ENV, "DR_EVAL_EVERY", "25")) +# Checkpoint-selection metric: "training" (historical; noisy at small batch +# sizes because a lucky scenario can look like a better policy) or "rollout" +# (deficit-free mean over the fixed held-out scenarios — the deployment metric). +save_metric = lowercase(get(ENV, "DR_SAVE_METRIC", "training")) +save_metric in ("training", "rollout") || + error("DR_SAVE_METRIC must be \"training\" or \"rollout\", got $save_metric") + +# ── Build strict subproblems (no deficit, no penalty) ──────────────────────── + +# Define the DiffOpt optimizer for subproblems +diff_optimizer = + () -> DiffOpt.diff_optimizer( + optimizer_with_attributes( + Ipopt.Optimizer, + "print_level" => 0, + "linear_solver" => "mumps", + ), + ) + +subproblems, state_params_in, state_params_out, uncertainty_samples, + initial_state, max_volume, hydro_meta = build_hydropowermodels( + joinpath(HydroPowerModels_dir, case_name), + formulation_file; + num_stages=num_stages, + optimizer=diff_optimizer, + strict=true, +) + +num_hydro = length(initial_state) + +# ── Logging ────────────────────────────────────────────────────────────────── + +lg = WandbLogger(; + project="RL", + name=save_file, + save_code=false, + config=Dict( + "layers" => layers, + "head_layers" => head_layers, + "activation" => "sigmoid (reachable)", + "optimizer" => string(optimizers), + "grad_clip" => grad_clip, + "training_method" => "subproblems-strict", + "penalty_schedule" => "none (strict)", + "num_stages" => num_stages, + "num_rollout_stages" => num_rollout_stages, + "num_epochs" => string(num_epochs), + "num_batches" => string(num_batches), + "num_train_per_batch" => string(_num_train_per_batch), + "pre_trained_model" => string(pre_trained_model), + "context_mode" => isempty(context_mode) ? "none" : context_mode, + "context_period" => context_period, + "context_horizon" => num_stages, + "n_context" => n_context, + "num_eval_scenarios" => num_eval_scenarios, + "eval_every" => eval_every, + "save_metric" => save_metric, + "lr" => lr_init, + "lr_final" => lr_final, + "lr_warmup" => lr_warmup, + ), +) + +# ── Build reachable policy ─────────────────────────────────────────────────── + +# HydroReachablePolicy: LSTM encoder over inflows + sigmoid feed-forward head +# over [encoded_inflow; reservoir_state], bounded to the one-stage reachable set. +base_model = hydro_reachable_policy( + hydro_meta, + layers; + combiner_layers=head_layers, + n_context=n_context, +) +models = isnothing(stage_context) ? base_model : ContextualPolicy(base_model, stage_context) +@info "Strict hydro policy context" context_mode=(isempty(context_mode) ? "none" : context_mode) context_period context_horizon=num_stages n_context + +# ── Load pretrained model (warmstart from non-strict training) ─────────────── + +if !isnothing(pre_trained_model) && pre_trained_model != "nothing" + model_save = JLD2.load(pre_trained_model) + model_state = model_save["model_state"] + # Load encoder/combiner weights; hydro bounds are preserved + load_policy_weights!(models, model_state) + @info "Loaded pretrained weights from $pre_trained_model" +end + +# ── Initial evaluation and callbacks ───────────────────────────────────────── + +Random.seed!(8788) +objective_values = [ + simulate_multistage( + subproblems, + state_params_in, + state_params_out, + initial_state, + DecisionRules.sample(uncertainty_samples), + models; + ) for _ in 1:2 +] +initial_training_obj = mean(objective_values) +convergence_criterium = StallingCriterium(num_epochs * num_batches, initial_training_obj, 0) + +# Fixed held-out scenarios, materialized once so every evaluation uses the same set. +# Use num_rollout_stages for evaluation (may differ from training num_stages). +Random.seed!(8789) +rollout_uncertainty = uncertainty_samples[1:num_rollout_stages] +eval_scenarios = [DecisionRules.sample(rollout_uncertainty) for _ in 1:num_eval_scenarios] +rollout_evaluation = RolloutEvaluation( + subproblems[1:num_rollout_stages], + state_params_in[1:num_rollout_stages], + state_params_out[1:num_rollout_stages], + initial_state, + eval_scenarios; + stride=eval_every, + policy_state=:realized, +) + +# Checkpoint-selection baseline. With save_metric == "rollout" the incumbent +# is the current model's held-out rollout cost, so a warmstarted run only saves +# checkpoints that genuinely improve on the loaded policy under the metric it +# is ultimately judged on. iter = eval_every satisfies the stride gate. +best_obj = if save_metric == "rollout" + rollout_evaluation(eval_every, models) + @info "Initial rollout evaluation (checkpoint baseline)" rollout_evaluation.last_objective_no_deficit rollout_evaluation.last_violation_share + rollout_evaluation.last_objective_no_deficit +else + initial_training_obj +end +model_path = joinpath(model_dir, save_file * ".jld2") +save_control = SaveBest(best_obj, model_path) + +# ── Train ──────────────────────────────────────────────────────────────────── + +# Total iteration budget, used by the learning-rate schedule. +total_iters = num_epochs * num_batches + +# No penalty schedule needed — strict mode has no deficit to penalize. +train_multistage( + models, + initial_state, + subproblems, + state_params_in, + state_params_out, + uncertainty_samples; + num_batches=total_iters, + num_train_per_batch=_num_train_per_batch, + optimizer=first(optimizers), + # Apply the learning-rate schedule through the optimizer state. With the + # default constant schedule adjust! is a no-op-equivalent every iteration. + adjust_hyperparameters=(iter, opt_state, ntpb) -> begin + Flux.Optimisers.adjust!(opt_state, lr_schedule(iter, total_iters)) + ntpb + end, + record=(sample_log, iter, model) -> begin + # In strict mode: objectives == objectives_no_deficit (no penalty term) + training_loss = mean(sample_log.objectives) + metrics = Dict( + "metrics/loss" => training_loss, + "metrics/training_loss" => training_loss, + "metrics/lr" => lr_schedule(iter, total_iters), + ) + rollout_evaluation(iter, model) + if iter % eval_every == 0 + metrics["metrics/rollout_objective_no_deficit"] = + rollout_evaluation.last_objective_no_deficit + metrics["metrics/rollout_target_violation_share"] = + rollout_evaluation.last_violation_share + end + Wandb.log(lg, metrics) + # Checkpoint selection: historical per-batch training loss, or the + # held-out rollout objective (only refreshed at eval iterations). + if save_metric == "rollout" + iter % eval_every == 0 && + save_control(iter, model, rollout_evaluation.last_objective_no_deficit) + else + save_control(iter, model, training_loss) + end + return convergence_criterium(iter, model, training_loss) + end, + penalty_schedule=nothing, +) + +# Finish the run +close(lg) diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl b/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl index 7428293..9fe9af0 100644 --- a/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl +++ b/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl @@ -59,7 +59,7 @@ diff_optimizer = ), ) -subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels( +subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, diff --git a/examples/HydroPowerModels/train_ldr_hydropowermodels.jl b/examples/HydroPowerModels/train_ldr_hydropowermodels.jl index c296ce2..bb7f4c5 100644 --- a/examples/HydroPowerModels/train_ldr_hydropowermodels.jl +++ b/examples/HydroPowerModels/train_ldr_hydropowermodels.jl @@ -84,7 +84,7 @@ diff_optimizer = "linear_solver" => "mumps", ), ) -subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, initial_state, max_volume = build_hydropowermodels( +subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, initial_state, max_volume, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, @@ -95,7 +95,7 @@ subproblems, state_params_in_sub, state_params_out_sub, uncertainty_samples_sub, # ── Build det-eq for training ──────────────────────────────────────────────── -subproblems_de, state_params_in, state_params_out, uncertainty_samples, _, _ = build_hydropowermodels( +subproblems_de, state_params_in, state_params_out, uncertainty_samples, _, _, _ = build_hydropowermodels( joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, diff --git a/examples/HydroPowerModels/validate_sddp_vs_jump.jl b/examples/HydroPowerModels/validate_sddp_vs_jump.jl new file mode 100644 index 0000000..34a5f5c --- /dev/null +++ b/examples/HydroPowerModels/validate_sddp_vs_jump.jl @@ -0,0 +1,307 @@ +# Validate that DecisionRules JuMP subproblems match SDDP's problem formulation. +# +# Reads SDDP trajectory data (states, inflows, stage costs) extracted by +# sddp/extract_sddp_trajectories.jl, then replays those trajectories through +# the JuMP subproblems built by load_hydropowermodels.jl. +# +# For each SDDP scenario and stage: +# 1. Set initial state = SDDP's reservoir_in +# 2. Set inflows = SDDP's inflow values +# 3. Set target = SDDP's reservoir_out (what SDDP realized) +# 4. Solve the JuMP subproblem +# 5. Compare: JuMP stage cost vs SDDP stage cost +# +# If the problems are equivalent, costs should match to solver tolerance. +# Any mismatch reveals structural differences (load scaling, constraints, etc.). +# +# Runs in the main DecisionRules project environment. +using DecisionRules +using JuMP, DiffOpt, Ipopt +using CSV, DataFrames +using JSON +using Statistics + +HydroPowerModels_dir = dirname(@__FILE__) +include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl")) + +case_name = "bolivia" +formulation = "ACPPowerModel" +case_dir = joinpath(HydroPowerModels_dir, case_name) +out_dir = joinpath(case_dir, formulation) + +println("=" ^ 60) +println("Validate SDDP vs JuMP subproblems") +println("=" ^ 60) + +# ── Load SDDP trajectory data ───────────────────────────────────────────── +traj_file = joinpath(out_dir, "sddp_trajectories_validate.csv") +meta_file = joinpath(out_dir, "sddp_validate_meta.csv") +demand_file = joinpath(out_dir, "sddp_validate_demands.csv") + +if !isfile(traj_file) + error("SDDP trajectory file not found: $traj_file\n" * + "Run sddp/extract_sddp_trajectories.jl first.") +end + +traj_df = CSV.read(traj_file, DataFrame) + +# Infer nHyd from trajectory columns (res_in_1, res_in_2, ...) +nhyd = count(n -> startswith(string(n), "res_in_"), names(traj_df)) +num_gen = count(n -> startswith(string(n), "pg_"), names(traj_df)) + +# Load metadata if available +if isfile(meta_file) + meta_df = CSV.read(meta_file, DataFrame) + meta = Dict(row.key => row.value for row in eachrow(meta_df)) + baseMVA = meta["baseMVA"] + load_scaler = meta["load_scaler"] +else + println(" WARNING: meta file not found, using defaults") + hydro_json = JSON.parsefile(joinpath(case_dir, "hydro.json")) + power_json = JSON.parsefile(joinpath(case_dir, "PowerModels.json")) + baseMVA = power_json["baseMVA"] + load_scaler = 0.6 +end + +scenarios = sort(unique(traj_df.scenario)) +num_scenarios = length(scenarios) +stages = sort(unique(traj_df.stage)) +num_stages = length(stages) + +println(" Trajectories: $num_scenarios scenarios × $num_stages stages") +println(" nHyd=$nhyd, num_gen=$num_gen, baseMVA=$baseMVA, load_scaler=$load_scaler") + +# ── Check demand consistency ─────────────────────────────────────────────── +# The MOF file has demands baked in. SDDP applies 0.6 scaling. +# Let's read the MOF file and compare demands. +println("\n--- Demand cross-check ---") +mof_model = JuMP.read_from_file( + joinpath(case_dir, formulation * ".mof.json"); use_nlp_block=false +) +# Find load constraints or demand parameters in MOF +# The demand is typically in the objective or power balance constraints. +# Let's check by looking at the pd/qd variables or fixed values. +mof_vars = JuMP.all_variables(mof_model) +demand_vars = filter(v -> occursin("deficit", JuMP.name(v)) || occursin("load", JuMP.name(v)), mof_vars) +println(" MOF demand-related variables: $(length(demand_vars))") +for v in demand_vars[1:min(3, length(demand_vars))] + println(" $(JuMP.name(v))") +end + +try + pb_cons = filter(c -> occursin("kcl", string(c)), JuMP.all_constraints(mof_model; include_variable_in_set_constraints=false)) + if !isempty(pb_cons) + println(" Found $(length(pb_cons)) power balance (kcl) constraints") + println(" First: $(pb_cons[1])") + end +catch e + println(" (could not inspect constraints: $e)") +end + +# ── Build JuMP subproblems ───────────────────────────────────────────────── +diff_optimizer = + () -> DiffOpt.diff_optimizer( + optimizer_with_attributes( + Ipopt.Optimizer, + "print_level" => 0, + "linear_solver" => "mumps", + ), + ) + +# Build with strict=true (hard equality targets, no deficit) +subproblems_strict, state_in_strict, state_out_strict, uncert_strict, + initial_state, max_volume, _ = build_hydropowermodels( + case_dir, formulation * ".mof.json"; + num_stages=num_stages, + optimizer=diff_optimizer, + strict=true, +) + +# Also build with penalty (to test non-strict) +subproblems_penalty, state_in_penalty, state_out_penalty, uncert_penalty, + _, _, _ = build_hydropowermodels( + case_dir, formulation * ".mof.json"; + num_stages=num_stages, + optimizer=diff_optimizer, + penalty_l1=:auto, penalty_l2=:auto, +) + +println("\n Built $num_stages subproblems (strict and penalty modes)") +println(" Initial state: $initial_state") +println(" SDDP initial: $([meta["initial_vol_$r"] for r in 1:nhyd])") + +# ── Cross-check inflow data ─────────────────────────────────────────────── +println("\n--- Inflow cross-check ---") +# DecisionRules reads inflows from the same CSV +nCen = length(uncert_strict[1]) # number of scenarios per stage +println(" DecisionRules nCen=$nCen") + +# Check that inflow values match +for s in scenarios[1:min(2, num_scenarios)] + stage_rows = filter(r -> r.scenario == s, traj_df) + for t in 1:min(3, num_stages) + row = stage_rows[t, :] + ω = row.noise_term + # DecisionRules uncertainty_samples[t][ω] should match SDDP inflows + dr_scenario = uncert_strict[t][ω] + for r in 1:nhyd + dr_inflow = dr_scenario[r][2] + sddp_inflow = row[Symbol("inflow_$r")] + if !isapprox(dr_inflow, sddp_inflow; atol=1e-10) + println(" MISMATCH: scenario=$s stage=$t hydro=$r: DR=$dr_inflow SDDP=$sddp_inflow") + end + end + end +end +println(" Inflow cross-check passed (spot-checked)") + +# ── Replay SDDP trajectories through JuMP subproblems ───────────────────── +println("\n--- Stage-by-stage cost comparison ---") + +results_rows = NamedTuple[] + +for s in scenarios + stage_rows = filter(r -> r.scenario == s, traj_df) + sddp_total = 0.0 + jump_strict_total = 0.0 + jump_penalty_total = 0.0 + max_cost_diff_strict = 0.0 + max_cost_diff_penalty = 0.0 + + for t in 1:num_stages + row = stage_rows[t, :] + ω = row.noise_term + sddp_cost = row.stage_objective + + # SDDP's realized states + sddp_res_in = [row[Symbol("res_in_$r")] for r in 1:nhyd] + sddp_res_out = [row[Symbol("res_out_$r")] for r in 1:nhyd] + sddp_inflows = [row[Symbol("inflow_$r")] for r in 1:nhyd] + + # ── Strict subproblem ── + sp = subproblems_strict[t] + # Set initial state + for (j, param) in enumerate(state_in_strict[t]) + set_parameter_value(param, sddp_res_in[j]) + end + # Set inflows + for (param, _) in uncert_strict[t][ω] + # Already the right values, but let's be explicit + end + dr_scenario = uncert_strict[t][ω] + for (param, val) in dr_scenario + set_parameter_value(param, val) + end + # Set target = SDDP's realized out state + for j in 1:nhyd + target_param = state_out_strict[t][j][1] + set_parameter_value(target_param, sddp_res_out[j]) + end + + optimize!(sp) + jump_strict_cost = objective_value(sp) + + # Read realized state from JuMP + jump_res_out = [value(state_out_strict[t][j][2]) for j in 1:nhyd] + + # ── Penalty subproblem ── + sp_p = subproblems_penalty[t] + for (j, param) in enumerate(state_in_penalty[t]) + set_parameter_value(param, sddp_res_in[j]) + end + for (param, val) in uncert_penalty[t][ω] + set_parameter_value(param, val) + end + for j in 1:nhyd + target_param = state_out_penalty[t][j][1] + set_parameter_value(target_param, sddp_res_out[j]) + end + + optimize!(sp_p) + jump_penalty_cost = objective_value(sp_p) + jump_penalty_res_out = [value(state_out_penalty[t][j][2]) for j in 1:nhyd] + # Penalty cost includes deficit — compute cost without deficit + # The deficit penalty is the diff between penalty and strict cost ideally + + # State comparison + state_diff = maximum(abs.(jump_res_out .- sddp_res_out)) + state_diff_penalty = maximum(abs.(jump_penalty_res_out .- sddp_res_out)) + + cost_diff_strict = jump_strict_cost - sddp_cost + cost_diff_penalty = jump_penalty_cost - sddp_cost + + sddp_total += sddp_cost + jump_strict_total += jump_strict_cost + jump_penalty_total += jump_penalty_cost + max_cost_diff_strict = max(max_cost_diff_strict, abs(cost_diff_strict)) + max_cost_diff_penalty = max(max_cost_diff_penalty, abs(cost_diff_penalty)) + + push!(results_rows, ( + scenario = s, + stage = t, + noise_term = ω, + sddp_cost = sddp_cost, + jump_strict_cost = jump_strict_cost, + jump_penalty_cost = jump_penalty_cost, + cost_diff_strict = cost_diff_strict, + cost_diff_penalty = cost_diff_penalty, + rel_diff_strict = abs(cost_diff_strict) / max(abs(sddp_cost), 1e-12), + max_state_diff_strict = state_diff, + max_state_diff_penalty = state_diff_penalty, + )) + + if t <= 3 || abs(cost_diff_strict) > 1.0 + println(" s=$s t=$t: SDDP=$(round(sddp_cost; digits=2)) " * + "JuMP_strict=$(round(jump_strict_cost; digits=2)) " * + "diff=$(round(cost_diff_strict; digits=4)) " * + "state_diff=$(round(state_diff; sigdigits=3))") + end + end + + println(" Scenario $s totals: " * + "SDDP=$(round(sddp_total; digits=1)) " * + "JuMP_strict=$(round(jump_strict_total; digits=1)) " * + "diff=$(round(jump_strict_total - sddp_total; digits=1)) " * + "max_stage_diff=$(round(max_cost_diff_strict; digits=4))") +end + +# ── Summary ──────────────────────────────────────────────────────────────── +println("\n" * "=" ^ 60) +println("Validation Summary") +println("=" ^ 60) + +df_results = DataFrame(results_rows) +println(" Total comparisons: $(nrow(df_results))") +println(" Max stage cost diff (strict): $(round(maximum(abs.(df_results.cost_diff_strict)); digits=4))") +println(" Mean stage cost diff (strict): $(round(mean(abs.(df_results.cost_diff_strict)); digits=4))") +println(" Max relative diff (strict): $(round(maximum(df_results.rel_diff_strict); sigdigits=3))") +println(" Max state diff (strict): $(round(maximum(df_results.max_state_diff_strict); sigdigits=3))") +println(" Max stage cost diff (penalty): $(round(maximum(abs.(df_results.cost_diff_penalty)); digits=4))") + +# Per-scenario total comparison +for s in scenarios + s_rows = filter(r -> r.scenario == s, df_results) + sddp_total = sum(s_rows.sddp_cost) + jump_total = sum(s_rows.jump_strict_cost) + println(" Scenario $s: SDDP=$(round(sddp_total; digits=1)) JuMP=$(round(jump_total; digits=1)) " * + "gap=$(round(jump_total - sddp_total; digits=1)) " * + "($(round((jump_total - sddp_total) / sddp_total * 100; digits=3))%)") +end + +tol = 1.0 +big_diffs = filter(r -> abs(r.cost_diff_strict) > tol, df_results) +if nrow(big_diffs) > 0 + println("\n WARNING: $(nrow(big_diffs)) stages have |cost diff| > $tol") + println(" Largest differences:") + sorted = sort(big_diffs, :cost_diff_strict; by=abs, rev=true) + for r in eachrow(sorted[1:min(10, nrow(sorted)), :]) + println(" s=$(r.scenario) t=$(r.stage): diff=$(round(r.cost_diff_strict; digits=4))") + end +else + println("\n ✓ All stage cost differences within $tol — problems are equivalent") +end + +# Save detailed results +results_file = joinpath(out_dir, "validate_sddp_vs_jump.csv") +CSV.write(results_file, df_results) +println("\nSaved: $results_file") diff --git a/examples/inventory_control/README.md b/examples/inventory_control/README.md index a4ae4be..0b97023 100644 --- a/examples/inventory_control/README.md +++ b/examples/inventory_control/README.md @@ -68,6 +68,52 @@ subproblems. Two strategies are available: For the relaxed formulation (no integer variables), `NoIntegerStrategy` is used — subproblems are solved and duals read as-is. +## Strict Mode (reachable policy, no penalty tuning) + +The `strict` and `strict_integer` variants replace the L1 target penalty with +the hard equality `s_mid == s_target` — no deficit variable, no penalty +hyperparameter, and the target-constraint dual is the pure shadow price +(the same construction as the hydro strict mode). + +Strict mode requires every target to be one-stage feasible. Here the exact +reachable set is a one-liner: with realized incoming inventory `s` and order +`q ∈ [0, Q_max]`, the order-up-to position satisfies + +``` +s_mid = s + q ∈ [s, s + Q_max] +``` + +and nothing else constrains it (`s_out` is free; the hold/backlog split is +always feasible). `InventoryReachablePolicy` maps the network output onto +exactly this interval. Three design points, each load-bearing: + +1. **Boundary attainment.** With fixed ordering cost `K`, "do not order" + (`q = 0`) is an essential decision at the interval boundary. The policy + uses `hardsigmoid` (attains 0 and 1 exactly on finite inputs) instead of + `sigmoid` (strictly interior) — a strict-sigmoid policy would be forced + to pay `K` every period. In the hydro case the boundary is economically + irrelevant and this distinction does not matter; here it is first-order. +2. **Pass-through state components.** The demand-history entries of the + state have singleton reachable sets (the observed demands): the policy + passes them through deterministically, and the stage models leave their + target parameters unconstrained. +3. **Stage-wise training only.** Strict variants train on the stage-wise + subproblems (closed loop). The target (`s_mid`, pre-demand) and the + carried state (`s_out = s_mid − d`, post-demand) are different + quantities, so the deterministic equivalent's target-feedback recursion + would hand the policy the wrong state for its reachable bounds — unlike + hydro, where target and state are the same reservoir volume and the + strict regular DE is safe by induction. + +Run them like any other variant (one tag per invocation): + +```bash +julia --project=examples/inventory_control \ + examples/inventory_control/train_dr_inventory.jl strict +julia --project=examples/inventory_control \ + examples/inventory_control/train_dr_inventory.jl strict_integer +``` + ## Score-Function Gradient Mixing `ScoreFunctionConfig` adds a REINFORCE-style correction to the dual diff --git a/examples/inventory_control/build_inventory_problem.jl b/examples/inventory_control/build_inventory_problem.jl index b6137b8..a18c579 100644 --- a/examples/inventory_control/build_inventory_problem.jl +++ b/examples/inventory_control/build_inventory_problem.jl @@ -152,9 +152,16 @@ Returns the five-tuple expected by `simulate_multistage` and - `T`, `K`, `c`, `h`, `p`, `Q_max`: problem parameters. - `I_0`: initial inventory. - `num_scenarios`: number of uncertainty samples per SGD batch. -- `penalty`: target-deficit penalty λ. +- `penalty`: target-deficit penalty λ (ignored when `strict = true`). - `seed`: RNG seed for demand sampling. - `integer`: whether to include binary setup variable z. +- `strict`: replace the L1 target penalty with the hard equality + `s_mid == s_target` (no deficit variable, no penalty hyperparameter). + Requires a policy whose targets are always one-stage reachable, i.e. + `s_target ∈ [s_in, s_in + Q_max]` — see [`InventoryReachablePolicy`](@ref). + The demand pass-through targets stay unconstrained in both modes: their + reachable sets are singletons (the observed demands), so constraining them + adds nothing and predicting them is not a decision. """ function build_inventory_subproblems(; T = INVENTORY_T, @@ -168,6 +175,7 @@ function build_inventory_subproblems(; penalty = INVENTORY_PENALTY, seed = 42, integer = true, + strict = false, ) # Fix the random seed so demand samples are reproducible. Random.seed!(seed) @@ -242,9 +250,19 @@ function build_inventory_subproblems(; # Split end-of-period inventory into holding and backlog parts. @constraint(m, inv_hold - back == s_out) - # L1 target-deficit penalty: λ · |s_mid - ŝ_target|. - _, deficit = create_deficit!(m, 1; penalty_l1=penalty) - @constraint(m, deficit[1] == s_mid - s_target) + if strict + # Strict mode: hard equality, no slack, no penalty term. The dual + # of this constraint is the pure shadow price ∂q_t/∂ŝ_target. + # Feasible iff s_target ∈ [s_in, s_in + Q_max] (exactly the + # one-stage reachable set of s_mid = s_in + q with q ∈ [0, Q_max]; + # s_out and the hold/backlog split are free, so no other + # constraint restricts s_mid). + @constraint(m, s_mid == s_target) + else + # L1 target-deficit penalty: λ · |s_mid - ŝ_target|. + _, deficit = create_deficit!(m, 1; penalty_l1=penalty) + @constraint(m, deficit[1] == s_mid - s_target) + end # Store the model and parameter mappings. subproblems[t] = m @@ -287,9 +305,21 @@ The penalty term is - `T`, `K`, `c`, `h`, `p`, `Q_max`: problem parameters. - `I_0`: initial inventory. - `num_scenarios`: number of uncertainty samples per SGD batch. -- `penalty`: target-deficit penalty ``\\lambda``. +- `penalty`: target-deficit penalty ``\\lambda`` (ignored when `strict = true`). - `seed`: RNG seed for demand sampling. - `integer`: whether to include binary setup variable z. +- `strict`: replace the L1 penalty with hard equalities + ``s^{mid}_t = \\hat{s}_t`` and drop the penalty term from the objective. + **Feasibility caveat**: unlike the hydro case, the target (order-up-to + position ``s^{mid}``) and the carried state (post-demand inventory + ``s^{out} = s^{mid} - d``) are *different quantities*, so the standard + target-feedback DE recursion hands the policy ``\\hat{s}^{mid}_{t-1}`` + where a reachable policy expects ``s^{out}_{t-1}`` — the induction + guarantee does NOT transfer naively. Either train stage-wise (what the + `strict` variants do; see `build_training_problem`) or reconstruct the + realized inventory inside the recursion as + ``s^{out}_{t-1} = \\hat{s}_{t-1} - d_{t-1}`` (target minus the demand + pass-through) before computing reachable bounds. # Examples ```julia @@ -311,6 +341,7 @@ function build_inventory_det_equivalent(; penalty = INVENTORY_PENALTY, seed = 42, integer = true, + strict = false, ) # Fix the random seed so demand samples are reproducible. Random.seed!(seed) @@ -366,22 +397,40 @@ function build_inventory_det_equivalent(; # Split end-of-period inventory into holding and backlog. @constraint(m, [t=1:T], inv_hold[t] - back[t] == s_out[t]) - # --- Target-deficit penalty via NormOneCone --- - # norm_deficit_arr[t] ≥ |s_mid[t] - s_target[t]| (L1 norm). - @variable(m, norm_deficit_arr[1:T] >= 0.0) - @variable(m, deficit_arr[1:T]) - @constraint(m, [t=1:T], deficit_arr[t] == s_mid[t] - s_target[t]) - @constraint(m, [t=1:T], [norm_deficit_arr[t]; deficit_arr[t:t]] in MOI.NormOneCone(2)) + if strict + # --- Strict targets: hard equalities, no deficit machinery --- + # Feasible for any target path with ŝ_t ∈ [s_{t-1}, s_{t-1} + Q_max] + # generated by a reachable policy: by induction, the equality pins + # s_mid[t] (hence s_out[t] = ŝ_t − d_t) to the value the policy + # planned from, so every later target stays reachable. + @constraint(m, [t=1:T], s_mid[t] == s_target[t]) - # --- Objective: operational cost + target penalty --- - if integer - @objective(m, Min, - sum(K * z[t] + c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T) + - penalty * sum(norm_deficit_arr)) + # --- Objective: pure operational cost --- + if integer + @objective(m, Min, + sum(K * z[t] + c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T)) + else + @objective(m, Min, + sum(c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T)) + end else - @objective(m, Min, - sum(c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T) + - penalty * sum(norm_deficit_arr)) + # --- Target-deficit penalty via NormOneCone --- + # norm_deficit_arr[t] ≥ |s_mid[t] - s_target[t]| (L1 norm). + @variable(m, norm_deficit_arr[1:T] >= 0.0) + @variable(m, deficit_arr[1:T]) + @constraint(m, [t=1:T], deficit_arr[t] == s_mid[t] - s_target[t]) + @constraint(m, [t=1:T], [norm_deficit_arr[t]; deficit_arr[t:t]] in MOI.NormOneCone(2)) + + # --- Objective: operational cost + target penalty --- + if integer + @objective(m, Min, + sum(K * z[t] + c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T) + + penalty * sum(norm_deficit_arr)) + else + @objective(m, Min, + sum(c * q[t] + h * inv_hold[t] + p * back[t] for t in 1:T) + + penalty * sum(norm_deficit_arr)) + end end # --- Build parameter mappings for DecisionRules interface --- @@ -623,3 +672,172 @@ function build_lstm_exante_policy(; seed=2024, hidden=16) return LSTMExAntePolicy(encoder, combiner, state) end + +# --------------------------------------------------------------------------- +# Reachable ex-ante policy (strict mode: always-feasible targets) +# --------------------------------------------------------------------------- + +@doc raw""" + InventoryReachablePolicy{E,C,S} + +Recurrent ex-ante policy whose targets are always one-stage reachable, +enabling **strict mode** (hard target equalities, no penalty hyperparameter). + +## Exact reachable set + +The only decision behind the target is the order quantity +``q \in [0, Q_{max}]`` with ``s^{mid} = s + q``, where ``s`` is the realized +incoming inventory. ``s^{mid}`` appears in no other restrictive constraint +(``s^{out}`` is free and the hold/backlog split is feasible for any sign), so +the one-stage reachable set of the target is **exactly** + +```math +\hat{s} \in [\, s,\; s + Q_{max} \,], +``` + +with no conservatism and no cross-component coupling. The binary setup +variable does not shrink this set (``z = 1`` admits every +``q \in [0, Q_{max}]``); it only reshapes the cost over it. + +## Boundary attainment (why `hardsigmoid`, not `sigmoid`) + +With a fixed ordering cost ``K > 0``, the endpoint ``q = 0`` ("do not +order") is an economically essential decision, not a measure-zero boundary. +A strict sigmoid keeps ``q = Q_{max}\,\sigma(z) > 0`` for every finite +``z``, which under strict equalities forces ``z_{setup} = 1`` and pays +``K`` every period — silently deleting the no-order action from the policy +class. The output therefore uses `hardsigmoid`, which attains 0 and 1 +exactly on finite inputs: + +```math +\hat{s} = s + Q_{max} \cdot \operatorname{hard\sigma}(z), \qquad +\operatorname{hard\sigma}(z) = \min(\max(z/6 + 1/2,\, 0),\, 1). +``` + +(Contrast with the hydro reachable policy, where the interval endpoints are +economically irrelevant and a strict sigmoid is harmless.) + +## Pass-through components + +The second and third state components (``d_{t-1}``, ``d_{t-2}`` histories) +are exogenous information states: their one-stage "reachable sets" are +singletons — the realized demand values. The policy passes them through +deterministically (no trainable parameters); the stage models leave their +target parameters unconstrained in both penalty and strict modes. + +## Information pattern and gradients + +Strictly **ex-ante**: the LSTM encoder consumes only the *lagged* demand +``d_{t-1}``; current demand ``d_t`` is never used for the order decision +(it appears in the output only as state pass-through plumbing). The bound +map ``\hat{s} = s + Q_{max} y`` is affine and exact, so — unlike the hydro +policy's physics bounds — it is left differentiable: the additive ``s`` +term carries an exact gradient path through the realized state. + +# Fields +- `encoder::E`: `Flux.LSTMCell` over the normalized lagged demand. +- `combiner::C`: `Dense` head mapping `[encoded; s/100; d_{t-2}/100]` to the + pre-activation order fraction. +- `state::S`: LSTM hidden state, threaded across stages; reset per scenario. +- `Q_max::Float32`: order capacity (the constant reachable-interval width). + +# Examples +```julia +policy = build_reachable_inventory_policy(; seed = 2024) +Flux.reset!(policy) +target = policy(Float32[d_t, inventory, d_lag1, d_lag2]) +@assert inventory <= target[1] <= inventory + INVENTORY_Q_MAX +``` +""" +mutable struct InventoryReachablePolicy{E,C,S} + encoder::E + combiner::C + state::S + Q_max::Float32 +end + +Functors.@functor InventoryReachablePolicy (encoder, combiner) + +function (policy::InventoryReachablePolicy)(x) + # Extract features from input: [d_t, inventory, d_{t-1}, d_{t-2}]. + # Only d_{t-1} (lagged) feeds the LSTM — d_t is NOT used (ex-ante). + inventory = Float32(x[2]) + last_demand = Float32(x[3]) + prev_demand = Float32(x[4]) + + # Match the element type of the LSTM state (Float32 during training). + T = eltype(first(policy.state)) + + # Feed the normalized lagged demand through the LSTM cell and thread + # the hidden state to the next stage call within this scenario. + encoded, new_state = policy.encoder(T[last_demand / 100], policy.state) + policy.state = new_state + + # Concatenate LSTM output with current inventory and prev demand. + combined = vcat(encoded, T[inventory / 100, prev_demand / 100]) + + # Map to the order fraction y ∈ [0, 1]; hardsigmoid attains both + # endpoints exactly (q = 0 and q = Q_max are real decisions). + y = Flux.hardsigmoid(policy.combiner(combined)[1]) + + # Affine map onto the exact reachable interval [s, s + Q_max], computed + # in Float64 FROM THE RAW INPUT STATE: the implied order quantity is then + # q = target − s = Q_max·y ≥ 0 bitwise. Using the Float32-cast state here + # would make the strict equality infeasible whenever y saturates at 0 and + # Float32(s) < s (the solver would need q ≈ −1e−5·s < 0) — a failure mode + # the integer variant hits constantly because the fixed cost K actively + # pushes the policy to the q = 0 boundary. + target_s_mid = Float64(x[2]) + Float64(policy.Q_max) * Float64(y) + + # Return [target_s_mid, d_t, d_{t-1}] — target + state pass-throughs. + return [target_s_mid, Float64(x[1]), Float64(last_demand)] +end + +""" + Flux.reset!(policy::InventoryReachablePolicy) -> Nothing + +Reset the LSTM hidden state to its initial value. Must be called at every +scenario boundary so demand-history memory does not leak across rollouts. +""" +function Flux.reset!(policy::InventoryReachablePolicy) + # Restore the LSTM hidden state to its fresh initial values. + policy.state = Flux.initialstates(policy.encoder) + return nothing +end + +""" + build_reachable_inventory_policy(; seed = 2024, hidden = 16, Q_max = INVENTORY_Q_MAX) + -> InventoryReachablePolicy + +Construct the reachable ex-ante policy used by the strict variants. + +Architecture matches [`build_lstm_exante_policy`](@ref) exactly — +LSTMCell(1 → hidden) encoder, Dense(hidden + 2 → 1) head — so +penalty-vs-strict comparisons isolate the output parameterization, not +model capacity. + +# Keyword Arguments +- `seed::Int`: random seed for weight initialization. +- `hidden::Int`: LSTM hidden dimension. +- `Q_max`: order capacity (reachable-interval width). + +# Examples +```julia +policy = build_reachable_inventory_policy(; seed = 2024, hidden = 16) +``` +""" +function build_reachable_inventory_policy(; seed=2024, hidden=16, Q_max=INVENTORY_Q_MAX) + # Fix the random seed for reproducible weight initialization. + Random.seed!(seed) + + # LSTM cell: 1 input (normalized lagged demand) → hidden state. + encoder = Flux.LSTMCell(1 => hidden) + + # Dense head: [LSTM output; inventory; prev_demand] → order fraction. + combiner = Dense(hidden + 2, 1) + + # Initialize the LSTM hidden state to its default zeros. + state = Flux.initialstates(encoder) + + return InventoryReachablePolicy(encoder, combiner, state, Float32(Q_max)) +end diff --git a/examples/inventory_control/compare_results.jl b/examples/inventory_control/compare_results.jl index f9d6f2e..b53671e 100644 --- a/examples/inventory_control/compare_results.jl +++ b/examples/inventory_control/compare_results.jl @@ -803,6 +803,7 @@ function relaxed_results() lstm_costs = optional_costs("relaxed_lstm", "dr") hp_costs = optional_costs("relaxed_hp", "dr") lstm_hp_costs = optional_costs("relaxed_lstm_hp", "dr") + strict_costs = optional_costs("strict", "dr") # Load scalar baseline metadata. base_stock_level = read_scalar(resolve_file("relaxed_basestock_S_star.txt")) @@ -820,6 +821,8 @@ function relaxed_results() push!(results, MethodResult("TS-DDR Relaxed (LSTM)", lstm_costs)) !isnothing(lstm_hp_costs) && push!(results, MethodResult("TS-DDR Relaxed (LSTM+HP)", lstm_hp_costs)) + !isnothing(strict_costs) && + push!(results, MethodResult("TS-DDR Strict (Reachable)", strict_costs)) # Append non-TS-DDR baselines. push!(results, MethodResult("SDDP (PAR)", sddp_costs)) @@ -834,6 +837,8 @@ function relaxed_results() push!(timing_tags, "relaxed_hp") !isnothing(resolve_file_optional("relaxed_lstm_hp_dr_timing.csv")) && push!(timing_tags, "relaxed_lstm_hp") + !isnothing(resolve_file_optional("strict_dr_timing.csv")) && + push!(timing_tags, "strict") return results, load_timing(timing_tags), base_stock_level, sddp_bound end @@ -862,6 +867,7 @@ function integer_results() hp_costs = optional_costs("integer_hp", "dr") lstm_costs = optional_costs("integer_lstm", "dr") lstm_sf_costs = optional_costs("integer_lstm_sf", "dr") + strict_integer_costs = optional_costs("strict_integer", "dr") # Load scalar baseline metadata. base_stock_level = read_scalar(resolve_file("integer_basestock_S_star.txt")) @@ -882,6 +888,8 @@ function integer_results() push!(results, MethodResult("TS-DDR (LSTM)", lstm_costs)) !isnothing(lstm_sf_costs) && push!(results, MethodResult("TS-DDR (LSTM+SF)", lstm_sf_costs)) + !isnothing(strict_integer_costs) && + push!(results, MethodResult("TS-DDR Strict (Reachable+Int)", strict_integer_costs)) # Append non-TS-DDR baselines. push!(results, MethodResult("SDDP (MIP fwd)", sddp_mip_forward_costs)) @@ -891,7 +899,7 @@ function integer_results() # Collect timing tags for all present variants. timing_tags = ["integer", "integer_cr"] - for tag in ["integer_sf", "integer_hp", "integer_lstm", "integer_lstm_sf"] + for tag in ["integer_sf", "integer_hp", "integer_lstm", "integer_lstm_sf", "strict_integer"] !isnothing(resolve_file_optional("$(tag)_dr_timing.csv")) && push!(timing_tags, tag) end diff --git a/examples/inventory_control/train_dr_inventory.jl b/examples/inventory_control/train_dr_inventory.jl index 7777e1b..bff6228 100644 --- a/examples/inventory_control/train_dr_inventory.jl +++ b/examples/inventory_control/train_dr_inventory.jl @@ -115,6 +115,23 @@ struct InventoryTrainingVariant penalty::Float64 policy_builder::Function penalty_schedule_fn::Function + # Strict mode: hard target equalities (no deficit, no penalty). Requires a + # reachable policy — see InventoryReachablePolicy and strict_variants(). + strict::Bool +end + +# Backward-compatible 11-argument constructor: every historical call site +# predates strict mode, so it defaults to the penalty formulation. +function InventoryTrainingVariant( + tag, integer, num_batches, train_per_batch, learning_rate, warmup_batches, + training_integer_strategy, score_function, penalty, policy_builder, + penalty_schedule_fn, +) + return InventoryTrainingVariant( + tag, integer, num_batches, train_per_batch, learning_rate, warmup_batches, + training_integer_strategy, score_function, penalty, policy_builder, + penalty_schedule_fn, false, + ) end function InventoryTrainingVariant( @@ -168,6 +185,20 @@ function penalty_schedule_for(variant::InventoryTrainingVariant) return [warmup_phase, full_penalty_phase] end +""" + no_penalty_schedule(variant::InventoryTrainingVariant) -> Nothing + +Return `nothing`: strict-mode models have no deficit variables, so there is +no penalty to schedule (`train_multistage` accepts +`penalty_schedule = nothing`). + +# Examples +```julia +schedule = no_penalty_schedule(variant) # nothing +``` +""" +no_penalty_schedule(::InventoryTrainingVariant) = nothing + """ method_label(variant::InventoryTrainingVariant) -> String @@ -184,6 +215,10 @@ label = method_label(variant) function method_label(variant::InventoryTrainingVariant) tag = variant.tag + # --- Strict variants (reachable policy, no penalty) --- + tag == "strict" && return "TS-DDR Strict (Reachable)" + tag == "strict_integer" && return "TS-DDR Strict (Reachable+Int)" + # --- Relaxed tuned variants --- tag == "relaxed_lstm" && return "TS-DDR Relaxed (LSTM)" tag == "relaxed_hp" && return "TS-DDR Relaxed (HighPenalty)" @@ -403,7 +438,27 @@ det_eq, state_in, state_out, sampler, initial_state = ``` """ function build_training_problem(variant::InventoryTrainingVariant) - # Training uses a deterministic equivalent so target-dual gradients are coupled. + if variant.strict + # Strict variants train STAGE-WISE (closed loop). In this problem the + # target (order-up-to position s_mid) and the carried state + # (post-demand inventory s_out = s_mid − d) are different quantities, + # so the deterministic equivalent's target-feedback recursion would + # hand the reachable policy s_mid where it expects s_out — computing + # the reachable interval from the wrong state and breaking the strict + # feasibility guarantee. Stage-wise training feeds realized states, + # for which the interval [s, s + Q_max] is exact. (Contrast with the + # hydro case, where target and carried state are the same reservoir + # volume and the strict regular DE is safe by induction.) + return build_inventory_subproblems(; + num_scenarios = N_TRAIN_SCENARIOS, + seed = 42, + integer = variant.integer, + strict = true, + ) + end + + # Penalty variants use a deterministic equivalent so target-dual gradients + # are coupled across stages. return build_inventory_det_equivalent(; num_scenarios = N_TRAIN_SCENARIOS, penalty = variant.penalty, @@ -433,6 +488,7 @@ function build_evaluation_problem(variant::InventoryTrainingVariant) penalty = variant.penalty, seed = 99, integer = variant.integer, + strict = variant.strict, ) end @@ -485,6 +541,37 @@ function estimate_initial_loss( ) end +# Stage-wise method (strict variants train on subproblems with realized-state +# feedback): rolls the policy in closed loop, matching the training semantics. +function estimate_initial_loss( + policy, + subproblems::Vector{JuMP.Model}, + state_params_in, + state_params_out, + uncertainty_sampler, + initial_state, + variant::InventoryTrainingVariant, +) + # Use a small fixed sample only to seed SaveBest with a finite baseline. + Random.seed!(111) + + return mean( + let uncertainty_sample = sample(uncertainty_sampler) + # Closed-loop rollout: each stage sees the realized state. + Flux.reset!(policy) + simulate_multistage( + subproblems, + state_params_in, + state_params_out, + initial_state, + uncertainty_sample, + policy; + integer_strategy = variant.training_integer_strategy, + ) + end for _ in 1:12 + ) +end + """ train_variant!(policy, variant, det_eq, state_params_in, state_params_out, uncertainty_sampler, initial_state, model_path, curve_path; @@ -528,7 +615,12 @@ train_variant!(policy, variant, det_eq, spi, spo, sampler, x0, function train_variant!( policy, variant::InventoryTrainingVariant, - det_eq::JuMP.Model, + # Penalty variants train on the deterministic equivalent (JuMP.Model); + # strict variants train stage-wise (Vector{JuMP.Model}, realized-state + # feedback) because the target (order-up-to position s_mid) and the + # carried state (post-demand inventory s_out) are different quantities — + # target feedback would hand the reachable policy the wrong state. + det_eq::Union{JuMP.Model,Vector{JuMP.Model}}, state_params_in, state_params_out, uncertainty_sampler, @@ -577,6 +669,12 @@ function train_variant!( # Fix optimizer randomness for repeatability. Random.seed!(2024) + # The score-function keyword exists only on the deterministic-equivalent + # overload of train_multistage; the stage-wise overload (strict variants) + # must not receive it. + score_function_kwargs = det_eq isa JuMP.Model ? + (score_function = variant.score_function,) : NamedTuple() + elapsed_seconds = @elapsed train_multistage( policy, initial_state, @@ -589,7 +687,7 @@ function train_variant!( optimizer = Flux.Adam(variant.learning_rate), integer_strategy = variant.training_integer_strategy, penalty_schedule = variant.penalty_schedule_fn(variant), - score_function = variant.score_function, + score_function_kwargs..., record = (sample_log, iteration, current_policy) -> begin loss = isempty(sample_log.objectives_no_deficit) ? NaN : @@ -1084,6 +1182,45 @@ function inventory_training_variants() ), # Variant B: LSTM with tuned score function lstm_score_function_variant(), + # --- Strict variants (hard target equalities, no penalty tuning) --- + # The reachable policy guarantees ŝ ∈ [s, s + Q_max] exactly, so the + # strict equality s_mid == ŝ is always feasible and its dual is the + # pure shadow price — no penalty hyperparameter, no annealing. + InventoryTrainingVariant( + "strict", + false, + 800, + 10, + 1.0e-3, + 120, + NoIntegerStrategy(), + nothing, + INVENTORY_PENALTY, # unused in strict mode + () -> build_reachable_inventory_policy(; seed = 2024), + no_penalty_schedule, + true, + ), + # Strict + binary setup: FixedDiscreteIntegerStrategy reads exact LP + # shadow prices at the fixed integer assignment. The K·z jump at + # ŝ = s (order/no-order switch) is invisible to that local dual; a + # mixed score-function gradient can be layered on later using + # PENALTY-mode rollout models (ScoreFunctionConfig owns its own + # subproblems), since strict rollouts would reject perturbed targets + # that leave the reachable interval. + InventoryTrainingVariant( + "strict_integer", + true, + 800, + 10, + 8.0e-4, + 120, + FixedDiscreteIntegerStrategy(), + nothing, + INVENTORY_PENALTY, # unused in strict mode + () -> build_reachable_inventory_policy(; seed = 2024), + no_penalty_schedule, + true, + ), ] end diff --git a/src/DecisionRules.jl b/src/DecisionRules.jl index 5f90a6a..882e934 100644 --- a/src/DecisionRules.jl +++ b/src/DecisionRules.jl @@ -17,6 +17,7 @@ export simulate_multistage, simulate_states, simulate_stage, dense_multilayer_nn, + dense_policy_head, variable_to_parameter, create_deficit!, default_annealed_schedule, @@ -32,6 +33,10 @@ export simulate_multistage, ContinuousRelaxationIntegerStrategy, StallingCriterium, policy_input_dim, + ContextualPolicy, + context_at, + stage_phase_context, + vcat_contexts, normalize_recur_state, StateConditionedPolicy, state_conditioned_policy, @@ -113,11 +118,53 @@ tests to ensure that controlled problems never silently produce zero gradients. """ struct ErrorGradientFallback <: AbstractGradientFallback end +""" + _zero_cotangents(n_in, n_out) + +Create a tuple of zero/no tangents compatible with the `get_next_state` rrule pullback signature. + +Used by [`handle_gradient_error`](@ref) to produce a safe, neutral gradient when the +solver or DiffOpt differentiation fails. The returned tuple matches the cotangent +layout expected by `ChainRulesCore.rrule` for `get_next_state`: +four `NoTangent()` entries (for the function itself and non-differentiable arguments), +followed by dense zero vectors for the state-in and state-out dimensions, and a +trailing `NoTangent()`. + +# Arguments +- `n_in::Int`: dimension of the incoming state vector. +- `n_out::Int`: dimension of the outgoing state vector. + +# Returns +A `Tuple` of `NoTangent` and `Vector{Float64}` elements that Zygote can propagate +without error. +""" _zero_cotangents(n_in, n_out) = ( NoTangent(), NoTangent(), NoTangent(), NoTangent(), zeros(n_in), zeros(n_out), NoTangent(), ) +""" + handle_gradient_error(fallback::AbstractGradientFallback, e, n_state_in, n_state_out) + +Handle an exception raised inside the `get_next_state` rrule pullback. + +This is the **rrule-level** extension point: it is called when the backward pass +through a single stage fails (e.g., the solver returned an infeasible status and +DiffOpt cannot differentiate). Concrete methods decide whether to absorb the +error or propagate it. + +# Arguments +- `fallback::AbstractGradientFallback`: dispatch tag controlling recovery behavior. +- `e`: the caught exception. +- `n_state_in::Int`: dimension of the incoming state vector (needed to build zero cotangents). +- `n_state_out::Int`: dimension of the outgoing state vector. + +# Returns +A cotangent tuple (same layout as [`_zero_cotangents`](@ref)) when the error is +absorbed, or does not return (re-throws) when the error is propagated. + +See [`AbstractGradientFallback`](@ref) for how to implement custom subtypes. +""" function handle_gradient_error(::ZeroGradientFallback, e, n_state_in, n_state_out) @warn "get_next_state pullback failed — returning zero gradients" exception=(e, catch_backtrace()) return _zero_cotangents(n_state_in, n_state_out) @@ -127,6 +174,27 @@ function handle_gradient_error(::ErrorGradientFallback, e, n_state_in, n_state_o rethrow(e) end +""" + handle_training_error(fallback::AbstractGradientFallback, e, iter) + +Handle an exception raised during a full training iteration (gradient computation +and parameter update). + +This is the **training-loop-level** extension point: it is called when +`Zygote.gradient` or the subsequent optimizer update throws (e.g., a DiffOpt +assertion error or a numerical issue in the loss computation). Unlike +[`handle_gradient_error`](@ref), which operates inside a single-stage rrule, +this handler wraps the entire forward-backward pass for one iteration. + +# Arguments +- `fallback::AbstractGradientFallback`: dispatch tag controlling recovery behavior. +- `e`: the caught exception. +- `iter::Int`: current training iteration index (used in log messages). + +# Returns +- `true` to skip this iteration and continue training. +- Does not return (re-throws) when the error should propagate. +""" function handle_training_error(::ZeroGradientFallback, e, iter) @warn "Gradient computation failed at iter $iter — skipping update" exception=(e, catch_backtrace()) return true @@ -136,6 +204,25 @@ function handle_training_error(::ErrorGradientFallback, e, iter) rethrow(e) end +""" + handle_rollout_error(fallback::AbstractGradientFallback, e, iter) + +Handle an exception raised during a rollout evaluation scenario. + +This is the **rollout-level** extension point: it is called when a single +out-of-sample scenario fails during [`RolloutEvaluation`](@ref) (e.g., solver +infeasibility on an unseen uncertainty sample). Absorbing the error skips that +scenario and lets the evaluation continue with the remaining samples. + +# Arguments +- `fallback::AbstractGradientFallback`: dispatch tag controlling recovery behavior. +- `e`: the caught exception. +- `iter::Int`: scenario index within the rollout batch. + +# Returns +- `true` to skip this scenario and continue the rollout. +- Does not return (re-throws) when the error should propagate. +""" function handle_rollout_error(::ZeroGradientFallback, e, iter) @warn "Rollout scenario failed at iter $iter — skipping" exception=(e, catch_backtrace()) return true diff --git a/src/dense_multilayer_nn.jl b/src/dense_multilayer_nn.jl index 113ffac..f19578c 100644 --- a/src/dense_multilayer_nn.jl +++ b/src/dense_multilayer_nn.jl @@ -1,17 +1,97 @@ using Functors using ChainRulesCore +raw""" + dense_policy_head(num_inputs, num_outputs, layers; activation=Flux.relu) + +Create the feed-forward target head used after recurrent uncertainty encoding. + +Unlike [`dense_multilayer_nn`](@ref), the final layer also uses `activation`. +This is useful for bounded policies where the head must emit normalized values +such as `sigmoid(z) ∈ [0, 1]`. + +Mathematically, when `layers = [h₁, …, h_L]`, the head computes + +```math +H_\theta(z) = +\sigma\left(W_{L+1} + \sigma\left(W_L \cdots \sigma(W_1 z + b_1) \cdots + b_L\right) + + b_{L+1}\right). +``` + +This helper exists so state-conditioned TS-DDR policies can make the map +`[encoded_uncertainty; state] → target` nonlinear while keeping recurrence +confined to the uncertainty encoder. + +# Arguments +- `num_inputs::Int`: number of combined features, usually `encoded_uncertainty` + concatenated with the previous state. +- `num_outputs::Int`: number of target outputs. +- `layers::AbstractVector{Int}`: hidden widths for the nonrecurrent head. +- `activation`: activation used by each hidden layer and by the output layer. + +# Returns +- A `Dense` layer when `layers` is empty, otherwise a `Chain` of dense layers. + +# Examples +```julia +head = dense_policy_head(16, 3, [32, 32]; activation=sigmoid) +``` + +See also: [`state_conditioned_policy`](@ref), [`dense_multilayer_nn`](@ref) +""" +function dense_policy_head( + num_inputs::Int, + num_outputs::Int, + layers::AbstractVector{Int}; + activation=Flux.relu, +) + isempty(layers) && return Dense(num_inputs => num_outputs, activation) + head_layers = Any[Dense(num_inputs => layers[1], activation)] + for i in 1:(length(layers) - 1) + push!(head_layers, Dense(layers[i] => layers[i + 1], activation)) + end + push!(head_layers, Dense(layers[end] => num_outputs, activation)) + return Chain(head_layers...) +end + """ dense_multilayer_nn(num_inputs, num_outputs, layers; activation=Flux.relu, dense=Dense) Create a multi-layer neural network with the specified architecture. +Given hidden layer widths ``(h_1, \\dots, h_L)`` and activation ``\\sigma``, +the resulting `Chain` computes + +```math +f(x) = W_{L+1} \\, \\sigma\\bigl(W_L \\cdots \\sigma(W_1 x + b_1) \\cdots + b_L\\bigr) + b_{L+1}, +``` + +where ``W_k \\in \\mathbb{R}^{h_k \\times h_{k-1}}``, ``h_0 = \\text{num\\_inputs}``, +and the final layer ``W_{L+1} \\in \\mathbb{R}^{\\text{num\\_outputs} \\times h_L}`` +has no activation. When `layers` is empty, there is no hidden/output +distinction: the returned single dense layer uses `activation`, preserving the +same constructor semantics as `Dense(num_inputs, num_outputs, activation)`. + +If `dense` is a recurrent type (`LSTM`, `GRU`, `RNN`), pair notation +`in => out` is used instead of `(in, out, σ)`, and the last layer omits the +activation so it can act as a linear projection. + # Arguments -- `num_inputs::Int`: Number of input features -- `num_outputs::Int`: Number of output features -- `layers::Vector{Int}`: Hidden layer sizes -- `activation`: Activation function (default: Flux.relu) -- `dense`: Layer type (Dense, LSTM, etc.) +- `num_inputs::Int`: number of input features ``h_0``. +- `num_outputs::Int`: number of output features. +- `layers::Vector{Int}`: hidden layer widths ``(h_1, \\dots, h_L)``. +- `activation`: element-wise activation ``\\sigma`` (default: `Flux.relu`). +- `dense`: layer constructor (`Dense`, `LSTM`, `GRU`, `RNN`). + +# Returns +- `Chain` (or a single layer when `layers` is empty). + +# Examples +```julia +mlp = dense_multilayer_nn(10, 3, [64, 32]) # 10 → 64 → 32 → 3 +rnn = dense_multilayer_nn(5, 2, [16]; dense=LSTM) # 5 → 16 → 2 (LSTM) +``` """ function dense_multilayer_nn( num_inputs::Int, @@ -20,9 +100,14 @@ function dense_multilayer_nn( activation=Flux.relu, dense=Dense, ) + # Recurrent layers use pair notation (in => out) and ignore activation. is_recurrent = dense in (LSTM, GRU, RNN) + + # Helper that builds one hidden layer with the appropriate constructor form. _make_layer(in_dim, out_dim) = is_recurrent ? dense(in_dim => out_dim) : dense(in_dim, out_dim, activation) + + # No hidden layers: preserve Dense(input, output, activation) semantics. if length(layers) == 0 return if is_recurrent dense(num_inputs => num_outputs) @@ -30,10 +115,17 @@ function dense_multilayer_nn( dense(num_inputs, num_outputs, activation) end end + + # Interior hidden layers: h_2 → h_3 → … → h_{L-1}. midlayers = [_make_layer(layers[i], layers[i + 1]) for i in 1:(length(layers) - 1)] + + # First hidden layer maps from the input dimension. first_layer = _make_layer(num_inputs, layers[1]) + + # Final layer omits activation so the output is an unrestricted linear map. last_layer = is_recurrent ? dense(layers[end] => num_outputs) : dense(layers[end], num_outputs) + return Chain(first_layer, midlayers..., last_layer) end @@ -42,62 +134,200 @@ end Compute the input dimension for a policy network. -Policy networks receive `[uncertainty..., previous_state...]` as input, -so the input dimension is `num_uncertainties + num_states`. +Policy networks receive ``[w_t;\\; x_{t-1}]`` as input, so the required +dimension is -This format is consistent between subproblems and deterministic equivalent -formulations, enabling warmstarting policies trained with det_eq for use +```math +d_{\\text{in}} = \\dim(w_t) + \\dim(x_{t-1}). +``` + +This format is consistent between subproblem and deterministic-equivalent +formulations, enabling warm-starting policies trained with det_eq for use with subproblems. # Arguments -- `num_uncertainties::Int`: Number of uncertainty parameters per stage -- `num_states::Int`: Number of state variables +- `num_uncertainties::Int`: dimensionality of the stage uncertainty ``w_t``. +- `num_states::Int`: dimensionality of the state ``x_{t-1}``. # Returns -- `Int`: Total input dimension for the policy network +- `Int`: total input dimension ``d_{\\text{in}}``. + +# Examples +```julia +d = policy_input_dim(5, 3) # 8 +``` """ function policy_input_dim(num_uncertainties::Int, num_states::Int) + # Input is the concatenation [w_t; x_{t-1}]. return num_uncertainties + num_states end +function policy_input_dim(num_uncertainties::Int, num_states::Int, num_context::Int) + # Input is the concatenation [context_t; w_t; x_{t-1}]. + return num_context + num_uncertainties + num_states +end + """ policy_input_dim(uncertainty_samples, initial_state) Compute the input dimension for a policy network from problem data. +Infers ``\\dim(w_t)`` from the first uncertainty sample and ``\\dim(x)`` +from the initial state vector, then delegates to +[`policy_input_dim(::Int, ::Int)`](@ref). + # Arguments -- `uncertainty_samples`: Uncertainty samples from problem construction -- `initial_state`: Initial state vector +- `uncertainty_samples::Vector`: uncertainty samples; `length(uncertainty_samples[1])` + gives ``\\dim(w_t)``. +- `initial_state::Vector`: initial state vector; `length(initial_state)` gives + ``\\dim(x)``. # Returns -- `Int`: Total input dimension for the policy network +- `Int`: total input dimension ``d_{\\text{in}}``. + +# Examples +```julia +d = policy_input_dim(uncertainty_samples, x0) +``` """ function policy_input_dim(uncertainty_samples::Vector, initial_state::Vector) + # Infer dimensions from the first sample and the initial state. num_uncertainties = length(uncertainty_samples[1]) num_states = length(initial_state) return policy_input_dim(num_uncertainties, num_states) end +""" + ContextualPolicy(policy, context) + +Wrap a stage policy so each call receives known exogenous context before the +usual policy input. The wrapped policy is called as +`policy(vcat(context_at(context, t), input))`, where `t` is advanced once per +call and reset by `Flux.reset!`. + +This keeps training and rollout loops unchanged: context is data attached to +the trajectory, while the inner policy remains the only trainable component. +Matrix contexts are interpreted as `d_context x T`; function contexts are +called as `context(t)`. +""" +mutable struct ContextualPolicy{P,C} + policy::P + context::C + t::Int +end + +ContextualPolicy(policy, context) = ContextualPolicy(policy, context, 0) + +Functors.@functor ContextualPolicy (policy,) + +""" + context_at(context, t) + +Return the context vector for one-based stage `t`. +""" +function context_at(context::AbstractMatrix, t::Integer) + 1 <= t <= size(context, 2) || + throw(BoundsError(context, (:, t))) + return view(context, :, t) +end + +context_at(context::Function, t::Integer) = context(t) + +function (m::ContextualPolicy)(input) + m.t += 1 + return m.policy(vcat(context_at(m.context, m.t), input)) +end + +function Flux.reset!(m::ContextualPolicy) + m.t = 0 + Flux.reset!(m.policy) + return nothing +end + +""" + stage_phase_context(T; period, include_progress=true) + +Build a `d x T` context matrix encoding known calendar/stage information. +Rows are `sin(2*pi*t/period)`, `cos(2*pi*t/period)`, and, when +`include_progress=true`, `t/T`. + +The sine/cosine pair represents a cyclic process without a discontinuity +between the final and first period positions. A raw index would put those +neighbors far apart; one-hot period features would be exact but high +dimensional and would not encode adjacency. The optional progress feature has +a different purpose: it tells the policy how close it is to the optimization +horizon endpoint. +""" +function stage_phase_context(T::Integer; period::Integer, include_progress::Bool=true) + T >= 1 || throw(ArgumentError("T must be positive")) + period >= 1 || throw(ArgumentError("period must be positive")) + nrows = include_progress ? 3 : 2 + ctx = Matrix{Float32}(undef, nrows, T) + for t in 1:T + θ = 2f0 * Float32(pi) * Float32(t) / Float32(period) + ctx[1, t] = sin(θ) + ctx[2, t] = cos(θ) + if include_progress + ctx[3, t] = Float32(t) / Float32(T) + end + end + return ctx +end + +""" + vcat_contexts(a, b, ...) + +Vertically concatenate context matrices after checking that they cover the +same number of stages. +""" +function vcat_contexts(contexts::AbstractMatrix...) + isempty(contexts) && return Matrix{Float32}(undef, 0, 0) + T = size(first(contexts), 2) + all(size(c, 2) == T for c in contexts) || + throw(ArgumentError("all contexts must have the same number of columns")) + return vcat(contexts...) +end + """ StateConditionedPolicy -A policy architecture that separates temporal encoding from state conditioning: -- `encoder`: a recurrent cell (`LSTMCell`/`GRUCell`/`RNNCell`, or a `Chain` of them) - that encodes only the uncertainty sequence (temporal dependencies) -- `combiner`: a `Dense` layer that combines the encoder output with the previous - state to produce the next state +A policy architecture that separates temporal encoding from state conditioning. -Flux's recurrent cells are stateless (Flux >= 0.16): each call returns -`(output, new_state)` instead of mutating an internal `Recur`. `StateConditionedPolicy` -therefore carries the encoder's recurrent state itself in `state`, threading it through -one call per stage. Call `Flux.reset!` to clear it (back to `Flux.initialstates`) at the -start of a rollout. +The encoder is a recurrent cell (`LSTMCell`/`GRUCell`/`RNNCell`, or a `Chain` +of cells) that processes only the uncertainty sequence to capture temporal +dependencies. The combiner is a feed-forward target head that merges the +encoder output with the previous state to predict the next state. + +Given uncertainty ``w_t`` and previous state ``x_{t-1}``, the forward pass is + +```math +h_t, s_t = \\text{LSTM}(w_t, s_{t-1}) \\\\ +\\hat{x}_t = f_{\\text{combine}}([h_t;\\; x_{t-1}]) +``` -Input format: [uncertainty..., previous_state...] +where ``s_t`` is the hidden recurrent state carried across stages and +``f_{\\text{combine}}`` is a `Dense` layer with activation ``\\sigma``. + +Flux's recurrent cells are stateless (Flux >= 0.16): each call returns +`(output, new_state)` instead of mutating an internal `Recur`. +`StateConditionedPolicy` therefore carries the encoder's recurrent state +itself in `state`, threading it through one call per stage. Call +`Flux.reset!` to clear it (back to `Flux.initialstates`) at the start of +a rollout. + +# Fields +- `encoder::E`: recurrent cell or `Chain` of cells encoding uncertainty. +- `combiner::C`: feed-forward head mapping ``[h_t;\\; x_{t-1}]`` to + ``\\hat{x}_t``. +- `state::S`: current recurrent state ``s_t``, carried across calls. +- `n_uncertainty::Int`: dimensionality of ``w_t``. +- `n_state::Int`: dimensionality of ``x_{t-1}``. + +Input format: `[uncertainty..., previous_state...]`. """ mutable struct StateConditionedPolicy{E,C,S} encoder::E # Recurrent cell, or Chain of cells, that processes uncertainty only - combiner::C # Dense that combines encoder output with previous state + combiner::C # Feed-forward head combining encoder output with state state::S # Encoder recurrent state, carried across calls n_uncertainty::Int # Number of uncertainty dimensions n_state::Int # Number of state dimensions @@ -108,41 +338,91 @@ end Functors.@functor StateConditionedPolicy (encoder, combiner) """ - materialize_tangent(x) - -Recursively convert ChainRulesCore tangent types (MutableTangent, Tangent) -to plain NamedTuples/Arrays that Flux.update! can handle. + materialize_tangent(x::Number) -> Number -This is needed because Zygote produces MutableTangent for mutable structs (like Flux.Recur), -but Flux.update!/Optimisers.jl expects plain NamedTuples. +Return numeric tangents unchanged (leaf values are already plain scalars). """ materialize_tangent(x::Number) = x + +""" + materialize_tangent(x::AbstractArray) -> AbstractArray + +Return array tangents unchanged (arrays are already `Flux.update!`-compatible). +""" materialize_tangent(x::AbstractArray) = x + +""" + materialize_tangent(::Nothing) -> Nothing + +Map `nothing` tangents (unused parameters) to `nothing`. +""" materialize_tangent(::Nothing) = nothing + +""" + materialize_tangent(::ChainRulesCore.NoTangent) -> Nothing + +Map `NoTangent` (structural zeros for non-differentiable fields) to `nothing`. +""" materialize_tangent(::ChainRulesCore.NoTangent) = nothing + +""" + materialize_tangent(::ChainRulesCore.ZeroTangent) -> Nothing + +Map `ZeroTangent` (additive identity in the tangent space) to `nothing`. +""" materialize_tangent(::ChainRulesCore.ZeroTangent) = nothing +""" + materialize_tangent(t::ChainRulesCore.MutableTangent) -> NamedTuple + +Unwrap a `MutableTangent` (produced by Zygote for mutable structs such as +`Flux.Recur`) by extracting its backing `NamedTuple` and recursing. +""" function materialize_tangent(t::ChainRulesCore.MutableTangent) - # MutableTangent stores values in RefValues, extract them + # MutableTangent stores values in RefValues; extract them via backing. backing = ChainRulesCore.backing(t) return materialize_tangent(backing) end +""" + materialize_tangent(t::ChainRulesCore.Tangent) -> NamedTuple + +Unwrap an immutable `Tangent` by extracting its backing and recursing. +""" function materialize_tangent(t::ChainRulesCore.Tangent) + # Tangent backing is already a NamedTuple; recurse to handle nested types. backing = ChainRulesCore.backing(t) return materialize_tangent(backing) end +""" + materialize_tangent(nt::NamedTuple{K}) -> NamedTuple{K} + +Recurse through every field of a `NamedTuple`, materializing each value. +""" function materialize_tangent(nt::NamedTuple{K}) where {K} + # Apply materialize_tangent element-wise and reconstruct the same keys. vals = map(materialize_tangent, values(nt)) return NamedTuple{K}(vals) end +""" + materialize_tangent(t::Tuple) -> Tuple + +Recurse through every element of a `Tuple`, materializing each value. +""" function materialize_tangent(t::Tuple) return map(materialize_tangent, t) end +""" + materialize_tangent(ref::Base.RefValue) + +Dereference a `RefValue` wrapper (used inside `MutableTangent` fields) +and recurse on the contained value. +""" function materialize_tangent(ref::Base.RefValue) + # RefValue wraps a single value; extract it before recursing. return materialize_tangent(ref[]) end @@ -153,6 +433,7 @@ Return the underlying recurrent cell of `layer`. `Flux.LSTM`/`GRU`/`RNN` wrap a (`LSTMCell`/`GRUCell`/`RNNCell`) in a `.cell` field; if `layer` has no such field it is already a cell and is returned unchanged. """ +# Unwrap .cell field if present (LSTM → LSTMCell); return unchanged otherwise. _as_cell(layer) = hasfield(typeof(layer), :cell) ? layer.cell : layer """ @@ -161,7 +442,9 @@ _as_cell(layer) = hasfield(typeof(layer), :cell) ? layer.cell : layer Return the initial recurrent state for `encoder`: `Flux.initialstates(encoder)` for a single cell, or a tuple of per-layer initial states for a `Chain` of cells. """ +# Single cell: return (h0, c0) or equivalent from Flux.initialstates. _init_recurrent_state(cell) = Flux.initialstates(cell) +# Chain of cells: one initial state per layer, returned as a tuple. _init_recurrent_state(chain::Chain) = map(_init_recurrent_state, chain.layers) """ @@ -171,38 +454,106 @@ Advance `encoder` by one step on input `x` from recurrent `state`, returning the output and the updated state. For a `Chain` of cells, each layer's output feeds the next and each layer's state is threaded independently. """ +# Single cell: one stateful call returns (output, new_state). _step_encoder(cell, x, state) = cell(x, state) function _step_encoder(chain::Chain, x, states::Tuple) + # Delegate to the recursive tuple-based implementation. return _step_encoder_layers(chain.layers, x, states) end +""" + _step_encoder_layers(layers, x, states) -> (output, new_states) + +Recursively advance a tuple of recurrent layers by one time step. + +Each layer receives the output of the previous layer as input and its own +independent recurrent state. The base case (`layers == ()`) returns the +input unchanged with an empty state tuple. + +# Arguments +- `layers::Tuple`: remaining recurrent cells to evaluate. +- `x`: current input (or output of the prior layer). +- `states::Tuple`: per-layer recurrent states, same length as `layers`. + +# Returns +- `output`: output of the last layer in `layers`. +- `new_states::Tuple`: updated recurrent states, one per layer. +""" _step_encoder_layers(::Tuple{}, x, ::Tuple{}) = x, () function _step_encoder_layers(layers::Tuple, x, states::Tuple) + # Advance the first layer with its own recurrent state. out, new_state = _step_encoder(first(layers), x, first(states)) + + # Recurse on remaining layers, feeding this layer's output as input. rest_out, rest_states = _step_encoder_layers(Base.tail(layers), out, Base.tail(states)) + + # Reassemble the full state tuple: this layer's state followed by the rest. return rest_out, (new_state, rest_states...) end +""" + _state_eltype(state) -> Type + +Return the scalar element type of a recurrent state. + +For nested tuple states (e.g. LSTM's `(h, c)` or a `Chain`'s tuple of +per-layer states) this recurses into the first element until it reaches an +`AbstractVector`, then returns `eltype(v)`. The result is used to cast +inputs to the encoder's precision before each step. + +# Arguments +- `state::Tuple`: nested recurrent state. +- `v::AbstractVector`: leaf state vector. + +# Returns +- `Type`: the scalar element type (e.g. `Float32`). +""" _state_eltype(state::Tuple) = _state_eltype(first(state)) _state_eltype(v::AbstractVector) = eltype(v) +""" + (m::StateConditionedPolicy)(x) -> AbstractVector + +Execute the forward pass of the state-conditioned policy. + +The input `x` is split into the uncertainty portion ``w_t`` and the previous +state ``x_{t-1}``. The forward pass computes + +```math +h_t, s_t = \\text{encoder}(w_t, s_{t-1}) \\\\ +\\hat{x}_t = f_{\\text{combine}}([h_t;\\; x_{t-1}]) +``` + +where ``s_t`` is the updated recurrent state (stored in `m.state` for the +next call) and ``f_{\\text{combine}}`` is a `Dense` layer. + +# Arguments +- `x::AbstractVector`: concatenated input `[w_t..., x_{t-1}...]` of length + `m.n_uncertainty + m.n_state`. + +# Returns +- `AbstractVector`: predicted next state ``\\hat{x}_t``. +""" function (m::StateConditionedPolicy)(x) - # Split input into uncertainty and previous state + # Split the concatenated input into uncertainty w_t and previous state x_{t-1}. uncertainty = x[1:m.n_uncertainty] prev_state = x[(m.n_uncertainty + 1):end] - # Encode uncertainty through the recurrent encoder, carrying state across calls. - # Cast to encoder precision to keep the recurrent state type stable - # (avoids a Zygote codegen bug with nested-tuple convert when solver - # feeds Float64 into a Float32 LSTM). + # Determine the encoder's scalar precision from its current recurrent state. + # Casting the input avoids a Zygote codegen bug with nested-tuple convert + # when an upstream solver feeds Float64 into a Float32 LSTM. T = _state_eltype(m.state) + + # Advance the recurrent encoder by one step: h_t, s_t = encoder(w_t, s_{t-1}). encoded, new_state = _step_encoder(m.encoder, T.(uncertainty), m.state) + + # Persist the new recurrent state so the next call starts from s_t. m.state = new_state - # Combine encoded uncertainty with previous state + # Concatenate encoder output h_t with previous state x_{t-1}. combined = vcat(encoded, prev_state) - # Output next state prediction + # Map the combined vector through the feed-forward combiner to produce x_hat_t. return m.combiner(combined) end @@ -213,29 +564,62 @@ Reset the encoder's recurrent state to `Flux.initialstates`, e.g. before startin new rollout. """ function Flux.reset!(m::StateConditionedPolicy) + # Reinitialize s_0 to the cell's default (zeros for LSTM h/c). m.state = _init_recurrent_state(m.encoder) return nothing end """ state_conditioned_policy(n_uncertainty, n_state, n_output, layers; - activation=Flux.relu, encoder_type=Flux.LSTM) + activation=Flux.relu, encoder_type=Flux.LSTM, + combiner_layers=Int[]) + +Create a [`StateConditionedPolicy`](@ref) with the specified architecture. + +The resulting policy computes + +```math +h_t, s_t = \\text{encoder}(w_t, s_{t-1}) \\\\ +\\hat{x}_t = \\sigma\\bigl(W [h_t;\\; x_{t-1}] + b\\bigr) +``` -Create a StateConditionedPolicy with the specified architecture. +where the encoder is a stack of recurrent cells of width +``(l_1, \\dots, l_L)`` and the combiner has input dimension +``l_L + n_{\\text{state}}``. `combiner_layers` adds optional hidden layers +to this nonrecurrent combiner, making the state-to-target map nonlinear while +keeping recurrence confined to the uncertainty sequence. This distinction is +important for TS-DDR hydro experiments: inflow history is recurrent, but the +current reservoir state is used only as contemporaneous conditioning. # Arguments -- `n_uncertainty::Int`: Number of uncertainty input dimensions -- `n_state::Int`: Number of state dimensions (both input and output) -- `n_output::Int`: Number of output dimensions (typically same as n_state) -- `layers::Vector{Int}`: Hidden layer sizes for the encoder -- `activation`: Activation function for dense layers (default: relu) -- `encoder_type`: Recurrent layer/cell type (`LSTM`, `GRU`, `RNN`, or their `*Cell` - variants; default: `Flux.LSTM`). Must support `Flux.initialstates` and the stateful - `(x, state) -> (output, new_state)` call (Flux >= 0.16). - -# Architecture -- Encoder: encoder_type(n_uncertainty => layers[1]) -> ... -> layers[end] -- Combiner: Dense(layers[end] + n_state => n_output) +- `n_uncertainty::Int`: dimensionality of the uncertainty input ``w_t``. +- `n_state::Int`: dimensionality of the state ``x_{t-1}`` (both input and output). +- `n_output::Int`: output dimension (typically equal to `n_state`). +- `layers::Vector{Int}`: hidden layer widths ``(l_1, \\dots, l_L)`` for the encoder. +- `activation`: activation ``\\sigma`` for the combiner `Dense` layer + (default: `Flux.relu`). +- `encoder_type`: recurrent layer/cell constructor (`LSTM`, `GRU`, `RNN`, or + their `*Cell` variants; default: `Flux.LSTM`). Must support + `Flux.initialstates` and the stateful `(x, state) -> (output, new_state)` + interface (Flux >= 0.16). +- `combiner_layers::Vector{Int}`: hidden widths for the state-conditioned + target head. The default `Int[]` preserves the original single Dense head. + +# Returns +- `StateConditionedPolicy`: ready-to-use policy with initialized recurrent state. + +# Examples +```julia +policy = state_conditioned_policy(5, 3, 3, [16, 16]) +x = randn(Float32, 8) # [uncertainty(5); state(3)] +y = policy(x) # predicted next state (length 3) + +deep_head = state_conditioned_policy( + 5, 3, 3, [16, 16]; + activation = sigmoid, + combiner_layers = [32, 32], +) +``` """ function state_conditioned_policy( n_uncertainty::Int, @@ -244,28 +628,38 @@ function state_conditioned_policy( layers::Vector{Int}; activation=Flux.relu, encoder_type=Flux.LSTM, + combiner_layers=Int[], ) - # Build encoder (stack of recurrent cells that process uncertainty) + # Build encoder: a stack of recurrent cells that process only the uncertainty + # input w_t. The number of hidden layers determines the encoder topology. if length(layers) == 0 + # No hidden layers: single cell maps uncertainty directly to n_state. encoder = _as_cell(encoder_type(n_uncertainty => n_state)) encoder_output_dim = n_state elseif length(layers) == 1 + # One hidden layer: single cell with the specified width. encoder = _as_cell(encoder_type(n_uncertainty => layers[1])) encoder_output_dim = layers[1] else + # Multiple hidden layers: chain of cells, first maps from n_uncertainty. encoder_layers = [_as_cell(encoder_type(n_uncertainty => layers[1]))] for i in 1:(length(layers) - 1) + # Each subsequent cell maps from the previous layer's width. push!(encoder_layers, _as_cell(encoder_type(layers[i] => layers[i + 1]))) end encoder = Chain(encoder_layers...) encoder_output_dim = layers[end] end - # Build combiner (Dense that combines encoder output with previous state) - # Input: [encoded_uncertainty, previous_state] - # Output: next_state - combiner = Dense(encoder_output_dim + n_state => n_output, activation) + # Build combiner: feed-forward [h_t; x_{t-1}] → x_hat_t with activation σ. + combiner = dense_policy_head( + encoder_output_dim + n_state, + n_output, + collect(Int, combiner_layers); + activation=activation, + ) + # Initialize the recurrent state to Flux.initialstates for the chosen cell(s). return StateConditionedPolicy( encoder, combiner, _init_recurrent_state(encoder), n_uncertainty, n_state ) diff --git a/src/multiple_shooting.jl b/src/multiple_shooting.jl index 322d174..5173c54 100644 --- a/src/multiple_shooting.jl +++ b/src/multiple_shooting.jl @@ -33,6 +33,16 @@ Assumptions: using Base: accumulate using Zygote: Zygote +""" + _print_window_status_and_params(window, status; context="") + +Print a diagnostic dump for a window solve that did not reach an optimal status. + +Outputs: a primal feasibility report for `window.model`, the solver `status`, +the `window.stage_range`, and a sorted listing of every JuMP parameter in the +window model with its current value. The optional `context` string is included +in the header for traceability. +""" function _print_window_status_and_params(window, status; context::AbstractString="") header = isempty(context) ? "solve_window status" : "solve_window status ($context)" println("=====") @@ -92,6 +102,13 @@ function extract_uncertainty_params(window_uncertainties) end end +""" + _param_init_value(src::JuMP.VariableRef) -> Float64 + +Return the current parameter value of `src` if it is a JuMP `MOI.Parameter`; +otherwise return `0.0`. Silently returns `0.0` if `parameter_value` throws +(e.g., when the parameter has been deleted from the model). +""" function _param_init_value(src::JuMP.VariableRef) if JuMP.is_parameter(src) try @@ -103,6 +120,13 @@ function _param_init_value(src::JuMP.VariableRef) return 0.0 end +""" + _as_float64_vec(val) -> Vector{Float64} + +Coerce `val` to a `Vector{Float64}`. If `val` is already a `Vector{Float64}`, +return it as-is; if it is another `AbstractVector`, convert element-wise; if it +is a scalar, wrap it in a single-element vector. +""" function _as_float64_vec(val) if val isa AbstractVector return val isa Vector{Float64} ? val : Float64.(val) @@ -110,6 +134,16 @@ function _as_float64_vec(val) return [Float64(val)] end +""" + _create_like_variable(m::JuMP.Model, src::JuMP.VariableRef, t::Int; + force_parameter=false) -> JuMP.VariableRef + +Create a new variable in `m` that mirrors `src`. If `src` is a JuMP parameter +(or `force_parameter` is `true`), the new variable is created as an +`MOI.Parameter` initialized to [`_param_init_value`](@ref)`(src)`; otherwise a +free decision variable is created. The variable is named via +[`var_set_name!`](@ref) with stage suffix `t`. +""" function _create_like_variable( m::JuMP.Model, src::JuMP.VariableRef, t::Int; force_parameter::Bool=false ) @@ -368,6 +402,15 @@ function solve_window( ) end +""" + _set_window_parameters!(window_state_in_params, window_state_out_params, + s_in, targets) -> Nothing + +Write numeric initial-state and target values into the JuMP `MOI.Parameter` +variables of a window model before solving. `s_in` is broadcast into +`window_state_in_params`; each element of `targets` is broadcast into the +corresponding stage's target parameters in `window_state_out_params`. +""" function _set_window_parameters!( window_state_in_params, window_state_out_params, diff --git a/src/parameter_duals.jl b/src/parameter_duals.jl index 581f8d5..21b6c58 100644 --- a/src/parameter_duals.jl +++ b/src/parameter_duals.jl @@ -33,8 +33,19 @@ model = Model(HiGHS.Optimizer) @constraint(model, con, x >= 2 * p) @objective(model, Min, 3 * x + p) optimize!(model) -dual_p = compute_parameter_dual(model, p) # Should be -2 * dual(con) + 1 +dual_p = compute_parameter_dual(model, p) # = 2 * dual(con) + 1 (constraint x - 2p >= 0 has coef -2 on p, contribution -coef*dual) +# Hand-check: at the optimum x* = 2p the objective is 7p, so d(obj)/dp = 7 = 2*3 + 1 with dual(con) = 3. ``` + +# Limitations +Only affine, quadratic, and vector-affine constraint functions are inspected: if the +parameter appears in any other constraint function type (e.g. a nonlinear +`ScalarNonlinearFunction` constraint), that contribution is skipped with a one-time +warning and is NOT captured in the returned sensitivity. + +The constraint-dual formula is validated for `MIN_SENSE` objectives; for `MAX_SENSE` +models JuMP's dual sign conventions differ and this function's constraint contribution +has not been validated. """ function compute_parameter_dual(model::JuMP.Model, param::JuMP.VariableRef) if !JuMP.is_parameter(param) @@ -73,6 +84,20 @@ function _get_dual_from_constraints(model::JuMP.Model, param::JuMP.VariableRef) dual_contribution += _get_dual_from_vector_affine_constraints( model, param, F, S ) + # Variable-in-set constraints (variable bounds, MOI.Parameter definitions) + # carry no parameter coefficient to extract: the parameter's own definition + # constraint contributes nothing here, and a parameter cannot appear inside + # another variable's bound constraint. Skip them silently. + elseif F <: JuMP.VariableRef + # Intentionally no contribution. + else + # Any other constraint function type (e.g. ScalarNonlinearFunction) is not + # inspected, so if the parameter appears there its sensitivity contribution + # is silently zero. Warn once so the gradient gap is visible, but do not + # error: constraints of these types often do not involve parameters at all. + @warn "compute_parameter_dual: constraints of type ($F, $S) are not " * + "inspected; parameter sensitivities from such constraints are not " * + "captured." maxlog = 1 end end @@ -94,15 +119,11 @@ function _get_dual_from_affine_constraints(model::JuMP.Model, param::JuMP.Variab # Check if parameter appears in this constraint coef = _get_parameter_coefficient(func, param) if !iszero(coef) - try - con_dual = JuMP.dual(con) - # The dual contribution is -coefficient * constraint_dual - # This follows from the Lagrangian: L = f(x) + λ*(g(x) - p*coef - b) - # ∂L/∂p = -λ * coef - dual_contribution -= coef * con_dual - catch - # If dual is not available, skip this constraint - end + con_dual = JuMP.dual(con) + # The dual contribution is -coefficient * constraint_dual. + # This follows from the Lagrangian: L = f(x) + λ*(g(x) - p*coef - b) + # ∂L/∂p = -λ * coef. + dual_contribution -= coef * con_dual end end @@ -132,12 +153,8 @@ function _get_dual_from_quadratic_constraints(model::JuMP.Model, param::JuMP.Var total_coef = coef + quad_coef if !iszero(total_coef) - try - con_dual = JuMP.dual(con) - dual_contribution -= total_coef * con_dual - catch - # If dual is not available, skip this constraint - end + con_dual = JuMP.dual(con) + dual_contribution -= total_coef * con_dual end end @@ -158,17 +175,12 @@ function _get_dual_from_vector_affine_constraints( con_obj = JuMP.constraint_object(con) func = con_obj.func # Vector of AffExpr - try + coefs = [_get_parameter_coefficient(expr, param) for expr in func] + if any(!iszero, coefs) con_dual = JuMP.dual(con) # Vector of duals - - for (i, expr) in enumerate(func) - coef = _get_parameter_coefficient(expr, param) - if !iszero(coef) - dual_contribution -= coef * con_dual[i] - end + for (coef, dual_entry) in zip(coefs, con_dual) + dual_contribution -= coef * dual_entry end - catch - # If dual is not available, skip this constraint end end diff --git a/src/score_function.jl b/src/score_function.jl index 120d392..b396fc0 100644 --- a/src/score_function.jl +++ b/src/score_function.jl @@ -58,6 +58,10 @@ and the mixed gradient is - `perturbation_std::Real`: Gaussian standard deviation ``\\sigma``. - `num_rollouts::Integer`: number of perturbed rollouts ``M`` per sample. - `baseline::Symbol`: either `:mean` for mean-centering costs or `:none`. + Mean-centering reduces variance but, because the baseline is computed from the + same ``M`` rollouts, it introduces a small ``O(1/M)`` bias (effectively scaling + the estimator by ``(M-1)/M``) that vanishes as `num_rollouts` grows; a + leave-one-out baseline would be exactly unbiased. # Examples ```julia @@ -509,7 +513,10 @@ function _center_rollout_costs( costs::AbstractVector{<:Real}, baseline::Symbol, ) - # A mean baseline reduces variance without changing the expected gradient. + # A mean baseline computed from the SAME rollouts reduces variance but adds a + # small O(1/M) bias (the estimator is effectively scaled by (M-1)/M); the bias + # vanishes as the rollout count grows. A leave-one-out baseline would be + # exactly unbiased; mean-centering is kept for its simplicity. baseline_value = baseline === :mean ? mean(costs) : 0.0 return Float64.(costs) .- baseline_value diff --git a/src/simulate_multistage.jl b/src/simulate_multistage.jl index a5116e7..2703b07 100644 --- a/src/simulate_multistage.jl +++ b/src/simulate_multistage.jl @@ -72,6 +72,31 @@ function simulate_stage( ) end +""" + _set_stage_parameters!(state_param_in, state_param_out, uncertainty, + state_in, state_out_target) -> Nothing + +Write MOI parameter values into a single-stage JuMP subproblem before solving. + +Three groups of parameters are set: + +1. **Incoming state** ``x_{t-1}``: each element of `state_param_in` receives + the corresponding entry of `state_in`. +2. **Uncertainty** ``w_t``: each `(parameter, value)` pair in `uncertainty` + is written directly. +3. **Outgoing target** ``\\hat{x}_t``: the first element of each tuple in + `state_param_out` (the target parameter) receives the corresponding + entry of `state_out_target`. + +After this call the subproblem is ready for `optimize!`. + +# Arguments +- `state_param_in`: JuMP parameter variables for the incoming state. +- `state_param_out`: `(target_parameter, realized_state_variable)` pairs. +- `uncertainty`: `(parameter, value)` pairs for stage uncertainty ``w_t``. +- `state_in::AbstractVector{<:Real}`: realized incoming state ``x_{t-1}``. +- `state_out_target::AbstractVector{<:Real}`: policy target ``\\hat{x}_t``. +""" function _set_stage_parameters!( state_param_in, state_param_out, @@ -79,17 +104,17 @@ function _set_stage_parameters!( state_in, state_out_target, ) - # Update state parameters + # Write the realized incoming state into the input-state parameters. for (i, state_var) in enumerate(state_param_in) set_parameter_value(state_var, state_in[i]) end - # Update uncertainty + # Write sampled exogenous values into the uncertainty parameters. for (uncertainty_param, uncertainty_value) in uncertainty set_parameter_value(uncertainty_param, uncertainty_value) end - # Update state parameters out + # Write policy targets into the output-state target parameters. for i in 1:length(state_param_out) state_var = state_param_out[i][1] set_parameter_value(state_var, state_out_target[i]) @@ -97,6 +122,27 @@ function _set_stage_parameters!( return nothing end +""" + _simulate_stage(subproblem, state_param_in, state_param_out, uncertainty, + state_in, state_out_target, integer_strategy) -> Float64 + +Forward-solve one stage of the multistage problem and return the objective. + +Sets all parameters via [`_set_stage_parameters!`](@ref), then solves +`subproblem` through [`with_sensitivity_solution`](@ref) (which applies the +`integer_strategy` for models with discrete variables). Returns the scalar +optimal objective value ``q_t(x_{t-1}, w_t; \\hat{x}_t)``. + +# Arguments +- `subproblem::JuMP.Model`: stage-``t`` JuMP model. +- `state_param_in`: incoming-state parameters. +- `state_param_out`: `(target_parameter, realized_state_variable)` pairs. +- `uncertainty`: `(parameter, value)` pairs for ``w_t``. +- `state_in`: realized incoming state ``x_{t-1}``. +- `state_out_target`: policy target ``\\hat{x}_t``. +- `integer_strategy::AbstractIntegerStrategy`: controls how discrete + variables are handled during the solve. +""" function _simulate_stage( subproblem::JuMP.Model, state_param_in, @@ -106,15 +152,57 @@ function _simulate_stage( state_out_target, integer_strategy::AbstractIntegerStrategy, ) + # Write all parameter values into the model before solving. _set_stage_parameters!( state_param_in, state_param_out, uncertainty, state_in, state_out_target ) + # Solve and extract the objective value inside the sensitivity wrapper. return with_sensitivity_solution(subproblem, integer_strategy) do sensitivity_model + # Cache the deficit-free objective while the model is clean. Integer + # strategies dirty the model on cleanup (restoring integer bounds), so + # a later logger call would otherwise find a dirty model with no cache + # and throw. Mirrors the deterministic-equivalent forward pass. + subproblem.ext[:_last_obj_no_deficit] = + get_objective_no_target_deficit(sensitivity_model) return objective_value(sensitivity_model) end end +""" + _simulate_stage_with_parameter_duals(subproblem, state_param_in, state_param_out, + uncertainty, state_in, state_out_target, + integer_strategy) + -> (objective, d_state_in, d_state_out_target) + +Forward-solve one stage and extract parameter duals for the rrule pullback. + +Like [`_simulate_stage`](@ref), this sets parameters and solves the stage +problem. In addition it reads the dual sensitivities via [`pdual`](@ref): + +```math +\\mu_t = \\frac{\\partial q_t}{\\partial x_{t-1}}, \\qquad +\\lambda_t = \\frac{\\partial q_t}{\\partial \\hat{x}_t}. +``` + +These duals are the preferred (closed-form) gradient path used by the +[`simulate_stage`](@ref) rrule when parameter duals are available from the +solver. + +# Arguments +- `subproblem`: stage-``t`` JuMP model. +- `state_param_in`: incoming-state parameters (yields ``\\mu_t``). +- `state_param_out`: target parameters (yields ``\\lambda_t``). +- `uncertainty`: `(parameter, value)` pairs for ``w_t``. +- `state_in`: realized incoming state ``x_{t-1}``. +- `state_out_target`: policy target ``\\hat{x}_t``. +- `integer_strategy::AbstractIntegerStrategy`: discrete-variable strategy. + +# Returns +- `objective::Float64`: optimal stage cost ``q_t``. +- `d_state_in::Vector{Float64}`: ``\\mu_t`` (sensitivities w.r.t. incoming state). +- `d_state_out_target::Vector{Float64}`: ``\\lambda_t`` (sensitivities w.r.t. target). +""" function _simulate_stage_with_parameter_duals( subproblem, state_param_in, @@ -124,12 +212,20 @@ function _simulate_stage_with_parameter_duals( state_out_target, integer_strategy::AbstractIntegerStrategy, ) + # Write all parameter values into the model before solving. _set_stage_parameters!( state_param_in, state_param_out, uncertainty, state_in, state_out_target ) return with_sensitivity_solution(subproblem, integer_strategy) do sensitivity_model + # Read the optimal objective value. objective = objective_value(sensitivity_model) + # Cache the deficit-free objective while the model is clean (integer + # strategies dirty the model on cleanup; see _simulate_stage). + subproblem.ext[:_last_obj_no_deficit] = + get_objective_no_target_deficit(sensitivity_model) + # Extract duals w.r.t. incoming state parameters (mu_t). d_state_in = pdual.(state_param_in) + # Extract duals w.r.t. target parameters (lambda_t). d_state_out_target = pdual.([s[1] for s in state_param_out]) return objective, d_state_in, d_state_out_target end @@ -336,29 +432,132 @@ function ChainRulesCore.rrule( return y, public_pullback end +""" + get_objective_no_target_deficit(subproblem::JuMP.Model; + norm_deficit="norm_deficit") -> Float64 + +Compute the operational cost of a solved subproblem, excluding the +target-deficit penalty. + +The full objective includes a penalty ``C_\\delta \\|\\delta_t\\|`` that +penalizes deviations between realized and target states. This function +strips those terms so that logged costs reflect true operational cost: + +```math +\\text{cost}_t = q_t - \\sum_{j \\in \\mathcal{D}} c_j \\, \\delta_j, +``` + +where ``\\mathcal{D}`` is the set of variables whose names contain +`norm_deficit`. + +If the model is dirty (parameters changed since last solve), returns the +cached value from a previous successful call (stored in +`subproblem.ext[:_last_obj_no_deficit]`). If no such value exists, or if the +objective shape is unsupported, throws instead of inventing a cost. + +# Arguments +- `subproblem::JuMP.Model`: a solved JuMP model. + +# Keywords +- `norm_deficit::AbstractString`: substring matched against variable names + to identify deficit-penalty terms. +""" function get_objective_no_target_deficit( subproblem::JuMP.Model; norm_deficit::AbstractString="norm_deficit" ) + # If parameters were changed after the last solve, return the cached value. if subproblem.is_model_dirty - return get(subproblem.ext, :_last_obj_no_deficit, 0.0) + if haskey(subproblem.ext, :_last_obj_no_deficit) + return subproblem.ext[:_last_obj_no_deficit] + end + error( + "Cannot read objective without target deficit: " * + "model is dirty and no cached value exists", + ) end - try - obj = JuMP.objective_function(subproblem) - objective_val = objective_value(subproblem) - for term in obj.terms - if occursin(norm_deficit, JuMP.name(term[1])) - objective_val -= term[2] * value(term[1]) - end + + obj = JuMP.objective_function(subproblem) + objective_val = + objective_value(subproblem) - _target_deficit_penalty_value(obj, norm_deficit) + subproblem.ext[:_last_obj_no_deficit] = objective_val + return objective_val +end + +""" + _target_deficit_penalty_value(obj, norm_deficit) -> Float64 + +Return the part of a JuMP objective expression that is attributed to target +deficit variables. A variable is treated as a target-deficit variable when its +JuMP name contains `norm_deficit`. + +Proof sketch for the affine case: if the solved objective is +`q(x) + sum_i c_i d_i`, and `d_i` are exactly the matched target-deficit +variables, then the operational cost is the solved objective value minus +`sum_i c_i value(d_i)`. + +Quadratic terms involving target-deficit variables are deliberately rejected. +For such an objective, subtracting only the affine coefficient would not remove +the whole penalty and would produce a silently biased operational cost. +""" +function _target_deficit_penalty_value( + obj::JuMP.GenericAffExpr, norm_deficit::AbstractString +) + penalty = 0.0 + for (variable, coefficient) in obj.terms + if occursin(norm_deficit, JuMP.name(variable)) + penalty += coefficient * JuMP.value(variable) end - return objective_val - catch - return get(subproblem.ext, :_last_obj_no_deficit, 0.0) end + return penalty +end + +# Quadratic objectives are allowed only when target-deficit variables appear in +# the affine part; otherwise the penalty shape is ambiguous to this helper. +function _target_deficit_penalty_value( + obj::JuMP.GenericQuadExpr, norm_deficit::AbstractString +) + for (pair, _) in obj.terms + if occursin(norm_deficit, JuMP.name(pair.a)) || + occursin(norm_deficit, JuMP.name(pair.b)) + error("Quadratic target-deficit penalty terms are unsupported") + end + end + return _target_deficit_penalty_value(obj.aff, norm_deficit) +end + +# A bare deficit variable has implicit coefficient one. +function _target_deficit_penalty_value( + obj::JuMP.VariableRef, norm_deficit::AbstractString +) + return occursin(norm_deficit, JuMP.name(obj)) ? JuMP.value(obj) : 0.0 end +_target_deficit_penalty_value(::Real, ::AbstractString) = 0.0 + +function _target_deficit_penalty_value(obj, ::AbstractString) + error("Unsupported objective type for target-deficit stripping: $(typeof(obj))") +end + +""" + get_objective_no_target_deficit(subproblems::Vector{JuMP.Model}; + norm_deficit="norm_deficit") -> Float64 + +Sum the deficit-free operational costs across all stage subproblems. + +Calls the single-model [`get_objective_no_target_deficit`](@ref) on each +element and returns the total. + +# Arguments +- `subproblems::Vector{JuMP.Model}`: one solved JuMP model per stage. + +# Keywords +- `norm_deficit::AbstractString`: substring matched against variable names + to identify deficit-penalty terms. +""" function get_objective_no_target_deficit( subproblems::Vector{JuMP.Model}; norm_deficit::AbstractString="norm_deficit" ) + # Accumulate deficit-free costs across all stages. total_objective = 0.0 for subproblem in subproblems total_objective += get_objective_no_target_deficit( @@ -368,7 +567,11 @@ function get_objective_no_target_deficit( return total_objective end -# define ChainRulesCore.rrule of get_objective_no_target_deficit +# NOTE: get_objective_no_target_deficit is intentionally NON-DIFFERENTIABLE. +# This rrule returns NoTangent() for every input, i.e. a hard-zero gradient. The +# value is a logging/metric quantity only (deficit-free operational cost read from +# an already-solved model), and it MUST NOT appear in a loss whose gradient matters: +# any dependence of the loss on the policy through this function is silently dropped. function ChainRulesCore.rrule( ::typeof(get_objective_no_target_deficit), subproblem; norm_deficit="norm_deficit" ) @@ -379,10 +582,40 @@ function ChainRulesCore.rrule( return objective_val, _pullback end +""" + apply_rule(stage::Int, decision_rule, uncertainty, state_in) -> Vector + +Apply a single (shared) policy to produce the target state for stage `stage`. + +The policy receives a concatenated input vector +``[w_t^{(1)}, \\ldots, w_t^{(n_w)}, x_{t-1}^{(1)}, \\ldots, x_{t-1}^{(n_x)}]`` +and returns the next target ``\\hat{x}_t = \\pi_\\theta(w_t, x_{t-1})``. + +# Arguments +- `stage::Int`: current stage index (unused when a single rule is shared). +- `decision_rule`: callable policy ``\\pi_\\theta``. +- `uncertainty`: `(parameter, value)` pairs for ``w_t``; values are extracted. +- `state_in`: realized incoming state ``x_{t-1}``. +""" function apply_rule(::Int, decision_rule::T, uncertainty, state_in) where {T} + # Concatenate uncertainty values and incoming state into the policy input. return decision_rule(vcat([uncertainty[i][2] for i in 1:length(uncertainty)], state_in)) end +""" + apply_rule(stage::Int, decision_rules::Vector, uncertainty, state_in) -> Vector + +Apply a stage-specific policy from a vector of per-stage decision rules. + +Dispatches to `apply_rule(stage, decision_rules[stage], uncertainty, state_in)`, +selecting the rule at index `stage`. + +# Arguments +- `stage::Int`: current stage index, used to select `decision_rules[stage]`. +- `decision_rules::Vector`: one callable policy per stage. +- `uncertainty`: `(parameter, value)` pairs for ``w_t``. +- `state_in`: realized incoming state ``x_{t-1}``. +""" function apply_rule(stage::Int, decision_rules::Vector{T}, uncertainty, state_in) where {T} return apply_rule(stage, decision_rules[stage], uncertainty, state_in) end @@ -467,6 +700,32 @@ function simulate_multistage( ) end +""" + _set_multistage_parameters!(state_params_in, state_params_out, + uncertainties, states) -> Nothing + +Write MOI parameter values into a deterministic-equivalent JuMP model +across all ``T`` stages before solving. + +For each stage ``t = 1, \\ldots, T``: + +- **Initial state** (``t = 1`` only): `state_params_in[1]` receives + `states[1]` (the initial state ``x_0``). +- **Uncertainty** ``w_t``: each `(parameter, value)` pair in + `uncertainties[t]` is written. +- **Target** ``\\hat{x}_t``: the target parameters in + `state_params_out[t]` receive `states[t + 1]`. + +Note that `state_params_in[t]` for ``t > 1`` is NOT set here because in the +deterministic equivalent the incoming state is an internal variable linked +by constraints, not a parameter. + +# Arguments +- `state_params_in`: per-stage vectors of incoming-state parameters. +- `state_params_out`: per-stage vectors of `(target_parameter, state_variable)`. +- `uncertainties`: per-stage `(parameter, value)` pairs for ``w_t``. +- `states`: length-``(T+1)`` target trajectory ``[x_0, \\hat{x}_1, \\ldots, \\hat{x}_T]``. +""" function _set_multistage_parameters!( state_params_in, state_params_out, @@ -475,19 +734,20 @@ function _set_multistage_parameters!( ) for t in 1:length(state_params_in) state = states[t] - # Update state parameters in + # Only the initial state (t=1) is set as a parameter; later incoming + # states are internal variables in the deterministic equivalent. if t == 1 for (i, state_var) in enumerate(state_params_in[t]) set_parameter_value(state_var, state[i]) end end - # Update uncertainty + # Write sampled exogenous values into this stage's uncertainty parameters. for (uncertainty_param, uncertainty_value) in uncertainties[t] set_parameter_value(uncertainty_param, uncertainty_value) end - # Update state parameters out + # Write policy targets into this stage's output-state target parameters. for i in 1:length(state_params_out[t]) state_var = state_params_out[t][i][1] set_parameter_value(state_var, states[t + 1][i]) @@ -496,6 +756,26 @@ function _set_multistage_parameters!( return nothing end +""" + _simulate_multistage_det(det_equivalent, state_params_in, state_params_out, + uncertainties, states, integer_strategy) -> Float64 + +Solve the deterministic-equivalent model for a single uncertainty trajectory +and return the optimal objective. + +Sets all parameters via [`_set_multistage_parameters!`](@ref), solves the +coupled full-horizon problem ``Q(w; \\theta)`` through +[`with_sensitivity_solution`](@ref), and caches both the full objective and +the deficit-free operational cost in `det_equivalent.ext` for logging. + +# Arguments +- `det_equivalent::JuMP.Model`: full-horizon coupled JuMP model. +- `state_params_in`: per-stage incoming-state parameters. +- `state_params_out`: per-stage `(target_parameter, state_variable)` pairs. +- `uncertainties`: per-stage `(parameter, value)` pairs for ``w_t``. +- `states`: length-``(T+1)`` target trajectory from the policy. +- `integer_strategy::AbstractIntegerStrategy`: discrete-variable strategy. +""" function _simulate_multistage_det( det_equivalent::JuMP.Model, state_params_in, @@ -504,10 +784,13 @@ function _simulate_multistage_det( states, integer_strategy::AbstractIntegerStrategy, ) + # Write the full target trajectory and uncertainty into the DE model. _set_multistage_parameters!(state_params_in, state_params_out, uncertainties, states) return with_sensitivity_solution(det_equivalent, integer_strategy) do sensitivity_model + # Read the optimal objective value after solving. obj = objective_value(sensitivity_model) + # Cache both the full and deficit-free objectives for logging. sensitivity_model.ext[:_last_obj] = obj sensitivity_model.ext[:_last_obj_no_deficit] = get_objective_no_target_deficit(sensitivity_model) @@ -515,6 +798,42 @@ function _simulate_multistage_det( end end +""" + _simulate_multistage_det_with_parameter_duals(det_equivalent, state_params_in, + state_params_out, uncertainties, + states, integer_strategy) + -> (objective, Δ_states) + +Solve the deterministic equivalent and extract parameter duals for the +rrule pullback. + +Like [`_simulate_multistage_det`](@ref), this sets parameters and solves the +full-horizon problem. In addition it reads the dual sensitivities via +[`pdual`](@ref) for each stage: + +```math +\\Delta_{\\text{states}}[1] = \\frac{\\partial Q}{\\partial x_0}, \\qquad +\\Delta_{\\text{states}}[t+1] = \\lambda_t = \\frac{\\partial Q}{\\partial \\hat{x}_t}, +\\quad t = 1, \\ldots, T. +``` + +These ``\\lambda_t`` duals are the envelope-theorem gradient used in the +TS-DDR training objective (arXiv:2405.14973, Eq. 1.2). + +# Arguments +- `det_equivalent`: full-horizon coupled JuMP model. +- `state_params_in`: per-stage incoming-state parameters. +- `state_params_out`: per-stage `(target_parameter, state_variable)` pairs. +- `uncertainties`: per-stage `(parameter, value)` pairs for ``w_t``. +- `states`: length-``(T+1)`` target trajectory from the policy. +- `integer_strategy::AbstractIntegerStrategy`: discrete-variable strategy. + +# Returns +- `objective::Float64`: optimal coupled objective ``Q(w; \\theta)``. +- `Δ_states::Vector{Vector{Float64}}`: length-``(T+1)`` vector of parameter + duals; `Δ_states[1]` holds ``\\partial Q / \\partial x_0`` and + `Δ_states[t+1]` holds ``\\lambda_t``. +""" function _simulate_multistage_det_with_parameter_duals( det_equivalent, state_params_in, @@ -523,13 +842,18 @@ function _simulate_multistage_det_with_parameter_duals( states, integer_strategy::AbstractIntegerStrategy, ) + # Write the full target trajectory and uncertainty into the DE model. _set_multistage_parameters!(state_params_in, state_params_out, uncertainties, states) return with_sensitivity_solution(det_equivalent, integer_strategy) do sensitivity_model + # Read the optimal objective value after solving. objective = objective_value(sensitivity_model) + # Cache both the full and deficit-free objectives for logging. sensitivity_model.ext[:_last_obj] = objective sensitivity_model.ext[:_last_obj_no_deficit] = get_objective_no_target_deficit(sensitivity_model) + # Build the per-stage dual vector: initial-state duals at index 1, + # target-constraint duals lambda_t at index t+1. Δ_states = similar(states) Δ_states[1] = pdual.(state_params_in[1]) for t in 1:length(state_params_out) @@ -641,7 +965,9 @@ function ChainRulesCore.rrule( integer_strategy, ) pdual_available = true - catch + catch err + # Surface the reason the fast path was skipped without altering the fallback. + @debug "pdual path failed; falling back to DiffOpt reverse differentiation" exception = (err, catch_backtrace()) y = _simulate_stage( subproblem, state_param_in, @@ -791,7 +1117,9 @@ function ChainRulesCore.rrule( integer_strategy, ) pdual_available = true - catch + catch err + # Surface the reason the fast path was skipped without altering the fallback. + @debug "pdual path failed; falling back to DiffOpt reverse differentiation" exception = (err, catch_backtrace()) y = _simulate_multistage_det( det_equivalent, state_params_in, @@ -1084,30 +1412,52 @@ Q(\theta; w) = \sum_{t=1}^{T} q_t(x_{t-1}, w_t; \hat{x}_t), ``` -where each realized ``x_t`` is read from the previous stage solve. The gradient -therefore contains both the target duals ``\lambda_t`` and the sensitivity of -later realized states with respect to earlier targets. In the notation of the -extension note, +where each realized ``x_t`` is read from the previous stage solve. In this +single-shooting rollout ``\theta`` influences later stage costs through two +couplings: the realized-state chain (the solver map +``x_t = X_t(x_{t-1}, \hat{x}_t; w_t)`` propagates earlier targets forward), +and the **policy-feedback path** (the policy input at stage ``t`` is the +realized state ``x_{t-1}``, which itself depends on earlier targets). The exact +gradient is the total-derivative recursion ```math -\nabla_\theta Q(\theta; w) +\frac{d Q}{d \theta} = \sum_{t=1}^{T} \left[ - \frac{\partial q_t}{\partial \hat{x}_t} + \frac{\partial q_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d \theta} + - \sum_{k=t+1}^{T} - \frac{\partial q_k}{\partial x_{k-1}} - \prod_{j=t+1}^{k-1} - \frac{\partial x_j}{\partial x_{j-1}} - \frac{\partial x_t}{\partial \hat{x}_t} -\right] -\nabla_\theta \pi_\theta(w_t, x_{t-1}). + \frac{\partial q_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d \theta} +\right], +``` + +with the coupled state and target recursions (and ``d x_0 / d\theta = 0``) + +```math +\frac{d x_t}{d \theta} += +\frac{\partial X_t}{\partial x_{t-1}} \frac{d x_{t-1}}{d \theta} ++ +\frac{\partial X_t}{\partial \hat{x}_t} \frac{d \hat{x}_t}{d \theta}, +\qquad +\frac{d \hat{x}_t}{d \theta} += +\nabla_\theta \pi_\theta(w_t, x_{t-1}) ++ +\frac{\partial \pi_\theta}{\partial x_{t-1}} \frac{d x_{t-1}}{d \theta}. ``` -The dual terms come from target and transition constraints; the state -sensitivities are computed through DiffOpt in the rrules for -[`simulate_stage`](@ref) and [`get_next_state`](@ref). +All derivative products must be read as **total** derivatives that include the +policy-feedback composition: a term such as +``\partial \pi_\theta / \partial x_{k-1} \cdot d x_{k-1} / d\theta`` is present +at every stage, so ``\theta`` reaches stage ``k`` not only through the +realized-state chain ``\partial x_j / \partial x_{j-1}`` but also through the +policy input ``x_{k-1}``. Reverse-mode AD (Zygote through the stage rrules) +computes exactly this full chain: the dual terms from target and transition +constraints supply ``\partial q_t / \partial \hat{x}_t`` and +``\partial q_t / \partial x_{t-1}``, while the state sensitivities are computed +through DiffOpt in the rrules for [`simulate_stage`](@ref) and +[`get_next_state`](@ref). # Arguments - `model`: differentiable Flux-compatible policy. It receives @@ -1231,9 +1581,11 @@ function train_multistage( return objective end catch e - if handle_training_error(gradient_fallback, e, iter) - nothing - end + # handle_training_error rethrows for ErrorGradientFallback and logs a + # warning + returns true (skip this iteration) for ZeroGradientFallback; + # either way no gradient is available, so the try-expression is nothing. + handle_training_error(gradient_fallback, e, iter) + nothing end record(sample_log, iter, model) && break @@ -1248,19 +1600,6 @@ function train_multistage( return model end -function sim_states(t, m, initial_state, uncertainty_sample_vec, prev_states) - # Input: [uncertainty, previous_predicted_state] - # For t=1: return initial_state (no prediction needed) - # For t>1: policy receives [uncertainty[t-1], prev_states[t-1]] - if t == 1 - return Float32.(initial_state) - else - uncertainties_t = uncertainty_sample_vec[t - 1] - prev_state = prev_states[t - 1] - return m(vcat(uncertainties_t, prev_state)) - end -end - @doc raw""" train_multistage(model, initial_state, det_equivalent::JuMP.Model, state_params_in, state_params_out, uncertainty_sampler; @@ -1487,9 +1826,11 @@ function train_multistage( return objective end catch e - if handle_training_error(gradient_fallback, e, iter) - nothing - end + # handle_training_error rethrows for ErrorGradientFallback and logs a + # warning + returns true (skip this iteration) for ZeroGradientFallback; + # either way no gradient is available, so the try-expression is nothing. + handle_training_error(gradient_fallback, e, iter) + nothing end record(sample_log, iter, model) && break diff --git a/src/utils.jl b/src/utils.jl index 6dc07e5..5d6ee1c 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -27,21 +27,48 @@ end Create deficit variables to penalize state deviations in a JuMP model. -Supports three modes: -- L1 norm only: Uses `MOI.NormOneCone` (default if no penalty specified) -- L2 squared norm only: Uses sum of squared deviations (solver-compatible alternative to SecondOrderCone) -- Both norms: Creates both constraints with separate penalties +Supports three modes controlled by the penalty keywords. Let +``d \\in \\mathbb{R}^n`` be the deficit vector (`len = n`). + +**L1 norm only** (default when no penalty keyword is given, or `penalty_l1` +alone): + +```math +\\text{norm\\_deficit} \\geq \\| d \\|_1 = \\sum_{i=1}^{n} |d_i|, +\\quad \\text{objective} \\mathrel{+}= \\lambda_1 \\cdot \\text{norm\\_deficit}. +``` + +Implemented via `MOI.NormOneCone(1 + n)`. + +**L2 squared norm only** (`penalty_l2` alone): + +```math +\\text{norm\\_deficit} \\geq \\| d \\|_2^2 = \\sum_{i=1}^{n} d_i^2, +\\quad \\text{objective} \\mathrel{+}= \\lambda_2 \\cdot \\text{norm\\_deficit}. +``` + +**Both norms** (`penalty_l1` and `penalty_l2`): + +```math +\\text{norm\\_deficit} + \\geq \\lambda_1 \\| d \\|_1 + \\lambda_2 \\| d \\|_2^2, +\\quad \\text{objective} \\mathrel{+}= 1 \\cdot \\text{norm\\_deficit}. +``` # Arguments -- `model`: The JuMP model to add deficit variables to -- `len`: Number of deficit variables (typically dimension of state) -- `penalty_l1`: Penalty coefficient for L1 norm (NormOneCone). If `nothing` and L1 is used, defaults to max objective coefficient. -- `penalty_l2`: Penalty coefficient for L2 squared norm (sum of squares). If `nothing` and L2 is used, defaults to max objective coefficient. -- `penalty`: Legacy argument. If provided and penalty_l1/penalty_l2 are both `nothing`, uses this for L1 norm only. +- `model`: The JuMP model to add deficit variables to. +- `len`: Number of deficit variables (typically dimension of state). +- `penalty_l1`: Penalty coefficient ``\\lambda_1`` for the L1 norm + (`NormOneCone`). Pass `:auto` to use `max |objective coefficients|`. +- `penalty_l2`: Penalty coefficient ``\\lambda_2`` for the L2 squared norm + (sum of squares). Pass `:auto` to use `max |objective coefficients|`. +- `penalty`: Legacy argument. If provided and `penalty_l1`/`penalty_l2` are + both `nothing`, uses this for L1 norm only. # Returns -- `norm_deficit`: Single variable representing total penalized deviation (for logging compatibility) -- `_deficit`: Vector of deficit variables for each state dimension +- `norm_deficit`: Single variable representing total penalized deviation (for + logging compatibility). +- `_deficit`: Vector of deficit variables for each state dimension. # Examples ```julia @@ -381,6 +408,14 @@ function normalize_recur_state(state) end end +""" + (callback::SaveBest)(iter, model, loss) -> Bool + +Compare `loss` against the incumbent `callback.best_loss`. When `loss` is +strictly smaller, copy `model` to CPU, normalize recurrent state via +[`normalize_recur_state`](@ref), and write the Flux state to +`callback.model_path` with JLD2. Always returns `false` (never stops training). +""" function (callback::SaveBest)(iter, model, loss) if loss < callback.best_loss m = cpu(model) @@ -392,11 +427,36 @@ function (callback::SaveBest)(iter, model, loss) return false end +""" + StallingCriterium(patience::Int, best_loss::Float64, stall_count::Int) + +Early-stopping callback that halts training when the loss stalls. + +Tracks the number of consecutive iterations without improvement. When +`stall_count` reaches `patience`, the callback returns `true` to signal that +training should stop. Use `best_loss = Inf` and `stall_count = 0` for a fresh +start. + +# Arguments +- `patience::Int`: maximum consecutive non-improving iterations before stopping. +- `best_loss::Float64`: incumbent best loss. Use `Inf` to accept the first value. +- `stall_count::Int`: current stall counter (typically initialized to `0`). +""" mutable struct StallingCriterium <: Function patience::Int best_loss::Float64 stall_count::Int end + +""" + (callback::StallingCriterium)(iter, model, loss) -> Bool + +Update the stall counter and return `true` when `stall_count >= patience`. + +If `loss < best_loss`, reset the counter to zero and update the incumbent. +Otherwise increment `stall_count`. Returns `true` (stop training) once the +patience budget is exhausted; `false` otherwise. +""" function (callback::StallingCriterium)(iter, model, loss) if loss < callback.best_loss callback.best_loss = loss @@ -448,9 +508,25 @@ function Base.empty!(sample_log::SampleLog) return sample_log end +""" + _reset_sample_log!(sample_log::SampleLog) -> SampleLog + _reset_sample_log!(sample_log) -> typeof(sample_log) + +Clear the per-batch cache of a [`SampleLog`](@ref) before the next training +batch. For a `SampleLog`, delegates to `Base.empty!`; for any other type the +call is a no-op and returns `sample_log` unchanged. +""" _reset_sample_log!(sample_log::SampleLog) = empty!(sample_log) _reset_sample_log!(sample_log) = sample_log +""" + _total_objective_value(model::JuMP.Model) -> Float64 + _total_objective_value(models::Vector{JuMP.Model}) -> Float64 + +Return the objective value of `model`, falling back to the cached value +`model.ext[:_last_obj]` (default `0.0`) when the model is dirty or +`objective_value` throws. The multi-model method sums across all models. +""" function _total_objective_value(model::JuMP.Model) if model.is_model_dirty return get(model.ext, :_last_obj, 0.0) @@ -476,6 +552,13 @@ function (sample_log::SampleLog)(s::Int, models) return nothing end +""" + _sequential_mean(values) -> Float64 + +Compute the arithmetic mean of `values` using sequential accumulation (left +fold). This preserves the same floating-point summation order as the historical +`loss += ...; loss /= n` pattern inside the training loops. +""" # Sequential accumulation keeps the same floating-point summation order as the # historical `loss += ...; loss /= n` pattern inside the training loops. function _sequential_mean(values) @@ -642,6 +725,19 @@ function RolloutEvaluation( ) end +""" + _simulate_multistage_target_feedback(subproblems, state_params_in, + state_params_out, initial_state, uncertainties, decision_rules, + integer_strategy) -> Float64 + +Run a stage-wise rollout where the **target** state (not the realized state) is +fed back into the policy at each stage, matching the deterministic-equivalent +target-generation semantics from [`simulate_states`](@ref). Each stage +subproblem is solved sequentially; the accumulated objective value (including +target-deficit penalties) is returned. + +Used by [`RolloutEvaluation`](@ref) when `policy_state == :target`. +""" function _simulate_multistage_target_feedback( subproblems::Vector{JuMP.Model}, state_params_in, @@ -684,6 +780,18 @@ function _simulate_multistage_target_feedback( return objective end +""" + (evaluation::RolloutEvaluation)(iter, model) -> Nothing + +Evaluate the policy on the held-out scenario set every `stride` iterations. + +On active iterations (`iter % stride == 0`), rolls `model` out over all fixed +scenarios using either closed-loop (`:realized`) or target-feedback (`:target`) +semantics. Prints `metrics/rollout_objective_no_deficit` and +`metrics/rollout_target_violation_share`, and caches the values in +`evaluation.last_objective_no_deficit` / `evaluation.last_violation_share`. +Scenarios that fail to solve are skipped with a warning if all fail. +""" function (evaluation::RolloutEvaluation)(iter, model) iter % evaluation.stride == 0 || return nothing total = 0.0 @@ -739,6 +847,13 @@ function (evaluation::RolloutEvaluation)(iter, model) return nothing end +""" + var_set_name!(src::JuMP.VariableRef, dest::JuMP.VariableRef, t::Int) -> Nothing + +Name `dest` after `src` with a `#t` stage suffix. If `src` has a JuMP name, the +result is `"#"`; otherwise the MOI variable index is used as fallback, +producing `"_[]#"`. +""" function var_set_name!(src::JuMP.VariableRef, dest::JuMP.VariableRef, t::Int) name = JuMP.name(src) if !isempty(name) @@ -751,6 +866,21 @@ function var_set_name!(src::JuMP.VariableRef, dest::JuMP.VariableRef, t::Int) end end +""" + add_child_model_vars!(model, subproblem, t, state_params_in, state_params_out, + initial_state, var_src_to_dest, skip_parameter_refs) + -> Dict{VariableRef,VariableRef} + +Copy decision variables from stage-`t` `subproblem` into the deterministic- +equivalent `model`, populating the source-to-destination mapping +`var_src_to_dest`. State-coupling variables (incoming parameters and outgoing +realized-state/target pairs) are handled specially: at `t == 1` they become +fresh parameters in `model`; at `t > 1` incoming state parameters are linked to +the previous stage's realized state variables. Each copied variable is renamed +via [`var_set_name!`](@ref) with a `#t` suffix. Mutates `state_params_in`, +`state_params_out`, `var_src_to_dest`, and `skip_parameter_refs` in place; +returns `var_src_to_dest`. +""" function add_child_model_vars!( model::JuMP.Model, subproblem::JuMP.Model, @@ -759,6 +889,7 @@ function add_child_model_vars!( state_params_out::Vector{Vector{Tuple{Any,VariableRef}}}, initial_state::Vector{Float64}, var_src_to_dest::Dict{VariableRef,VariableRef}, + skip_parameter_refs::Set{VariableRef}, ) allvars = all_variables(subproblem) allvars = setdiff(allvars, state_params_in[t]) @@ -799,23 +930,30 @@ function add_child_model_vars!( for (i, src) in enumerate(state_params_in[t]) if src isa VariableRef var_src_to_dest[src] = state_params_out[t - 1][i][2] + push!(skip_parameter_refs, src) end state_params_in[t][i] = state_params_out[t - 1][i][2] - # delete parameter constraint associated with src - if src isa VariableRef - for con in - JuMP.all_constraints(subproblem, VariableRef, MOI.Parameter{Float64}) - obj = JuMP.constraint_object(con) - if obj.func == src - JuMP.delete(subproblem, con) - end - end - end end end return var_src_to_dest end +""" + copy_and_replace_variables(src, map::Dict{VariableRef,VariableRef}) + +Deep-copy a JuMP expression `src`, substituting every `VariableRef` key in +`map` with its destination value. Dispatches on the concrete expression type: + +- `Vector`: element-wise recursive call. +- `Real`: returned as-is (no variables to replace). +- `VariableRef`: direct lookup in `map`. +- `GenericAffExpr`: rebuild with remapped variable keys and same coefficients. +- `GenericQuadExpr`: rebuild affine part and remapped `UnorderedPair` keys. +- `GenericNonlinearExpr`: recursively remap arguments, then reconstruct via + `@expression`. + +Throws an error for unrecognized expression types. +""" function copy_and_replace_variables( src::Vector, map::Dict{JuMP.VariableRef,JuMP.VariableRef} ) @@ -875,6 +1013,18 @@ function copy_and_replace_variables(src::Any, ::Dict{JuMP.VariableRef,JuMP.Varia ) end +""" + create_constraint(model, obj, var_src_to_dest) -> ConstraintRef + +Add a constraint to `model` whose function is `obj.func` with all `VariableRef` +keys replaced via `var_src_to_dest` (see [`copy_and_replace_variables`](@ref)). + +Four methods handle different constraint types: +- Generic `ScalarConstraint`: uses `@constraint(model, new_func in obj.set)`. +- `ScalarConstraint{NonlinearExpr, MOI.EqualTo}`: `new_func == obj.set.value`. +- `ScalarConstraint{NonlinearExpr, MOI.LessThan}`: `new_func <= obj.set.upper`. +- `ScalarConstraint{NonlinearExpr, MOI.GreaterThan}`: `new_func >= obj.set.lower`. +""" function create_constraint(model, obj, var_src_to_dest) new_func = copy_and_replace_variables(obj.func, var_src_to_dest) return @constraint(model, new_func in obj.set) @@ -901,10 +1051,26 @@ function create_constraint( return @constraint(model, new_func >= obj.set.lower) end +""" + add_child_model_exps!(model, subproblem, var_src_to_dest, skip_parameter_refs, + state_params_out, state_params_in, t) -> Dict + +Copy all constraints and the objective contribution from stage-`t` `subproblem` +into the deterministic-equivalent `model`, remapping variables through +`var_src_to_dest` via [`create_constraint`](@ref) and +[`copy_and_replace_variables`](@ref). Constraint-based state parameters and +input parameters at `t == 1` are updated to point to the new model's +constraint refs. Parameter constraints for incoming states that are linked to +the previous realized state are skipped via `skip_parameter_refs`; all other +parameter constraints are copied. The subproblem objective is added to +`model`'s existing objective. Returns a `Dict` mapping source `ConstraintRef` +to destination `ConstraintRef`. +""" function add_child_model_exps!( model::JuMP.Model, subproblem::JuMP.Model, var_src_to_dest::Dict{VariableRef,VariableRef}, + skip_parameter_refs::Set{VariableRef}, state_params_out, state_params_in, t, @@ -914,6 +1080,11 @@ function add_child_model_exps!( cons_to_cons = Dict() for con in JuMP.all_constraints(subproblem; include_variable_in_set_constraints=true) #, F, S) obj = JuMP.constraint_object(con) + if obj.func isa VariableRef && + obj.set isa MOI.Parameter && + obj.func in skip_parameter_refs + continue + end c = create_constraint(model, obj, var_src_to_dest) cons_to_cons[con] = c if (state_params_out[t][1][1] isa ConstraintRef) @@ -961,6 +1132,9 @@ stage `t` with the incoming state parameter of stage `t+1`. Returns `(model, uncertainties_new)` where `uncertainties_new` has the same format as the input but with variable refs remapped to the deterministic-equivalent model. +This function also remaps `state_params_in` and `state_params_out` in place. +Copy those arrays before calling if you need to keep the original stage-wise +references for rollout evaluation or later model construction. """ function deterministic_equivalent!( model::JuMP.Model, @@ -972,6 +1146,7 @@ function deterministic_equivalent!( ) set_objective_sense(model, objective_sense(subproblems[1])) var_src_to_dest = Dict{VariableRef,VariableRef}() + skip_parameter_refs = Set{VariableRef}() for t in 1:length(subproblems) DecisionRules.add_child_model_vars!( model, @@ -981,13 +1156,20 @@ function deterministic_equivalent!( state_params_out, initial_state, var_src_to_dest, + skip_parameter_refs, ) end cons_to_cons = Vector{Dict}(undef, length(subproblems)) for t in 1:length(subproblems) cons_to_cons[t] = DecisionRules.add_child_model_exps!( - model, subproblems[t], var_src_to_dest, state_params_out, state_params_in, t + model, + subproblems[t], + var_src_to_dest, + skip_parameter_refs, + state_params_out, + state_params_in, + t, ) end @@ -1039,6 +1221,14 @@ function _remap_uncertainties( ] end +""" + find_variables(model::JuMP.Model, variable_name_parts::Vector{<:AbstractString}) + +Return variables from `model` whose JuMP name contains **all** substrings in +`variable_name_parts`. When the initial filter yields more than one variable, +results are reordered by matching `"[i]"` for `i = 1, 2, ...` to +produce a consistently indexed vector. +""" function find_variables(model::JuMP.Model, variable_name_parts::Vector{S}) where {S} all_vars = all_variables(model) interest_vars = all_vars[findall( diff --git a/test/runtests.jl b/test/runtests.jl index 84814a4..016e25d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -317,6 +317,13 @@ include("test_score_function.jl") uncertainty_samples = [[(uncertainty_1, [2.0])], [(uncertainty_2, [1.0])]] initial_state = [5.0] + n_param_cons_before = length( + JuMP.all_constraints(subproblem2, VariableRef, MOI.Parameter{Float64}) + ) + has_state_in_param_before = any( + con -> JuMP.constraint_object(con).func == state_in_2, + JuMP.all_constraints(subproblem2, VariableRef, MOI.Parameter{Float64}), + ) det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent!( quiet_nonlinear_ipopt_model(), subproblems, @@ -325,6 +332,14 @@ include("test_score_function.jl") initial_state, uncertainty_samples, ) + @test length( + JuMP.all_constraints(subproblem2, VariableRef, MOI.Parameter{Float64}) + ) == n_param_cons_before + @test has_state_in_param_before + @test any( + con -> JuMP.constraint_object(con).func == state_in_2, + JuMP.all_constraints(subproblem2, VariableRef, MOI.Parameter{Float64}), + ) obj_val = DecisionRules.simulate_multistage( det_equivalent, @@ -448,6 +463,26 @@ include("test_score_function.jl") optimize!(model5) @test compute_parameter_dual(model5, state_in5) ≈ -30.0 rtol=1.0e-1 @test compute_parameter_dual(model5, state_out5) ≈ 30.0 rtol=1.0e-1 + + model6 = Model() + @variable(model6, x6) + @variable(model6, p6 in MOI.Parameter(1.0)) + @constraint(model6, x6 - p6 >= 0) + @objective(model6, Min, x6) + @test_throws Exception compute_parameter_dual(model6, p6) + + # Test 7: documented docstring example (MIN sense) + # min 3x + p s.t. x >= 2*p, x >= 0 + # At optimality x* = 2p, so the objective is 7p and ∂obj/∂p = 7. + # The constraint normalizes to x - 2p >= 0 (coef of p is -2), dual = 3: + # contribution -(-2) * 3 = 6, plus the objective coefficient 1 → 7. + model7 = quiet_ipopt_model() + @variable(model7, x7 >= 0) + @variable(model7, p7 in MOI.Parameter(1.0)) + @constraint(model7, con7, x7 >= 2 * p7) + @objective(model7, Min, 3 * x7 + p7) + optimize!(model7) + @test compute_parameter_dual(model7, p7) ≈ 7.0 rtol=1.0e-2 end @testset "create_deficit!" begin @@ -1040,6 +1075,19 @@ include("test_score_function.jl") @test policy.n_uncertainty == n_uncertainty @test policy.n_state == n_state + policy_deep_head = state_conditioned_policy( + n_uncertainty, + n_state, + n_output, + layers; + activation=sigmoid, + encoder_type=Flux.LSTM, + combiner_layers=[7, 5], + ) + @test policy_deep_head.combiner isa Flux.Chain + Flux.reset!(policy_deep_head) + @test length(policy_deep_head(rand(Float32, n_uncertainty + n_state))) == n_output + # Test forward pass Flux.reset!(policy) input = rand(Float32, n_uncertainty + n_state) @@ -2128,6 +2176,9 @@ include("test_score_function.jl") # Empty layers (single layer) m_empty = dense_multilayer_nn(3, 2, Int[]; activation=relu, dense=Dense) @test size(m_empty(rand(Float32, 3))) == (2,) + m_empty.weight .= -1 + m_empty.bias .= 0 + @test all(m_empty(ones(Float32, 3)) .== 0) # Empty layers LSTM m_empty_lstm = dense_multilayer_nn(3, 2, Int[]; dense=LSTM) @@ -2147,12 +2198,36 @@ include("test_score_function.jl") @testset "policy_input_dim" begin @test policy_input_dim(5, 3) == 8 @test policy_input_dim(0, 4) == 4 + @test policy_input_dim(5, 3, 2) == 10 uncertainty_samples = [[(nothing, [1.0, 2.0]), (nothing, [3.0])]] initial_state = [0.0, 0.0, 0.0] @test policy_input_dim(uncertainty_samples, initial_state) == 5 end + @testset "ContextualPolicy" begin + ctx = stage_phase_context(4; period=4) + @test size(ctx) == (3, 4) + @test ctx[:, 1] == context_at(ctx, 1) + @test_throws BoundsError context_at(ctx, 5) + + ctx2 = fill(Float32(2), 1, 4) + @test size(vcat_contexts(ctx, ctx2)) == (4, 4) + @test_throws ArgumentError vcat_contexts(ctx, fill(Float32(0), 1, 3)) + + seen = Vector{Float32}[] + inner = x -> (push!(seen, Float32.(x)); Float32[x[1] + x[end]]) + policy = ContextualPolicy(inner, Float32[10 20; 30 40]) + + @test policy(Float32[1, 2]) == Float32[12] + @test policy(Float32[3, 4]) == Float32[24] + @test seen[1] == Float32[10, 30, 1, 2] + @test seen[2] == Float32[20, 40, 3, 4] + Flux.reset!(policy) + @test policy.t == 0 + @test context_at(t -> Float32[t, t + 1], 3) == Float32[3, 4] + end + @testset "normalize_recur_state" begin plain = (a=1.0, b=[2.0, 3.0]) @test normalize_recur_state(plain) == plain @@ -2271,6 +2346,24 @@ include("test_score_function.jl") DecisionRules.get_objective_no_target_deficit(sp1) + DecisionRules.get_objective_no_target_deficit(sp2) @test total ≈ indiv + + quad_model = quiet_ipopt_model() + @variable(quad_model, x) + @variable(quad_model, norm_deficit >= 0) + @constraint(quad_model, x == 2) + @constraint(quad_model, norm_deficit == 3) + @objective(quad_model, Min, x^2 + 10 * norm_deficit) + optimize!(quad_model) + @test DecisionRules.get_objective_no_target_deficit(quad_model) ≈ 4.0 atol=1.0e-4 + + bad_quad = quiet_ipopt_model() + @variable(bad_quad, z) + @variable(bad_quad, norm_deficit_bad >= 0) + @constraint(bad_quad, z == 1) + @constraint(bad_quad, norm_deficit_bad == 2) + @objective(bad_quad, Min, z^2 + norm_deficit_bad^2) + optimize!(bad_quad) + @test_throws ErrorException DecisionRules.get_objective_no_target_deficit(bad_quad) end @testset "materialize_tangent edge cases" begin