diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 9b09538..1b23ebc 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -9,3 +9,9 @@ updates: directory: "/" schedule: interval: "weekly" + ignore: + # The v5→v7 bump silently broke coverage uploads in the sibling + # DecisionRules.jl repo. 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 d6b04cb..be1a6df 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -17,6 +17,7 @@ jobs: permissions: actions: write contents: read + id-token: write # OIDC token for tokenless Codecov uploads (see codecov step) strategy: fail-fast: false matrix: @@ -37,8 +38,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: codecov-action@v7 silently stopped uploading coverage in + # the sibling DecisionRules.jl repo (Codecov "Missing Head Commit" on + # PRs); v5 is the last version verified to upload from this workflow + # shape. Before re-bumping, migrate deliberately 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 AND this repository to be activated on + # codecov.io (it currently is not); if uploads fail with an OIDC error, + # either install the app + activate the repo, 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 so upload breakage is visible instead of silent. + fail_ci_if_error: true diff --git a/.gitignore b/.gitignore index a432b01..d337d8a 100644 --- a/.gitignore +++ b/.gitignore @@ -42,3 +42,4 @@ logs/ # Slurm batch scripts (user-specific, not part of the package) *.sbatch *.sh +examples/HydroPowerModels/results/ diff --git a/Project.toml b/Project.toml index 42ea84c..dec18e2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,33 +1,42 @@ name = "DecisionRulesExa" uuid = "7c3e91a4-d8f2-4b6a-9e15-a2c4f7b80d53" -authors = ["Andrew Rosemberg and contributors"] version = "0.1.0" +authors = ["Andrew Rosemberg and contributors"] [deps] -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] +CSV = "0.10.16" CUDA = "6" ChainRulesCore = "1.26" ExaModels = "0.11" Flux = "0.16" +JSON = "1.6.1" MadNLP = "0.10" MadNLPGPU = "0.10" NLPModels = "0.21" +Tables = "1.13.0" Zygote = "0.7" julia = "1.10, 1.11, 1.12" [extras] +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "Statistics"] diff --git a/README.md b/README.md index a06eca1..aedb525 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,16 @@ train_tsddr( ) ``` +> **Note on the uncertainty parameter**: `train_tsddr` writes the full sampled +> trajectory (length `T * nw`) into `p_uncertainty` with +> `ExaModels.set_parameter!`, which enforces an exact size match (ExaModels ≥ +> 0.11). The `p_w` built by `build_deterministic_equivalent` / +> `build_linear_tracking_problem` holds only the `(T - 1) * nw` dynamics +> entries, so for `train_tsddr` your NLP needs an uncertainty parameter of +> length `T * nw` (as the Hydro example's `p_inflow` is). See the +> `"train_tsddr open-loop smoke test"` testset in `test/runtests.jl` for a +> minimal full-length variant of the problem above. + For GPU, replace `backend = nothing` with `backend = CUDABackend()` and add `linear_solver = CUDSSSolver` to `madnlp_kwargs`. ## What you need to provide @@ -73,6 +83,58 @@ For a custom problem you need: The package provides `build_deterministic_equivalent` for generic problems and `build_linear_tracking_problem` as a ready-made demo. For domain-specific models (power systems, robotics), build the ExaModels NLP directly — see `examples/HydroPowerModels/` for a complete AC-OPF example. +## Strict reachable target equality + +The usual TS-DDR deterministic equivalent uses slack-penalized target +constraints, + +```text +x_t - pi_theta(w_t, x_{t-1}) = delta_t, +objective += rho * penalty(delta_t). +``` + +This is the right default for open-loop target trajectories and for cases where +the policy can request states that are not reachable from the previous realized +state. The target multipliers are then gradients of the penalized projection +problem, so their quality depends on the penalty calibration. + +For policies whose output is guaranteed to lie in a one-stage reachable state +set, a stricter formulation is possible: + +```text +x_t = pi_theta(w_t, x_{t-1}), pi_theta(w_t, x_{t-1}) in R(w_t, x_{t-1}). +``` + +In that case the deterministic equivalent does not need target slack variables +or target penalties. The multiplier on the equality is the local envelope +sensitivity of the true stage problem with respect to the policy-imposed next +state, not the sensitivity of a penalized approximation. This is useful when: + +- users can define a differentiable or piecewise differentiable map into a + subset of the one-stage reachable set; +- total recourse is guaranteed by the model for every state produced by that + map. + +There are two strict hydro paths: + +- **Embedded strict DE** evaluates the policy inside the NLP against realized + reservoir states. This is the usual way to make strict mode safe because the + policy sees the state from which its next target must be reachable. +- **Regular strict DE with reachable rollout** computes targets before solving + the NLP, but starts from the true initial state and feeds the previous target + back to the reachable policy. If + `x̂_t ∈ R(x̂_{t-1}, w_t)` and `x̂_0 = x_0`, the full target path is feasible by + induction. The strict equality then forces the realized path to equal that + reachable target path. + +Do not use strict equality for a generic open-loop target policy. For +unreachable targets, the slack-penalty formulation is the robust fallback. + +The hydro reachable policy keeps recurrence over inflows only. Optional +`combiner_layers` / `DR_HEAD_LAYERS` add a nonlinear feed-forward map from +`[encoded_inflow; reservoir_state]` to targets without adding recurrence over +the state input. + ## Parallel GPU solves When training samples are independent, multiple NLP instances can be solved concurrently on the same GPU. Pass a `problem_pool` of independent ExaModels problem copies to `train_tsddr`: @@ -217,7 +279,24 @@ Choose DecisionRules.jl when: - [`examples/end_to_end_cpu.jl`](examples/end_to_end_cpu.jl) — minimal CPU demo with a linear tracking problem - [`examples/end_to_end_gpu.jl`](examples/end_to_end_gpu.jl) — same demo on GPU with CUDSS -- [`examples/HydroPowerModels/`](examples/HydroPowerModels/) — full multi-stage hydrothermal scheduling with DC and AC OPF +- [`examples/HydroPowerModels/`](examples/HydroPowerModels/) — full multi-stage hydrothermal scheduling with DC and AC OPF (open-loop DE, embedded closed-loop, strict targets, critic control variate) + +## Repository Map + +| Path | Purpose | +|---|---| +| `src/DecisionRulesExa.jl` | Module entrypoint and public exports | +| `src/policy.jl` | MLP, state-conditioned LSTM policies, bounded target policies, nonlinear target heads | +| `src/deterministic_equivalent.jl` | Generic open-loop deterministic-equivalent builder and solve helpers | +| `src/embedded_deterministic_equivalent.jl` | Generic embedded-policy deterministic equivalent with nonlinear oracle | +| `src/training.jl` | `train_tsddr`, embedded training, solver retry/warm-start handling | +| `src/rollout.jl` | Stage-wise rollout evaluation for ExaModels problems | +| `src/critic_control_variate.jl` | Scalar critic/control-variate helpers | +| `src/utils.jl` | Indexing and small shared utilities | +| `examples/end_to_end_cpu.jl` | Minimal CPU training demo | +| `examples/end_to_end_gpu.jl` | Minimal GPU training demo | +| `examples/HydroPowerModels/` | Bolivia hydrothermal scheduling examples | +| `test/runtests.jl` | Unit and smoke tests | ## Citation diff --git a/examples/HydroPowerModels/Project.toml b/examples/HydroPowerModels/Project.toml index 30e1a40..975a23d 100644 --- a/examples/HydroPowerModels/Project.toml +++ b/examples/HydroPowerModels/Project.toml @@ -1,4 +1,5 @@ [deps] +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" diff --git a/examples/HydroPowerModels/README.md b/examples/HydroPowerModels/README.md index 65f91a6..d6d3a41 100644 --- a/examples/HydroPowerModels/README.md +++ b/examples/HydroPowerModels/README.md @@ -1,16 +1,25 @@ # HydroPowerModels Example -Multi-stage hydrothermal scheduling using DecisionRulesExa.jl with DC or AC OPF formulations. +Multi-stage hydrothermal scheduling using DecisionRulesExa.jl with DC or AC +OPF formulations. This is the GPU-accelerated counterpart of the +[DecisionRules.jl hydro example](https://github.com/LearningToOptimize/DecisionRules.jl/tree/main/examples/HydroPowerModels). ## Problem description -A hydro-dominated power system (Bolivia test case) is operated over a planning horizon of up to 96 stages. At each stage, the operator must decide generator dispatch, reservoir outflows, and spillage subject to: +A hydro-dominated power system (Bolivia test case) is operated over a +planning horizon of up to 126 stages (96 operational + 30 look-ahead). +At each stage, the operator must decide generator dispatch, reservoir +outflows, and spillage subject to: - **Power flow constraints** (DC linearization or full AC polar OPF) - **Reservoir dynamics** (water balance with stochastic inflows) - **Generator and transmission limits** -The TS-DDR policy (an LSTM network) predicts target reservoir levels at each stage. The deterministic-equivalent NLP projects these targets onto the feasible set via slack-penalized target constraints. Training uses envelope-theorem gradients: dual multipliers on the target constraints give the policy gradient without differentiating through the solver. +The TS-DDR policy (an LSTM network) predicts target reservoir levels at +each stage. The deterministic-equivalent NLP projects these targets onto +the feasible set. Training uses envelope-theorem gradients: dual +multipliers on the target constraints give the policy gradient without +differentiating through the solver. ## Formulations @@ -28,7 +37,8 @@ The `bolivia/` directory contains: - `PowerModels.json` — power system topology (39 buses, 55 branches, 19 generators) - `hydro.json` — hydro unit parameters (7 reservoirs) - `inflows.csv` — historical inflow scenarios (144 stages x 200 scenarios x 7 reservoirs) -- `_demand.csv` — per-stage bus demand scaling +- `_demand.csv` — inactive per-stage demand candidate; rename to `demand.csv` + when you intentionally want training scripts to use it Pre-solved deterministic-equivalent references (MOF format) are provided for validation: - `DCPPowerModel.mof.json` @@ -38,19 +48,67 @@ Pre-solved deterministic-equivalent references (MOF format) are provided for val | File | Description | |---|---| -| `train_hydro_exa.jl` | Main training script with penalty scheduling, parallel GPU solves, and W&B logging | -| `train_hydro_exa_critic.jl` | Critic/control-variate variant of the main training script; uses normalized hydro features, a replay buffer, and cheap critic rollouts | -| `hydro_power_data.jl` | Data parsing (PowerModels JSON, hydro JSON, inflows CSV) | -| `hydro_power_exa.jl` | ExaModels problem builder for DC and AC OPF formulations | +| `Project.toml` | Example-specific environment with CUDA, CUDSS, MadNLPGPU, W&B, and JLD2 | +| `README.md` | This guide | +| `bolivia/PowerModels.json` | Bolivia network topology and generator data | +| `bolivia/hydro.json` | Hydro unit limits, initial volumes, cascade metadata, and stage duration | +| `bolivia/inflows.csv` | Historical inflow scenarios | +| `bolivia/_demand.csv` | Inactive per-stage demand candidate; ignored unless renamed to `demand.csv` | +| `bolivia/DCPPowerModel.mof.json` | Pre-exported DC OPF stage template | +| `bolivia/ACPPowerModel.mof.json` | Pre-exported AC polar OPF stage template | +| `bolivia/SOCWRConicPowerModel.mof.json` | Convex relaxation template used for validation/baselines | +| `hydro_power_data.jl` | Data parsing for PowerModels JSON, hydro JSON, inflows, and demand | +| `hydro_power_exa.jl` | Regular open-loop ExaModels DE builder; supports `strict_targets=true` | +| `hydro_reachable_policy.jl` | Reachable hydro target policy, cascade clamps, and reachability proof sketches | +| `hydro_power_exa_embedded.jl` | Embedded-policy DE builder and nonlinear oracle | +| `hydro_training_utils.jl` | Shared training-entrypoint helpers such as environment layer parsing | +| `train_hydro_exa.jl` | Open-loop DE training with penalty scheduling, parallel GPU solves, and W&B logging | +| `train_hydro_exa_embedded.jl` | Embedded (closed-loop) DE training; supports `DR_STRICT_EMBEDDED_TARGETS=true` for penalty-free strict mode | +| `train_hydro_exa_strict.jl` | Regular strict DE training using reachable-policy target rollout and no target slack penalty | +| `train_hydro_exa_critic.jl` | Critic/control-variate variant; adds a scalar critic with replay buffer and cheap rollout samples | | `eval_exa_de.jl` | Validation script comparing ExaModels results against JuMP reference | -| `Project.toml` | Example-specific dependencies (W&B, JLD2, CUDA, etc.) | ## Running -### GPU training (recommended) +### Strict regular DE training (recommended for the current experiments) -```julia -# From this directory: +```bash +DR_ENCODER_LAYERS=128,128 \ +DR_HEAD_LAYERS=128,128 \ +julia --project -t auto train_hydro_exa_strict.jl +``` + +This path removes all target slack penalties from the regular deterministic +equivalent. The target trajectory is generated before the solve, but it is not a +generic open-loop trajectory: `HydroReachablePolicy` starts from the known +initial reservoir state and feeds each previous target into the next policy call. +Because every emitted target is one-stage reachable from the previous target, +the full strict DE target path is feasible by induction. The strict solve then +forces the realized reservoir path to equal that reachable target path. + +Use: + +- `DR_ENCODER_LAYERS`: recurrent LSTM layers over inflows only. +- `DR_HEAD_LAYERS`: optional feed-forward hidden layers over + `[encoded_inflow; reservoir_state]`. This is the nonlinear state-to-target + map; it does not add recurrence over the state input. + +### Strict embedded training + +```bash +DR_STRICT_EMBEDDED_TARGETS=true julia --project -t auto train_hydro_exa_embedded.jl +``` + +Strict mode embeds the policy inside the NLP and enforces hard equality +targets — no penalty tuning needed. Requires a reachable-set policy +(built automatically from the hydro data). The current embedded oracle manually +codes the reachable-policy Jacobian for the default single Dense target head; +use the regular strict DE script for multilayer state-conditioned heads unless +the embedded oracle is extended. + +### Open-loop DE training (GPU) + +```bash julia --project -t auto train_hydro_exa.jl ``` @@ -58,21 +116,17 @@ Set `USE_GPU = true` in `train_hydro_exa.jl` (default). Requires a CUDA-capable ### GPU training with critic control variate -```julia -# From this directory: +```bash julia --project -t auto train_hydro_exa_critic.jl ``` The critic script keeps the dual-multiplier actor update but adds a damped control variate (`critic_cv_weight = 0.5`) trained on the stage-wise rollout -objective without target penalty. Its default critic rollout uses -`policy_state = :target`; set `CRITIC_POLICY_STATE = :realized` for closed-loop -critic labels. Deterministic-equivalent critic fitting remains available as an -ablation through `DeterministicEquivalentCriticTarget()`. +objective without target penalty. ### CPU training -Set `USE_GPU = false` in `train_hydro_exa.jl`, then run the same command. +Set `USE_GPU = false` in the training script, then run the same command. ### Configuration @@ -86,12 +140,165 @@ Key parameters in `train_hydro_exa.jl`: | `NUM_BATCHES` | 100 | Gradient steps per epoch | | `NUM_WORKERS` | 4 | Parallel GPU solver instances | | `LAYERS` | `[128, 128]` | LSTM hidden layer sizes | +| `DR_HEAD_LAYERS` | `""` | Optional nonrecurrent state-conditioned target-head hidden sizes | | `LR` | 1e-3 | Learning rate | | `DEFICIT_COST` | 1e5 | Load-shedding penalty ($/pu) | +| `DR_TARGET_PENALTY_MULT` | `8.0` | Bolivia hydro default multiplier applied to the `:auto` target-penalty coefficients | +| `DR_PENALTY_SCHEDULE` | `const` | Penalty schedule; `const` uses `DR_TARGET_PENALTY_MULT` at every stage | + +### Bolivia target-penalty calibration + +The hydro builders keep `target_penalty = :auto`; `:auto` is the package helper +that computes the base L2/L1 target-slack penalty coefficients from the case +data. For the Bolivia AC hydro case, the case scripts then multiply those base +coefficients by `DR_TARGET_PENALTY_MULT`, which defaults to `8.0`. + +This default came from the SDDP bridge diagnostic in +`compare_sddp_policy_rollout.jl`. We loaded the trained SDDP cuts, simulated one +scenario, took the SDDP stage-1 reservoir `out` states as targets, and solved the +corresponding one-stage Exa target-penalty problem. The first multiplier that +enforced the SDDP targets to solver tolerance was `8.0`; larger multipliers did +not materially improve target matching and started to worsen numerical behavior. + +Representative stage-1 sweep: + +| Multiplier | Exa no-target obj. | Gap vs SDDP | Target violation share | Max target slack | +|---:|---:|---:|---:|---:| +| 1 | 1568.300 | -1599.767 | 1.62e-1 | 2.018e-2 | +| 3 | 1648.143 | -1519.924 | 3.25e-1 | 1.995e-2 | +| 5 | 1669.268 | -1498.799 | 4.38e-1 | 1.994e-2 | +| 6 | 2727.777 | -440.291 | 1.29e-1 | 7.877e-3 | +| 7 | 3082.389 | -85.679 | 2.53e-2 | 1.345e-3 | +| 8 | 3163.132 | -4.936 | 5.38e-7 | 1.159e-8 | +| 10 | 3163.129 | -4.939 | 2.77e-7 | 3.105e-9 | + +The same SDDP-target bridge was also run as a horizon-10 deterministic +equivalent: all ten SDDP reservoir `out` states were injected as Exa targets and +the Exa problem was solved once over the ten-stage horizon. With multiplier +`8.0` and no target-penalty discount (`gamma = 1.0`), the target trajectory +matched to numerical tolerance and the Exa no-target objective was within about +`49.4` of the SDDP rollout objective over a `31668.1` objective. Mild discounting +at `gamma = 0.99` was similar; stronger discounting degraded target enforcement. + +| Horizon | Multiplier | Penalty gamma | Exa no-target obj. | Gap vs SDDP | Target violation share | Max target slack | +|---:|---:|---:|---:|---:|---:|---:| +| 10 | 8 | 1.00 | 31618.626 | -49.434 | 7.85e-8 | 1.18e-8 | +| 10 | 8 | 0.99 | 31618.629 | -49.431 | 1.87e-7 | 6.37e-8 | +| 10 | 8 | 0.97 | 31617.250 | -50.810 | 3.77e-5 | 2.27e-5 | +| 10 | 8 | 0.95 | 31127.120 | -540.940 | 1.39e-2 | 1.02e-2 | + +The full 126-stage SDDP-target deterministic equivalent gives the same +directional result. Flat multiplier `8.0` enforces the complete target +trajectory to numerical tolerance; discounting the target penalty creates +increasing target leakage and worsens the no-target objective gap. + +| Horizon | Multiplier | Penalty gamma | Exa no-target obj. | Gap vs SDDP | Target violation share | Max target slack | +|---:|---:|---:|---:|---:|---:|---:| +| 126 | 8 | 1.00 | 375724.588 | -615.045 | 6.65e-8 | 6.32e-9 | +| 126 | 8 | 0.99 | 375721.696 | -617.938 | 7.54e-6 | 1.45e-4 | +| 126 | 8 | 0.97 | 375477.644 | -861.990 | 2.86e-4 | 4.47e-3 | +| 126 | 8 | 0.95 | 375175.342 | -1164.292 | 4.89e-4 | 1.55e-2 | + +### Strict reachable hydro targets + +The hydro builders support strict target mode when paired with +`HydroReachablePolicy`: + +```julia +policy = hydro_reachable_policy(hydro_data, LAYERS; + activation = sigmoid, + encoder_type = Flux.LSTM, + combiner_layers = HEAD_LAYERS, +) + +prob = build_embedded_hydro_de(policy, power_data, hydro_data, T; + formulation = :ac_polar, + strict_targets = true, +) +``` + +The training script exposes the same path with: + +```bash +export DR_STRICT_EMBEDDED_TARGETS=true +julia --project -t auto train_hydro_exa_embedded.jl +``` + +In embedded strict mode the oracle enforces + +```text +reservoir[t+1, r] = pi_theta(inflow[t], reservoir[t])[r] +``` + +directly. No `delta_pos`, `delta_neg`, L2 target penalty, or L1 target penalty +is created for the target equality. This is intended for closed-loop embedded +training: the policy is evaluated with the realized previous reservoir state, so +it can produce a physically reachable next reservoir by construction. + +In regular strict DE, `train_hydro_exa_strict.jl` instead constructs the target +trajectory externally: + +```text +target[0] = initial_volume +target[t] = pi_theta(inflow[t], target[t-1]) +``` + +Since `pi_theta` maps into the one-stage reachable set from its input state, the +target trajectory is feasible by induction. This is the specific reason strict +mode is valid for regular DE here, despite regular DE target generation usually +being open-loop after the initial state. + +For the current hydro water balance, + +```text +x[t+1, r] = x[t, r] - K*outflow[t, r] - spill[t, r] + + K*inflow[t, r] + upstream_contributions[t, r], +``` + +the helper maps the network's normalized output `y in [0, 1]` into a conservative +one-stage reachable interval: + +```text +upper_r = min(max_vol_r, x_r + K*inflow_r - K*min_turn_r) + +lower_r = min_vol_r +``` + +The lower bound is `min_vol` because this data model has nonnegative spill with +no finite upper bound. If a user supplies finite spill caps through +`spill_max`, the helper instead uses: + +```text +lower_r = max(min_vol_r, + x_r + K*inflow_r - K*max_turn_r - spill_max_r) +``` + +and still emits: + +```text +target_r = lower_r + (upper_r - lower_r) * y_r. +``` + +This is a smooth affine map in the network output, so incoming gradients are not +destroyed by post-hoc clipping. During the actor update, the reachability bounds +are treated as constants because they depend on realized state, inflow, and case +limits rather than trainable weights; the incoming adjoint still reaches the +network as `(upper_r - lower_r) * adjoint_r`. The bound derivatives with respect +to the previous reservoir state are included in the embedded oracle Jacobian and +VJP for the NLP solve. + +The helper is deliberately conservative for cascaded reservoirs: it does not +rely on optional upstream releases to make a target feasible. If a new hydro +case has finite spill limits, nonlinear transition physics, travel times, or +other coupled state-transition constraints, the strict option should only be +used after providing a case-specific reachable-set map. Otherwise keep the +slack-penalty builder and tune the target penalty. ### Training features -- **Penalty scheduling**: target penalty multiplier ramps through phases (0.1 -> 1.0 -> 10.0 -> 30.0) over training +- **Penalty scheduling**: default is constant multiplier `8.0` on top of the + `:auto` base penalty; annealed schedules remain available as explicit + ablations - **Sample scheduling**: `num_train_per_batch` increases from `NUM_WORKERS` to `8 * NUM_WORKERS` - **Evaluation scheduling**: rollout evaluation starts with 4 scenarios and ramps to 32 at halfway - **Parallel solves**: independent NLP copies solved concurrently via `Threads.@spawn` worker pool diff --git a/examples/HydroPowerModels/compare_paired_evals.jl b/examples/HydroPowerModels/compare_paired_evals.jl new file mode 100644 index 0000000..8456476 --- /dev/null +++ b/examples/HydroPowerModels/compare_paired_evals.jl @@ -0,0 +1,259 @@ +# compare_paired_evals.jl +# +# Verdict report for the cross-package equivalence check between +# DecisionRules.jl (MAIN, JuMP/Ipopt) and DecisionRulesExa.jl (EXA, +# ExaModels/MadNLP) on the 100 paired scenarios. +# +# Loads: +# 1. MAIN ground truth: paired_strict_rollout.jld2 (eval_paired_tsddr.jl) +# 2. MAIN policy reference: paired_policy_reference.jld2 (dump_paired_policy_reference.jl) +# 3. EXA evaluation: paired_exa_strict.jld2 (eval_paired_exa_strict.jl) +# 4. MAIN paired_costs.csv (for the SDDP column) +# +# and prints: +# (a) policy-parity gate results, +# (b) per-scenario cost deltas MAIN-stagewise vs EXA-stagewise, +# (c) EXA-stagewise vs EXA-DE per-scenario deltas, +# (d) trajectory deltas (vol / gen / policy targets), +# (e) mean costs of every leg side by side with SDDP, +# (f) the explicit list of known structural differences with measured +# activity/inertness. +# +# Agreement classification is honest and threshold-based: +# EXACT : relative deviation < 1e-6 +# TIGHT : relative deviation < 1e-3 (solver-tolerance regime) +# DISCREPANT : anything larger +# +# Usage: +# julia --project compare_paired_evals.jl + +using JLD2 +using Statistics +using CSV, Tables + +const SCRIPT_DIR = dirname(@__FILE__) +const MAIN_HPM_DIR = "/storage/scratch1/9/arosemberg3/DecisionRules.jl/examples/HydroPowerModels" + +const MAIN_RESULTS = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "results", + "paired_strict_rollout.jld2") +const MAIN_REFERENCE = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "results", + "paired_policy_reference.jld2") +const MAIN_COSTS_CSV = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "paired_costs.csv") +const EXA_RESULTS = joinpath(SCRIPT_DIR, "bolivia", "ACPPowerModel", "results", + "paired_exa_strict.jld2") + +""" + classify(rel::Real) -> String + +Classify a relative deviation: `"EXACT"` below `1e-6`, `"TIGHT"` below `1e-3` +(solver-tolerance regime), `"DISCREPANT"` otherwise (`"NaN"` if not finite). +""" +function classify(rel::Real) + isfinite(rel) || return "NaN" + rel < 1e-6 && return "EXACT" + rel < 1e-3 && return "TIGHT" + return "DISCREPANT" +end + +""" + report_deltas(label, a, b) -> Float64 + +Print mean/max absolute and relative deviations between two equal-length +vectors (NaN entries dropped pairwise), classify the max relative deviation, +and return it. Relative deviation is `|a - b| / max(|a|, 1)` so near-zero +values do not explode the ratio. +""" +function report_deltas(label, a, b) + mask = .!isnan.(a) .& .!isnan.(b) + n = count(mask) + if n == 0 + println(" $label: no comparable entries (all NaN)") + return NaN + end + av, bv = Float64.(a[mask]), Float64.(b[mask]) + absd = abs.(av .- bv) + reld = absd ./ max.(abs.(av), 1.0) + println(" $label [n=$n]") + println(" mean |Δ| = $(mean(absd)) max |Δ| = $(maximum(absd))") + println(" mean rel = $(mean(reld)) max rel = $(maximum(reld)) → $(classify(maximum(reld)))") + return maximum(reld) +end + +# ── Load everything ──────────────────────────────────────────────────────────── + +main = JLD2.load(MAIN_RESULTS) +ref = JLD2.load(MAIN_REFERENCE) +exa = JLD2.load(EXA_RESULTS) + +main_costs = Float64.(main["costs"]) # [S] +main_vol = Float64.(main["vol_trajectories"]) # [T × S] +main_gen = Float64.(main["gen_trajectories"]) # [T × S] +main_idx = Int.(main["scenario_indices"]) # [T × S] + +exa_costs = Float64.(exa["costs"]) # [S] +exa_vol = Float64.(exa["vol_trajectories"]) # [T × S] +exa_gen = Float64.(exa["gen_trajectories"]) # [T × S] +exa_ok = Bool.(exa["scenario_ok"]) +de_costs = Float64.(exa["de_costs"]) # [N] +de_ok = Bool.(exa["de_ok"]) +n_de = Int(exa["num_de_scenarios"]) + +T, S = size(main_vol) + +# SDDP column from MAIN's paired_costs.csv (quoted headers contain commas, so +# CSV.jl is required for parsing). +sddp_costs = Float64[] +if isfile(MAIN_COSTS_CSV) + cols = Tables.columntable(CSV.File(MAIN_COSTS_CSV)) + sddp_key = findfirst(k -> occursin("SDDP", String(k)), collect(keys(cols))) + if sddp_key !== nothing + global sddp_costs = Float64.(collect(cols[collect(keys(cols))[sddp_key]])) + end +end + +println("=" ^ 72) +println("CROSS-PACKAGE PAIRED EVALUATION — VERDICT REPORT") +println(" MAIN results: $MAIN_RESULTS") +println(" MAIN reference: $MAIN_REFERENCE") +println(" EXA results: $EXA_RESULTS") +println(" T=$T stages, S=$S scenarios, DE scenarios N=$n_de") +println("=" ^ 72) + +# Sanity: both sides must have evaluated the same scenario index matrix. +ref_idx = Int.(ref["scenario_indices"]) +idx_match = main_idx == ref_idx +println("\nScenario-index matrices identical (MAIN results vs reference): $idx_match") +idx_match || println(" !! The two evaluations did NOT use the same scenarios — all cost/trajectory comparisons below are void.") + +# ── (a) Policy-parity gate ───────────────────────────────────────────────────── +println("\n(a) POLICY PARITY GATE (from eval_paired_exa_strict.jl)") +println(" checkpoint loader path: $(exa["gate_loader_mode"])", + exa["gate_loader_mode"] == "stock" ? "" : + " ← EXA's stock loader cannot load MAIN checkpoints (LSTM wrapper vs bare LSTMCell state)") +println(" probe parity (single calls): ", + exa["gate_probe_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_probe_max_dev"])") +println(" open-loop, EXA policy as-is: ", + exa["gate_openloop_asis_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_openloop_asis_max_dev"])") +println(" open-loop, state-threaded: ", + exa["gate_openloop_threaded_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_openloop_threaded_max_dev"])") +if exa["gate_probe_pass"] && !exa["gate_openloop_asis_pass"] && exa["gate_openloop_threaded_pass"] + println(" → INTERPRETATION: weights load correctly; the divergence is the EXA") + println(" policy's MEMORYLESS LSTM encoding (Flux 0.16 restarts recurrent") + println(" state on every call; Flux.reset! is a no-op) vs MAIN's explicit") + println(" state threading across stages. The two packages evaluate") + println(" DIFFERENT functions of the inflow history with the same weights.") +end +println(" inflow indexing/units: ", + exa["gate_inflow_pass"] ? "PASS (exact)" : "FAIL", + " max|Δ| = $(exa["gate_inflow_recon_max_dev"])") +println(" metadata devs: K=$(exa["gate_dev_K"]) min_vol=$(exa["gate_dev_min_vol"]) " * + "max_vol=$(exa["gate_dev_max_vol"]) min_turn=$(exa["gate_dev_min_turn"]) " * + "max_turn=$(exa["gate_dev_max_turn"]) upstream_max=$(exa["gate_dev_upstream_max"])") + +# ── (b) MAIN stage-wise vs EXA stage-wise costs ──────────────────────────────── +println("\n(b) PER-SCENARIO COSTS: MAIN stage-wise (Ipopt) vs EXA stage-wise (MadNLP)") +println(" NOTE: if gate (a) shows the policies compute different targets, part of") +println(" this delta is the policy-function difference, not the subproblem builder.") +report_deltas("total cost per scenario", main_costs, exa_costs) +n_fail = count(.!exa_ok) +n_fail > 0 && println(" EXA stage-wise scenarios failed (excluded): $n_fail") + +# ── (c) EXA stage-wise vs EXA DE ─────────────────────────────────────────────── +println("\n(c) EXA STAGE-WISE vs EXA STRICT FULL-HORIZON DE (first $n_de scenarios)") +println(" Strict-mode claim: identical targets pin the same reservoir path, so the") +println(" DE objective must equal the stage-wise sum for the same scenario.") +report_deltas("total cost per scenario", exa_costs[1:n_de], de_costs) +# Per-stage decomposition of the DE objective vs the stage-wise stage costs. +sc = Float64.(exa["stage_costs"])[:, 1:n_de] +dsc = Float64.(exa["de_stage_costs"]) +report_deltas("per-stage costs (all t, s ≤ N)", vec(sc), vec(dsc)) +n_de_fail = count(.!de_ok) +n_de_fail > 0 && println(" EXA DE scenarios failed (excluded): $n_de_fail") + +# ── (d) Trajectory deltas ────────────────────────────────────────────────────── +println("\n(d) TRAJECTORY DELTAS across all (t, s)") +report_deltas("vol_trajectories (Σ_r volume/0.0036, MW) MAIN vs EXA-stagewise", + vec(main_vol), vec(exa_vol)) +report_deltas("gen_trajectories (Σ_thermal pg·baseMVA) MAIN vs EXA-stagewise", + vec(main_gen), vec(exa_gen)) +# Policy target paths: MAIN open-loop reference vs EXA realized targets. +xhat_ref = Float64.(ref["xhat_trajectories"]) # [T × nHyd × S] +tgt_exa = Float64.(exa["target_trajectories"]) # [nHyd × T × S] +tgt_exa_perm = permutedims(tgt_exa, (2, 1, 3)) # → [T × nHyd × S] +report_deltas("policy target trajectories (all t, r, s) MAIN-ref vs EXA", + vec(xhat_ref), vec(tgt_exa_perm)) +report_deltas("DE vol_trajectories vs EXA-stagewise (s ≤ N)", + vec(Float64.(exa["vol_trajectories"])[:, 1:n_de]), + vec(Float64.(exa["de_vol_trajectories"]))) +report_deltas("DE gen_trajectories vs EXA-stagewise (s ≤ N)", + vec(Float64.(exa["gen_trajectories"])[:, 1:n_de]), + vec(Float64.(exa["de_gen_trajectories"]))) + +# ── (e) Mean costs side by side ──────────────────────────────────────────────── +println("\n(e) MEAN COSTS ($S scenarios; DE over first $n_de)") +_mean_ok(v) = (m = v[.!isnan.(v)]; isempty(m) ? NaN : mean(m)) +println(" MAIN stage-wise strict (Ipopt): $(round(_mean_ok(main_costs); digits=1))") +println(" EXA stage-wise strict (MadNLP): $(round(_mean_ok(exa_costs); digits=1))") +println(" EXA strict DE (MadNLP, N=$n_de): $(round(_mean_ok(de_costs); digits=1))") +println(" MAIN stage-wise mean over s ≤ $n_de: $(round(_mean_ok(main_costs[1:n_de]); digits=1))") +if !isempty(sddp_costs) + println(" SDDP (paired_costs.csv): $(round(_mean_ok(sddp_costs); digits=1))") +else + println(" SDDP column not found in $MAIN_COSTS_CSV") +end +println(" rollout_tsddr cross-check (s=1): $(exa["rollout_tsddr_check_objective"]) vs custom loop $(exa_costs[1])") + +# ── (f) Known structural differences, with measured activity ────────────────── +println("\n(f) KNOWN STRUCTURAL DIFFERENCES (EXA builder vs MAIN mof.json subproblem)") + +max_def_sw = maximum(filter(!isnan, Float64.(exa["max_deficit"])); init = -Inf) +max_def_de = maximum(filter(!isnan, Float64.(exa["de_max_deficit"])); init = -Inf) +max_dq_sw = maximum(filter(!isnan, Float64.(exa["max_abs_deficit_q"])); init = -Inf) +max_dq_de = maximum(filter(!isnan, Float64.(exa["de_max_abs_deficit_q"])); init = -Inf) + +println(""" + 1. DEFICIT COST: MAIN objective uses 6000·Σ deficit (= cost_deficit 60 × baseMVA + 100, scaled at mof export). This evaluation used deficit_cost = + $(exa["deficit_cost_used"]) to match. EXA TRAINING uses 1e5 instead. + Measured max deficit: stage-wise = $max_def_sw pu, DE = $max_def_de pu. + → the 1e5-vs-6000 difference is $(max(max_def_sw, max_def_de) <= 1e-8 ? "INERT (deficit never activates)" : "ACTIVE — deficit occurs, costs are NOT comparable to training runs"). + 2. REACTIVE SLACK: EXA adds a FREE zero-cost deficit_q to every reactive KCL; + MAIN's mof.json has hard reactive balance. Not disableable via kwargs. + Measured max |deficit_q|: stage-wise = $max_dq_sw pu, DE = $max_dq_de pu. + → $(max(max_dq_sw, max_dq_de) <= 1e-6 ? "INERT on these scenarios" : "ACTIVE — the EXA AC feasible set is genuinely relaxed vs MAIN (reactive balance violated at zero cost); expect lower EXA costs"). + 3. THERMAL LIMITS: MAIN enforces quadratic p²+q² ≤ rate_a² per branch end + (62 quadratic constraints in mof.json); EXA box-bounds p_fr/q_fr/p_to/q_to + in [−rate_a, rate_a] — a relaxation in the (|p|,|q|) corners. Not measured + directly here; shows up as cost differences when branch limits bind. + 4. MIN-VIOLATION SLACKS: mof.json has zero-cost min_outflow/min_volume + violation slacks; EXA enforces outflow ≥ min_turn hard. Bolivia has + min_turn ≡ 0 and min_vol ≡ 0 → inert. + 5. GEN COSTS: identical by construction (linear c1 per pu, c2 = 0, from the + same PowerModels.json; verified against mof.json coefficients). + Turbine coupling identical: baseMVA·pg = pf·outflow. + 6. DEMAND: mof.json bakes in 0.6× PowerModels loads (pd and qd); EXA used + load_scaler = $(exa["load_scaler_used"]) → matched. + 7. SPILL UNITS: both sides use spill with coefficient 1 (volume units) in the + water balance and K = 0.0036 on inflow/outflow → identical dynamics. + 8. POLICY RECURRENCE: see gate (a). EXA's encoder is memoryless per stage + (Flux 0.16 LSTM restarts from initialstates each call); MAIN threads LSTM + state across stages. Same weights, different function. This affects EXA + TRAINING and evaluation alike: the EXA pipeline optimizes/evaluates a + policy without inflow memory. + 9. CHECKPOINT LOADER: loader path used = "$(exa["gate_loader_mode"])". If not + "stock", DecisionRulesExa's load_stateconditioned_policy! cannot ingest + MAIN checkpoints (MAIN saves bare LSTMCell states; EXA wraps cells in + Flux.LSTM) — cross-package warmstarts silently depend on a custom loader. + 10. SOLVERS/TOLERANCES: MAIN = Ipopt (mumps, default tol); EXA = MadNLP + (tol = $(exa["solver_tol"])). Agreement at TIGHT (<1e-3 rel) is the + expected ceiling for cost comparisons even with identical models. + 11. INITIAL STATE: EXA training clamps x0 into [min_vol, max_vol]; this + evaluation used MAIN's raw initial_state (denormal ≈ 1e-316 values ≈ 0); + difference ≤ 1e-315 hm³ → inert. +""") +println("=" ^ 72) +println("END OF REPORT") +println("=" ^ 72) diff --git a/examples/HydroPowerModels/eval_paired_exa_strict.jl b/examples/HydroPowerModels/eval_paired_exa_strict.jl new file mode 100644 index 0000000..593f739 --- /dev/null +++ b/examples/HydroPowerModels/eval_paired_exa_strict.jl @@ -0,0 +1,992 @@ +# eval_paired_exa_strict.jl +# +# Cross-package equivalence check against DecisionRules.jl (the JuMP/Ipopt MAIN +# repo). Evaluates the SAME stage-wise-strict TS-DDR checkpoint on the SAME 100 +# paired scenarios that MAIN's eval_paired_tsddr.jl evaluated, using +# +# (a) the stage-wise strict ExaModels rollout (1-stage strict problems, +# closed-loop realized-state feedback), and +# (b) the strict regular full-horizon deterministic equivalent (DE), +# +# and records per-scenario objectives and operative solution values for +# comparison against MAIN's ground truth (paired_strict_rollout.jld2). +# +# DISCREPANCIES ARE THE DELIVERABLE. Nothing here is tuned to force agreement; +# every structural difference is measured and saved. Known structural +# differences handled/documented here (see also compare_paired_evals.jl): +# +# 1. DEFICIT COST. MAIN's ACPPowerModel.mof.json objective is +# Σ_g c1_g·pg_g + 6000·Σ_b deficit_b (linear; c2 = 0 for all g) +# where 6000 = cost_deficit (60 $/MWh) × baseMVA (100): HydroPowerModels +# scales the deficit cost by baseMVA when exporting the subproblem. The +# EXA builder uses `deficit_cost` per pu directly, so this script passes +# deficit_cost = power_data.cost_deficit * power_data.baseMVA = 6000.0. +# NOTE: train_hydro_exa_strict.jl trains with DEFICIT_COST = 1e5 instead — +# inert iff deficit never activates (max deficit is recorded per scenario). +# 2. DEMAND. The mof.json bakes in 0.6 × PowerModels.json loads (see MAIN's +# export_subproblem_mof.jl, which scales pd AND qd by 0.6). The EXA builder +# reproduces this with load_scaler = 0.6 and demand_matrix = nothing. +# 3. REACTIVE SLACK. By default the EXA AC builder adds a FREE, zero-cost +# `deficit_q` variable to every reactive KCL; MAIN's mof.json has no +# reactive slack (hard reactive balance). The builder kwarg +# `reactive_deficit_cost` now controls this; this script reads the +# DR_REACTIVE_DEFICIT env var (default "hard" → Inf, i.e. NO reactive +# slack — MAIN-faithful; "free" → nothing reproduces the historical +# relaxation; a number → linear |deficit_q| penalty at that cost) and +# passes it to BOTH the stage problem and the DE leg. max |deficit_q| is +# still recorded per scenario (exact zeros in hard mode). +# 4. THERMAL LIMITS. mof.json enforces quadratic branch limits +# p² + q² ≤ rate_a² (62 ScalarQuadraticFunction ≤ constraints); the EXA +# builder only box-bounds each of p_fr, q_fr, p_to, q_to in +# [−rate_a, rate_a] — a superset of the disk (relaxation in the corners). +# 5. MIN-VIOLATION SLACKS. mof.json has min_outflow_violation / +# min_volume_violation slack variables with ZERO objective cost (they only +# relax the ≥ min bounds); the EXA builder enforces outflow ≥ min_turn +# hard. For Bolivia min_turn ≡ 0, so both are inert. +# 6. POLICY RECURRENCE (measured by the parity gate below). RESOLVED: the +# EXA HydroReachablePolicy now threads the LSTM state across stages +# explicitly (`DecisionRulesExa._step_encoder`, mirroring MAIN's +# `DecisionRules._step_encoder`), and `Flux.reset!(policy)` is a REAL +# reset back to `Flux.initialstates`. The as-is open-loop gate below is +# therefore expected to PASS with max|Δ| = 0.0 against MAIN. The +# independent manually-threaded diagnostic (gate 4) is retained as a +# cross-check: if as-is and threaded ever disagree, the policy forward +# pass has regressed from MAIN's semantics. +# +# Requires: MAIN's paired_policy_reference.jld2 (produced by +# dump_paired_policy_reference.jl in the MAIN repo) — run that job first. +# +# Environment variables: +# DR_DE_SCENARIOS = "10" (number of scenarios for the full-horizon DE leg) +# DR_MAX_ITER = "9000" +# DR_REACTIVE_DEFICIT = "hard" ("hard" → Inf = no reactive slack (MAIN-faithful); +# "free" → nothing = historical free slack; +# a number → linear |deficit_q| cost) +# DR_CHECKPOINT = (checkpoint to evaluate; default: the +# MAIN reference checkpoint hardcoded below) +# DR_CHECKPOINT_KIND = "main" ("main" → apply ALL MAIN-reference parity gates; +# "exa" → EXA-trained checkpoint: SKIP the +# probe-parity and open-loop-trajectory gates +# (only meaningful for the exact reference +# weights) but KEEP the policy-independent +# inflow-indexing gate; scenario inflow values +# still come from the reference JLD2) +# DR_ENCODER_LAYERS = "128,128" (LSTM encoder widths; must match checkpoint; +# DR_LAYERS is the legacy alias, as in +# train_hydro_exa_strict.jl) +# DR_HEAD_LAYERS = "" (state-conditioned head hidden widths; must +# match checkpoint; "" → linear head) +# DR_CONTEXT = "" (""/"none", "phase", or "phase+progress"; +# defaults to the reference file metadata +# when present) +# DR_CONTEXT_HORIZON = "126" (denominator/horizon used for progress +# context; defaults to reference metadata) +# DR_REFERENCE_FILE = (paired_policy_reference*.jld2 from MAIN) +# DR_OUTPUT_TAG = "" ("" → save to results/paired_exa_strict.jld2; +# "" → results/paired_exa_strict_.jld2 +# with checkpoint path + all knobs recorded) +# +# Usage: +# julia --project -t auto eval_paired_exa_strict.jl + +using DecisionRulesExa +using ExaModels +using MadNLP +using Flux +using Statistics, Random +using JLD2 + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) # parse_layers +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = "ACPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") + +# Absolute paths into the MAIN (DecisionRules.jl) repository. +const MAIN_HPM_DIR = "/storage/scratch1/9/arosemberg3/DecisionRules.jl/examples/HydroPowerModels" +const DEFAULT_REFERENCE_FILE = joinpath( + MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "results", "paired_policy_reference.jld2" +) +const REFERENCE_FILE = get(ENV, "DR_REFERENCE_FILE", DEFAULT_REFERENCE_FILE) +# Default checkpoint: the MAIN reference checkpoint against which the parity +# gates below were designed. DR_CHECKPOINT overrides it with any other +# checkpoint (MAIN- or EXA-trained). +const DEFAULT_MODEL_PATH = joinpath( + MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "models", + "bolivia-ACPPowerModel-h126-r96-subproblems-strict-2026-07-01T09:41:53.026.jld2", +) +const MODEL_PATH = get(ENV, "DR_CHECKPOINT", DEFAULT_MODEL_PATH) +# Checkpoint kind: "main" (DecisionRules.jl-trained; MAIN-reference parity +# gates apply) or "exa" (train_hydro_exa_strict.jl-trained; the probe-parity +# and open-loop-trajectory gates are SKIPPED because the reference probe +# outputs / trajectories were produced by the specific reference checkpoint — +# comparing independently trained weights against them is meaningless. The +# inflow-indexing gate is policy-independent and is KEPT, and scenario inflow +# values are still sourced from the reference JLD2 as authoritative data). +const CHECKPOINT_KIND = lowercase(strip(get(ENV, "DR_CHECKPOINT_KIND", "main"))) +CHECKPOINT_KIND in ("main", "exa") || + error("DR_CHECKPOINT_KIND must be \"main\" or \"exa\", got \"$CHECKPOINT_KIND\"") +const MAIN_PARITY_GATES = CHECKPOINT_KIND == "main" + +# Architecture knobs — parsed exactly as train_hydro_exa_strict.jl parses them +# (including the DR_LAYERS legacy alias); they MUST match the checkpoint. +const ENCODER_LAYERS = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) + +# Output tag: "" keeps the historical filename results/paired_exa_strict.jld2; +# a non-empty tag saves to results/paired_exa_strict_.jld2 so evaluating a +# new checkpoint never clobbers the reference results. +const OUTPUT_TAG = String(strip(get(ENV, "DR_OUTPUT_TAG", ""))) +const OUT_SUFFIX = isempty(OUTPUT_TAG) ? "" : "_$(OUTPUT_TAG)" +# Record the new provenance knobs in the JLD2 whenever any of them was +# explicitly set; with all of them unset the saved file keeps exactly the +# historical key set (byte-identical default behavior). +const RECORD_KNOBS = any(haskey.(Ref(ENV), + ("DR_CHECKPOINT", "DR_CHECKPOINT_KIND", "DR_ENCODER_LAYERS", "DR_LAYERS", + "DR_HEAD_LAYERS", "DR_CONTEXT", "DR_CONTEXT_HORIZON", "DR_REFERENCE_FILE", + "DR_OUTPUT_TAG"))) +# mof.json demand = 0.6 × PowerModels.json pd/qd (export_subproblem_mof.jl). +const LOAD_SCALER = 0.6 +const NUM_DE_SCENARIOS = parse(Int, get(ENV, "DR_DE_SCENARIOS", "10")) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +# Reactive-slack control (header item 3): "hard" → Inf (no deficit_q, +# MAIN-faithful), "free" → nothing (historical free slack), number → linear +# |deficit_q| penalty at that cost. Passed to both the stage problem and the +# DE leg builders. +const REACTIVE_DEFICIT_RAW = lowercase(strip(get(ENV, "DR_REACTIVE_DEFICIT", "hard"))) +const REACTIVE_DEFICIT_COST = REACTIVE_DEFICIT_RAW == "hard" ? Inf : + REACTIVE_DEFICIT_RAW == "free" ? nothing : + parse(Float64, REACTIVE_DEFICIT_RAW) +# CPU MadNLP: same pattern as train_hydro_exa_strict.jl's SOLVER_KWARGS, but +# this evaluation runs on a CPU node (backend = nothing everywhere). +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen) baseMVA=$(power_data.baseMVA) cost_deficit=$(power_data.cost_deficit)" + +# Deficit cost matching MAIN's mof.json objective coefficient (see header, item 1): +# 6000 = cost_deficit ($/MWh) × baseMVA, applied per pu of shed load. +const DEFICIT_COST = power_data.cost_deficit * power_data.baseMVA +@info " deficit_cost used for equivalence: $DEFICIT_COST (training used 1e5)" + +@info "Loading hydro data (num_stages = full inflow history)..." +# num_stages = nothing keeps the RAW rows of inflows.csv (47 for Bolivia — +# fewer than the 96 eval stages). MAIN's read_inflow tiles those rows +# vertically for longer horizons (load_hydropowermodels.jl:13-21), so stage t +# maps to raw row mod1(t, nrows); the reconstruction gate below applies the +# same cyclic convention. +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; num_stages = nothing) +nHyd = hydro_data.nHyd +@info " nHyd=$nHyd nScenarios=$(hydro_data.nScenarios) K=$(hydro_data.K)" + +# ── Load the MAIN policy reference ──────────────────────────────────────────── + +isfile(REFERENCE_FILE) || error( + "Missing reference file $REFERENCE_FILE — run dump_paired_policy_reference.jl " * + "in the MAIN repo first." +) +ref = JLD2.load(REFERENCE_FILE) +const T_EVAL = Int(ref["num_eval_stages"]) +const NUM_SCEN = Int(ref["num_scenarios"]) +inflow_ref = ref["inflow_values"] # [T × nHyd × S] authoritative values +xhat_ref = ref["xhat_trajectories"] # [T × nHyd × S] MAIN open-loop targets +scen_idx = Int.(ref["scenario_indices"]) # [T × S] +x0_ref = Float64.(ref["initial_state"]) +probe_in = Float32.(ref["probe_inputs"]) # [2 nHyd × 3] +probe_out_ref = ref["probe_outputs"] # [nHyd × 3] +@info "Loaded reference: T=$T_EVAL, S=$NUM_SCEN from $REFERENCE_FILE" +@assert size(inflow_ref) == (T_EVAL, nHyd, NUM_SCEN) +@assert length(x0_ref) == nHyd + +ref_string(key::AbstractString, default::AbstractString) = + haskey(ref, key) ? string(ref[key]) : default +ref_int(key::AbstractString, default::Int) = + haskey(ref, key) ? Int(ref[key]) : default + +const CONTEXT_MODE = canonical_context_mode( + get(ENV, "DR_CONTEXT", ref_string("context_mode", "")) +) +const CONTEXT_PERIOD = ref_int("context_period", countlines(INFLOW_FILE)) +const CONTEXT_HORIZON = parse( + Int, + get(ENV, "DR_CONTEXT_HORIZON", string(ref_int("context_horizon", 126))), +) +CONTEXT_HORIZON >= T_EVAL || + error("DR_CONTEXT_HORIZON=$CONTEXT_HORIZON must cover reference T_EVAL=$T_EVAL") +const STAGE_CONTEXT = build_stage_context(CONTEXT_MODE, CONTEXT_HORIZON, CONTEXT_PERIOD) +const N_CONTEXT = isnothing(STAGE_CONTEXT) ? 0 : size(STAGE_CONTEXT, 1) +@info "Policy context" context_mode=(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) CONTEXT_PERIOD CONTEXT_HORIZON N_CONTEXT + +# ── Build the EXA policy and load the requested checkpoint ──────────────────── +# Constructor call mirrors train_hydro_exa_strict.jl exactly (sigmoid activation +# and Flux.LSTM encoder are the constructor defaults); combiner_layers must +# match the checkpoint's head architecture. + +Random.seed!(42) +base_policy = hydro_reachable_policy( + hydro_data, + ENCODER_LAYERS; + combiner_layers = HEAD_LAYERS, + n_context = N_CONTEXT, +) +policy = isnothing(STAGE_CONTEXT) ? base_policy : ContextualPolicy(base_policy, STAGE_CONTEXT) +policy_core(policy) = policy isa ContextualPolicy ? policy.policy : policy +stage_encoder_input(policy, t::Int, w) = + policy isa ContextualPolicy ? vcat(Float32.(context_at(policy.context, t)), w) : w +@info "Policy built" encoder_layers = ENCODER_LAYERS head_layers = HEAD_LAYERS checkpoint_kind = CHECKPOINT_KIND context_mode=(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) +isfile(MODEL_PATH) || error("Checkpoint not found: $MODEL_PATH (set DR_CHECKPOINT)") + +""" + load_main_checkpoint!(policy, model_state) -> String + +Load a DecisionRules.jl (MAIN) checkpoint into the EXA `HydroReachablePolicy`, +returning a string describing which loading path succeeded. + +# Arguments +- `policy`: EXA `HydroReachablePolicy` (encoder is a `Chain` of `Flux.LSTM` + wrapper layers, each holding a `.cell`). +- `model_state`: the checkpoint's `Flux.state`, whose `encoder` field was saved + from MAIN's `Chain` of BARE `LSTMCell`s (MAIN's `_as_cell` strips the `LSTM` + wrapper), i.e. `state.encoder.layers[i] = (Wi = …, Wh = …, bias = …)`. + +# Returns +- `"stock"` if the package loader `load_stateconditioned_policy!` succeeded. +- `"cell-by-cell"` if the stock loader failed and the weights were injected + directly into each `LSTM.cell` plus the combiner. + +# Notes +`load_stateconditioned_policy!` now includes the cell-by-cell MAIN-checkpoint +fallback natively (`DecisionRulesExa._load_encoder_state!`), so the stock path +is expected to SUCCEED and report `"stock"`. EXA-trained checkpoints +(train_hydro_exa_strict.jl saves `Flux.state(cpu(m))` of the same +`HydroReachablePolicy` type, encoder stored as `Flux.LSTM` wrappers) also load +through the stock path. This wrapper's own cell-by-cell +branch is retained as belt-and-braces; reaching it would indicate a loader +regression and is recorded in the output. The cell-by-cell path performs the +mathematically identical weight injection: the wrapper's `cell` has exactly +the fields `(Wi, Wh, bias)` MAIN saved. +""" +function load_main_checkpoint!(policy, model_state) + try + load_stateconditioned_policy!(policy, model_state) + return "stock" + catch err + @warn "Stock EXA loader failed on the MAIN checkpoint (expected: MAIN saves bare LSTMCell states, EXA wraps cells in Flux.LSTM). Falling back to cell-by-cell injection." exception = err + core = policy_core(policy) + inner_state = hasproperty(model_state, :policy) ? getproperty(model_state, :policy) : model_state + enc_state = getproperty(inner_state, :encoder) + layer_states = getproperty(enc_state, :layers) + length(layer_states) == length(core.encoder.layers) || + error("Encoder depth mismatch: checkpoint has $(length(layer_states)) layers, policy has $(length(core.encoder.layers))") + for (layer, lstate) in zip(core.encoder.layers, layer_states) + # MAIN layer state keys (Wi, Wh, bias) match LSTMCell's fields. + Flux.loadmodel!(layer.cell, lstate) + end + Flux.loadmodel!(core.combiner, getproperty(inner_state, :combiner)) + Flux.reset!(policy) + return "cell-by-cell" + end +end + +model_state = JLD2.load(MODEL_PATH, "model_state") +loader_mode = load_main_checkpoint!(policy, model_state) +@info "Checkpoint loaded via: $loader_mode ($MODEL_PATH)" + +# ── Policy-parity gate ──────────────────────────────────────────────────────── + +""" + _max_abs_dev(a, b) -> Float64 + +Maximum absolute element-wise deviation `max_i |a_i − b_i|` between two arrays +of equal size. +""" +_max_abs_dev(a, b) = maximum(abs.(Float64.(a) .- Float64.(b))) + +gate = Dict{String, Any}("loader_mode" => loader_mode) +core_policy = policy_core(policy) + +# (0) Frozen hydro-metadata parity: the bounds the two policies scale into. +gate["dev_K"] = abs(core_policy.K - Float64(ref["policy_K"])) +gate["dev_min_vol"] = _max_abs_dev(core_policy.min_vol, ref["policy_min_vol"]) +gate["dev_max_vol"] = _max_abs_dev(core_policy.max_vol, ref["policy_max_vol"]) +gate["dev_min_turn"] = _max_abs_dev(core_policy.min_turn, ref["policy_min_turn"]) +gate["dev_max_turn"] = _max_abs_dev(core_policy.max_turn, ref["policy_max_turn"]) +gate["dev_upstream_max"] = _max_abs_dev(core_policy.upstream_max_inflow, ref["policy_upstream_max"]) +@info "Metadata deviations" gate["dev_K"] gate["dev_min_vol"] gate["dev_max_vol"] gate["dev_min_turn"] gate["dev_max_turn"] gate["dev_upstream_max"] + +# (1) Probe parity: single policy calls from the reset state. Tests weight +# loading independent of any recurrence-threading semantics. ONLY meaningful +# when evaluating the exact MAIN reference checkpoint the probe outputs were +# generated from — SKIPPED for EXA-trained checkpoints. +probe_out_exa = fill(NaN, nHyd, 3) +if MAIN_PARITY_GATES + for p in 1:3 + Flux.reset!(policy) # REAL reset: each probe starts from initialstates + probe_out_exa[:, p] = Float64.(policy(probe_in[:, p])) + end + gate["probe_outputs_exa"] = probe_out_exa + gate["probe_max_dev"] = _max_abs_dev(probe_out_exa, probe_out_ref) + # Float32 tolerance: relative to the target scale (max_vol up to ~138). + probe_pass = gate["probe_max_dev"] <= 1e-5 * max(1.0, maximum(abs.(probe_out_ref))) + gate["probe_pass"] = probe_pass + println(probe_pass ? "GATE probe parity: PASS" : "GATE probe parity: FAIL", + " (max abs dev = $(gate["probe_max_dev"]))") +else + # The reference probe outputs are the REFERENCE checkpoint's responses; + # an independently trained EXA checkpoint has different weights, so a + # weight-parity comparison would fail by construction and prove nothing. + gate["probe_max_dev"] = NaN + gate["probe_pass"] = "skipped" + println("GATE probe parity: SKIPPED (DR_CHECKPOINT_KIND=exa — reference " * + "probe outputs only characterize the MAIN reference checkpoint)") +end + +# (2) Inflow reconstruction: rebuild w[t, r, s] from the EXA loader's +# scenario_inflows and the reference scenario indices; compare against the +# reference values. This validates scenario indexing and units across loaders +# (the historically buggy spot). +inflow_recon_dev = 0.0 +first_mismatch = nothing +# MAIN's read_inflow tiles the raw inflow rows vertically when num_stages +# exceeds the file length (load_hydropowermodels.jl:13-21), so stage t maps to +# raw row mod1(t, nrows). The EXA loader keeps the raw (untiled) matrix; apply +# the same cyclic convention here so both sides index identical physical data. +n_inflow_rows = size(hydro_data.scenario_inflows[1], 1) +for s in 1:NUM_SCEN, t in 1:T_EVAL, r in 1:nHyd + v_exa = hydro_data.scenario_inflows[r][mod1(t, n_inflow_rows), scen_idx[t, s]] + dev = abs(v_exa - inflow_ref[t, r, s]) + if dev > inflow_recon_dev + global inflow_recon_dev = dev + global first_mismatch = (t = t, r = r, s = s, exa = v_exa, ref = inflow_ref[t, r, s]) + end +end +gate["inflow_recon_max_dev"] = inflow_recon_dev +gate["inflow_pass"] = inflow_recon_dev == 0.0 +if gate["inflow_pass"] + println("GATE inflow indexing: PASS (exact reconstruction)") +else + println("GATE inflow indexing: FAIL (max abs dev = $inflow_recon_dev at $first_mismatch)") + # Diagnose common failure patterns at the first mismatching coordinate. + # All probes use cyclic row indexing so t beyond the raw row count cannot + # itself throw while diagnosing an indexing mismatch. + t, r, s = first_mismatch.t, first_mismatch.r, first_mismatch.s + ω = scen_idx[t, s] + tr = mod1(t, n_inflow_rows) + println(" diagnostics at (t=$t → raw row $tr, r=$r, s=$s, ω=$ω):") + println(" ref value = $(inflow_ref[t, r, s])") + println(" exa [tr, ω] = $(hydro_data.scenario_inflows[r][tr, ω])") + ω <= size(hydro_data.scenario_inflows[r], 1) && tr <= size(hydro_data.scenario_inflows[r], 2) && + println(" exa transposed [ω, tr] = $(hydro_data.scenario_inflows[r][ω, tr])") + println(" exa row-offset [tr+1, ω] = $(hydro_data.scenario_inflows[r][mod1(tr + 1, n_inflow_rows), ω])") + println(" exa unit ratio (exa/ref) = $(hydro_data.scenario_inflows[r][tr, ω] / inflow_ref[t, r, s])") + println(" CONTINUING with the REFERENCE inflow values as authoritative scenario data.") +end + +# (3) Open-loop trajectory with the EXA policy AS-IS (state-threaded forward +# pass — see header item 6). Expected to MATCH MAIN exactly (max|Δ| = 0.0). +""" + open_loop_targets_asis(policy, x0, w_mat) -> Matrix{Float64} + +Open-loop target recursion `x̂_t = π(w_t, x̂_{t-1})`, `x̂_0 = x0`, using the EXA +policy exactly as its own training/eval pipeline calls it. The policy forward +pass now threads the LSTM recurrent state across stages (DecisionRules.jl +semantics), so this trajectory is expected to reproduce MAIN's exactly. + +# Arguments +- `policy`: EXA `HydroReachablePolicy`. +- `x0`: initial reservoir state (length nHyd). +- `w_mat`: `[T × nHyd]` inflow values. + +# Returns +- `[T × nHyd]` matrix of open-loop targets. +""" +function open_loop_targets_asis(policy, x0, w_mat) + Flux.reset!(policy) # REAL reset: start from initialstates + T, nH = size(w_mat) + prev = Float32.(x0) + out = zeros(Float64, T, nH) + for t in 1:T + target = policy(vcat(Float32.(w_mat[t, :]), prev)) + out[t, :] = Float64.(target) + prev = Float32.(target) # open-loop: previous target as next state + end + return out +end + +# (4) Open-loop trajectory with MANUALLY threaded recurrent state, replicating +# MAIN's `_step_encoder` semantics with the SAME loaded weights. Retained as an +# independent cross-check of (3): both are now expected to pass; if (3) fails +# while (4) passes, the policy forward pass has regressed from MAIN's +# threading semantics (weight parity still holds). +""" + open_loop_targets_threaded(policy, x0, w_mat) -> Matrix{Float64} + +Open-loop target recursion with the LSTM state carried across stages, +mirroring DecisionRules.jl's `HydroReachablePolicy` forward pass: + +```math +(h_t^{(l)}, c_t^{(l)}) = \\mathrm{LSTMCell}^{(l)}(h_t^{(l-1)}, (h_{t-1}^{(l)}, c_{t-1}^{(l)})), +\\qquad \\hat{x}_t = \\mathrm{lower} + (\\mathrm{upper} - \\mathrm{lower}) \\cdot \\sigma(\\mathrm{combiner}([h_t; \\hat{x}_{t-1}])) +``` + +followed by the same cascade clamp as the EXA forward pass. Uses +`layer.cell(x, state)` directly (Flux 0.16 `LSTMCell` returns +`(h, (h, c))`). + +# Arguments +- `policy`: EXA `HydroReachablePolicy` with MAIN weights loaded. +- `x0`: initial reservoir state. +- `w_mat`: `[T × nHyd]` inflow values. + +# Returns +- `[T × nHyd]` matrix of open-loop targets under MAIN's recurrence semantics. +""" +function open_loop_targets_threaded(policy, x0, w_mat) + core = policy_core(policy) + cells = [layer.cell for layer in core.encoder.layers] + # Zero initial recurrent state per layer, as Flux.initialstates gives. + states = Any[Flux.initialstates(c) for c in cells] + T, nH = size(w_mat) + prev = Float32.(x0) + out = zeros(Float64, T, nH) + for t in 1:T + w = Float32.(w_mat[t, :]) + # Thread the recurrent state layer by layer across stages. + h = stage_encoder_input(policy, t, w) + for (i, c) in enumerate(cells) + h, states[i] = c(h, states[i]) + end + # Same head + reachable-bounds scaling + cascade clamp as the EXA + # forward pass (these helpers come from hydro_reachable_policy.jl). + y = core.combiner(vcat(h, prev)) + lower, upper = _hydro_reachable_bounds(core, w, prev, y) + raw = lower .+ (upper .- lower) .* y + target = isempty(core.cascade) ? raw : + min.(raw, _cascade_upper_bounds(core, raw, w, prev)) + out[t, :] = Float64.(target) + prev = Float32.(target) + end + return out +end + +if MAIN_PARITY_GATES + w_s1 = inflow_ref[:, :, 1] # scenario 1 inflows [T × nHyd] + xhat_s1_ref = xhat_ref[:, :, 1] # MAIN open-loop trajectory + global xhat_s1_asis = open_loop_targets_asis(policy, x0_ref, w_s1) + global xhat_s1_threaded = open_loop_targets_threaded(policy, x0_ref, w_s1) + gate["openloop_asis_max_dev"] = _max_abs_dev(xhat_s1_asis, xhat_s1_ref) + gate["openloop_threaded_max_dev"] = _max_abs_dev(xhat_s1_threaded, xhat_s1_ref) + tol_traj = 1e-5 * max(1.0, maximum(abs.(xhat_s1_ref))) + gate["openloop_asis_pass"] = gate["openloop_asis_max_dev"] <= tol_traj + gate["openloop_threaded_pass"] = gate["openloop_threaded_max_dev"] <= tol_traj + println("GATE open-loop (EXA policy as-is): ", + gate["openloop_asis_pass"] ? "PASS" : "FAIL", + " (max abs dev = $(gate["openloop_asis_max_dev"]))") + println("GATE open-loop (state-threaded diag): ", + gate["openloop_threaded_pass"] ? "PASS" : "FAIL", + " (max abs dev = $(gate["openloop_threaded_max_dev"]))") +else + # The reference open-loop trajectories (`xhat_trajectories`) were rolled + # out by the MAIN reference checkpoint; comparing an independently trained + # EXA checkpoint's trajectory against them is meaningless (its targets + # SHOULD differ). Its own open-loop targets are still saved via the DE + # leg's de_target_trajectories. + global xhat_s1_asis = zeros(Float64, 0, 0) + global xhat_s1_threaded = zeros(Float64, 0, 0) + gate["openloop_asis_max_dev"] = NaN + gate["openloop_threaded_max_dev"] = NaN + gate["openloop_asis_pass"] = "skipped" + gate["openloop_threaded_pass"] = "skipped" + println("GATE open-loop trajectories: SKIPPED (DR_CHECKPOINT_KIND=exa — " * + "reference trajectories only characterize the MAIN reference checkpoint)") +end + +# ── Stage problem and callbacks (copied from train_hydro_exa_strict.jl) ─────── +# demand_matrix = nothing: the builder bakes in load_scaler × default demand for +# its single stage, matching the mof.json's 0.6-scaled loads at every stage. + +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = nothing, # CPU node — no GPU backend + float_type = Float64, + formulation = FORMULATION, + target_penalty = :auto, # irrelevant in strict mode + deficit_cost = DEFICIT_COST, # 6000 = MAIN mof.json coefficient + demand_matrix = nothing, + load_scaler = LOAD_SCALER, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, # header item 3 + ) +end +@info "Building 1-stage strict ExaModels stage problem (CPU, formulation=$FORMULATION)..." +rollout_prob = _build_rollout_de() + +# Same callback as train_hydro_exa_strict.jl's set_hydro_rollout_stage!, minus +# the demand update (demand_mat === nothing here). +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +# Realized state: in strict mode hydro_solution reads the (cascade-clamped) +# reservoir parameter trajectory, so the realized state equals the clamped +# target — the strict analogue of MAIN reading value(reservoir_out). +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +# Thermal generator positions, derived with the same hydro.json/PowerModels.json +# logic as MAIN's eval_paired_tsddr.jl lines 74-80: a generator is thermal iff +# its grid index is not any hydro unit's index_grid. +hydro_grid_idx = Set(power_data.gens[h.gen_pos].idx for h in hydro_data.units) +thermal_pos = [pos for (pos, g) in enumerate(power_data.gens) if !(g.idx in hydro_grid_idx)] +@info "Thermal generators: $(length(thermal_pos)) of $(power_data.nGen)" + +""" + volume_to_mw(volume; k = 0.0036) -> Float64 + +Convert a reservoir volume (hm³) to the MW-equivalent used by MAIN's +`eval_paired_tsddr.jl`: `volume / k` with `k = 0.0036`. +""" +volume_to_mw(volume; k = 0.0036) = volume / k + +""" + recompute_stage_costs(sol, power_data, deficit_cost) -> Vector{Float64} + +Recompute per-stage objective values from a `hydro_solution` NamedTuple: + +```math +c_t = \\sum_g (c_{2,g}\\, pg_{g,t}^2 + c_{1,g}\\, pg_{g,t}) + c_d \\sum_b \\mathrm{deficit}_{b,t} +``` + +which is the complete strict-mode EXA objective (no other terms exist). Used +both to split the full-horizon DE objective into stages and as an internal +consistency check on 1-stage solves. + +# Arguments +- `sol`: NamedTuple from [`hydro_solution`](@ref) (`pg` is `[nGen × T]`, + `deficit` is `[nBus × T]`). +- `power_data::PowerData`: generator cost coefficients. +- `deficit_cost::Real`: load-shedding cost per pu (6000 for this check). + +# Returns +- `Vector{Float64}` of length `T` with per-stage objective values. +""" +function recompute_stage_costs(sol, power_data, deficit_cost) + T = size(sol.pg, 2) + costs = zeros(Float64, T) + for t in 1:T + c = 0.0 + for (gpos, g) in enumerate(power_data.gens) + # Quadratic + linear generation cost (c2 = 0 for all Bolivia gens). + c += g.cost2 * sol.pg[gpos, t]^2 + g.cost1 * sol.pg[gpos, t] + end + # Per-bus active load shedding at the deficit cost. + c += deficit_cost * sum(@view sol.deficit[:, t]) + costs[t] = c + end + return costs +end + +# ── Stage-wise leg: closed-loop strict rollout on all paired scenarios ──────── +# Mirrors DecisionRulesExa.rollout_tsddr's fresh-solver path (reuse_solver = +# false, warmstart irrelevant for fresh solvers, retry_on_failure = true) while +# additionally extracting per-stage solution components that rollout_tsddr does +# not expose (pg, deficit, spill, outflow). A rollout_tsddr cross-check on +# scenario 1 verifies the custom loop matches the package machinery. + +@info "Stage-wise leg: $NUM_SCEN scenarios × $T_EVAL stages (cold MadNLP solves)..." + +costs_stagewise = fill(NaN, NUM_SCEN) # per-scenario total objective +stage_costs = fill(NaN, T_EVAL, NUM_SCEN) # per-stage objective +vol_trajectories = fill(NaN, T_EVAL, NUM_SCEN) # Σ_r volume_to_mw(state_r) after stage t +gen_trajectories = fill(NaN, T_EVAL, NUM_SCEN) # Σ_thermal pg × baseMVA at stage t +reservoir_traj = fill(NaN, nHyd, T_EVAL + 1, NUM_SCEN) # realized states incl. x0 +outflow_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) +spill_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) +target_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) # raw policy targets +max_deficit = fill(NaN, NUM_SCEN) # max_b,t deficit (pu) +sum_deficit = fill(NaN, NUM_SCEN) # Σ_b,t deficit (pu) +max_abs_deficit_q = fill(NaN, NUM_SCEN) # max_b,t |deficit_q| (pu) — EXA-only slack +scenario_ok = falses(NUM_SCEN) +n_retries_total = 0 +max_objective_recompute_dev = 0.0 # internal consistency check + +baseMVA = power_data.baseMVA + +for s in 1:NUM_SCEN + Flux.reset!(policy) # REAL reset at the scenario boundary + state = copy(x0_ref) # closed-loop realized state (Float64) + reservoir_traj[:, 1, s] = state + total = 0.0 + failed = false + for t in 1:T_EVAL + # Authoritative inflow values from the MAIN reference (see gate item 2). + w_t = inflow_ref[t, :, s] + # Policy call with Float32 inputs, exactly as MAIN's closed loop does. + x_hat = Float64.(policy(vcat(Float32.(w_t), Float32.(state)))) + target_traj[:, t, s] = x_hat + + # Write x_{t-1}, w_t, x̂_t into the strict stage problem. + set_hydro_rollout_stage!(rollout_prob, state, w_t, x_hat, t) + + # Cold one-shot solve (rollout_tsddr's fresh-solver path). + result = MadNLP.madnlp(rollout_prob.model; SOLVER_KWARGS...) + if !solve_succeeded(result) || !isfinite(result.objective) + # Single cold retry, mirroring rollout_tsddr's retry_on_failure. + global n_retries_total += 1 + result = MadNLP.madnlp(rollout_prob.model; SOLVER_KWARGS...) + end + if !solve_succeeded(result) || !isfinite(result.objective) + @warn "Stage solve failed after retry" scenario = s stage = t status = result.status + failed = true + break + end + + sol = hydro_solution(rollout_prob, result) + total += result.objective + stage_costs[t, s] = result.objective + # Internal consistency: the strict stage objective must equal the + # recomputed gen+deficit cost exactly (same terms, same data). + recomputed = recompute_stage_costs(sol, power_data, DEFICIT_COST)[1] + global max_objective_recompute_dev = + max(max_objective_recompute_dev, abs(recomputed - result.objective)) + + # Advance the closed loop on the realized state (= clamped target). + state = Float64.(sol.reservoir[:, end]) + reservoir_traj[:, t + 1, s] = state + outflow_traj[:, t, s] = sol.outflow[:, 1] + spill_traj[:, t, s] = sol.spill[:, 1] + + # Operative metrics with the IDENTICAL formulas as MAIN's evaluation. + vol_trajectories[t, s] = sum(volume_to_mw(state[r]) for r in 1:nHyd) + gen_trajectories[t, s] = sum(sol.pg[g, 1] * baseMVA for g in thermal_pos) + + # Deficit activity (quantifies whether cost/slack differences are inert). + def_max = maximum(sol.deficit[:, 1]) + def_sum = sum(sol.deficit[:, 1]) + dq_max = maximum(abs.(sol.deficit_q[:, 1])) + max_deficit[s] = isnan(max_deficit[s]) ? def_max : max(max_deficit[s], def_max) + sum_deficit[s] = isnan(sum_deficit[s]) ? def_sum : sum_deficit[s] + def_sum + max_abs_deficit_q[s] = isnan(max_abs_deficit_q[s]) ? dq_max : max(max_abs_deficit_q[s], dq_max) + end + if !failed + costs_stagewise[s] = total + scenario_ok[s] = true + end + if s % 10 == 0 || s == NUM_SCEN + ok_costs = costs_stagewise[1:s][scenario_ok[1:s]] + running_mean = isempty(ok_costs) ? NaN : mean(ok_costs) + println(" [$s/$NUM_SCEN] cost = $(round(total; digits=1)), running mean = $(round(running_mean; digits=1)), retries = $n_retries_total") + end +end +@info "Stage-wise leg done" n_ok = count(scenario_ok) n_retries_total max_objective_recompute_dev + +# ── rollout_tsddr cross-check on scenario 1 ─────────────────────────────────── +# Runs the actual package machinery (same callbacks) to verify the custom loop +# above reproduces it. Float32 initial state so the policy sees Float32 inputs; +# the solver interface buffers are Float64 regardless. +w_flat_s1 = vec(permutedims(inflow_ref[:, :, 1])) # stage-major [w_1; w_2; …] +rollout_check = rollout_tsddr( + policy, + Float32.(x0_ref), + rollout_prob, + Float32.(w_flat_s1); + horizon = T_EVAL, + n_uncertainty = nHyd, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + policy_state = :realized, + retry_on_failure = true, +) +rollout_check_obj = rollout_check === nothing ? NaN : rollout_check.objective +@info "rollout_tsddr cross-check (scenario 1)" custom_loop = costs_stagewise[1] rollout_tsddr = rollout_check_obj + +# ── DE leg: strict full-horizon deterministic equivalent ────────────────────── +# Strict-mode claim under test: with the target trajectory generated by the +# open-loop recursion (previous target as next policy state — exactly +# train_hydro_exa_strict.jl's rollout_reachable_targets), the T-stage strict DE +# objective equals the stage-wise sum for the same scenario, because the +# reservoir path is pinned to the same targets and stages decouple. + +@info "DE leg: building $T_EVAL-stage strict DE and solving $NUM_DE_SCENARIOS scenarios..." +de_prob = build_hydro_de(power_data, hydro_data, T_EVAL; + backend = nothing, + float_type = Float64, + formulation = FORMULATION, + target_penalty = :auto, + deficit_cost = DEFICIT_COST, + demand_matrix = nothing, + load_scaler = LOAD_SCALER, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, # header item 3 +) + +""" + rollout_reachable_targets(policy, x0, w_flat, T, nHyd) -> Vector{Float64} + +Roll out the reachable-policy target trajectory for the strict regular DE +(verbatim semantics of train_hydro_exa_strict.jl): start from the feasible +`x0`, feed `[w_t; previous_target]` to the policy, and record each target as +the next previous state, so every target is one-stage reachable from the prior +target by induction. + +# Arguments +- `policy`: reachable hydro policy with input `[inflow; previous_state]`. +- `x0`: initial reservoir state. +- `w_flat`: stage-major flat inflow vector of length `T * nHyd`. +- `T::Int`: number of stages. +- `nHyd::Int`: number of hydro reservoirs. + +# Returns +- `Vector{Float64}`: stage-major target trajectory for + `ExaModels.set_parameter!(prob.core, prob.p_target, targets)`. +""" +function rollout_reachable_targets(policy, x0, w_flat, T, nHyd) + Flux.reset!(policy) + prev = Float32.(x0) + targets = Vector{Vector{Float32}}(undef, T) + for t in 1:T + wt = Float32.(view(w_flat, ((t - 1) * nHyd + 1):(t * nHyd))) + target = policy(vcat(wt, prev)) + targets[t] = Float32.(target) + prev = targets[t] + end + return Float64.(vcat(targets...)) +end + +de_costs = fill(NaN, NUM_DE_SCENARIOS) +de_stage_costs = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_vol = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_gen = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_reservoir = fill(NaN, nHyd, T_EVAL + 1, NUM_DE_SCENARIOS) +de_outflow = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_spill = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_targets = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_max_deficit = fill(NaN, NUM_DE_SCENARIOS) +de_max_abs_deficit_q = fill(NaN, NUM_DE_SCENARIOS) +de_ok = falses(NUM_DE_SCENARIOS) + +for s in 1:NUM_DE_SCENARIOS + # Stage-major flat inflow vector for scenario s from the reference values. + w_flat = vec(permutedims(inflow_ref[:, :, s])) + # Open-loop target trajectory with the EXA policy (previous target as state). + targets = rollout_reachable_targets(policy, x0_ref, w_flat, T_EVAL, nHyd) + de_targets[:, :, s] = reshape(targets, nHyd, T_EVAL) + + # Write parameters and pin the strict reservoir path (prepare_solve! also + # applies the Float64 cascade clamp and enforces p_reservoir[1:nHyd] = x0). + ExaModels.set_parameter!(de_prob.core, de_prob.p_x0, x0_ref) + ExaModels.set_parameter!(de_prob.core, de_prob.p_inflow, w_flat) + ExaModels.set_parameter!(de_prob.core, de_prob.p_target, targets) + prepare_solve!(de_prob, x0_ref, w_flat, targets) + + result = MadNLP.madnlp(de_prob.model; SOLVER_KWARGS...) + if !solve_succeeded(result) || !isfinite(result.objective) + @warn "DE solve failed" scenario = s status = result.status + continue + end + de_ok[s] = true + de_costs[s] = result.objective + + sol = hydro_solution(de_prob, result) + de_stage_costs[:, s] = recompute_stage_costs(sol, power_data, DEFICIT_COST) + de_reservoir[:, :, s] = sol.reservoir + de_outflow[:, :, s] = sol.outflow + de_spill[:, :, s] = sol.spill + for t in 1:T_EVAL + # Operative metrics with the identical MAIN formulas, taken from the + # post-stage reservoir state and per-stage thermal generation. + de_vol[t, s] = sum(volume_to_mw(sol.reservoir[r, t + 1]) for r in 1:nHyd) + de_gen[t, s] = sum(sol.pg[g, t] * baseMVA for g in thermal_pos) + end + de_max_deficit[s] = maximum(sol.deficit) + de_max_abs_deficit_q[s] = maximum(abs.(sol.deficit_q)) + println(" DE [$s/$NUM_DE_SCENARIOS] objective = $(round(result.objective; digits=1)), stagewise total = $(round(costs_stagewise[s]; digits=1))") +end +@info "DE leg done" n_ok = count(de_ok) + +# ── Save everything ──────────────────────────────────────────────────────────── +out_dir = joinpath(SCRIPT_DIR, CASE_NAME, FORM_LABEL, "results") +mkpath(out_dir) +# DR_OUTPUT_TAG suffixes the filename so non-reference evaluations never +# clobber the untagged reference results file. +out_file = joinpath(out_dir, "paired_exa_strict$(OUT_SUFFIX).jld2") +# Provenance knobs, recorded whenever any of the new env vars was set (the +# all-defaults run keeps exactly the historical key set). +knob_extras = RECORD_KNOBS ? ( + checkpoint_path = MODEL_PATH, + checkpoint_kind = CHECKPOINT_KIND, + encoder_layers = ENCODER_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, + output_tag = OUTPUT_TAG, + main_parity_gates_applied = MAIN_PARITY_GATES, +) : (;) +jldsave(out_file; + knob_extras..., + # Gate results (policy parity, inflow indexing, metadata) + gate_loader_mode = loader_mode, + gate_probe_max_dev = gate["probe_max_dev"], + gate_probe_pass = gate["probe_pass"], + gate_probe_outputs_exa = probe_out_exa, + gate_inflow_recon_max_dev = gate["inflow_recon_max_dev"], + gate_inflow_pass = gate["inflow_pass"], + gate_openloop_asis_max_dev = gate["openloop_asis_max_dev"], + gate_openloop_asis_pass = gate["openloop_asis_pass"], + gate_openloop_threaded_max_dev = gate["openloop_threaded_max_dev"], + gate_openloop_threaded_pass = gate["openloop_threaded_pass"], + gate_dev_K = gate["dev_K"], + gate_dev_min_vol = gate["dev_min_vol"], + gate_dev_max_vol = gate["dev_max_vol"], + gate_dev_min_turn = gate["dev_min_turn"], + gate_dev_max_turn = gate["dev_max_turn"], + gate_dev_upstream_max = gate["dev_upstream_max"], + xhat_s1_asis = xhat_s1_asis, + xhat_s1_threaded = xhat_s1_threaded, + # Stage-wise leg + costs = costs_stagewise, + stage_costs = stage_costs, + vol_trajectories = vol_trajectories, + gen_trajectories = gen_trajectories, + reservoir_trajectories = reservoir_traj, + outflow_trajectories = outflow_traj, + spill_trajectories = spill_traj, + target_trajectories = target_traj, + max_deficit = max_deficit, + sum_deficit = sum_deficit, + max_abs_deficit_q = max_abs_deficit_q, + scenario_ok = collect(scenario_ok), + n_retries_total = n_retries_total, + max_objective_recompute_dev = max_objective_recompute_dev, + rollout_tsddr_check_objective = rollout_check_obj, + # DE leg + de_costs = de_costs, + de_stage_costs = de_stage_costs, + de_vol_trajectories = de_vol, + de_gen_trajectories = de_gen, + de_reservoir_trajectories = de_reservoir, + de_outflow_trajectories = de_outflow, + de_spill_trajectories = de_spill, + de_target_trajectories = de_targets, + de_max_deficit = de_max_deficit, + de_max_abs_deficit_q = de_max_abs_deficit_q, + de_ok = collect(de_ok), + # Configuration provenance + deficit_cost_used = DEFICIT_COST, + load_scaler_used = LOAD_SCALER, + reactive_deficit_setting = REACTIVE_DEFICIT_RAW, + reactive_deficit_cost_used = + REACTIVE_DEFICIT_COST === nothing ? "free" : Float64(REACTIVE_DEFICIT_COST), + solver_tol = SOLVER_KWARGS.tol, + solver_max_iter = SOLVER_KWARGS.max_iter, + model_path = MODEL_PATH, + reference_file = REFERENCE_FILE, + num_eval_stages = T_EVAL, + num_scenarios = NUM_SCEN, + num_de_scenarios = NUM_DE_SCENARIOS, +) +println("Saved: $out_file") + +# ── Final summary ───────────────────────────────────────────────────────────── + +""" + _split_csv_line(line) -> Vector{String} + +Split one CSV line into fields with minimal double-quote awareness: commas +inside `"…"` do not delimit (needed for the MAIN header column +`"TS-DDR (strict, paired)"`). No escape handling beyond quote toggling. +""" +function _split_csv_line(line::AbstractString) + fields = String[] # accumulated fields + buf = IOBuffer() # current field characters + inq = false # inside a quoted region? + for c in line + if c == '"' + inq = !inq # toggle quoting; quotes are dropped + elseif c == ',' && !inq + push!(fields, String(take!(buf))) # unquoted comma ends the field + else + write(buf, c) + end + end + push!(fields, String(take!(buf))) # trailing field + return fields +end + +""" + _mean_csv_column(path, needle) -> Float64 + +Mean of the numeric column whose header contains `needle` in the CSV at +`path`: + +```math +\\bar{c} = \\frac{1}{n} \\sum_{i=1}^{n} c_i +``` + +Returns `NaN` when the file is missing, the column is not found, or no row +parses — the summary then simply reports `NaN` for that baseline. +""" +function _mean_csv_column(path::AbstractString, needle::AbstractString) + isfile(path) || return NaN # ground truth absent + lines = readlines(path) + length(lines) >= 2 || return NaN # header + ≥1 data row + header = _split_csv_line(lines[1]) # quote-aware header parse + col = findfirst(h -> occursin(needle, h), header) # column by substring + col === nothing && return NaN + vals = Float64[] + for ln in lines[2:end] + fs = split(ln, ',') # data rows are plain numeric + length(fs) == length(header) || continue # skip malformed rows + v = tryparse(Float64, strip(fs[col])) + v === nothing || push!(vals, v) + end + return isempty(vals) ? NaN : mean(vals) +end + +# Successful-scenario cost vectors for the two legs of THIS evaluation. +ok_costs = costs_stagewise[collect(scenario_ok)] +de_ok_costs = de_costs[collect(de_ok)] + +# Untagged ground-truth baselines from the MAIN repo (comparability anchors): +# the .026 reference checkpoint's stage-wise mean (results/paired_strict_rollout.jld2) +# and the SDDP paired mean (SDDP-SOC column of paired_costs.csv). NaN if absent. +main_gt_file = joinpath(MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "results", "paired_strict_rollout.jld2") +main_gt_mean = isfile(main_gt_file) ? mean(Float64.(JLD2.load(main_gt_file, "costs"))) : NaN +sddp_mean = _mean_csv_column( + joinpath(MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "paired_costs.csv"), "SDDP") + +println("\n" * "=" ^ 64) +println("SUMMARY — paired EXA strict evaluation") +println(" Checkpoint: $MODEL_PATH") +println(" Kind: $CHECKPOINT_KIND (encoder=$(ENCODER_LAYERS), head=$(HEAD_LAYERS))") +println(" Context: $(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) (period=$CONTEXT_PERIOD, horizon=$CONTEXT_HORIZON)") +println(" Reference: $REFERENCE_FILE") +println(" Output tag: $(isempty(OUTPUT_TAG) ? "(none)" : OUTPUT_TAG)") +println(" Stage-wise mean: $(round(mean(ok_costs); digits=1)) over $(length(ok_costs))/$NUM_SCEN scenarios") +println(" Stage-wise std: $(round(std(ok_costs); digits=1))") +println(" DE-leg mean: $(round(mean(de_ok_costs); digits=1)) over $(length(de_ok_costs))/$NUM_DE_SCENARIOS scenarios") +println(" MAIN .026 mean (GT): $(round(main_gt_mean; digits=1))") +println(" SDDP mean (GT): $(round(sddp_mean; digits=1))") +println("=" ^ 64) diff --git a/examples/HydroPowerModels/hydro_power_exa.jl b/examples/HydroPowerModels/hydro_power_exa.jl index 40feee9..1df8d9d 100644 --- a/examples/HydroPowerModels/hydro_power_exa.jl +++ b/examples/HydroPowerModels/hydro_power_exa.jl @@ -15,12 +15,22 @@ # reservoir dynamics (initial condition, water balance, turbine coupling) # p_target[t,r] − reservoir[t+1,r] − δ[t,r] = 0 ← ADDED LAST # -# Target constraints are added last so result.multipliers[target_con_range] -# gives ∇_{x̂} Q directly (envelope theorem for policy gradient). +# Target constraints are added last in non-strict mode so +# target_multipliers(prob, result) gives ∇_{x̂} Q. In strict mode the +# reservoir trajectory is a parameter and target_multipliers transforms +# water-balance duals into target sensitivities. +# +# Strict-mode invariant: with reservoir a parameter, the initial condition +# reservoir[1,r] − p_x0[r] = 0 would be a parameter-only constraint (an +# all-zero Jacobian row), so the builders omit it and the initial condition is +# enforced purely by data: prepare_solve! must keep p_reservoir[1:nHyd] equal +# to the initial state x0 (it writes reservoir_vals = vcat(init, xhat)). Any +# code that updates p_reservoir must preserve p_reservoir[1:nH] == x0. using ExaModels using MadNLP using LinearAlgebra +import DecisionRulesExa: prepare_solve!, target_multipliers # ── Index helpers ───────────────────────────────────────────────────────────── # All arrays are flat, stage-major: index (t, i) → (t-1)*n + i @@ -30,6 +40,15 @@ using LinearAlgebra @inline _bri(nBR, t, br) = (t-1)*nBR + br # branch / pf @inline _ri(nH, t, r) = (t-1)*nH + r # hydro: reservoir / outflow / spill / delta +# ── 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 + # ── Problem struct ──────────────────────────────────────────────────────────── """ @@ -65,8 +84,22 @@ struct HydroExaDEProblem horizon::Int # formulation formulation::Symbol # :dc or :ac_polar - # range into result.multipliers for target constraints + # range into result.multipliers for target constraints (or water balance in strict mode) target_con_range::UnitRange{Int} + strict_targets::Bool + # strict mode: reservoir is a parameter, not a variable (length (T+1)*nHyd) + p_reservoir # nothing when !strict_targets + strict_reservoir_values::Vector{Float64} + # cascade data for target clamping (empty if no upstream connections) + cascade::Vector{CascadeLink} + K::Float64 + min_turn::Vector{Float64} + max_vol::Vector{Float64} + # Reactive-slack mode of the AC builder (:none for DC): + # :free — single FREE zero-cost deficit_q variable per bus/stage + # :penalized — deficit_q = δq⁺ − δq⁻ (δq± ≥ 0) with linear cost c·Σ(δq⁺+δq⁻) + # :hard — no reactive slack (hard reactive balance, MAIN-faithful) + reactive_deficit_mode::Symbol end # ── AC branch coefficient helper ────────────────────────────────────────────── @@ -146,7 +179,9 @@ end backend=nothing, float_type=Float64, formulation=:dc, target_penalty=:auto, target_penalty_l1=:auto, demand_matrix=nothing, - reactive_demand_matrix=nothing, deficit_cost=nothing) + reactive_demand_matrix=nothing, deficit_cost=nothing, + load_scaler=1.0, strict_targets=false, + reactive_deficit_cost=nothing) -> HydroExaDEProblem Build the T-stage hydro-power deterministic equivalent. @@ -170,6 +205,22 @@ Pass `:auto` (default) to use `2 × max_gen_cost`, matching JuMP's `penalty_l2 = `target_penalty_l1` sets the L1 coefficient for the `λ·|δ|` target slack penalty. Pass `:auto` (default) to use the same value as L2 ρ. Pass `nothing` to disable L1. The L1 term is reformulated as `λ·(δ⁺ + δ⁻)` with `δ = δ⁺ − δ⁻`, `δ⁺,δ⁻ ≥ 0`. + +`reactive_deficit_cost` (AC only) controls the reactive slack `deficit_q` in the +reactive KCL: +- `nothing` (default) — current behavior: one FREE, zero-cost `deficit_q` + variable per bus/stage (relaxes the reactive balance; model is byte-identical + to builds preceding this kwarg). +- finite `c ≥ 0` — `|deficit_q|` is penalized linearly: the slack is split into + `deficit_q = δq⁺ − δq⁻` with `δq⁺, δq⁻ ≥ 0` and objective `c·Σ(δq⁺ + δq⁻)`. + Constraint count and ORDER are unchanged (the split variables enter the same + reactive-KCL rows), so `target_con_range` is unaffected in both strict and + non-strict modes; only variable count changes (+T·nBus vs the default). +- `Inf` — omit `deficit_q` entirely: hard reactive balance, matching MAIN's + ACPPowerModel.mof.json (−T·nBus variables vs the default; constraint + count/order again unchanged). +Passing a non-`nothing` value with `formulation = :dc` throws (DC has no +reactive balance). """ function build_hydro_de(power_data::PowerData, hydro_data::HydroData, @@ -182,19 +233,24 @@ function build_hydro_de(power_data::PowerData, demand_matrix = nothing, reactive_demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 1.0, + strict_targets::Bool = false, + reactive_deficit_cost::Union{Nothing,Real} = nothing) formulation in (:dc, :ac_polar) || error("formulation must be :dc or :ac_polar, got :$formulation") if formulation === :dc + reactive_deficit_cost === nothing || + error("reactive_deficit_cost applies only to formulation = :ac_polar (DC has no reactive balance)") return _build_dc_hydro_de(power_data, hydro_data, T; backend=backend, float_type=float_type, target_penalty=target_penalty, target_penalty_l1=target_penalty_l1, demand_matrix=demand_matrix, deficit_cost=deficit_cost, - load_scaler=load_scaler) + load_scaler=load_scaler, + strict_targets=strict_targets) else return _build_ac_hydro_de(power_data, hydro_data, T; backend=backend, float_type=float_type, @@ -203,8 +259,30 @@ function build_hydro_de(power_data::PowerData, demand_matrix=demand_matrix, reactive_demand_matrix=reactive_demand_matrix, deficit_cost=deficit_cost, - load_scaler=load_scaler) + load_scaler=load_scaler, + strict_targets=strict_targets, + reactive_deficit_cost=reactive_deficit_cost) + end +end + +function _build_cascade_links(hydro_data::HydroData) + K = Float64(hydro_data.K) + spill_dests = Dict{Int,Set{Int}}() + for conn in hydro_data.upstream_spills + push!(get!(spill_dests, conn.upstream_pos, Set{Int}()), conn.downstream_pos) end + cascade = CascadeLink[] + for conn in hydro_data.upstream_turns + d, u = conn.downstream_pos, conn.upstream_pos + has_spill = haskey(spill_dests, u) && d in spill_dests[u] + push!(cascade, CascadeLink(d, u, !has_spill, Float32(K * hydro_data.units[u].max_turn))) + end + for conn in hydro_data.upstream_spills + d, u = conn.downstream_pos, conn.upstream_pos + already = any(c -> c.downstream == d && c.upstream == u, cascade) + already || push!(cascade, CascadeLink(d, u, false, Float32(K * hydro_data.units[u].max_turn))) + end + return cascade end # ── DC builder ──────────────────────────────────────────────────────────────── @@ -218,7 +296,8 @@ function _build_dc_hydro_de(power_data::PowerData, target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 1.0, + strict_targets::Bool = false) nBus = power_data.nBus nGen = power_data.nGen @@ -258,9 +337,17 @@ function _build_dc_hydro_de(power_data::PowerData, deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) # Reservoir levels: (T+1)*nHyd - res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) - res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) - reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + # In strict mode, reservoir = [x0; targets] is predetermined → parameter. + # Eliminates (T+1)*nHyd variables and nHyd + T*nHyd constraints. + if strict_targets + p_reservoir = ExaModels.parameter(core, zeros(float_type, (T+1) * nHyd)) + reservoir = p_reservoir + else + p_reservoir = nothing + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + end # Turbine outflow: T*nHyd out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) @@ -271,8 +358,10 @@ function _build_dc_hydro_de(power_data::PowerData, spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) # Target slack: δ = δ⁺ − δ⁻ with δ⁺,δ⁻ ≥ 0 (L1+L2 Lagrangian penalty) - delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) - delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + if !strict_targets + delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + end # ── Parameters ──────────────────────────────────────────────────────────── @@ -306,19 +395,22 @@ function _build_dc_hydro_de(power_data::PowerData, for item in def_cost_items ) - # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² - delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - ExaModels.objective(core, - p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 - for item in delta_items - ) + if !strict_targets + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - # L1 penalty: λ·(δ⁺ + δ⁻) - if use_l1 + # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² ExaModels.objective(core, - p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 for item in delta_items ) + + # L1 penalty: λ·(δ⁺ + δ⁻) + if use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end end # ── Constraints ─────────────────────────────────────────────────────────── @@ -391,15 +483,21 @@ function _build_dc_hydro_de(power_data::PowerData, ) n_con += T * nBus - # 5. Initial reservoir condition - ic_items = [(r = r,) for r in 1:nHyd] - ExaModels.constraint(core, - reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] - for item in ic_items - ) - n_con += nHyd + # 5. Initial reservoir condition (skip in strict: reservoir is a parameter, + # so reservoir[1,r] − p_x0[r] = 0 would be parameter-only — an all-zero + # Jacobian row. The initial condition is instead maintained by the invariant + # that prepare_solve! writes p_reservoir[1:nHyd] = x0; see file-top comment.) + if !strict_targets + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + end - # 6. Water balance + # 6. Water balance (reservoir is a parameter in strict mode — same expression) + wb_con_start = n_con + 1 wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), out_idx = _ri(nHyd, t, r), @@ -447,26 +545,40 @@ function _build_dc_hydro_de(power_data::PowerData, n_con += T * nHyd # ── TARGET CONSTRAINTS (ADDED LAST) ─────────────────────────────────────── - # x̂ − x − (δ⁺ − δ⁻) = 0 - target_items = [(param_idx = _ri(nHyd, t, r), - res_idx = _ri(nHyd, t+1, r), - delta_idx = _ri(nHyd, t, r)) - for t in 1:T for r in 1:nHyd] - ExaModels.constraint(core, - p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] - for item in target_items - ) - target_con_range = (n_con + 1):(n_con + T * nHyd) + if strict_targets + # Strict: reservoir is a parameter. target_multipliers transforms these + # water-balance duals into ∇_{x̂} Q by the adjacent-stage chain rule. + target_con_range = wb_con_start:(wb_con_start + T * nHyd - 1) + else + # x̂ − x − (δ⁺ − δ⁻) = 0 + target_items = [(param_idx = _ri(nHyd, t, r), + res_idx = _ri(nHyd, t+1, r), + delta_idx = _ri(nHyd, t, r)) + for t in 1:T for r in 1:nHyd] + ExaModels.constraint(core, + p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] + for item in target_items + ) + target_con_range = (n_con + 1):(n_con + T * nHyd) + end model = ExaModels.ExaModel(core) + cascade = _build_cascade_links(hydro_data) + return HydroExaDEProblem( core, model, p_demand, nothing, p_x0, p_inflow, p_target, p_penalty_half, Float64(ρ / 2), p_penalty_l1, Float64(ρ_l1), nHyd, nBus, nGen, nBranch, T, - :dc, target_con_range, + :dc, target_con_range, strict_targets, + p_reservoir, + strict_targets ? zeros(Float64, (T + 1) * nHyd) : Float64[], + cascade, Float64(K), + Float64.([h.min_turn for h in hydro_data.units]), + Float64.([h.max_vol for h in hydro_data.units]), + :none, # DC has no reactive balance, hence no reactive slack ) end @@ -482,7 +594,9 @@ function _build_ac_hydro_de(power_data::PowerData, demand_matrix = nothing, reactive_demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 1.0, + strict_targets::Bool = false, + reactive_deficit_cost::Union{Nothing,Real} = nothing) nBus = power_data.nBus nGen = power_data.nGen @@ -501,6 +615,20 @@ function _build_ac_hydro_de(power_data::PowerData, baseMVA = float_type(power_data.baseMVA) cd = float_type(deficit_cost !== nothing ? deficit_cost : power_data.cost_deficit) + # Reactive-slack mode (see build_hydro_de docstring): + # nothing → :free (default; byte-identical to builds preceding the kwarg), + # finite c ≥ 0 → :penalized (|deficit_q| costed linearly via δq⁺/δq⁻), + # Inf → :hard (no reactive slack — MAIN-faithful hard reactive balance). + rq_mode = if reactive_deficit_cost === nothing + :free + elseif isinf(Float64(reactive_deficit_cost)) + :hard + else + (isnan(Float64(reactive_deficit_cost)) || reactive_deficit_cost < 0) && + error("reactive_deficit_cost must be nothing, a finite cost ≥ 0, or Inf; got $reactive_deficit_cost") + :penalized + end + core = ExaModels.ExaCore(float_type; backend = backend) # ── Variables ───────────────────────────────────────────────────────────── @@ -538,13 +666,28 @@ function _build_ac_hydro_de(power_data::PowerData, # Active deficit (load shedding): T*nBus (non-negative) deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) - # Reactive deficit (free; allows reactive balance at zero cost) - deficit_q = ExaModels.variable(core, T * nBus) + # Reactive deficit — declared in place of the historical free variable so + # the variable ORDER of the :free mode is byte-identical to older builds. + if rq_mode === :free + # Free, zero-cost slack (relaxes the reactive balance). + deficit_q = ExaModels.variable(core, T * nBus) + elseif rq_mode === :penalized + # Split slack deficit_q = δq⁺ − δq⁻ with δq⁺, δq⁻ ≥ 0; the linear cost + # c·Σ(δq⁺ + δq⁻) is added in the objective section below. + deficit_q_pos = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + deficit_q_neg = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + end # :hard → no reactive slack variable at all # Reservoir levels: (T+1)*nHyd - res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) - res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) - reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + if strict_targets + p_reservoir = ExaModels.parameter(core, zeros(float_type, (T+1) * nHyd)) + reservoir = p_reservoir + else + p_reservoir = nothing + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + end # Turbine outflow: T*nHyd out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) @@ -555,8 +698,10 @@ function _build_ac_hydro_de(power_data::PowerData, spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) # Target slack: δ = δ⁺ − δ⁻ with δ⁺,δ⁻ ≥ 0 (L1+L2 Lagrangian penalty) - delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) - delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + if !strict_targets + delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + end # ── Parameters ──────────────────────────────────────────────────────────── @@ -601,22 +746,43 @@ function _build_ac_hydro_de(power_data::PowerData, for item in def_cost_items ) - # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² - delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - ExaModels.objective(core, - p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 - for item in delta_items - ) + # Linear reactive-slack penalty c·Σ(δq⁺ + δq⁻) = c·Σ|deficit_q| (only in + # :penalized mode; :free and :hard add no objective term here, keeping the + # default model byte-identical). + if rq_mode === :penalized + cq = float_type(reactive_deficit_cost) + rq_cost_items = [(idx = _bi(nBus, t, b), c = cq) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * (deficit_q_pos[item.idx] + deficit_q_neg[item.idx]) + for item in rq_cost_items + ) + end - # L1 penalty: λ·(δ⁺ + δ⁻) - if use_l1 + if !strict_targets + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + + # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² ExaModels.objective(core, - p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 for item in delta_items ) + + # L1 penalty: λ·(δ⁺ + δ⁻) + if use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end end # ── Constraints ─────────────────────────────────────────────────────────── + # NOTE (target_con_range audit): the reactive-slack modes differ ONLY in + # variables and objective terms. Constraint COUNT and ORDER are identical + # in all three modes (:penalized adds constraint! terms to EXISTING + # reactive-KCL rows; :hard omits the slack term from those same rows), so + # the n_con accounting below and the resulting target_con_range are + # unaffected in both strict and non-strict modes. n_con = 0 # 1. Reference angle: va[t, ref] = 0 @@ -764,21 +930,41 @@ function _build_ac_hydro_de(power_data::PowerData, item.brow => q_to[item.bcol] for item in kcl_qto_items ) - ExaModels.constraint!(core, c_kcl_q, - item.brow => -deficit_q[item.dcol] - for item in kcl_def_items - ) + # Reactive slack term: modifies the EXISTING reactive-KCL rows only — + # constraint count/order identical across all rq_mode values. + if rq_mode === :free + ExaModels.constraint!(core, c_kcl_q, + item.brow => -deficit_q[item.dcol] + for item in kcl_def_items + ) + elseif rq_mode === :penalized + # deficit_q = δq⁺ − δq⁻ enters the KCL as −δq⁺ + δq⁻. + ExaModels.constraint!(core, c_kcl_q, + item.brow => -deficit_q_pos[item.dcol] + for item in kcl_def_items + ) + ExaModels.constraint!(core, c_kcl_q, + item.brow => deficit_q_neg[item.dcol] + for item in kcl_def_items + ) + end # :hard → no slack term: hard reactive balance n_con += T * nBus - # 9. Initial reservoir condition - ic_items = [(r = r,) for r in 1:nHyd] - ExaModels.constraint(core, - reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] - for item in ic_items - ) - n_con += nHyd + # 9. Initial reservoir condition (skip in strict: reservoir is a parameter, + # so reservoir[1,r] − p_x0[r] = 0 would be parameter-only — an all-zero + # Jacobian row. The initial condition is instead maintained by the invariant + # that prepare_solve! writes p_reservoir[1:nHyd] = x0; see file-top comment.) + if !strict_targets + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + end - # 10. Water balance + # 10. Water balance (reservoir is a parameter in strict mode — same expression) + wb_con_start = n_con + 1 wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), out_idx = _ri(nHyd, t, r), @@ -826,26 +1012,40 @@ function _build_ac_hydro_de(power_data::PowerData, n_con += T * nHyd # ── TARGET CONSTRAINTS (ADDED LAST) ─────────────────────────────────────── - # x̂ − x − (δ⁺ − δ⁻) = 0 - target_items = [(param_idx = _ri(nHyd, t, r), - res_idx = _ri(nHyd, t+1, r), - delta_idx = _ri(nHyd, t, r)) - for t in 1:T for r in 1:nHyd] - ExaModels.constraint(core, - p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] - for item in target_items - ) - target_con_range = (n_con + 1):(n_con + T * nHyd) + if strict_targets + # Strict: reservoir is a parameter. target_multipliers transforms these + # water-balance duals into ∇_{x̂} Q by the adjacent-stage chain rule. + target_con_range = wb_con_start:(wb_con_start + T * nHyd - 1) + else + # x̂ − x − (δ⁺ − δ⁻) = 0 + target_items = [(param_idx = _ri(nHyd, t, r), + res_idx = _ri(nHyd, t+1, r), + delta_idx = _ri(nHyd, t, r)) + for t in 1:T for r in 1:nHyd] + ExaModels.constraint(core, + p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] + for item in target_items + ) + target_con_range = (n_con + 1):(n_con + T * nHyd) + end model = ExaModels.ExaModel(core) + cascade = _build_cascade_links(hydro_data) + return HydroExaDEProblem( core, model, p_demand, p_reactive_demand, p_x0, p_inflow, p_target, p_penalty_half, Float64(ρ / 2), p_penalty_l1, Float64(ρ_l1), nHyd, nBus, nGen, nBranch, T, - :ac_polar, target_con_range, + :ac_polar, target_con_range, strict_targets, + p_reservoir, + strict_targets ? zeros(Float64, (T + 1) * nHyd) : Float64[], + cascade, Float64(K), + Float64.([h.min_turn for h in hydro_data.units]), + Float64.([h.max_vol for h in hydro_data.units]), + rq_mode, ) end @@ -877,6 +1077,57 @@ function set_inflows!(prob::HydroExaDEProblem, w::AbstractVector) return prob end +function prepare_solve!(prob::HydroExaDEProblem, init_state, w_flat, xhat_flat) + if prob.p_reservoir !== nothing + nH = prob.nHyd + T = prob.horizon + init = Float64.(vec(Array(init_state))) + inflow = Float64.(vec(Array(w_flat))) + xhat = Float64.(vec(Array(xhat_flat))) + if !isempty(prob.cascade) + K = prob.K + for t in 1:T + x_prev = t == 1 ? init : view(xhat, (t-2)*nH+1:(t-1)*nH) + targets_t = view(xhat, (t-1)*nH+1:t*nH) + inflow_t = view(inflow, (t-1)*nH+1:t*nH) + for conn in prob.cascade + u, d = conn.upstream, conn.downstream + R_u = K * inflow_t[u] + x_prev[u] - targets_t[u] + if conn.turn_only + max_contrib = min(Float64(conn.K_max_turn), max(0.0, R_u)) + else + max_contrib = max(0.0, R_u) + end + true_upper = x_prev[d] + K * inflow_t[d] - K * prob.min_turn[d] + max_contrib + true_upper = min(prob.max_vol[d], true_upper) + targets_t[d] = min(targets_t[d], true_upper) + end + end + end + # Strict-mode invariant: p_reservoir[1:nH] must equal x0, because the + # builders omit the (parameter-only) initial-condition constraint in + # strict mode — see file-top comment. vcat(init, xhat) guarantees it. + reservoir_vals = vcat(init, xhat) + copyto!(prob.strict_reservoir_values, reservoir_vals) + ExaModels.set_parameter!(prob.core, prob.p_reservoir, reservoir_vals) + end + return nothing +end + +function target_multipliers(prob::HydroExaDEProblem, result) + λ = result.multipliers[prob.target_con_range] + prob.strict_targets || return λ + + nH = prob.nHyd + T = prob.horizon + raw = vec(Array(λ)) + out = copy(raw) + if T > 1 + out[1:(T - 1) * nH] .-= raw[(nH + 1):(T * nH)] + end + return out +end + # ── Post-processing ─────────────────────────────────────────────────────────── """ @@ -903,12 +1154,20 @@ function hydro_solution(prob::HydroExaDEProblem, result) pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG pf_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + if prob.strict_targets + res_sol = reshape(copy(prob.strict_reservoir_values), nH, T+1) + else + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + end out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - delta_sol = dp_sol .- dn_sol + if prob.strict_targets + delta_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + delta_sol = dp_sol .- dn_sol + end return (va=va_sol, pg=pg_sol, pf=pf_sol, deficit=def_sol, reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) else # :ac_polar @@ -921,13 +1180,33 @@ function hydro_solution(prob::HydroExaDEProblem, result) p_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR q_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - def_q_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + # Reactive slack layout depends on the builder's reactive_deficit_mode: + # :free — one free deficit_q block + # :penalized — δq⁺ then δq⁻ blocks; deficit_q = δq⁺ − δq⁻ + # :hard — no slack variable; report exact zeros + if prob.reactive_deficit_mode === :penalized + dqp_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + dqn_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + def_q_sol = dqp_sol .- dqn_sol + elseif prob.reactive_deficit_mode === :hard + def_q_sol = zeros(eltype(sol), nB, T) + else + def_q_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + end + if prob.strict_targets + res_sol = reshape(copy(prob.strict_reservoir_values), nH, T+1) + else + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + end out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - delta_sol = dp_sol .- dn_sol + if prob.strict_targets + delta_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + delta_sol = dp_sol .- dn_sol + end return (va=va_sol, vm=vm_sol, pg=pg_sol, qg=qg_sol, p_fr=p_fr_sol, q_fr=q_fr_sol, p_to=p_to_sol, q_to=q_to_sol, deficit=def_sol, deficit_q=def_q_sol, diff --git a/examples/HydroPowerModels/hydro_power_exa_embedded.jl b/examples/HydroPowerModels/hydro_power_exa_embedded.jl new file mode 100644 index 0000000..c23c934 --- /dev/null +++ b/examples/HydroPowerModels/hydro_power_exa_embedded.jl @@ -0,0 +1,1251 @@ +# hydro_power_exa_embedded.jl +# +# Embedded-NN deterministic equivalent for the Hydro power system. +# Target constraints from hydro_power_exa.jl are replaced by a +# VectorNonlinearOracle that evaluates the Flux policy inline. +# +# Slack target constraint (matching regular DE sign convention): +# π_θ(inflow_t, reservoir_t) − reservoir_{t+1,r} − δ⁺_{t,r} + δ⁻_{t,r} = 0 +# +# Strict target constraint: +# π_θ(inflow_t, reservoir_t) − reservoir_{t+1,r} = 0 +# +# Recurrent encoder handling: the per-stage encoder outputs h_t are cached in +# h_cache, computed by threading the recurrent state across stages +# (DecisionRulesExa._step_encoder — DecisionRules.jl semantics; Flux ≥ 0.16 +# cells are stateless so the Chain cannot be called directly per stage). Since +# the encoder reads only inflows, h_t depends on w_{1..t} and never on +# reservoir states; the cache is invalidated when inflows or policy parameters +# change (invalidate_policy_cache!/set_inflows!) and the oracle's analytic +# per-stage Jacobian ∂π_t/∂x_t through the combiner remains the FULL +# derivative — state threading adds no ∂h_t/∂x dependence. +# +# Depends on: hydro_power_data.jl, hydro_power_exa.jl, hydro_reachable_policy.jl + +using Flux +using LinearAlgebra: I +import DecisionRulesExa: set_x0!, set_uncertainty!, set_targets!, invalidate_policy_cache! + +# Keep legacy `include("hydro_power_exa_embedded.jl")` scripts working while +# letting training entrypoints include the policy explicitly. +if !isdefined(@__MODULE__, :HydroReachablePolicy) + include(joinpath(@__DIR__, "hydro_reachable_policy.jl")) +end + +# ── Problem struct ────────────────────────────────────────────────────────────── + +struct EmbeddedHydroExaDEProblem{P, VT <: AbstractVector{Float64}} + core + model + p_demand + p_reactive_demand + p_x0 + p_inflow + p_penalty_half + base_penalty_half::Float64 + p_penalty_l1 + base_penalty_l1::Float64 + policy::P + nHyd::Int + nx::Int + nw::Int + nBus::Int + nGen::Int + nBranch::Int + horizon::Int + formulation::Symbol + target_con_range::UnitRange{Int} + _res_start::Int + _dp_start::Int + _dn_start::Int + _nvar::Int + _inflow_buf::VT + _x0_buf::VT + _h_cache_dirty::Ref{Bool} + strict_targets::Bool +end + +# ── Interface (duck-typing for train_tsddr_embedded) ──────────────────────────── + +""" + set_x0!(prob::EmbeddedHydroExaDEProblem, x0) + +Set the initial reservoir state for the embedded hydro DE. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `x0::AbstractVector`: initial reservoir volumes, one value per hydro unit. + +# Returns +- `prob`. + +# Notes +The oracle closure reads `_x0_buf` directly when evaluating the first-stage +policy input, so this method updates both the ExaModels parameter and the +oracle-side buffer. +""" +function set_x0!(prob::EmbeddedHydroExaDEProblem, x0::AbstractVector) + length(x0) == prob.nHyd || error("x0 length must be nHyd=$(prob.nHyd)") + ExaModels.set_parameter!(prob.core, prob.p_x0, x0) + copyto!(prob._x0_buf, Float64.(x0)) + return prob +end + +""" + set_inflows!(prob::EmbeddedHydroExaDEProblem, w) + +Set the full inflow trajectory parameter. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `w::AbstractVector`: stage-major vector with length `prob.horizon * prob.nHyd`. + +# Returns +- `prob`. + +# Notes +Updating inflows invalidates cached recurrent encoder states because each stage +encoder input has changed. +""" +function set_inflows!(prob::EmbeddedHydroExaDEProblem, w::AbstractVector) + expected = prob.horizon * prob.nHyd + length(w) == expected || error("w must have length T*nHyd=$expected") + ExaModels.set_parameter!(prob.core, prob.p_inflow, w) + copyto!(prob._inflow_buf, Float64.(w)) + prob._h_cache_dirty[] = true + return prob +end + +""" + set_uncertainty!(prob::EmbeddedHydroExaDEProblem, w) + +Set the uncertainty trajectory for generic embedded training code. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `w::AbstractVector`: stage-major inflow trajectory. + +# Returns +- `prob`. + +# Notes +In the hydro example, uncertainty is exactly the inflow trajectory, so this +delegates to [`set_inflows!`](@ref). +""" +function set_uncertainty!(prob::EmbeddedHydroExaDEProblem, w::AbstractVector) + set_inflows!(prob, w) +end + +""" + set_targets!(prob::EmbeddedHydroExaDEProblem, target) + +Ignore external targets for embedded hydro DEs. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: ignored. +- `target::AbstractVector`: ignored. + +# Returns +- `nothing`. + +# Notes +Embedded hydro DEs generate targets inside the NLP through the policy oracle, so +there is no target parameter to update. +""" +function set_targets!(::EmbeddedHydroExaDEProblem, ::AbstractVector) + return nothing +end + +""" + invalidate_policy_cache!(prob::EmbeddedHydroExaDEProblem) + +Mark cached recurrent encoder states dirty after policy parameters change. The +oracle recomputes those states lazily on the next function/Jacobian evaluation. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. + +# Returns +- `prob`. +""" +function invalidate_policy_cache!(prob::EmbeddedHydroExaDEProblem) + prob._h_cache_dirty[] = true + return prob +end + +""" + set_demand!(prob::EmbeddedHydroExaDEProblem, demand_matrix) + +Update active power demand parameters. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `demand_matrix::AbstractMatrix`: `T x nBus` matrix of active demand values. + +# Returns +- `prob`. + +# Notes +The matrix is flattened in stage-major order to match the ExaModels parameter +layout. +""" +function set_demand!(prob::EmbeddedHydroExaDEProblem, demand_matrix::AbstractMatrix) + T, nB = size(demand_matrix) + T == prob.horizon || error("demand_matrix must have T=$(prob.horizon) rows") + nB == prob.nBus || error("demand_matrix must have nBus=$(prob.nBus) cols") + flat = [demand_matrix[t, b] for t in 1:T for b in 1:nB] + ExaModels.set_parameter!(prob.core, prob.p_demand, flat) + return prob +end + +""" + embedded_hydro_realized_states(prob, result) + +Extract realized reservoir states from an embedded hydro solve. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `result`: MadNLP result with a flat primal solution. + +# Returns +- A flat vector containing stages `1:T`, excluding the initial state. +""" +function embedded_hydro_realized_states(prob::EmbeddedHydroExaDEProblem, result) + T = prob.horizon + nH = prob.nHyd + sol = result.solution + return sol[prob._res_start + nH : prob._res_start + (T + 1) * nH - 1] +end + +""" + hydro_solution(prob::EmbeddedHydroExaDEProblem, result) -> NamedTuple + +Reshape the flat NLP solution into named physical blocks. This mirrors the +regular hydro DE post-processing helper so rollout diagnostics can read +reservoirs, generation, deficits, outflows, spills, and target slacks by name. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `result`: MadNLP result with a flat primal solution. + +# Returns +- A `NamedTuple` of solution arrays. + +# Notes +For strict targets, slack arrays are returned as zeros because no slack +variables are present in the strict NLP. +""" +function hydro_solution(prob::EmbeddedHydroExaDEProblem, result) + T = prob.horizon + nH = prob.nHyd + nG = prob.nGen + nBR = prob.nBranch + nB = prob.nBus + sol = result.solution + off = 0 + + if prob.formulation === :dc + va_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + pf_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + if prob.strict_targets + dp_sol = zeros(eltype(sol), nH, T) + dn_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + end + delta_sol = dp_sol .- dn_sol + return (va=va_sol, pg=pg_sol, pf=pf_sol, deficit=def_sol, + reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) + else + va_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + vm_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + qg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + p_fr_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + q_fr_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + p_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + q_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + def_q_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + if prob.strict_targets + dp_sol = zeros(eltype(sol), nH, T) + dn_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + end + delta_sol = dp_sol .- dn_sol + return (va=va_sol, vm=vm_sol, pg=pg_sol, qg=qg_sol, + p_fr=p_fr_sol, q_fr=q_fr_sol, p_to=p_to_sol, q_to=q_to_sol, + deficit=def_sol, deficit_q=def_q_sol, + reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) + end +end + +# ── Oracle builder helper ─────────────────────────────────────────────────────── + +function _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets::Bool = false) + + _res(s) = res_start + (s-1)*nHyd : res_start + s*nHyd - 1 + _dp(t) = dp_start + (t-1)*nHyd : dp_start + t*nHyd - 1 + _dn(t) = dn_start + (t-1)*nHyd : dn_start + t*nHyd - 1 + _inflow(t) = (t-1)*nHyd+1 : t*nHyd + _crow(t) = (t-1)*nHyd+1 : t*nHyd + + encoder = policy.encoder + combiner = policy.combiner + combiner isa Flux.Dense || + throw(ArgumentError("embedded hydro DE currently requires the default single Dense reachable-policy head; use combiner_layers=Int[] or regular strict DE for multilayer state heads")) + n_h = size(combiner.weight, 2) - policy.n_state + reachable_policy = policy isa HydroReachablePolicy + has_spill_cap = reachable_policy && policy.spill_max !== nothing + + jac_r = Int[] + jac_c = Int[] + for t in 1:T, r in 1:nHyd + row = (t-1)*nHyd + r + push!(jac_r, row); push!(jac_c, res_start + t*nHyd + r - 1) + if !strict_targets + push!(jac_r, row); push!(jac_c, dp_start + (t-1)*nHyd + r - 1) + push!(jac_r, row); push!(jac_c, dn_start + (t-1)*nHyd + r - 1) + end + if t > 1 + for j in 1:nHyd + push!(jac_r, row); push!(jac_c, res_start + (t-1)*nHyd + j - 1) + end + end + end + nnzj = length(jac_r) + + const_jac_cpu = zeros(Float64, nnzj) + nn_jac_ranges_flat = Vector{UnitRange{Int}}(undef, T * nHyd) + k = 0 + for t in 1:T, r in 1:nHyd + const_jac_cpu[k+1] = -1.0 + k += 1 + if !strict_targets + const_jac_cpu[k+1] = -1.0 + const_jac_cpu[k+2] = 1.0 + k += 2 + end + if t > 1 + nn_jac_ranges_flat[(t-1)*nHyd + r] = (k+1):(k+nHyd) + k += nHyd + end + end + const_jac_dev = similar(inflow_buf, Float64, nnzj) + copyto!(const_jac_dev, const_jac_cpu) + + W_state = @view combiner.weight[:, n_h+1:end] + act = combiner.σ + has_output_bounds = getfield(policy, :output_lower) !== nothing + output_lower_f32 = similar(x0_buf, Float32, nHyd) + output_scale_f32 = similar(x0_buf, Float32, nHyd) + if has_output_bounds + output_lower_f32 .= getfield(policy, :output_lower) + output_scale_f32 .= getfield(policy, :output_scale) + else + fill!(output_lower_f32, 0f0) + fill!(output_scale_f32, 1f0) + end + min_vol_f32 = similar(x0_buf, Float32, nHyd) + max_vol_f32 = similar(x0_buf, Float32, nHyd) + min_turn_f32 = similar(x0_buf, Float32, nHyd) + max_turn_f32 = similar(x0_buf, Float32, nHyd) + spill_max_f32 = similar(x0_buf, Float32, nHyd) + upstream_max_f32 = similar(x0_buf, Float32, nHyd) + if reachable_policy + min_vol_f32 .= policy.min_vol + max_vol_f32 .= policy.max_vol + min_turn_f32 .= policy.min_turn + max_turn_f32 .= policy.max_turn + upstream_max_f32 .= policy.upstream_max_inflow + if policy.spill_max !== nothing + spill_max_f32 .= policy.spill_max + else + fill!(spill_max_f32, Float32(Inf)) + end + else + fill!(upstream_max_f32, 0f0) + end + + function _act_deriv!(σ_prime, output) + if act === NNlib.sigmoid || act === NNlib.sigmoid_fast + σ_prime .= output .* (one(eltype(output)) .- output) + elseif act === Base.tanh || act === NNlib.tanh_fast + σ_prime .= one(eltype(output)) .- output .* output + elseif act === identity + fill!(σ_prime, one(eltype(σ_prime))) + else + σ_prime .= output .* (one(eltype(output)) .- output) + end + return σ_prime + end + + function _activation!(dest, z) + if act === NNlib.sigmoid || act === NNlib.sigmoid_fast + dest .= inv.(one(eltype(dest)) .+ exp.(-z)) + elseif act === Base.tanh || act === NNlib.tanh_fast + dest .= tanh.(z) + elseif act === identity + dest .= z + else + dest .= act.(z) + end + return dest + end + + # Pre-allocated buffers — all on the same device as x0_buf + x0_f32 = similar(x0_buf, Float32, nHyd) + x_prev_f32 = similar(x0_buf, Float32, nHyd) + infl_f32 = similar(x0_buf, Float32, nHyd) + _comb_in = similar(x0_buf, Float32, n_h + nHyd) + z_buf = similar(x0_buf, Float32, nHyd) + nn_out_f32 = similar(x0_buf, Float32, nHyd) + y_norm_f32 = similar(x0_buf, Float32, nHyd) + lower_f32 = similar(x0_buf, Float32, nHyd) + upper_f32 = similar(x0_buf, Float32, nHyd) + scale_f32 = similar(x0_buf, Float32, nHyd) + dlower_dx_f32 = similar(x0_buf, Float32, nHyd) + dupper_dx_f32 = similar(x0_buf, Float32, nHyd) + dbound_dx_f32 = similar(x0_buf, Float32, nHyd) + σ_prime_buf = similar(x0_buf, Float32, nHyd) + λ_f32_buf = similar(x0_buf, Float32, nHyd) + d_xprev_buf = similar(x0_buf, Float32, nHyd) + J_buf = similar(x0_buf, Float64, nHyd, nHyd) + dbound_dx_f64 = similar(x0_buf, Float64, nHyd) + diag_mask = similar(x0_buf, Float64, nHyd, nHyd) + copyto!(diag_mask, Matrix{Float64}(I, nHyd, nHyd)) + h_cache = similar(x0_buf, Float32, n_h, T) + h_cache_dirty = Ref(true) + + function _reachable_bounds!(inflow, x_prev) + K32 = Float32(policy.K) + upper_f32 .= x_prev .+ K32 .* inflow .- K32 .* min_turn_f32 .+ upstream_max_f32 + dupper_dx_f32 .= ifelse.(upper_f32 .<= max_vol_f32, 1f0, 0f0) + upper_f32 .= min.(max_vol_f32, upper_f32) + + if has_spill_cap + lower_f32 .= x_prev .+ K32 .* inflow .- K32 .* max_turn_f32 .- spill_max_f32 + dlower_dx_f32 .= ifelse.(lower_f32 .>= min_vol_f32, 1f0, 0f0) + lower_f32 .= max.(min_vol_f32, lower_f32) + else + lower_f32 .= min_vol_f32 + fill!(dlower_dx_f32, 0f0) + end + + dupper_dx_f32 .= ifelse.(upper_f32 .< lower_f32, dlower_dx_f32, dupper_dx_f32) + upper_f32 .= max.(upper_f32, lower_f32) + scale_f32 .= upper_f32 .- lower_f32 + return nothing + end + + function _combiner_fwd!(nn_out, h, x_prev) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev) + nn_out .= lower_f32 .+ scale_f32 .* y_norm_f32 + else + nn_out .= output_lower_f32 .+ output_scale_f32 .* y_norm_f32 + end + return nn_out + end + + function _combiner_jac!(J_buf, σ_prime_buf, h, x_prev) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + _act_deriv!(σ_prime_buf, y_norm_f32) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev) + σ_prime_buf .*= scale_f32 + else + σ_prime_buf .*= output_scale_f32 + end + J_buf .= reshape(σ_prime_buf, :, 1) .* W_state + if reachable_policy + dbound_dx_f32 .= dlower_dx_f32 .+ y_norm_f32 .* (dupper_dx_f32 .- dlower_dx_f32) + dbound_dx_f64 .= dbound_dx_f32 + J_buf .+= reshape(dbound_dx_f64, :, 1) .* diag_mask + end + return nothing + end + + function _populate_h_cache!() + # Thread the recurrent encoder state across stages explicitly. Flux + # ≥ 0.16 cells are stateless — calling the Chain of LSTM wrappers + # directly (`encoder(x)`) would restart from `initialstates` on every + # stage, making the encoder memoryless. `_step_encoder` advances the + # underlying cells one stage at a time from a fresh initial state, + # exactly like DecisionRules.jl's `_step_encoder` and the threaded + # policy forward pass. h_t therefore depends on the inflow history + # w_{1..t} only (never on reservoir states), which is what makes this + # per-stage cache — and the oracle's per-stage direct Jacobian — + # exact for the embedded policy. + enc_state = DecisionRulesExa._init_recurrent_state(encoder) + for t in 1:T + infl_f32 .= view(inflow_buf, _inflow(t)) + h, enc_state = DecisionRulesExa._step_encoder(encoder, infl_f32, enc_state) + view(h_cache, :, t) .= h + end + h_cache_dirty[] = false + return nothing + end + + function _ensure_h_cache!() + h_cache_dirty[] && _populate_h_cache!() + return nothing + end + + function oracle_f!(c, xv) + _ensure_h_cache!() + x0_f32 .= view(x0_buf, 1:nHyd) + for t in 1:T + infl_f32 .= view(inflow_buf, _inflow(t)) + if t == 1 + x_prev_f32 .= x0_f32 + else + x_prev_f32 .= view(xv, _res(t)) + end + h = view(h_cache, :, t) + _combiner_fwd!(nn_out_f32, h, x_prev_f32) + + c_t = view(c, _crow(t)) + r_v = view(xv, _res(t+1)) + if strict_targets + c_t .= nn_out_f32 .- r_v + else + dp_v = view(xv, _dp(t)) + dn_v = view(xv, _dn(t)) + c_t .= nn_out_f32 .- r_v .- dp_v .+ dn_v + end + end + return nothing + end + + function oracle_jac!(vals, xv) + _ensure_h_cache!() + copyto!(vals, const_jac_dev) + for t in 2:T + infl_f32 .= view(inflow_buf, _inflow(t)) + x_prev_f32 .= view(xv, _res(t)) + h = view(h_cache, :, t) + _combiner_jac!(J_buf, σ_prime_buf, h, x_prev_f32) + for r in 1:nHyd + vals[nn_jac_ranges_flat[(t-1)*nHyd + r]] .= view(J_buf, r, :) + end + end + return nothing + end + + function oracle_vjp!(Jtv, xv, λ) + _ensure_h_cache!() + fill!(Jtv, 0.0) + for t in 1:T + λ_f64 = view(λ, _crow(t)) + view(Jtv, _res(t+1)) .-= λ_f64 + if !strict_targets + view(Jtv, _dp(t)) .-= λ_f64 + view(Jtv, _dn(t)) .+= λ_f64 + end + if t > 1 + h = view(h_cache, :, t) + infl_f32 .= view(inflow_buf, _inflow(t)) + x_prev_f32 .= view(xv, _res(t)) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev_f32 + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + _act_deriv!(σ_prime_buf, y_norm_f32) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev_f32) + σ_prime_buf .*= scale_f32 + dbound_dx_f32 .= dlower_dx_f32 .+ + y_norm_f32 .* (dupper_dx_f32 .- dlower_dx_f32) + else + σ_prime_buf .*= output_scale_f32 + fill!(dbound_dx_f32, 0f0) + end + λ_f32_buf .= λ_f64 + σ_prime_buf .*= λ_f32_buf + mul!(d_xprev_buf, W_state', σ_prime_buf) + if reachable_policy + d_xprev_buf .+= dbound_dx_f32 .* λ_f32_buf + end + view(Jtv, _res(t)) .+= d_xprev_buf + end + end + return nothing + end + + oracle = ExaModels.VectorNonlinearOracle( + nvar = nvar_total, + ncon = T * nHyd, + nnzj = nnzj, + jac_rows = jac_r, + jac_cols = jac_c, + lcon = zeros(T * nHyd), + ucon = zeros(T * nHyd), + f! = oracle_f!, + jac! = oracle_jac!, + vjp! = oracle_vjp!, + ) + return oracle, h_cache_dirty +end + +# ── DC builder ────────────────────────────────────────────────────────────────── + +function _build_embedded_dc_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + target_penalty::Union{Real,Symbol} = :auto, + target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, + demand_matrix = nothing, + deficit_cost::Union{Nothing,Real} = nothing, + load_scaler::Real = 1.0, + strict_targets::Bool = false, +) + nBus = power_data.nBus + nGen = power_data.nGen + nBranch = power_data.nBranch + nHyd = hydro_data.nHyd + K = float_type(hydro_data.K) + ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) + ρ_l1 = if target_penalty_l1 === :auto + ρ + elseif target_penalty_l1 === nothing + zero(float_type) + else + float_type(target_penalty_l1) + end + use_l1 = ρ_l1 > 0 + baseMVA = float_type(power_data.baseMVA) + cd = float_type(deficit_cost !== nothing ? deficit_cost : power_data.cost_deficit) + + core = ExaModels.ExaCore(float_type; backend = backend) + + # ── Variables (track offsets) ───────────────────────────────────────────── + var_offset = 0 + + va = ExaModels.variable(core, T * nBus) + var_offset += T * nBus + + pg_lb = float_type.(repeat([g.pmin for g in power_data.gens], T)) + pg_ub = float_type.(repeat([g.pmax for g in power_data.gens], T)) + pg = ExaModels.variable(core, T * nGen; lvar = pg_lb, uvar = pg_ub) + var_offset += T * nGen + + pf_lb = float_type.(repeat([-b.rate_a for b in power_data.branches], T)) + pf_ub = float_type.(repeat([ b.rate_a for b in power_data.branches], T)) + pf = ExaModels.variable(core, T * nBranch; lvar = pf_lb, uvar = pf_ub) + var_offset += T * nBranch + + deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + var_offset += T * nBus + + res_start = var_offset + 1 + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + var_offset += (T+1) * nHyd + + out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) + out_ub = float_type.(repeat([h.max_turn for h in hydro_data.units], T)) + outflow = ExaModels.variable(core, T * nHyd; lvar = out_lb, uvar = out_ub) + var_offset += T * nHyd + + spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += T * nHyd + + dp_start = strict_targets ? 0 : var_offset + 1 + delta_pos = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + dn_start = strict_targets ? 0 : var_offset + 1 + delta_neg = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + nvar_total = var_offset + + # ── Parameters ──────────────────────────────────────────────────────────── + init_demand = if demand_matrix !== nothing + float_type.(load_scaler .* [demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_demand, T)) + end + + p_demand = ExaModels.parameter(core, init_demand) + p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) + p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) + p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) + + # ── Objective ───────────────────────────────────────────────────────────── + gen_cost_items = [(t = t, g = g_pos, + c1 = float_type(g.cost1), c2 = float_type(g.cost2)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.objective(core, + item.c2 * pg[_gi(nGen, item.t, item.g)]^2 + + item.c1 * pg[_gi(nGen, item.t, item.g)] + for item in gen_cost_items + ) + + def_cost_items = [(t = t, b = b, c = cd) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * deficit[_bi(nBus, item.t, item.b)] + for item in def_cost_items + ) + + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + if !strict_targets + ExaModels.objective(core, + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 + for item in delta_items + ) + end + + if !strict_targets && use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end + + # ── Constraints 1–7 ────────────────────────────────────────────────────── + n_con = 0 + + nRef = length(power_data.ref_buses) + ref_items = [(t = t, ref = ref) for t in 1:T for ref in power_data.ref_buses] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.ref)] + for item in ref_items + ) + n_con += T * nRef + + ohm_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + b = float_type(br.b)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + item.b * (va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - pf[_bri(nBranch, item.t, item.br)] + for item in ohm_items + ) + n_con += T * nBranch + + ang_lb = float_type.(repeat([br.angmin for br in power_data.branches], T)) + ang_ub = float_type.(repeat([br.angmax for br in power_data.branches], T)) + ang_items = [(t = t, f = br.f_bus, tb = br.t_bus) + for t in 1:T for br in power_data.branches] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)] + for item in ang_items; + lcon = ang_lb, ucon = ang_ub, + ) + n_con += T * nBranch + + kcl_init_items = [(t = t, b = b) for t in 1:T for b in 1:nBus] + c_kcl = ExaModels.constraint(core, + p_demand[_bi(nBus, item.t, item.b)] + for item in kcl_init_items + ) + kcl_gen_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl, + item.brow => -pg[item.gcol] + for item in kcl_gen_items + ) + kcl_fr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl, + item.brow => pf[item.bcol] + for item in kcl_fr_items + ) + kcl_to_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl, + item.brow => -pf[item.bcol] + for item in kcl_to_items + ) + kcl_def_items = [(t = t, brow = _bi(nBus, t, b), dcol = _bi(nBus, t, b)) + for t in 1:T for b in 1:nBus] + ExaModels.constraint!(core, c_kcl, + item.brow => -deficit[item.dcol] + for item in kcl_def_items + ) + n_con += T * nBus + + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + + wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), + out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), + inflow_p = _ri(nHyd, t, r), K = K) + for t in 1:T for r in 1:nHyd] + c_wb = ExaModels.constraint(core, + reservoir[item.res_next] - reservoir[item.res_curr] + + item.K * outflow[item.out_idx] + + spill[item.spill_idx] + - item.K * p_inflow[item.inflow_p] + for item in wb_items + ) + if !isempty(hydro_data.upstream_turns) + wb_turn_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos), K = K) + for t in 1:T for conn in hydro_data.upstream_turns] + ExaModels.constraint!(core, c_wb, + item.row => -item.K * outflow[item.col] + for item in wb_turn_items + ) + end + if !isempty(hydro_data.upstream_spills) + wb_spill_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos)) + for t in 1:T for conn in hydro_data.upstream_spills] + ExaModels.constraint!(core, c_wb, + item.row => -spill[item.col] + for item in wb_spill_items + ) + end + n_con += T * nHyd + + tc_items = [(gen_col = _gi(nGen, t, h.gen_pos), out_col = _ri(nHyd, t, h.pos), + baseMVA = baseMVA, pf_h = float_type(h.pf)) + for t in 1:T for h in hydro_data.units] + ExaModels.constraint(core, + item.baseMVA * pg[item.gen_col] - item.pf_h * outflow[item.out_col] + for item in tc_items + ) + n_con += T * nHyd + + # ── Oracle (replaces target constraints) ────────────────────────────────── + inflow_buf = backend === nothing ? + zeros(Float64, T * nHyd) : + KernelAbstractions.zeros(backend, Float64, T * nHyd) + x0_buf = backend === nothing ? + zeros(Float64, nHyd) : + KernelAbstractions.zeros(backend, Float64, nHyd) + + oracle, h_cache_dirty = _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets = strict_targets) + ExaModels.constraint(core, oracle) + target_con_range = (n_con + 1):(n_con + T * nHyd) + + model = ExaModels.ExaModel(core) + + return EmbeddedHydroExaDEProblem( + core, model, + p_demand, nothing, p_x0, p_inflow, + p_penalty_half, Float64(ρ / 2), + p_penalty_l1, Float64(ρ_l1), + policy, nHyd, nHyd, nHyd, + nBus, nGen, nBranch, T, + :dc, target_con_range, + res_start, dp_start, dn_start, nvar_total, + inflow_buf, x0_buf, h_cache_dirty, + strict_targets, + ) +end + +# ── AC polar builder ──────────────────────────────────────────────────────────── + +function _build_embedded_ac_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + target_penalty::Union{Real,Symbol} = :auto, + target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, + demand_matrix = nothing, + reactive_demand_matrix = nothing, + deficit_cost::Union{Nothing,Real} = nothing, + load_scaler::Real = 1.0, + strict_targets::Bool = false, +) + nBus = power_data.nBus + nGen = power_data.nGen + nBranch = power_data.nBranch + nHyd = hydro_data.nHyd + K = float_type(hydro_data.K) + ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) + ρ_l1 = if target_penalty_l1 === :auto + ρ + elseif target_penalty_l1 === nothing + zero(float_type) + else + float_type(target_penalty_l1) + end + use_l1 = ρ_l1 > 0 + baseMVA = float_type(power_data.baseMVA) + cd = float_type(deficit_cost !== nothing ? deficit_cost : power_data.cost_deficit) + + core = ExaModels.ExaCore(float_type; backend = backend) + + # ── Variables ───────────────────────────────────────────────────────────── + var_offset = 0 + + va = ExaModels.variable(core, T * nBus) + var_offset += T * nBus + + vm_lb = float_type.(repeat([b.vmin for b in power_data.buses], T)) + vm_ub = float_type.(repeat([b.vmax for b in power_data.buses], T)) + vm = ExaModels.variable(core, T * nBus; lvar = vm_lb, uvar = vm_ub, + start = ones(float_type, T * nBus)) + var_offset += T * nBus + + pg_lb = float_type.(repeat([g.pmin for g in power_data.gens], T)) + pg_ub = float_type.(repeat([g.pmax for g in power_data.gens], T)) + pg = ExaModels.variable(core, T * nGen; lvar = pg_lb, uvar = pg_ub) + var_offset += T * nGen + + qg_lb = float_type.(repeat([isfinite(g.qmin) ? g.qmin : -1e4 for g in power_data.gens], T)) + qg_ub = float_type.(repeat([isfinite(g.qmax) ? g.qmax : 1e4 for g in power_data.gens], T)) + qg = ExaModels.variable(core, T * nGen; lvar = qg_lb, uvar = qg_ub) + var_offset += T * nGen + + p_fr_lb = float_type.(repeat([-b.rate_a for b in power_data.branches], T)) + p_fr_ub = float_type.(repeat([ b.rate_a for b in power_data.branches], T)) + p_fr = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + q_fr = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + p_to = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + q_to = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + var_offset += 4 * T * nBranch + + deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + var_offset += T * nBus + + deficit_q = ExaModels.variable(core, T * nBus) + var_offset += T * nBus + + res_start = var_offset + 1 + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + var_offset += (T+1) * nHyd + + out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) + out_ub = float_type.(repeat([h.max_turn for h in hydro_data.units], T)) + outflow = ExaModels.variable(core, T * nHyd; lvar = out_lb, uvar = out_ub) + var_offset += T * nHyd + + spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += T * nHyd + + dp_start = strict_targets ? 0 : var_offset + 1 + delta_pos = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + dn_start = strict_targets ? 0 : var_offset + 1 + delta_neg = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + nvar_total = var_offset + + # ── Parameters ──────────────────────────────────────────────────────────── + init_demand = if demand_matrix !== nothing + float_type.(load_scaler .* [demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_demand, T)) + end + init_reactive_demand = if reactive_demand_matrix !== nothing + float_type.(load_scaler .* [reactive_demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_reactive_demand, T)) + end + + p_demand = ExaModels.parameter(core, init_demand) + p_reactive_demand = ExaModels.parameter(core, init_reactive_demand) + p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) + p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) + p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) + + br_ac = [_ac_branch_coeffs(br, float_type) for br in power_data.branches] + + # ── Objective ───────────────────────────────────────────────────────────── + gen_cost_items = [(t = t, g = g_pos, + c1 = float_type(g.cost1), c2 = float_type(g.cost2)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.objective(core, + item.c2 * pg[_gi(nGen, item.t, item.g)]^2 + + item.c1 * pg[_gi(nGen, item.t, item.g)] + for item in gen_cost_items + ) + + def_cost_items = [(t = t, b = b, c = cd) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * deficit[_bi(nBus, item.t, item.b)] + for item in def_cost_items + ) + + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + if !strict_targets + ExaModels.objective(core, + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 + for item in delta_items + ) + end + if !strict_targets && use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end + + # ── Constraints ─────────────────────────────────────────────────────────── + n_con = 0 + + # 1. Reference angle + nRef = length(power_data.ref_buses) + ref_items = [(t = t, ref = ref) for t in 1:T for ref in power_data.ref_buses] + ExaModels.constraint(core, va[_bi(nBus, item.t, item.ref)] for item in ref_items) + n_con += T * nRef + + # 2. AC from-end active power flow + pfr_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c3 = br_ac[br_pos].c3, c4 = br_ac[br_pos].c4, c5 = br_ac[br_pos].c5) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + p_fr[_bri(nBranch, item.t, item.br)] + - item.c5 * vm[_bi(nBus, item.t, item.f)]^2 + - item.c3 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * cos(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - item.c4 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * sin(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + for item in pfr_items + ) + n_con += T * nBranch + + # 3. AC from-end reactive power flow + qfr_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c3 = br_ac[br_pos].c3, c4 = br_ac[br_pos].c4, c6 = br_ac[br_pos].c6) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + q_fr[_bri(nBranch, item.t, item.br)] + + item.c6 * vm[_bi(nBus, item.t, item.f)]^2 + + item.c4 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * cos(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - item.c3 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * sin(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + for item in qfr_items + ) + n_con += T * nBranch + + # 4. AC to-end active power flow + pto_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c1 = br_ac[br_pos].c1, c2 = br_ac[br_pos].c2, c7 = br_ac[br_pos].c7) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + p_to[_bri(nBranch, item.t, item.br)] + - item.c7 * vm[_bi(nBus, item.t, item.tb)]^2 + - item.c1 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * cos(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + - item.c2 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * sin(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + for item in pto_items + ) + n_con += T * nBranch + + # 5. AC to-end reactive power flow + qto_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c1 = br_ac[br_pos].c1, c2 = br_ac[br_pos].c2, c8 = br_ac[br_pos].c8) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + q_to[_bri(nBranch, item.t, item.br)] + + item.c8 * vm[_bi(nBus, item.t, item.tb)]^2 + + item.c2 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * cos(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + - item.c1 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * sin(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + for item in qto_items + ) + n_con += T * nBranch + + # 6. Phase angle difference + ang_lb = float_type.(repeat([br.angmin for br in power_data.branches], T)) + ang_ub = float_type.(repeat([br.angmax for br in power_data.branches], T)) + ang_items = [(t = t, f = br.f_bus, tb = br.t_bus) for t in 1:T for br in power_data.branches] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)] + for item in ang_items; + lcon = ang_lb, ucon = ang_ub, + ) + n_con += T * nBranch + + # 7. Active KCL + kcl_p_init = [(t = t, b = b, gs = float_type(power_data.buses[b].gs)) + for t in 1:T for b in 1:nBus] + c_kcl_p = ExaModels.constraint(core, + p_demand[_bi(nBus, item.t, item.b)] + + item.gs * vm[_bi(nBus, item.t, item.b)]^2 + for item in kcl_p_init + ) + kcl_pg_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => -pg[item.gcol] for item in kcl_pg_items) + kcl_pfr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => p_fr[item.bcol] for item in kcl_pfr_items) + kcl_pto_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => p_to[item.bcol] for item in kcl_pto_items) + kcl_def_items = [(t = t, brow = _bi(nBus, t, b), dcol = _bi(nBus, t, b)) + for t in 1:T for b in 1:nBus] + ExaModels.constraint!(core, c_kcl_p, + item.brow => -deficit[item.dcol] for item in kcl_def_items) + n_con += T * nBus + + # 8. Reactive KCL + kcl_q_init = [(t = t, b = b, bs = float_type(power_data.buses[b].bs)) + for t in 1:T for b in 1:nBus] + c_kcl_q = ExaModels.constraint(core, + p_reactive_demand[_bi(nBus, item.t, item.b)] + - item.bs * vm[_bi(nBus, item.t, item.b)]^2 + for item in kcl_q_init + ) + kcl_qg_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => -qg[item.gcol] for item in kcl_qg_items) + kcl_qfr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => q_fr[item.bcol] for item in kcl_qfr_items) + kcl_qto_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => q_to[item.bcol] for item in kcl_qto_items) + ExaModels.constraint!(core, c_kcl_q, + item.brow => -deficit_q[item.dcol] for item in kcl_def_items) + n_con += T * nBus + + # 9. Initial reservoir + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + + # 10. Water balance + wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), + out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), + inflow_p = _ri(nHyd, t, r), K = K) + for t in 1:T for r in 1:nHyd] + c_wb = ExaModels.constraint(core, + reservoir[item.res_next] - reservoir[item.res_curr] + + item.K * outflow[item.out_idx] + + spill[item.spill_idx] + - item.K * p_inflow[item.inflow_p] + for item in wb_items + ) + if !isempty(hydro_data.upstream_turns) + wb_turn_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos), K = K) + for t in 1:T for conn in hydro_data.upstream_turns] + ExaModels.constraint!(core, c_wb, + item.row => -item.K * outflow[item.col] for item in wb_turn_items) + end + if !isempty(hydro_data.upstream_spills) + wb_spill_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos)) + for t in 1:T for conn in hydro_data.upstream_spills] + ExaModels.constraint!(core, c_wb, + item.row => -spill[item.col] for item in wb_spill_items) + end + n_con += T * nHyd + + # 11. Turbine coupling + tc_items = [(gen_col = _gi(nGen, t, h.gen_pos), out_col = _ri(nHyd, t, h.pos), + baseMVA = baseMVA, pf_h = float_type(h.pf)) + for t in 1:T for h in hydro_data.units] + ExaModels.constraint(core, + item.baseMVA * pg[item.gen_col] - item.pf_h * outflow[item.out_col] + for item in tc_items + ) + n_con += T * nHyd + + # ── Oracle ──────────────────────────────────────────────────────────────── + inflow_buf = backend === nothing ? + zeros(Float64, T * nHyd) : + KernelAbstractions.zeros(backend, Float64, T * nHyd) + x0_buf = backend === nothing ? + zeros(Float64, nHyd) : + KernelAbstractions.zeros(backend, Float64, nHyd) + + oracle, h_cache_dirty = _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets = strict_targets) + ExaModels.constraint(core, oracle) + target_con_range = (n_con + 1):(n_con + T * nHyd) + + model = ExaModels.ExaModel(core) + + return EmbeddedHydroExaDEProblem( + core, model, + p_demand, p_reactive_demand, p_x0, p_inflow, + p_penalty_half, Float64(ρ / 2), + p_penalty_l1, Float64(ρ_l1), + policy, nHyd, nHyd, nHyd, + nBus, nGen, nBranch, T, + :ac_polar, target_con_range, + res_start, dp_start, dn_start, nvar_total, + inflow_buf, x0_buf, h_cache_dirty, + strict_targets, + ) +end + +# ── Dispatcher ────────────────────────────────────────────────────────────────── + +function build_embedded_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + formulation::Symbol = :dc, + kwargs..., +) + formulation in (:dc, :ac_polar) || + error("formulation must be :dc or :ac_polar, got :$formulation") + + if formulation === :dc + return _build_embedded_dc_hydro_de(policy, power_data, hydro_data, T; kwargs...) + else + return _build_embedded_ac_hydro_de(policy, power_data, hydro_data, T; kwargs...) + end +end diff --git a/examples/HydroPowerModels/hydro_reachable_policy.jl b/examples/HydroPowerModels/hydro_reachable_policy.jl new file mode 100644 index 0000000..ee57be4 --- /dev/null +++ b/examples/HydroPowerModels/hydro_reachable_policy.jl @@ -0,0 +1,469 @@ +# hydro_reachable_policy.jl +# +# Hydro-specific feasible target policy for strict regular and embedded DEs. +# This file owns the policy architecture and reachability/cascade bounds; the +# ExaModels problem builders live in hydro_power_exa.jl and +# hydro_power_exa_embedded.jl. + +using Flux +using Zygote +import DecisionRulesExa: load_stateconditioned_policy! + +""" + hydro_reachable_policy(hydro_data, layers; activation=sigmoid, encoder_type=Flux.LSTM, + spill_max=nothing, combiner_layers=Int[]) + +Build a state-conditioned hydro policy whose outputs are one-stage reachable +reservoir targets. + +The recurrent encoder reads inflows. The nonrecurrent combiner reads +`[encoded_inflow; reservoir_state]`, emits normalized targets in `[0, 1]`, and +the wrapper maps them into the reachability interval implied by the current +inflow and previous reservoir state. + +# Arguments +- `hydro_data::HydroData`: hydro limits, initial volumes, stage duration, and + cascade metadata. +- `layers::AbstractVector{Int}`: recurrent encoder widths over inflows. + +# Keywords +- `activation`: activation used by the target head. It must be sigmoid-style so + normalized targets remain in `[0, 1]`. +- `encoder_type`: Flux recurrent layer constructor, usually `Flux.LSTM`. +- `spill_max`: optional finite spill cap used to tighten lower bounds. +- `combiner_layers`: hidden widths in the nonrecurrent state-conditioned head. + +# Returns +- `HydroReachablePolicy`: a Flux-compatible policy with trainable encoder and + combiner parameters plus fixed hydro reachability metadata. + +# Notes +For strict regular deterministic equivalents, rolling this policy from the true +initial state and feeding each previous target into the next policy call gives a +feasible target path by induction. Embedded strict DEs use realized reservoir +states inside the NLP. Cascade links are handled by clamping downstream targets +to the upper bound implied by same-stage upstream targets. + +The recurrent encoder state is threaded across stages explicitly (Flux ≥ 0.16 +cells are stateless, so calling the `LSTM` wrapper directly would restart from +`initialstates` every stage): the forward pass advances `state` by one cell +step per call, mirroring DecisionRules.jl's `HydroReachablePolicy` exactly. +Call `Flux.reset!(policy)` at scenario boundaries to restore the initial state. +""" +mutable struct HydroReachablePolicy{E,C,RS,V,S,I} + encoder::E + combiner::C + state::RS # Encoder recurrent state, threaded across stages + n_context::Int # Number of context dimensions prepended before inflow + n_uncertainty::Int + n_state::Int + min_vol::V + max_vol::V + min_turn::V + max_turn::V + spill_max::S + upstream_max_inflow::V + K::Float64 + output_lower::Nothing + output_scale::Nothing + cascade::Vector{CascadeLink} + cascade_upstream::I + cascade_downstream::I + cascade_turn_only::V + cascade_k_max_turn::V + cascade_reservoir_ids::I +end + +Flux.@layer HydroReachablePolicy trainable=(encoder, combiner) + +""" + _hydro_adapt_bound(x, ref) + +Return `x` as a vector with the same array family and element type as `ref`. + +# Arguments +- `x::AbstractVector`: hydro metadata stored on the policy. +- `ref::AbstractVector`: vector whose device family and element type should be + matched. + +# Returns +- A vector with `length(x)` whose storage is compatible with `ref`. + +# Notes +This keeps metadata such as `min_vol` and `max_vol` on the same device as the +policy forward pass: CPU inputs stay on CPU, GPU inputs stay on GPU. +""" +function _hydro_adapt_bound(x::AbstractVector, ref::AbstractVector) + typeof(x) === typeof(ref) && return x + y = similar(ref, length(x)) + copyto!(y, convert.(eltype(ref), x)) + return y +end + +""" + _hydro_adapt_index(x, ref) + +Return integer indices stored on the same device family as `ref`. + +# Arguments +- `x::AbstractVector`: integer indices stored in ordinary Julia memory. +- `ref::AbstractVector`: vector whose device family should be matched. + +# Returns +- An integer vector compatible with `ref`. + +# Notes +This helper is used before GPU gathers such as `inflow[upstream]`; copying to +`similar(ref, Int, ...)` avoids host-side scalar indexing during policy +evaluation. +""" +function _hydro_adapt_index(x::AbstractVector, ref::AbstractVector) + y = similar(ref, Int, length(x)) + copyto!(y, x) + return y +end + +""" + _hydro_reachable_bounds(policy, inflow, x_prev, ref) -> (lower, upper) + +Compute one-stage reservoir target bounds for the hydro water balance. + +# Arguments +- `policy::HydroReachablePolicy`: policy carrying fixed hydro metadata. +- `inflow`: current-stage inflow vector. +- `x_prev`: previous reservoir-state vector. +- `ref`: vector used to choose element type and device family. + +# Returns +- `(lower, upper)`: vectors defining the closed interval into which normalized + policy outputs are mapped. + +# Notes +For each reservoir, the simplified balance is + +```text +x_next = x_prev + K * inflow - K * turbine_out - spill + upstream_contrib. +``` + +The upper bound uses the smallest required turbine outflow (`min_turn`) and zero +spill. The lower bound is the physical minimum volume unless `spill_max` is +finite; with finite spill, the lowest reachable storage uses `max_turn` and +maximum spill. + +Proof sketch: every feasible turbine/spill choice satisfies the same balance +equation and variable bounds. Substituting extremal admissible outflow/spill +values gives an interval containing all one-stage reachable reservoir states. +Mapping a sigmoid output into this interval therefore produces a reachable +target in the relaxed one-stage balance. +""" +function _hydro_reachable_bounds(policy::HydroReachablePolicy, inflow, x_prev, ref) + min_vol = _hydro_adapt_bound(policy.min_vol, ref) + max_vol = _hydro_adapt_bound(policy.max_vol, ref) + min_turn = _hydro_adapt_bound(policy.min_turn, ref) + max_turn = _hydro_adapt_bound(policy.max_turn, ref) + upstream = _hydro_adapt_bound(policy.upstream_max_inflow, ref) + K = convert(eltype(ref), policy.K) + + upper_raw = x_prev .+ K .* inflow .- K .* min_turn .+ upstream + upper = min.(max_vol, upper_raw) + + lower = if policy.spill_max === nothing + min_vol + else + spill_max = _hydro_adapt_bound(policy.spill_max, ref) + lower_raw = x_prev .+ K .* inflow .- K .* max_turn .- spill_max + max.(min_vol, lower_raw) + end + + upper = max.(upper, lower) + return lower, upper +end +Zygote.@nograd _hydro_reachable_bounds + +""" + _cascade_upper_bounds(policy, target, inflow, x_prev) -> upper + +Tighten downstream reservoir upper bounds using same-stage upstream targets. + +# Arguments +- `policy::HydroReachablePolicy`: policy carrying cascade metadata. +- `target`: raw target vector before cascade clamping. +- `inflow`: current-stage inflow vector. +- `x_prev`: previous reservoir-state vector. + +# Returns +- A vector of per-reservoir upper bounds induced by incoming cascade links. + +# Notes +For a cascade link `u -> d`, the upstream target determines the upstream +release implied by the balance: + +```text +release_u = K * inflow_u + x_prev_u - target_u. +``` + +Proof sketch: downstream storage is increasing in upstream contribution. The +largest physically consistent contribution from an upstream target is exactly +the positive release implied by that target, possibly turbine-capped. Therefore +clamping the downstream target to this bound cannot remove any feasible target +that respects the upstream target and cascade balance. + +# Assumptions +- **Single-level cascades.** The implied release + `release_u = K * inflow_u + x_prev_u - target_u` omits the upstream unit's + own incoming cascade contribution: if `u` itself receives water from a unit + further upstream, its true release can be larger than computed here. The + omission is conservative — it can only under-estimate the available + downstream contribution, so the clamp is never infeasible, merely possibly + over-tight for multi-level chains. Bolivia's three links + (COR→SIS turbine-only, ZON→CHU turbine+spill, TAQ1→TAQ2 turbine+spill) are + all single-level, so the bound is exact for that case. +- **No gradient through the clamp.** `Zygote.@nograd` on this function means + that when the cascade clamp binds, the true dependence of the downstream + target on the upstream target carries no gradient — a deliberate + approximation that keeps the policy pullback cheap and well-defined. +- **Physically infeasible edge case.** If the cascade upper bound falls below + the reachable lower bound of `_hydro_reachable_bounds`, the clamped target + can fall below `lower`. There is no policy-level remedy: the stage is + genuinely infeasible, and downstream slack/deficit handling must absorb it. +""" +function _cascade_upper_bounds(policy::HydroReachablePolicy, target, inflow, x_prev) + cascade = policy.cascade + T = eltype(target) + K = T(policy.K) + n = length(target) + isempty(cascade) && return _hydro_adapt_bound(fill(T(Inf), n), target) + + upstream = _hydro_adapt_index(policy.cascade_upstream, target) + downstream = _hydro_adapt_index(policy.cascade_downstream, target) + reservoir_ids = _hydro_adapt_index(policy.cascade_reservoir_ids, target) + turn_only = _hydro_adapt_bound(policy.cascade_turn_only, target) + k_max_turn = _hydro_adapt_bound(policy.cascade_k_max_turn, target) + min_turn = _hydro_adapt_bound(policy.min_turn, target) + max_vol = _hydro_adapt_bound(policy.max_vol, target) + + # Release implied by asking the upstream reservoir to end at `target`. + release = K .* inflow[upstream] .+ x_prev[upstream] .- target[upstream] + positive_release = max.(zero(T), release) + max_contrib = ifelse.(turn_only .> zero(T), min.(k_max_turn, positive_release), positive_release) + + # One upper bound per cascade link, expressed for the downstream reservoir. + link_upper = x_prev[downstream] .+ + K .* inflow[downstream] .- + K .* min_turn[downstream] .+ + max_contrib + link_upper = min.(max_vol[downstream], link_upper) + + # Reduce link-wise bounds to one bound per reservoir. Reservoirs without + # incoming cascade links receive Inf and are left unchanged by `min`. + link_by_reservoir = ifelse.( + reshape(downstream, :, 1) .== reshape(reservoir_ids, 1, :), + reshape(link_upper, :, 1), + T(Inf), + ) + return vec(minimum(link_by_reservoir; dims = 1)) +end +Zygote.@nograd _cascade_upper_bounds + +""" + (policy::HydroReachablePolicy)(input) -> target + +Evaluate the reachable hydro policy. + +# Arguments +- `input`: concatenated vector `[context_t; inflow_t; x_{t-1}]`. + +# Returns +- A reservoir target vector in the one-stage reachable set. + +# Notes +The encoder reads `[context_t; inflow_t]`, threading its recurrent state across +calls (one cell step per stage, stored in `policy.state` — DecisionRules.jl +semantics). Reachability bounds and cascade clamps slice out only the true +inflow entries, so prepended context never changes physical feasibility logic. +The combiner reads both encoded inflow/context and previous reservoir state, +emits normalized targets, and those targets are mapped into the reachability +interval before cascade clamping. + +The cascade clamp inherits the assumptions documented on +[`_cascade_upper_bounds`](@ref): + +- Cascades are treated as single-level — the upstream release omits that unit's + own incoming cascade contribution, which is conservative (never infeasible, + possibly over-tight) on multi-level chains. Bolivia's three links + (COR→SIS turn-only, ZON→CHU turn+spill, TAQ1→TAQ2 turn+spill) are all + single-level. +- Both `_hydro_reachable_bounds` and `_cascade_upper_bounds` are + `Zygote.@nograd`, so when the cascade clamp binds, the downstream target's + true dependence on the upstream target carries no gradient (deliberate + approximation). +- In the physically infeasible edge case where the cascade upper bound is + below the reachable lower bound, the returned target can fall below `lower`; + the stage is genuinely infeasible and no policy-level remedy exists. +""" +function (m::HydroReachablePolicy)(input) + # Split input: optional context first, then true inflow, then previous state. + # Physical reachability 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 + inflow = input[w_start:w_end] + x_prev = input[w_end+1:end] + encoder_input = if c_end == 0 + inflow + else + vcat(input[1:c_end], inflow) + end + + # Encode inflow through the recurrent encoder, threading state across calls + # (mirrors DecisionRules.jl: encoded, s_t = _step_encoder(enc, T.(w_t), s_{t-1})). + # Cast to encoder precision for type stability (avoids Zygote codegen bugs). + T = DecisionRulesExa._state_eltype(m.state) + h, new_state = DecisionRulesExa._step_encoder(m.encoder, T.(encoder_input), m.state) + # Thread the recurrent state to the next call. + m.state = new_state + + y = m.combiner(vcat(h, x_prev)) + lower, upper = _hydro_reachable_bounds(m, inflow, x_prev, y) + # Because `y` is sigmoid-bounded, this affine map stays in [lower, upper]. + raw_target = lower .+ (upper .- lower) .* y + 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!(policy::HydroReachablePolicy) -> Nothing + +Reset the inflow encoder's recurrent state to `Flux.initialstates`, e.g. at +scenario boundaries. + +# Notes +The combiner is feed-forward and does not carry recurrent state. The recurrent +state is re-derived from the (possibly device-moved) encoder weights on every +reset, so the state always matches the encoder's device and element type. +""" +function Flux.reset!(m::HydroReachablePolicy) + # Reinitialize the recurrent state from the encoder's initial states. + m.state = DecisionRulesExa._init_recurrent_state(m.encoder) + return nothing +end + +""" + load_stateconditioned_policy!(policy::HydroReachablePolicy, state) + +Load Flux parameters into a reachable hydro policy. + +# Arguments +- `policy::HydroReachablePolicy`: target policy to update in place. +- `state`: checkpoint object accepted by `Flux.loadmodel!`. + +# Returns +- `policy`. + +# Notes +- If the checkpoint contains only `encoder` and `combiner` fields, those + trainable parts are loaded while hydro reachability metadata from the current + case is preserved. +- Checkpoints trained BEFORE recurrent-state threading (memoryless-encoder era) + have the SAME weight structure and load unchanged — only the runtime + semantics differ (the encoder now carries memory across stages). +- DecisionRules.jl (MAIN) checkpoints save the encoder as a `Chain` of BARE + `LSTMCell`s (layer state `(Wi, Wh, bias)` instead of `(cell = …,)`); those + load through the documented cell-by-cell fallback in + `DecisionRulesExa._load_encoder_state!`. +- The recurrent state is reset after loading so the next rollout starts from + `Flux.initialstates` of the loaded weights. +""" +function load_stateconditioned_policy!(policy::HydroReachablePolicy, state) + try + Flux.loadmodel!(policy, state) + Flux.reset!(policy) + return policy + catch err + hasproperty(state, :encoder) && hasproperty(state, :combiner) || rethrow(err) + @warn "Full HydroReachablePolicy checkpoint load failed; loading encoder/combiner only and keeping hydro reachability bounds" exception=(err, catch_backtrace()) + DecisionRulesExa._load_encoder_state!(policy.encoder, getproperty(state, :encoder)) + Flux.loadmodel!(policy.combiner, getproperty(state, :combiner)) + Flux.reset!(policy) + return policy + end +end + +function hydro_reachable_policy( + hydro_data::HydroData, + layers::AbstractVector{Int}; + activation = sigmoid, + encoder_type = Flux.LSTM, + spill_max = nothing, + combiner_layers = Int[], + n_context::Int = 0, +) + (activation === sigmoid || activation === NNlib.sigmoid || activation === NNlib.sigmoid_fast) || + throw(ArgumentError("hydro_reachable_policy requires a sigmoid-style activation so normalized targets stay in [0, 1]")) + nHyd = hydro_data.nHyd + n_context >= 0 || throw(ArgumentError("n_context must be nonnegative")) + enc_sizes = vcat(nHyd + n_context, layers) + enc_layers = [encoder_type(enc_sizes[i] => enc_sizes[i+1]) + for i in 1:length(layers)] + encoder = Flux.Chain(enc_layers...) + encoder_width = isempty(layers) ? nHyd + n_context : layers[end] + combiner = DecisionRulesExa._dense_policy_head( + encoder_width + nHyd, + nHyd, + collect(Int, combiner_layers); + activation = activation, + ) + spill_vec = spill_max === nothing ? nothing : Float32.(collect(spill_max)) + if spill_vec !== nothing && length(spill_vec) != nHyd + throw(ArgumentError("spill_max length must be nHyd=$nHyd")) + end + + K = Float64(hydro_data.K) + upstream_max = zeros(Float32, nHyd) + for conn in hydro_data.upstream_turns + upstream_max[conn.downstream_pos] += Float32(K * hydro_data.units[conn.upstream_pos].max_turn) + end + + spill_dests = Dict{Int,Set{Int}}() + for conn in hydro_data.upstream_spills + push!(get!(spill_dests, conn.upstream_pos, Set{Int}()), conn.downstream_pos) + end + cascade = CascadeLink[] + for conn in hydro_data.upstream_turns + d, u = conn.downstream_pos, conn.upstream_pos + has_spill = haskey(spill_dests, u) && d in spill_dests[u] + push!(cascade, CascadeLink(d, u, !has_spill, Float32(K * hydro_data.units[u].max_turn))) + end + for conn in hydro_data.upstream_spills + d, u = conn.downstream_pos, conn.upstream_pos + already = any(c -> c.downstream == d && c.upstream == u, cascade) + already || push!(cascade, CascadeLink(d, u, false, Float32(K * hydro_data.units[u].max_turn))) + end + + return HydroReachablePolicy( + encoder, combiner, + DecisionRulesExa._init_recurrent_state(encoder), # initial recurrent state + n_context, nHyd, nHyd, + Float32.([h.min_vol for h in hydro_data.units]), + Float32.([h.max_vol for h in hydro_data.units]), + Float32.([h.min_turn for h in hydro_data.units]), + Float32.([h.max_turn for h in hydro_data.units]), + spill_vec, + upstream_max, + K, + nothing, nothing, + cascade, + # Typed comprehensions guarantee concrete Vector{Int}/Vector{Float32} + # element types: `getfield.(links, :field)` can infer as Vector{Real} + # (the field symbol does not always constant-propagate through fused + # broadcast), which breaks the struct's I/V type-parameter binding. + Int[c.upstream for c in cascade], + Int[c.downstream for c in cascade], + Float32[c.turn_only for c in cascade], + Float32[c.K_max_turn for c in cascade], + collect(1:nHyd), + ) +end diff --git a/examples/HydroPowerModels/hydro_training_utils.jl b/examples/HydroPowerModels/hydro_training_utils.jl new file mode 100644 index 0000000..b1f8ae6 --- /dev/null +++ b/examples/HydroPowerModels/hydro_training_utils.jl @@ -0,0 +1,50 @@ +# Shared helpers for HydroPowerModels training entrypoints. + +""" + parse_layers(s::AbstractString) -> Vector{Int} + +Parse a comma-separated hidden-layer specification. + +Empty or whitespace-only strings return `Int[]`, which lets environment +variables represent "no hidden layers" without a separate flag. + +# Arguments +- `s::AbstractString`: comma-separated layer widths, with optional whitespace. + +# Returns +- `Vector{Int}`: parsed hidden-layer widths. + +# Examples +```julia +parse_layers("128, 64") == [128, 64] +parse_layers("") == Int[] +parse_layers(" ") == Int[] +``` +""" +function parse_layers(s::AbstractString) + return isempty(strip(s)) ? + Int[] : + [parse(Int, strip(x)) for x in split(s, ",") if !isempty(strip(x))] +end + +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 + throw(ArgumentError("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 DecisionRulesExa.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 diff --git a/examples/HydroPowerModels/train_hydro_exa.jl b/examples/HydroPowerModels/train_hydro_exa.jl index 6ec71e9..4da967a 100644 --- a/examples/HydroPowerModels/train_hydro_exa.jl +++ b/examples/HydroPowerModels/train_hydro_exa.jl @@ -17,6 +17,7 @@ using MadNLPGPU, KernelAbstractions, CUDA using CUDSS, CUDSS_jll, cuDNN const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) @@ -32,33 +33,50 @@ const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") -const LAYERS = [128, 128] +const LAYERS = parse_layers(get(ENV, "DR_LAYERS", "128,128")) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) const ACTIVATION = sigmoid const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) const NUM_BATCHES = 100 const NUM_TRAIN_PER_BATCH = 1 const NUM_EVAL_SCENARIOS = 4 -const EVAL_EVERY = 25 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) const LR = 1f-3 -const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "10")) +const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "0")) const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) const DEFICIT_COST = 1e5 const USE_GPU = true const load_scaler = 0.6 const NUM_WORKERS = 1 -const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "annealed") +const DISCOUNT_GAMMA = parse(Float64, get(ENV, "DR_DISCOUNT_GAMMA", "1.0")) +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" [ - (1, div(NUM_EPOCHS * NUM_BATCHES, 4), 0.1), - (div(NUM_EPOCHS * NUM_BATCHES, 4) + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 2, 1.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 2 + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 3, 10.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 3 + 1, NUM_EPOCHS * NUM_BATCHES, 30.0), + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +elseif _PENALTY_MODE == "annealed_discount" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), ] else - [(1, NUM_EPOCHS * NUM_BATCHES, 1.0)] + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] end # Optional: ramp num_train_per_batch and eval scenarios over training. @@ -66,11 +84,21 @@ end const NUM_TRAIN_SCHEDULE = nothing # e.g. [(1,500,1),(501,2000,4),(2001,4000,8)] const EVAL_SCHEDULE = nothing # e.g. [(1,2000,4),(2001,4000,32)] -const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 9000) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) const _CLIP_TAG = GRAD_CLIP > 0 ? "-clip$(Int(GRAD_CLIP))" : "" -const _SCHED_TAG = _PENALTY_MODE == "annealed" ? "-anneal" : "-const" -const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-deteq-gpu$(_CLIP_TAG)$(_SCHED_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const _DISC_TAG = DISCOUNT_GAMMA < 1.0 ? "-disc$(replace(string(DISCOUNT_GAMMA), "." => ""))" : "" +const _SCHED_TAG = if _PENALTY_MODE == "annealed" + "-anneal" +elseif _PENALTY_MODE == "annealed_discount" + "-anndisc" +else + "-const" +end +const _LAYER_TAG = LAYERS == [128, 128] ? "" : "-L$(join(LAYERS, "_"))" +const _HEAD_TAG = isempty(HEAD_LAYERS) ? "" : "-H$(join(HEAD_LAYERS, "_"))" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-gpu$(_CLIP_TAG)$(_SCHED_TAG)$(_DISC_TAG)$(_LAYER_TAG)$(_HEAD_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") mkpath(MODEL_DIR) const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") @@ -85,9 +113,10 @@ power_data = load_power_data(PM_FILE) @info "Loading hydro data..." hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; - num_stages = NUM_STAGES * 10) + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) nHyd = hydro_data.nHyd T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES @info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" demand_mat = if isfile(DEMAND_FILE) @@ -103,6 +132,7 @@ resolved_pen = TARGET_PEN_ARG === :auto ? auto_target_penalty(power_data, hydro_data) : Float64(TARGET_PEN_ARG) @info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : (@info "Using CPU backend"; nothing) @@ -134,6 +164,8 @@ x0_init = Float32.([clamp(hydro_data.initial_volumes[r], hydro_data.units[r].min_vol, hydro_data.units[r].max_vol) for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) # ── Smoke test ──────────────────────────────────────────────────────────────── @@ -149,15 +181,29 @@ solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding resolved_pen_l1 = prob.base_penalty_l1 +const _discount_weights = Float64[DISCOUNT_GAMMA^(t-1) for t in 1:T for _ in 1:nHyd] +if DISCOUNT_GAMMA < 1.0 + @info "Discount γ=$(DISCOUNT_GAMMA): stage 1 weight=1.0, stage $T weight=$(round(DISCOUNT_GAMMA^(T-1); sigdigits=4))" +end + # ── Policy ──────────────────────────────────────────────────────────────────── -policy = StateConditionedPolicy(nHyd, nHyd, nHyd, LAYERS; - activation = ACTIVATION, - encoder_type = Flux.LSTM) +policy_active_mask = trues(nHyd) +policy = bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = policy_active_mask, + combiner_layers = HEAD_LAYERS) if !isnothing(PRE_TRAINED) @info "Loading pre-trained model from $(PRE_TRAINED)..." - Flux.loadmodel!(policy, JLD2.load(PRE_TRAINED, "model_state")) + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) +end + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" end # ── W&B logging ─────────────────────────────────────────────────────────────── @@ -170,10 +216,13 @@ lg = WandbLogger( "case" => CASE_NAME, "formulation" => FORM_LABEL, "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, "layers" => LAYERS, + "head_layers" => HEAD_LAYERS, "activation" => string(ACTIVATION), "target_penalty" => "auto=$(round(resolved_pen; digits=2))", "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, "deficit_cost" => DEFICIT_COST, "num_epochs" => NUM_EPOCHS, "num_batches" => NUM_BATCHES, @@ -185,6 +234,7 @@ lg = WandbLogger( "backend" => USE_GPU ? "GPU" : "CPU", "load_scaler" => load_scaler, "penalty_schedule" => string(PENALTY_SCHEDULE), + "discount_gamma" => DISCOUNT_GAMMA, "num_train_schedule" => string(something(NUM_TRAIN_SCHEDULE, "fixed")), "eval_schedule" => string(something(EVAL_SCHEDULE, "fixed")), "num_workers" => NUM_WORKERS, @@ -201,10 +251,10 @@ epoch_losses = Float64[] stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] function _build_rollout_de() build_hydro_de(power_data, hydro_data, 1; - backend = nothing, + backend = backend, float_type = Float64, formulation = FORMULATION, - target_penalty = TARGET_PEN_ARG, + target_penalty = resolved_pen * HYDRO_TARGET_PENALTY_MULT, deficit_cost = DEFICIT_COST, demand_matrix = stage_demand, load_scaler = load_scaler, @@ -212,8 +262,8 @@ function _build_rollout_de() end rollout_prob = _build_rollout_de() n_rollout_pool = max(NUM_WORKERS, NUM_EVAL_SCENARIOS) -rollout_pool = [_build_rollout_de() for _ in 1:n_rollout_pool] -@info "Rollout pool ready: $(n_rollout_pool) CPU stage-problem copies" +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:n_rollout_pool] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(n_rollout_pool) stage-problem copies)" : "sequential")" function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) @@ -226,49 +276,43 @@ function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) end hydro_realized_state(stage_prob, result) = - Array(hydro_solution(stage_prob, result).reservoir[:, end]) + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +const _rollout_pen = resolved_pen * HYDRO_TARGET_PENALTY_MULT +const _rollout_pen_l1 = _rollout_pen function hydro_objective_no_target_penalty(stage_prob, result) sol = hydro_solution(stage_prob, result) - delta = Array(sol.delta) - penalty_l2_cost = (resolved_pen / 2) * sum(abs2, delta) - penalty_l1_cost = resolved_pen_l1 * sum(abs, delta) + delta = sol.delta + penalty_l2_cost = (_rollout_pen / 2) * sum(abs2, delta) + penalty_l1_cost = _rollout_pen_l1 * sum(abs, delta) return result.objective - penalty_l2_cost - penalty_l1_cost end Random.seed!(8789) -eval_scenarios = [sample_scenario(hydro_data, T) for _ in 1:NUM_EVAL_SCENARIOS] +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:NUM_EVAL_SCENARIOS] rollout_evaluation = RolloutEvaluation( rollout_prob, x0_init, eval_scenarios; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, stride = EVAL_EVERY, policy_state = :target, stage_problem_pool = rollout_pool, + retry_on_failure = true, active_scenarios = NUM_EVAL_SCENARIOS, -) -realized_rollout_evaluation = RolloutEvaluation( - rollout_prob, - x0_init, - eval_scenarios; - horizon = T, - n_uncertainty = nHyd, - set_stage_parameters! = set_hydro_rollout_stage!, - realized_state = hydro_realized_state, - objective_no_target_penalty = hydro_objective_no_target_penalty, - madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, - stride = EVAL_EVERY, - policy_state = :realized, - stage_problem_pool = rollout_pool, - active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), ) Random.seed!(8788) @@ -281,6 +325,25 @@ function _schedule_value(schedule, iter, default) end current_penalty_mult = Ref(NaN) +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end train_tsddr( policy, @@ -300,48 +363,51 @@ train_tsddr( madnlp_kwargs = SOLVER_KWARGS, warmstart = true, problem_pool = problem_pool, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, adjust_hyperparameters = (iter, opt_state, n) -> begin mult = _schedule_value(PENALTY_SCHEDULE, iter, last(PENALTY_SCHEDULE)[3]) if mult != current_penalty_mult[] current_penalty_mult[] = mult ρ_half_scaled = prob.base_penalty_half * mult ρ_l1_scaled = prob.base_penalty_l1 * mult - penalty_vals = fill(ρ_half_scaled, T * nHyd) - penalty_l1_vals = fill(ρ_l1_scaled, T * nHyd) + penalty_vals = ρ_half_scaled .* _discount_weights + penalty_l1_vals = ρ_l1_scaled .* _discount_weights for (p, _, _, _) in problem_pool ExaModels.set_parameter!(p.core, p.p_penalty_half, penalty_vals) ExaModels.set_parameter!(p.core, p.p_penalty_l1, penalty_l1_vals) end - @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)))" + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)), γ=$DISCOUNT_GAMMA)" end if !isnothing(EVAL_SCHEDULE) n_eval = _schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) rollout_evaluation.active_scenarios = n_eval - realized_rollout_evaluation.active_scenarios = n_eval end return isnothing(NUM_TRAIN_SCHEDULE) ? n : _schedule_value(NUM_TRAIN_SCHEDULE, iter, n) end, record_loss = (iter, m, loss, tag) -> begin metrics = Dict{String, Any}(tag => loss, "batch" => iter) + _merge_batch_stats!(metrics, last_batch_stats[]) isfinite(loss) && push!(epoch_losses, loss) if iter % EVAL_EVERY == 0 rollout_evaluation(iter, m) - realized_rollout_evaluation(iter, m) metrics["metrics/rollout_objective_no_target_penalty"] = rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_objective_no_deficit"] = rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_target_violation_share"] = rollout_evaluation.last_violation_share - metrics["metrics/rollout_realized_objective_no_target_penalty"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_objective_no_deficit"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_target_violation_share"] = - realized_rollout_evaluation.last_violation_share metrics["metrics/rollout_n_ok"] = - realized_rollout_evaluation.last_n_ok + rollout_evaluation.last_n_ok end if !isnan(current_penalty_mult[]) diff --git a/examples/HydroPowerModels/train_hydro_exa_critic.jl b/examples/HydroPowerModels/train_hydro_exa_critic.jl index 3d9a5a0..ea93110 100644 --- a/examples/HydroPowerModels/train_hydro_exa_critic.jl +++ b/examples/HydroPowerModels/train_hydro_exa_critic.jl @@ -34,11 +34,12 @@ const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") const LAYERS = [128, 128] const ACTIVATION = sigmoid -const NUM_STAGES = 96 +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) const NUM_EPOCHS = 80 const NUM_BATCHES = 100 const MAX_EVAL_SCENARIOS = 32 -const EVAL_EVERY = 25 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) const EVAL_SCHEDULE = [ (1, div(NUM_EPOCHS * NUM_BATCHES, 2), 4), @@ -56,21 +57,32 @@ const CRITIC_BUFFER_SIZE = 512 const CRITIC_BATCH_SIZE = 32 const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) const DEFICIT_COST = 1e5 const USE_GPU = true const load_scaler = 0.6 const NUM_WORKERS = 4 const CRITIC_ROLLOUT_SAMPLES_PER_BATCH = 0 # eval rollouts feed the critic via external_critic_samples const CRITIC_POLICY_STATE = :target # set to :realized for closed-loop critic targets +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) const CRITIC_ROLLOUT_OBJECTIVE = :objective const NUM_CHEAP_CRITIC_SAMPLES_PER_BATCH = 4 * NUM_WORKERS -const PENALTY_SCHEDULE = [ - (1, div(NUM_EPOCHS * NUM_BATCHES, 4), 0.1), - (div(NUM_EPOCHS * NUM_BATCHES, 4) + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 2, 1.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 2 + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 3, 10.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 3 + 1, NUM_EPOCHS * NUM_BATCHES, 30.0), -] +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) +const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +else + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] +end const NUM_TRAIN_SCHEDULE = [ (1, div(NUM_EPOCHS * NUM_BATCHES, 5), NUM_WORKERS), @@ -80,9 +92,10 @@ const NUM_TRAIN_SCHEDULE = [ (div(NUM_EPOCHS * NUM_BATCHES, 5) * 4 + 1, NUM_EPOCHS * NUM_BATCHES, 8 * NUM_WORKERS), ] -const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 9000) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) -const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-deteq-gpu-critic-cv-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-gpu-critic-cv-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") mkpath(MODEL_DIR) const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") @@ -97,9 +110,10 @@ power_data = load_power_data(PM_FILE) @info "Loading hydro data..." hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; - num_stages = NUM_STAGES * 10) + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) nHyd = hydro_data.nHyd T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES @info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" demand_mat = if isfile(DEMAND_FILE) @@ -115,6 +129,7 @@ resolved_pen = TARGET_PEN_ARG === :auto ? auto_target_penalty(power_data, hydro_data) : Float64(TARGET_PEN_ARG) @info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : (@info "Using CPU backend"; nothing) @@ -133,6 +148,7 @@ end @info "Building $(T)-stage ExaModels DE (formulation=$FORMULATION)..." prob = _build_de() +resolved_pen_l1 = prob.base_penalty_l1 @info "Building $(NUM_WORKERS)-worker problem pool..." problem_pool = [(prob, prob.p_x0, prob.p_target, prob.p_inflow)] @@ -146,6 +162,8 @@ x0_init = Float32.([clamp(hydro_data.initial_volumes[r], hydro_data.units[r].min_vol, hydro_data.units[r].max_vol) for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) # ── Critic/control variate ─────────────────────────────────────────────────── @@ -196,13 +214,27 @@ solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding # ── Policy ──────────────────────────────────────────────────────────────────── -policy = StateConditionedPolicy(nHyd, nHyd, nHyd, LAYERS; - activation = ACTIVATION, - encoder_type = Flux.LSTM) +policy_active_mask = trues(nHyd) +policy = bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = policy_active_mask) if !isnothing(PRE_TRAINED) @info "Loading pre-trained model from $(PRE_TRAINED)..." - Flux.loadmodel!(policy, JLD2.load(PRE_TRAINED, "model_state")) + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) +end + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + control_variate = ScalarCriticControlVariate( + CUDA.cu(control_variate.critic); + featurizer = control_variate.featurizer, + value_loss_weight = control_variate.value_loss_weight, + gradient_loss_weight = control_variate.gradient_loss_weight, + ) + @info "Policy, critic, and x0 moved to GPU" end # ── W&B logging ─────────────────────────────────────────────────────────────── @@ -215,9 +247,12 @@ lg = WandbLogger( "case" => CASE_NAME, "formulation" => FORM_LABEL, "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, "layers" => LAYERS, "activation" => string(ACTIVATION), "target_penalty" => "auto=$(round(resolved_pen; digits=2))", + "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, "deficit_cost" => DEFICIT_COST, "num_epochs" => NUM_EPOCHS, "num_batches" => NUM_BATCHES, @@ -257,7 +292,7 @@ epoch_losses = Float64[] stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] function _build_rollout_de() build_hydro_de(power_data, hydro_data, 1; - backend = nothing, + backend = backend, float_type = Float64, formulation = FORMULATION, target_penalty = TARGET_PEN_ARG, @@ -267,8 +302,8 @@ function _build_rollout_de() ) end rollout_prob = _build_rollout_de() -rollout_pool = [_build_rollout_de() for _ in 1:NUM_WORKERS] -@info "Rollout pool ready: $(NUM_WORKERS) CPU stage-problem copies" +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:NUM_WORKERS] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(NUM_WORKERS) stage-problem copies)" : "sequential")" function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) @@ -281,59 +316,53 @@ function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) end hydro_realized_state(stage_prob, result) = - Array(hydro_solution(stage_prob, result).reservoir[:, end]) + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols function hydro_objective_no_target_penalty(stage_prob, result) sol = hydro_solution(stage_prob, result) - return result.objective - (resolved_pen / 2) * sum(abs2, Array(sol.delta)) + penalty_l2_cost = (resolved_pen / 2) * sum(abs2, sol.delta) + penalty_l1_cost = resolved_pen_l1 * sum(abs, sol.delta) + return result.objective - penalty_l2_cost - penalty_l1_cost end Random.seed!(8789) -eval_scenarios = [sample_scenario(hydro_data, T) for _ in 1:MAX_EVAL_SCENARIOS] +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:MAX_EVAL_SCENARIOS] rollout_evaluation = RolloutEvaluation( rollout_prob, x0_init, eval_scenarios; - horizon = T, - n_uncertainty = nHyd, - set_stage_parameters! = set_hydro_rollout_stage!, - realized_state = hydro_realized_state, - objective_no_target_penalty = hydro_objective_no_target_penalty, - madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, - stride = EVAL_EVERY, - policy_state = :target, - stage_problem_pool = rollout_pool, - active_scenarios = 4, -) -realized_rollout_evaluation = RolloutEvaluation( - rollout_prob, - x0_init, - eval_scenarios; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, stride = EVAL_EVERY, policy_state = :realized, stage_problem_pool = rollout_pool, + retry_on_failure = true, active_scenarios = 4, + state_bounds = (_min_vols_dev, _max_vols_dev), ) critic_training_target = RolloutCriticTarget( rollout_prob; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, policy_state = CRITIC_POLICY_STATE, objective_value = CRITIC_ROLLOUT_OBJECTIVE, + state_bounds = (_min_vols_dev, _max_vols_dev), ) Random.seed!(8788) @@ -381,15 +410,17 @@ train_tsddr( if mult != current_penalty_mult[] current_penalty_mult[] = mult ρ_half_scaled = prob.base_penalty_half * mult + ρ_l1_scaled = prob.base_penalty_l1 * mult penalty_vals = fill(ρ_half_scaled, T * nHyd) + penalty_l1_vals = fill(ρ_l1_scaled, T * nHyd) for (p, _, _, _) in problem_pool ExaModels.set_parameter!(p.core, p.p_penalty_half, penalty_vals) + ExaModels.set_parameter!(p.core, p.p_penalty_l1, penalty_l1_vals) end - @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)))" + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)))" end n_eval = _schedule_value(EVAL_SCHEDULE, iter, MAX_EVAL_SCENARIOS) rollout_evaluation.active_scenarios = n_eval - realized_rollout_evaluation.active_scenarios = n_eval return _schedule_value(NUM_TRAIN_SCHEDULE, iter, n) end, record_loss = (iter, m, loss, tag) -> begin @@ -398,7 +429,6 @@ train_tsddr( if iter % EVAL_EVERY == 0 rollout_evaluation(iter, m) - realized_rollout_evaluation(iter, m) append!(shared_critic_samples, critic_samples_from_evaluation( rollout_evaluation; @@ -410,14 +440,8 @@ train_tsddr( rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_target_violation_share"] = rollout_evaluation.last_violation_share - metrics["metrics/rollout_realized_objective_no_target_penalty"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_objective_no_deficit"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_target_violation_share"] = - realized_rollout_evaluation.last_violation_share metrics["metrics/rollout_n_ok"] = - realized_rollout_evaluation.last_n_ok + rollout_evaluation.last_n_ok end if !isnan(current_penalty_mult[]) diff --git a/examples/HydroPowerModels/train_hydro_exa_embedded.jl b/examples/HydroPowerModels/train_hydro_exa_embedded.jl new file mode 100644 index 0000000..5d8c788 --- /dev/null +++ b/examples/HydroPowerModels/train_hydro_exa_embedded.jl @@ -0,0 +1,457 @@ +# train_hydro_exa_embedded.jl +# +# Embedded-NN hydro training with ExaModels + MadNLP (AC or DC OPF). +# Policy is embedded directly in the NLP via VectorNonlinearOracle. +# Gradient: envelope theorem (multiplier-weighted policy Jacobian). +# +# 4 configurations via environment variables: +# DR_PENALTY_SCHEDULE = "const" | "annealed" +# DR_TARGET_PENALTY_MULT = 8.0 (Bolivia default multiplier on :auto penalties) +# DR_PRETRAIN_ITERS = 0 | 500 (regular TSDDR warmup before embedded training) +# DR_PRETRAIN_PENALTY_MULT = 0.1 (penalty multiplier during pretrain, default 0.1) + +using DecisionRulesExa +using ExaModels +using Flux +using Statistics, Random, Dates +using Wandb, Logging +using JLD2 +using MadNLP, MadNLPGPU +using CUDA, CUDSS, KernelAbstractions + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa_embedded.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = FORMULATION === :ac_polar ? "ACPPowerModel" : "DCPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") +const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") + +const LAYERS = parse_layers(get(ENV, "DR_LAYERS", "128,128")) +const ACTIVATION = sigmoid +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) +const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) +const NUM_BATCHES = 100 +const NUM_TRAIN_PER_BATCH = 1 +const NUM_EVAL_SCENARIOS = 4 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) +const LR = 1f-3 + +const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) +const DEFICIT_COST = 1e5 +const load_scaler = 0.6 + +const PRETRAIN_ITERS = parse(Int, get(ENV, "DR_PRETRAIN_ITERS", "0")) +const PRETRAIN_PENALTY_MULT = parse(Float64, get(ENV, "DR_PRETRAIN_PENALTY_MULT", "0.1")) +const STRICT_EMBEDDED_TARGETS = parse(Bool, get(ENV, "DR_STRICT_EMBEDDED_TARGETS", "false")) + +const DISCOUNT_GAMMA = parse(Float64, get(ENV, "DR_DISCOUNT_GAMMA", "1.0")) +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) +const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +elseif _PENALTY_MODE == "annealed_discount" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +else + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] +end + +const USE_GPU = parse(Bool, get(ENV, "DR_USE_GPU", "true")) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +const _DISC_TAG = DISCOUNT_GAMMA < 1.0 ? "-disc$(replace(string(DISCOUNT_GAMMA), "." => ""))" : "" +const _SCHED_TAG = if _PENALTY_MODE == "annealed" + "-anneal" +elseif _PENALTY_MODE == "annealed_discount" + "-anndisc" +else + "-const" +end +const _PRETRAIN_TAG = PRETRAIN_ITERS > 0 ? "-pre$(PRETRAIN_ITERS)" : "" +const _GPU_TAG = USE_GPU ? "-gpu" : "" +const _LAYER_TAG = LAYERS == [128, 128] ? "" : "-L$(join(LAYERS, "_"))" +const _STRICT_TAG = STRICT_EMBEDDED_TARGETS ? "-strict" : "" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-embedded$(_GPU_TAG)$(_SCHED_TAG)$(_DISC_TAG)$(_LAYER_TAG)$(_PRETRAIN_TAG)$(_STRICT_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") +mkpath(MODEL_DIR) +const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen)" + +@info "Loading hydro data..." +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) +nHyd = hydro_data.nHyd +T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES +@info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" + +demand_mat = if isfile(DEMAND_FILE) + @info "Loading demand from $(DEMAND_FILE)..." + load_demand(DEMAND_FILE, power_data; T = T) +else + nothing +end + +resolved_pen = TARGET_PEN_ARG === :auto ? + auto_target_penalty(power_data, hydro_data) : + Float64(TARGET_PEN_ARG) +@info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" + +backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : + (@info "Using CPU backend"; nothing) + +x0_init = Float32.([clamp(hydro_data.initial_volumes[r], + hydro_data.units[r].min_vol, + hydro_data.units[r].max_vol) + for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) + +const _discount_weights = Float64[DISCOUNT_GAMMA^(t-1) for t in 1:T for _ in 1:nHyd] +if DISCOUNT_GAMMA < 1.0 + @info "Discount γ=$(DISCOUNT_GAMMA): stage 1 weight=1.0, stage $T weight=$(round(DISCOUNT_GAMMA^(T-1); sigdigits=4))" +end + +# ── Policy ──────────────────────────────────────────────────────────────────── + +Random.seed!(42) +policy = if STRICT_EMBEDDED_TARGETS + @info "Using strict embedded hydro policy with one-stage reachable target bounds" + hydro_reachable_policy(hydro_data, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM) +else + bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = trues(nHyd)) +end + +# ── Optional pretrain with regular TSDDR ────────────────────────────────────── + +if PRETRAIN_ITERS > 0 + @info "Building regular DE for pretrain ($(PRETRAIN_ITERS) iters)..." + prob_reg = build_hydro_de(power_data, hydro_data, T; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + load_scaler = load_scaler, + ) + + if PRETRAIN_PENALTY_MULT != 1.0 + ρ_half_pre = prob_reg.base_penalty_half * PRETRAIN_PENALTY_MULT + ρ_l1_pre = prob_reg.base_penalty_l1 * PRETRAIN_PENALTY_MULT + ExaModels.set_parameter!(prob_reg.core, prob_reg.p_penalty_half, + fill(ρ_half_pre, T * nHyd)) + ExaModels.set_parameter!(prob_reg.core, prob_reg.p_penalty_l1, + fill(ρ_l1_pre, T * nHyd)) + @info " Pretrain penalty mult=$(PRETRAIN_PENALTY_MULT) (ρ/2=$(round(ρ_half_pre; digits=2)), l1=$(round(ρ_l1_pre; digits=2)))" + end + + @info "Pretraining policy with regular TSDDR..." + train_tsddr( + policy, x0_init, prob_reg, + prob_reg.p_x0, prob_reg.p_target, prob_reg.p_inflow, + () -> sample_scenario(hydro_data, T); + num_batches = PRETRAIN_ITERS, + num_train_per_batch = 1, + optimizer = Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + if iter % 50 == 0 + @info " Pretrain iter $iter/$PRETRAIN_ITERS: loss=$(round(loss; digits=2))" + end + return false + end, + ) + @info "Pretrain done." +end + +# ── Move policy + x0 to GPU (BEFORE embedded DE build so oracle captures GPU policy) ── + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" +end + +# ── Build embedded DE ───────────────────────────────────────────────────────── + +@info "Building embedded $(T)-stage $(FORMULATION) hydro DE..." +prob_emb = build_embedded_hydro_de(policy, power_data, hydro_data, T; + backend = backend, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + load_scaler = load_scaler, + strict_targets = STRICT_EMBEDDED_TARGETS, +) +@info " nvar=$(prob_emb._nvar) oracle_cons=$(length(prob_emb.target_con_range)) strict_targets=$(prob_emb.strict_targets)" + +resolved_pen_l1 = prob_emb.base_penalty_l1 + +# ── Smoke test ──────────────────────────────────────────────────────────────── + +w_mean = mean_inflow(hydro_data, T) +set_x0!(prob_emb, x0_init) +set_inflows!(prob_emb, w_mean) +@info "Smoke test: solving embedded DE with mean inflows..." +result0 = MadNLP.madnlp(prob_emb.model; SOLVER_KWARGS..., print_level = MadNLP.WARN) +@info " Status: $(result0.status) Obj: $(round(result0.objective; digits=4))" +solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding anyway" + +# ── W&B logging ─────────────────────────────────────────────────────────────── + +lg = WandbLogger( + project = "RL", + name = RUN_NAME, + save_code = false, + config = Dict( + "case" => CASE_NAME, + "formulation" => FORM_LABEL, + "method" => "embedded_nn", + "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, + "layers" => LAYERS, + "activation" => string(ACTIVATION), + "target_penalty" => "auto=$(round(resolved_pen; digits=2))", + "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, + "deficit_cost" => DEFICIT_COST, + "num_epochs" => NUM_EPOCHS, + "num_batches" => NUM_BATCHES, + "num_train_per_batch" => NUM_TRAIN_PER_BATCH, + "num_eval_scenarios" => NUM_EVAL_SCENARIOS, + "eval_every" => EVAL_EVERY, + "lr" => LR, + "penalty_schedule" => string(PENALTY_SCHEDULE), + "discount_gamma" => DISCOUNT_GAMMA, + "pretrain_iters" => PRETRAIN_ITERS, + "pretrain_penalty_mult" => PRETRAIN_PENALTY_MULT, + "strict_embedded_targets" => STRICT_EMBEDDED_TARGETS, + ), +) + +# ── Rollout evaluation ──────────────────────────────────────────────────────── + +stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = stage_demand, + load_scaler = load_scaler, + strict_targets = STRICT_EMBEDDED_TARGETS, + ) +end + +rollout_prob = _build_rollout_de() +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:NUM_EVAL_SCENARIOS] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(NUM_EVAL_SCENARIOS) stage-problem copies)" : "sequential")" + +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + if demand_mat !== nothing + set_demand!(stage_prob, load_scaler .* demand_mat[stage:stage, :]) + end + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +function hydro_objective_no_target_penalty(stage_prob, result) + if STRICT_EMBEDDED_TARGETS + return result.objective + end + sol = hydro_solution(stage_prob, result) + delta = sol.delta + penalty_l2_cost = (resolved_pen / 2) * sum(abs2, delta) + penalty_l1_cost = resolved_pen_l1 * sum(abs, delta) + return result.objective - penalty_l2_cost - penalty_l1_cost +end + +Random.seed!(8789) +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:NUM_EVAL_SCENARIOS] + +rollout_evaluation = RolloutEvaluation( + rollout_prob, x0_init, eval_scenarios; + horizon = T_ROLLOUT, n_uncertainty = nHyd, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + objective_no_target_penalty = hydro_objective_no_target_penalty, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + stride = EVAL_EVERY, + policy_state = :realized, + stage_problem_pool = rollout_pool, + retry_on_failure = true, + active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), +) + +# ── Training ────────────────────────────────────────────────────────────────── + +Random.seed!(8788) + +best_obj = Inf +epoch_losses = Float64[] + +function _schedule_value(schedule, iter, default) + for (lo, hi, val) in schedule + lo <= iter <= hi && return val + end + return default +end + +current_penalty_mult = Ref(NaN) +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end + +@info "Starting embedded training: $(NUM_EPOCHS) epochs × $(NUM_BATCHES) batches" + +train_tsddr_embedded( + policy, x0_init, prob_emb, + () -> sample_scenario(hydro_data, T); + num_batches = NUM_EPOCHS * NUM_BATCHES, + num_train_per_batch = NUM_TRAIN_PER_BATCH, + optimizer = Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + get_realized_states = embedded_hydro_realized_states, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Embedded training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Embedded training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, + adjust_hyperparameters = (iter, opt_state, n) -> begin + mult = _schedule_value(PENALTY_SCHEDULE, iter, last(PENALTY_SCHEDULE)[3]) + if mult != current_penalty_mult[] + current_penalty_mult[] = mult + if STRICT_EMBEDDED_TARGETS + @info "Strict target equality active; target slack penalties are not used" + else + ρ_half_scaled = prob_emb.base_penalty_half * mult + ρ_l1_scaled = prob_emb.base_penalty_l1 * mult + penalty_vals = ρ_half_scaled .* _discount_weights + penalty_l1_vals = ρ_l1_scaled .* _discount_weights + ExaModels.set_parameter!(prob_emb.core, prob_emb.p_penalty_half, penalty_vals) + ExaModels.set_parameter!(prob_emb.core, prob_emb.p_penalty_l1, penalty_l1_vals) + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)), γ=$DISCOUNT_GAMMA)" + end + end + return n + end, + record_loss = (iter, m, loss, tag) -> begin + metrics = Dict{String, Any}(tag => loss, "batch" => iter) + _merge_batch_stats!(metrics, last_batch_stats[]) + isfinite(loss) && push!(epoch_losses, loss) + + if iter % EVAL_EVERY == 0 + rollout_evaluation(iter, m) + metrics["metrics/rollout_objective_no_target_penalty"] = + rollout_evaluation.last_objective_no_target_penalty + metrics["metrics/rollout_target_violation_share"] = + rollout_evaluation.last_violation_share + metrics["metrics/rollout_n_ok"] = + rollout_evaluation.last_n_ok + end + + if !isnan(current_penalty_mult[]) + metrics["metrics/target_penalty_multiplier"] = current_penalty_mult[] + end + + batch_in_epoch = (iter - 1) % NUM_BATCHES + 1 + if batch_in_epoch == NUM_BATCHES + epoch = (iter - 1) ÷ NUM_BATCHES + 1 + mean_loss = isempty(epoch_losses) ? NaN : mean(epoch_losses) + n_ok = length(epoch_losses) + empty!(epoch_losses) + Wandb.log(lg, Dict("metrics/epoch_objective" => mean_loss, "epoch" => epoch)) + @info "Epoch $epoch/$NUM_EPOCHS mean=$(round(mean_loss; digits=2)) ok=$n_ok/$NUM_BATCHES" + if isfinite(mean_loss) && mean_loss < best_obj + global best_obj = mean_loss + jldsave(MODEL_PATH; model_state = Flux.state(cpu(m))) + @info " → New best: $(round(mean_loss; digits=4)) — saved $MODEL_PATH" + end + end + Wandb.log(lg, metrics) + return false + end, +) + +close(lg) +@info "Done. Best model saved to: $(MODEL_PATH)" diff --git a/examples/HydroPowerModels/train_hydro_exa_strict.jl b/examples/HydroPowerModels/train_hydro_exa_strict.jl new file mode 100644 index 0000000..c7620d6 --- /dev/null +++ b/examples/HydroPowerModels/train_hydro_exa_strict.jl @@ -0,0 +1,705 @@ +# train_hydro_exa_strict.jl +# +# Strict-mode regular DE training with ExaModels + MadNLP (AC polar OPF). +# Uses HydroReachablePolicy (sigmoid-bounded to one-stage reachable set) +# with strict_targets=true (delta variables fixed to zero, no target penalty). +# +# Key insight: ordinary regular DE target generation is open-loop after x0, so +# strict equality is usually unsafe: the optimizer may discover realized states +# different from the target path, and later targets were not computed from those +# realized states. Here the reachable policy is rolled out from the true x0 and +# uses the previous target as the next policy state. Since every target is +# one-stage reachable from the previous target, the full target trajectory is +# feasible by induction and strict regular DE is valid. +# +# Environment variables: +# DR_ENCODER_LAYERS = "128,128" (LSTM encoder layer sizes) +# DR_HEAD_LAYERS = "" (state-conditioned target-head hidden sizes) +# DR_LAYERS = "128,128" (legacy alias for DR_ENCODER_LAYERS) +# DR_NUM_STAGES = "126" +# DR_NUM_ROLLOUT_STAGES = "96" +# DR_NUM_EPOCHS = "80" +# DR_NUM_BATCHES = "100" +# DR_NUM_TRAIN_PER_BATCH = "1" (sampled trajectories per gradient step) +# DR_NUM_TRAIN_SCHEDULE = "" (optional "lo:hi:n,..." schedule) +# DR_CONTEXT = "" (""/"none", "phase", or "phase+progress") +# DR_NUM_EVAL_SCENARIOS = "4" (fixed held-out rollout-selection scenarios) +# DR_EVAL_SCHEDULE = "" (optional "lo:hi:n,..." active-eval schedule) +# DR_EVAL_EVERY = "50" +# DR_SAVE_METRIC = "training" ("training" or "rollout") +# DR_LR = "0.001" +# DR_LR_FINAL = DR_LR (cosine-decay final learning rate) +# DR_LR_WARMUP = "0" (linear warmup iterations from DR_LR/100) +# DR_PRETRAINED_MODEL = "" (optional diagnostic warmstart checkpoint) +# DR_REACTIVE_DEFICIT = "free" ("free", "hard", or a finite penalty) +# DR_GRAD_CLIP = "0" +# DR_MAX_ITER = "9000" +# +# Reproducible recipes: +# +# Historical baseline (all defaults, byte-compatible training semantics): +# julia --project -t auto train_hydro_exa_strict.jl +# +# Fresh fast/fair candidate (headline timing starts at this command): +# DR_REACTIVE_DEFICIT=hard DR_SAVE_METRIC=rollout \ +# DR_NUM_EVAL_SCENARIOS=8 DR_EVAL_EVERY=100 \ +# DR_NUM_TRAIN_PER_BATCH=2 \ +# DR_LR=5e-4 DR_LR_FINAL=5e-5 DR_LR_WARMUP=100 \ +# DR_NUM_EPOCHS=25 DR_NUM_BATCHES=100 \ +# DR_HEAD_LAYERS=128,128 \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Progressive-sampling Phase B (spend samples only after coarse movement): +# DR_REACTIVE_DEFICIT=hard DR_SAVE_METRIC=rollout \ +# DR_NUM_EVAL_SCENARIOS=12 DR_EVAL_SCHEDULE=1:800:4,801:1400:8,1401:1800:12 \ +# DR_NUM_TRAIN_PER_BATCH=1 DR_NUM_TRAIN_SCHEDULE=1:800:1,801:1400:2,1401:1800:4 \ +# DR_LR=7e-4 DR_LR_FINAL=2e-5 DR_LR_WARMUP=100 \ +# DR_NUM_EPOCHS=18 DR_NUM_BATCHES=100 \ +# DR_HEAD_LAYERS=256,256 \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Diagnostic warmstart (not a clean headline unless parent time is counted): +# DR_PRETRAINED_MODEL= DR_SAVE_METRIC=rollout \ +# DR_NUM_TRAIN_PER_BATCH=4 DR_LR=1e-4 DR_LR_FINAL=1e-5 DR_LR_WARMUP=50 \ +# DR_NUM_EPOCHS=10 DR_HEAD_LAYERS= \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Usage: +# julia --project -t auto train_hydro_exa_strict.jl + +using DecisionRulesExa +using StableRNGs +using ExaModels +using Flux +using Statistics, Random, Dates +using Wandb, Logging +using JLD2 +using MadNLP +using MadNLPGPU, KernelAbstractions, CUDA +using CUDSS, CUDSS_jll, cuDNN + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = FORMULATION === :ac_polar ? "ACPPowerModel" : "DCPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") +const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") + +function parse_reactive_deficit_cost(raw::AbstractString) + s = lowercase(strip(raw)) + if s == "free" + return nothing + elseif s == "hard" + return Inf + end + cost = parse(Float64, s) + (isnan(cost) || cost < 0) && + throw(ArgumentError("DR_REACTIVE_DEFICIT must be free, hard, or a finite nonnegative cost; got $raw")) + return cost +end + +function value_tag(x) + return replace(replace(string(x), "." => "p"), "-" => "m") +end + +function parse_int_schedule(raw::AbstractString, name::AbstractString) + s = strip(raw) + if isempty(s) || lowercase(s) in ("fixed", "none", "nothing") + return nothing + end + + schedule = Tuple{Int, Int, Int}[] + for item in split(s, r"[,;]") + token = strip(item) + isempty(token) && continue + + parts = split(token, ":") + local lo::Int + local hi::Int + local value::Int + if length(parts) == 3 + lo = parse(Int, strip(parts[1])) + hi = parse(Int, strip(parts[2])) + value = parse(Int, strip(parts[3])) + elseif length(parts) == 2 && occursin("-", parts[1]) + bounds = split(parts[1], "-") + length(bounds) == 2 || + throw(ArgumentError("$name schedule entry '$token' must be lo:hi:value or lo-hi:value")) + lo = parse(Int, strip(bounds[1])) + hi = parse(Int, strip(bounds[2])) + value = parse(Int, strip(parts[2])) + else + throw(ArgumentError("$name schedule entry '$token' must be lo:hi:value or lo-hi:value")) + end + + lo >= 1 || throw(ArgumentError("$name schedule lower bound must be >= 1 in '$token'")) + hi >= lo || throw(ArgumentError("$name schedule upper bound must be >= lower bound in '$token'")) + value >= 1 || throw(ArgumentError("$name schedule value must be >= 1 in '$token'")) + push!(schedule, (lo, hi, value)) + end + + isempty(schedule) && return nothing + sort!(schedule; by = first) + last_hi = 0 + for (lo, hi, _) in schedule + lo > last_hi || throw(ArgumentError("$name schedule has overlapping entries near iteration $lo")) + last_hi = hi + end + return schedule +end + +function schedule_value(schedule, iter::Int, default::Int) + isnothing(schedule) && return default + for (lo, hi, value) in schedule + lo <= iter <= hi && return value + end + return default +end + +function schedule_tag(schedule, prefix::AbstractString) + isnothing(schedule) && return "" + vals = unique(last.(schedule)) + return "-$(prefix)$(join(vals, "_"))" +end + +const ENCODER_LAYERS = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +const ACTIVATION = sigmoid +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) +const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) +const NUM_BATCHES = parse(Int, get(ENV, "DR_NUM_BATCHES", "100")) +const NUM_TRAIN_PER_BATCH = parse(Int, get(ENV, "DR_NUM_TRAIN_PER_BATCH", "1")) +const NUM_TRAIN_SCHEDULE = parse_int_schedule(get(ENV, "DR_NUM_TRAIN_SCHEDULE", ""), "DR_NUM_TRAIN_SCHEDULE") +const CONTEXT_MODE = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +const CONTEXT_PERIOD = countlines(INFLOW_FILE) +const CONTEXT_HORIZON = NUM_STAGES +const STAGE_CONTEXT = build_stage_context(CONTEXT_MODE, CONTEXT_HORIZON, CONTEXT_PERIOD) +const N_CONTEXT = isnothing(STAGE_CONTEXT) ? 0 : size(STAGE_CONTEXT, 1) +const NUM_EVAL_SCENARIOS = parse(Int, get(ENV, "DR_NUM_EVAL_SCENARIOS", "4")) +const EVAL_SCHEDULE = parse_int_schedule(get(ENV, "DR_EVAL_SCHEDULE", ""), "DR_EVAL_SCHEDULE") +if !isnothing(EVAL_SCHEDULE) && maximum(last.(EVAL_SCHEDULE)) > NUM_EVAL_SCENARIOS + throw(ArgumentError("DR_EVAL_SCHEDULE cannot exceed DR_NUM_EVAL_SCENARIOS=$(NUM_EVAL_SCENARIOS)")) +end +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) +const SAVE_METRIC = lowercase(strip(get(ENV, "DR_SAVE_METRIC", "training"))) +SAVE_METRIC in ("training", "rollout") || + throw(ArgumentError("DR_SAVE_METRIC must be training or rollout; got $SAVE_METRIC")) +const LR = parse(Float32, get(ENV, "DR_LR", "0.001")) +const LR_FINAL = parse(Float32, get(ENV, "DR_LR_FINAL", string(LR))) +const LR_WARMUP = parse(Int, get(ENV, "DR_LR_WARMUP", "0")) +const PRE_TRAINED = strip(get(ENV, "DR_PRETRAINED_MODEL", "")) +const HAS_PRETRAINED = !(isempty(PRE_TRAINED) || lowercase(PRE_TRAINED) == "nothing") +const REACTIVE_DEFICIT_RAW = get(ENV, "DR_REACTIVE_DEFICIT", "free") +const REACTIVE_DEFICIT_COST = parse_reactive_deficit_cost(REACTIVE_DEFICIT_RAW) +const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "0")) + +const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = 8.0 +const DEFICIT_COST = 1e5 +const USE_GPU = true +const load_scaler = 0.6 +const NUM_WORKERS = 1 + +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +const _CLIP_TAG = GRAD_CLIP > 0 ? "-clip$(Int(GRAD_CLIP))" : "" +const _ENC_TAG = ENCODER_LAYERS == [128, 128] ? "" : "-E$(join(ENCODER_LAYERS, "_"))" +const _HEAD_TAG = isempty(HEAD_LAYERS) ? "-Hlinear" : "-H$(join(HEAD_LAYERS, "_"))" +const _NT_TAG = isnothing(NUM_TRAIN_SCHEDULE) ? + (NUM_TRAIN_PER_BATCH > 1 ? "-nt$(NUM_TRAIN_PER_BATCH)" : "") : + schedule_tag(NUM_TRAIN_SCHEDULE, "ntsch") +const _CTX_TAG = context_run_tag(CONTEXT_MODE) +const _EV_TAG = schedule_tag(EVAL_SCHEDULE, "evsch") +const _SAVE_TAG = SAVE_METRIC == "rollout" ? "-rollout" : "" +const _WARM_TAG = HAS_PRETRAINED ? "-warm" : "" +const _RQ_TAG = REACTIVE_DEFICIT_COST === nothing ? "" : + isinf(Float64(REACTIVE_DEFICIT_COST)) ? "-rqhard" : + "-rq$(value_tag(REACTIVE_DEFICIT_COST))" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-strict-gpu$(_CLIP_TAG)$(_ENC_TAG)$(_HEAD_TAG)$(_NT_TAG)$(_CTX_TAG)$(_EV_TAG)$(_SAVE_TAG)$(_WARM_TAG)$(_RQ_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") +mkpath(MODEL_DIR) +const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") +const TOTAL_ITERS = NUM_EPOCHS * NUM_BATCHES + +function lr_schedule(iter::Int, total_iters::Int) + if iter <= LR_WARMUP + return LR * (0.01f0 + 0.99f0 * Float32(iter) / Float32(max(LR_WARMUP, 1))) + end + ρ = clamp((iter - LR_WARMUP) / max(total_iters - LR_WARMUP, 1), 0.0, 1.0) + return LR_FINAL + 0.5f0 * (LR - LR_FINAL) * (1f0 + cos(Float32(pi * ρ))) +end + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen)" + +@info "Loading hydro data..." +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) +nHyd = hydro_data.nHyd +T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES +@info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" + +demand_mat = if isfile(DEMAND_FILE) + @info "Loading demand from $(DEMAND_FILE)..." + load_demand(DEMAND_FILE, power_data; T = T) +else + nothing +end + +# ── Build ExaModels DE (strict: delta variables fixed to zero) ──────────────── + +resolved_pen = TARGET_PEN_ARG === :auto ? + auto_target_penalty(power_data, hydro_data) : + Float64(TARGET_PEN_ARG) +@info "Auto target penalty: ρ=$(round(resolved_pen; digits=2)) (not used — strict mode)" + +backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : + (@info "Using CPU backend"; nothing) + +function _build_de() + build_hydro_de(power_data, hydro_data, T; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + load_scaler = load_scaler, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, + ) +end + +@info "Building strict $(T)-stage ExaModels DE (formulation=$FORMULATION)..." +prob = _build_de() + +@info "Building $(NUM_WORKERS)-worker problem pool..." +problem_pool = [(prob, prob.p_x0, prob.p_target, prob.p_inflow)] +for i in 2:NUM_WORKERS + p = _build_de() + push!(problem_pool, (p, p.p_x0, p.p_target, p.p_inflow)) +end +@info " Pool ready: $(NUM_WORKERS) independent strict DE instances on GPU" + +x0_init = Float32.([clamp(hydro_data.initial_volumes[r], + hydro_data.units[r].min_vol, + hydro_data.units[r].max_vol) + for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) + +# ── Policy (HydroReachablePolicy — one-stage reachable sigmoid bounds) ──────── + +Random.seed!(42) +base_policy = hydro_reachable_policy(hydro_data, ENCODER_LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + combiner_layers = HEAD_LAYERS, + n_context = N_CONTEXT) +policy = isnothing(STAGE_CONTEXT) ? base_policy : ContextualPolicy(base_policy, STAGE_CONTEXT) + +if HAS_PRETRAINED + @info "Loading pre-trained policy checkpoint" PRE_TRAINED + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) + Flux.reset!(policy) +end + +""" + rollout_reachable_targets(policy, x0, w_flat, T, nHyd) -> Vector{Float64} + +Roll out a [`HydroReachablePolicy`] target trajectory before solving the strict +regular deterministic equivalent. + +The regular DE receives an external target vector, so it cannot query the policy +inside the NLP. This helper constructs that vector in a way that preserves +strict-mode feasibility: it starts from the known feasible initial state `x0`, +feeds `[w_t; previous_target]` to the policy, and stores each reachable target as +the next previous state. By induction, every target in the returned trajectory is +reachable from the prior target under the sampled inflow path. + +# Arguments +- `policy`: reachable hydro policy with input `[inflow; previous_state]`. +- `x0`: initial reservoir state. +- `w_flat`: stage-major flat inflow vector of length `T * nHyd`. +- `T::Int`: number of stages. +- `nHyd::Int`: number of hydro reservoir state components. + +# Returns +- `Vector{Float64}`: stage-major target trajectory suitable for + `ExaModels.set_parameter!(prob.core, prob.p_target, targets)`. + +# Examples +```julia +targets = rollout_reachable_targets(policy, x0_init, mean_inflow(hydro_data, T), T, nHyd) +``` +""" +function rollout_reachable_targets(policy, x0, w_flat, T, nHyd) + Flux.reset!(policy) + prev = x0 + targets = Vector{Vector{Float32}}(undef, T) + for t in 1:T + wt = Float32.(view(w_flat, ((t - 1) * nHyd + 1):(t * nHyd))) + target = policy(vcat(wt, prev)) + targets[t] = Float32.(target) + prev = targets[t] + end + return Float64.(vcat(targets...)) +end + +# ── Smoke test ──────────────────────────────────────────────────────────────── + +w_mean = mean_inflow(hydro_data, T) +xhat_mean = rollout_reachable_targets(policy, x0_init, w_mean, T, nHyd) +ExaModels.set_parameter!(prob.core, prob.p_x0, x0_init) +ExaModels.set_parameter!(prob.core, prob.p_inflow, w_mean) +ExaModels.set_parameter!(prob.core, prob.p_target, xhat_mean) +prepare_solve!(prob, x0_init, w_mean, xhat_mean) +@info "Smoke test: solving strict DE with mean inflows and reachable policy targets..." +result0 = MadNLP.madnlp(prob.model; SOLVER_KWARGS..., print_level = MadNLP.WARN) +@info " Status: $(result0.status) Objective: $(round(result0.objective; digits=4))" +isfinite(result0.objective) || error("Smoke test returned non-finite objective") +solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding anyway" + +Flux.reset!(policy) + +if USE_GPU + policy = policy isa ContextualPolicy ? + ContextualPolicy(CUDA.cu(policy.policy), CUDA.cu(policy.context)) : + CUDA.cu(policy) + Flux.reset!(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" +end + +@info "Strict Exa training config" RUN_NAME NUM_STAGES NUM_ROLLOUT_STAGES NUM_EPOCHS NUM_BATCHES NUM_TRAIN_PER_BATCH NUM_TRAIN_SCHEDULE CONTEXT_MODE CONTEXT_PERIOD CONTEXT_HORIZON N_CONTEXT NUM_EVAL_SCENARIOS EVAL_SCHEDULE EVAL_EVERY SAVE_METRIC LR LR_FINAL LR_WARMUP PRE_TRAINED REACTIVE_DEFICIT_RAW REACTIVE_DEFICIT_COST GRAD_CLIP MAX_ITER + +function checkpoint_policy_state(m) + if m isa ContextualPolicy + return Flux.state(ContextualPolicy(cpu(m.policy), Array(m.context))) + end + return Flux.state(cpu(m)) +end + +# ── W&B logging ─────────────────────────────────────────────────────────────── + +lg = WandbLogger( + project = "RL", + name = RUN_NAME, + save_code = false, + config = Dict( + "case" => CASE_NAME, + "formulation" => FORM_LABEL, + "method" => "deteq-strict", + "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, + "encoder_layers" => ENCODER_LAYERS, + "head_layers" => HEAD_LAYERS, + "activation" => string(ACTIVATION), + "target_penalty" => "strict (disabled)", + "deficit_cost" => DEFICIT_COST, + "reactive_deficit" => REACTIVE_DEFICIT_RAW, + "reactive_deficit_cost" => string(REACTIVE_DEFICIT_COST), + "num_epochs" => NUM_EPOCHS, + "num_batches" => NUM_BATCHES, + "num_train_per_batch" => NUM_TRAIN_PER_BATCH, + "num_train_schedule" => string(something(NUM_TRAIN_SCHEDULE, "fixed")), + "context_mode" => isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE, + "context_period" => CONTEXT_PERIOD, + "context_horizon" => CONTEXT_HORIZON, + "n_context" => N_CONTEXT, + "num_eval_scenarios" => NUM_EVAL_SCENARIOS, + "eval_schedule" => string(something(EVAL_SCHEDULE, "fixed")), + "eval_every" => EVAL_EVERY, + "save_metric" => SAVE_METRIC, + "lr" => LR, + "lr_final" => LR_FINAL, + "lr_warmup" => LR_WARMUP, + "pre_trained_model" => PRE_TRAINED, + "grad_clip" => GRAD_CLIP, + "backend" => USE_GPU ? "GPU" : "CPU", + "load_scaler" => load_scaler, + "strict_targets" => true, + "policy_type" => N_CONTEXT == 0 ? "HydroReachablePolicy" : "ContextualPolicy{HydroReachablePolicy}", + "num_workers" => NUM_WORKERS, + ), +) + +# ── Training ────────────────────────────────────────────────────────────────── + +Random.seed!(8788) + +best_obj = Inf +epoch_losses = Float64[] + +stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = stage_demand, + load_scaler = load_scaler, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, + ) +end +rollout_prob = _build_rollout_de() +n_rollout_pool = max(NUM_WORKERS, NUM_EVAL_SCENARIOS) +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:n_rollout_pool] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(n_rollout_pool) stage-problem copies)" : "sequential")" + +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + if demand_mat !== nothing + set_demand!(stage_prob, load_scaler .* demand_mat[stage:stage, :]) + end + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +hydro_objective_no_target_penalty(stage_prob, result) = result.objective + +# Held-out evaluation scenarios. Two modes: +# +# 1. DR_EVAL_PROTOCOL_IDS set (comma-separated column ids of the seeded paired +# protocol): the eval set is those exact columns of the 126×500 protocol +# matrix (StableRNG seed 20260706 — identical generation to the paired +# evaluation scripts). With the representative subset found by searching +# 300k candidate subsets against 7 evaluated policies — +# ids 2,39,81,119,130,156,200,206,378,493 (subset seed 77142) — the +# 10-scenario training-time evaluation tracks the full 500-scenario paired +# mean within ~75 cost units and preserves paired differences vs SDDP +# within ~35, so SaveBest selects on (a faithful proxy of) the deployment +# metric. +# 2. Unset (historical): random draws from the inflow process, seed 8789. +# NOTE: arbitrary small draws carry offsets of hundreds-to-thousands of +# cost units vs the paired protocol; never compare their values across +# runs with different eval sets. +const EVAL_PROTOCOL_IDS = let raw = strip(get(ENV, "DR_EVAL_PROTOCOL_IDS", "")) + isempty(raw) ? Int[] : [parse(Int, strip(x)) for x in split(raw, ",")] +end + +""" + protocol_eval_scenario(hydro_data, T, protocol_indices, s) -> Vector{Float64} + +Flat `T × nHyd` inflow vector for paired-protocol scenario column `s`: +stage `t` realizes joint inflow scenario `protocol_indices[t, s]`, with the +cyclic raw-row mapping shared by every paired evaluation script. +""" +function protocol_eval_scenario(hydro_data::HydroData, T::Int, protocol_indices, s::Int) + nHyd = hydro_data.nHyd + w = Vector{Float64}(undef, T * nHyd) + for t in 1:T + t_row = mod1(t, hydro_data.nStagesSample) + j = protocol_indices[t, s] + for r in 1:nHyd + w[(t-1)*nHyd + r] = hydro_data.scenario_inflows[r][t_row, j] + end + end + return w +end + +eval_scenarios = if isempty(EVAL_PROTOCOL_IDS) + Random.seed!(8789) + [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:NUM_EVAL_SCENARIOS] +else + # Same fixed-shape generation as the paired protocol (126 rows × 500 + # columns; see load_hydropowermodels.jl in the MAIN repo). + protocol_indices = rand(StableRNG(20260706), 1:hydro_data.nScenarios, 126, 500) + @assert T_ROLLOUT <= 126 && all(1 .<= EVAL_PROTOCOL_IDS .<= 500) + @assert length(EVAL_PROTOCOL_IDS) == NUM_EVAL_SCENARIOS "DR_NUM_EVAL_SCENARIOS must match the id count" + @info "Eval set = paired-protocol columns $(EVAL_PROTOCOL_IDS)" + [protocol_eval_scenario(hydro_data, T_ROLLOUT, protocol_indices, s) for s in EVAL_PROTOCOL_IDS] +end +rollout_evaluation = RolloutEvaluation( + rollout_prob, + x0_init, + eval_scenarios; + horizon = T_ROLLOUT, + n_uncertainty = nHyd, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + objective_no_target_penalty = hydro_objective_no_target_penalty, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + stride = EVAL_EVERY, + policy_state = :realized, + stage_problem_pool = rollout_pool, + retry_on_failure = true, + active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), +) + +current_num_train = Ref(NUM_TRAIN_PER_BATCH) +current_eval_scenarios = Ref(schedule_value(EVAL_SCHEDULE, 1, NUM_EVAL_SCENARIOS)) +rollout_evaluation.active_scenarios = current_eval_scenarios[] + +if SAVE_METRIC == "rollout" + rollout_evaluation(EVAL_EVERY, policy) + best_obj = rollout_evaluation.last_objective_no_target_penalty + @info "Initial rollout evaluation (checkpoint baseline)" best_obj rollout_evaluation.last_violation_share rollout_evaluation.last_n_ok +end + +Random.seed!(8788) + +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end + +train_tsddr( + policy, + x0_init, + prob, + prob.p_x0, + prob.p_target, + prob.p_inflow, + () -> sample_scenario(hydro_data, T); + num_batches = TOTAL_ITERS, + num_train_per_batch = NUM_TRAIN_PER_BATCH, + optimizer = GRAD_CLIP > 0 ? + Flux.Optimisers.OptimiserChain( + Flux.Optimisers.ClipGrad(GRAD_CLIP), + Flux.Adam(LR), + ) : Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + problem_pool = problem_pool, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, + adjust_hyperparameters = (LR_WARMUP == 0 && LR_FINAL == LR) ? + ((iter, opt_state, n) -> begin + n_next = schedule_value(NUM_TRAIN_SCHEDULE, iter, n) + if n_next != current_num_train[] + current_num_train[] = n_next + @info "num_train_per_batch → $n_next" + end + n_eval = schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) + if n_eval != current_eval_scenarios[] + current_eval_scenarios[] = n_eval + rollout_evaluation.active_scenarios = n_eval + @info "rollout active scenarios → $n_eval" + end + n_next + end) : + ((iter, opt_state, n) -> begin + Flux.Optimisers.adjust!(opt_state, lr_schedule(iter, TOTAL_ITERS)) + n_next = schedule_value(NUM_TRAIN_SCHEDULE, iter, n) + if n_next != current_num_train[] + current_num_train[] = n_next + @info "num_train_per_batch → $n_next" + end + n_eval = schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) + if n_eval != current_eval_scenarios[] + current_eval_scenarios[] = n_eval + rollout_evaluation.active_scenarios = n_eval + @info "rollout active scenarios → $n_eval" + end + n_next + end), + record_loss = (iter, m, loss, tag) -> begin + metrics = Dict{String, Any}( + tag => loss, + "batch" => iter, + "metrics/lr" => lr_schedule(iter, TOTAL_ITERS), + "metrics/num_train_per_batch" => current_num_train[], + "metrics/active_eval_scenarios" => current_eval_scenarios[], + ) + _merge_batch_stats!(metrics, last_batch_stats[]) + isfinite(loss) && push!(epoch_losses, loss) + + if iter % EVAL_EVERY == 0 + rollout_evaluation(iter, m) + metrics["metrics/rollout_objective_no_target_penalty"] = + rollout_evaluation.last_objective_no_target_penalty + metrics["metrics/rollout_objective_no_deficit"] = + rollout_evaluation.last_objective_no_target_penalty + metrics["metrics/rollout_target_violation_share"] = + rollout_evaluation.last_violation_share + metrics["metrics/rollout_n_ok"] = + rollout_evaluation.last_n_ok + if SAVE_METRIC == "rollout" + rollout_score = rollout_evaluation.last_objective_no_target_penalty + if isfinite(rollout_score) && rollout_score < best_obj + global best_obj = rollout_score + jldsave(MODEL_PATH; model_state = checkpoint_policy_state(m)) + @info " -> New best rollout: $(round(rollout_score; digits=4)) -- saved $MODEL_PATH" + end + end + end + + batch_in_epoch = (iter - 1) % NUM_BATCHES + 1 + if batch_in_epoch == NUM_BATCHES + epoch = (iter - 1) ÷ NUM_BATCHES + 1 + mean_loss = isempty(epoch_losses) ? NaN : mean(epoch_losses) + n_ok = length(epoch_losses) + empty!(epoch_losses) + Wandb.log(lg, Dict("metrics/epoch_objective" => mean_loss, "epoch" => epoch)) + @info "Epoch $epoch/$NUM_EPOCHS mean=$(round(mean_loss; digits=2)) ok=$n_ok/$NUM_BATCHES" + if SAVE_METRIC == "training" && isfinite(mean_loss) && mean_loss < best_obj + global best_obj = mean_loss + jldsave(MODEL_PATH; model_state = checkpoint_policy_state(m)) + @info " → New best: $(round(mean_loss; digits=4)) — saved $MODEL_PATH" + end + end + Wandb.log(lg, metrics) + return false + end, +) + +close(lg) +@info "Done. Best model saved to: $(MODEL_PATH)" diff --git a/src/DecisionRulesExa.jl b/src/DecisionRulesExa.jl index 4e23112..ad622bf 100644 --- a/src/DecisionRulesExa.jl +++ b/src/DecisionRulesExa.jl @@ -1,3 +1,34 @@ +""" + DecisionRulesExa + +GPU-accelerated companion to DecisionRules.jl for Two-Stage Deep Decision Rules +(TS-DDR) training with ExaModels and MadNLP. + +DecisionRulesExa implements the same target-projection workflow as +DecisionRules.jl: + +1. a policy predicts target states, +2. an NLP projects those targets onto the feasible set, and +3. target-constraint multipliers provide the policy-gradient signal. + +The package formulates inner optimization problems as `ExaModels.ExaModel` +instances solved by MadNLP, enabling GPU-native solves and warm-started repeated +training solves. + +# Main Types +- [`DeterministicEquivalentProblem`](@ref): deterministic equivalent with + explicit target parameters. +- [`EmbeddedDeterministicEquivalentProblem`](@ref): deterministic equivalent + whose target policy is embedded with `VectorNonlinearOracle`. +- [`StateConditionedPolicy`](@ref): recurrent policy for sequential target + rollout. +- [`MLPPolicy`](@ref): stateless policy for full-horizon target prediction. + +# Main Training APIs +- [`train_tsddr`](@ref): open-loop target-parameter training. +- [`train_tsddr_embedded`](@ref): embedded-policy training. +- [`rollout_tsddr`](@ref): stage-wise deployment-style evaluation. +""" module DecisionRulesExa using ExaModels @@ -12,6 +43,7 @@ using ChainRulesCore include("utils.jl") include("deterministic_equivalent.jl") +include("embedded_deterministic_equivalent.jl") include("policy.jl") include("critic_control_variate.jl") include("training.jl") @@ -39,9 +71,23 @@ export # Policies MLPPolicy, StateConditionedPolicy, + ContextualPolicy, + context_at, + stage_phase_context, + vcat_contexts, + ConstantStatePolicy, + FixedOutputPolicy, + bounded_state_policy, + load_stateconditioned_policy!, + + # Embedded-NN deterministic equivalent + EmbeddedDeterministicEquivalentProblem, + build_embedded_deterministic_equivalent, + invalidate_policy_cache!, # Training solve_succeeded, + prepare_solve!, materialize_tangent, _all_finite_gradient, AbstractCriticControlVariate, @@ -60,6 +106,7 @@ export critic_samples_from_evaluation, simulate_tsddr, train_tsddr, + train_tsddr_embedded, # Stage-wise rollout evaluation rollout_tsddr, diff --git a/src/critic_control_variate.jl b/src/critic_control_variate.jl index a1953d0..0341984 100644 --- a/src/critic_control_variate.jl +++ b/src/critic_control_variate.jl @@ -9,32 +9,82 @@ abstract type AbstractCriticTrainingTarget end """ DeterministicEquivalentCriticTarget() -Train critic value targets from the full deterministic-equivalent objective. -This is useful for ablations and for pure DE control-variate experiments. +Select critic value targets from the full deterministic-equivalent training +objective. + +# Returns +- `DeterministicEquivalentCriticTarget`: a target selector for deterministic- + equivalent critic supervision. + +# Notes +Use this target for ablations or for control-variate experiments tied to the +training surrogate rather than to deployed rollout performance. """ struct DeterministicEquivalentCriticTarget <: AbstractCriticTrainingTarget end """ - RolloutCriticTarget(stage_problem; kwargs...) - -Train critic value targets from stage-wise rollout evaluation. This is the -preferred target when the critic is meant to guide convergence of the deployed -rollout objective rather than the deterministic-equivalent surrogate. - -Required keyword callbacks match `rollout_tsddr`: -- `set_stage_parameters!` -- `realized_state` - -By default `policy_state = :target`, matching the differentiable target -recurrence used by the actor. Set `policy_state = :realized` to train on -closed-loop realized-state rollout targets. - -By default `objective_value = :objective`, so critic value targets include the -same target-penalty contribution that appears in the dual actor signal. Set -`objective_value = :objective_no_target_penalty` to train on the rollout -objective with target-slack penalties removed. + RolloutCriticTarget( + stage_problem; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + policy_state::Symbol = :target, + reuse_solver::Bool = false, + objective_value::Symbol = :objective, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + ) -> RolloutCriticTarget + +Select critic value targets from stage-wise rollout evaluation. + +# Arguments +- `stage_problem`: single-stage optimization problem template used by + [`rollout_tsddr`](@ref). + +# Keywords +- `horizon::Int`: number of rollout stages. +- `n_uncertainty::Int`: dimension of each per-stage uncertainty slice. +- `set_stage_parameters!::Function`: callback that writes state, + uncertainty, target, and stage index into `stage_problem`. +- `realized_state::Function`: callback that extracts the realized next state + from a solved stage. +- `objective_no_target_penalty::Function`: callback returning the stage + objective with target-tracking penalties removed. +- `madnlp_kwargs`: keyword arguments forwarded to the MadNLP solver wrapper. +- `warmstart::Bool`: whether rollout solves should warm-start from previous + stage information. +- `policy_state::Symbol`: `:target` trains against the differentiable target + recurrence; `:realized` trains against closed-loop realized states. +- `reuse_solver::Bool`: whether rollout evaluation may reuse one solver + object across stages. +- `objective_value::Symbol`: `:objective` uses the full rollout objective; + `:objective_no_target_penalty` removes target-slack penalties. +- `state_bounds`: optional `(lower, upper)` projection bounds for realized + states. +- `project_state`: optional custom projection callback for realized states. +- `retry_on_failure::Bool`: whether failed stage solves should be retried with + a cold-start solver. + +# Returns +- `RolloutCriticTarget`: a target selector carrying the rollout callbacks and + options. + +# Throws +- Throws an error if `policy_state` is not `:target` or `:realized`. +- Throws an error if `objective_value` is not `:objective` or + `:objective_no_target_penalty`. + +# Notes +This is the preferred target when the critic is meant to guide convergence of +the deployed rollout objective rather than the deterministic-equivalent +surrogate. """ -struct RolloutCriticTarget{S,R,O,M} <: AbstractCriticTrainingTarget +struct RolloutCriticTarget{S,R,O,M,B,P} <: AbstractCriticTrainingTarget stage_problem horizon::Int n_uncertainty::Int @@ -46,6 +96,9 @@ struct RolloutCriticTarget{S,R,O,M} <: AbstractCriticTrainingTarget policy_state::Symbol reuse_solver::Bool objective_value::Symbol + state_bounds::B + project_state::P + retry_on_failure::Bool end function RolloutCriticTarget( @@ -60,6 +113,9 @@ function RolloutCriticTarget( policy_state::Symbol = :target, reuse_solver::Bool = false, objective_value::Symbol = :objective, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, ) policy_state in (:target, :realized) || error("policy_state must be :target or :realized") @@ -77,32 +133,66 @@ function RolloutCriticTarget( policy_state, reuse_solver, objective_value, + state_bounds, + project_state, + retry_on_failure, ) end """ NoCriticControlVariate() -Default no-op critic configuration. Passing this to `train_tsddr` recovers the -original dual-multiplier actor update. +Construct the no-op critic configuration. + +# Returns +- `NoCriticControlVariate`: a sentinel that disables critic/control-variate + terms. + +# Notes +Passing this value to `train_tsddr` recovers the original dual-multiplier actor +update. """ struct NoCriticControlVariate <: AbstractCriticControlVariate end """ - ScalarCriticControlVariate(critic; featurizer=default_critic_featurizer, - value_loss_weight=0.1, - gradient_loss_weight=1.0) - -Wrap a scalar Flux-compatible critic `C(w, xhat)` for optional TS-DDR -control-variate training. The critic is called as `critic(features)`, where -`features = featurizer(initial_state, uncertainty, xhat)`. - -The critic loss is - - value_loss_weight * mse(C, objective) - + gradient_loss_weight * mse(gradient(xhat -> C, xhat), target_multipliers) - -Either loss weight may be zero. + ScalarCriticControlVariate( + critic; + featurizer = default_critic_featurizer, + value_loss_weight::Real = 0.1, + gradient_loss_weight::Real = 1.0, + ) -> ScalarCriticControlVariate + +Wrap a scalar Flux-compatible critic for optional TS-DDR control-variate +training. + +# Arguments +- `critic`: callable scalar model evaluated as `critic(features)`. + +# Keywords +- `featurizer`: callable + `featurizer(initial_state, uncertainty, xhat) -> features`. +- `value_loss_weight::Real`: nonnegative weight on objective-value matching. +- `gradient_loss_weight::Real`: nonnegative weight on target-gradient + matching. + +# Returns +- `ScalarCriticControlVariate`: critic configuration with loss weights stored + as `Float64`. + +# Throws +- Throws an error if either loss weight is negative. + +# Notes +For each [`CriticSample`](@ref), the critic loss is + +```math +w_v |C(f) - J|^2 ++ w_g \\frac{1}{n}\\|\\nabla_{\\hat{x}} C(f) - \\lambda_{\\hat{x}}\\|_2^2, +``` + +where `f = featurizer(initial_state, uncertainty, xhat)`, `J` is the scalar +objective target, and ``\\lambda_{\\hat{x}}`` is the target multiplier array. +Either weight may be zero. """ struct ScalarCriticControlVariate{C,F} <: AbstractCriticControlVariate critic::C @@ -128,11 +218,33 @@ function ScalarCriticControlVariate( end """ - CriticSample(initial_state, uncertainty, xhat, objective_value, - target_multipliers; metadata=nothing) + CriticSample( + initial_state, + uncertainty, + xhat, + objective_value::Real, + target_multipliers; + metadata = nothing, + ) -> CriticSample + +Store one already-solved TS-DDR scenario as scalar-critic supervision. + +# Arguments +- `initial_state`: initial state used for the scenario. +- `uncertainty`: scenario uncertainty trajectory. +- `xhat`: policy target trajectory. +- `objective_value::Real`: scalar value target for the critic. +- `target_multipliers`: multiplier-like target with the same shape as `xhat`. + +# Keywords +- `metadata`: optional payload retained with the sample. + +# Returns +- `CriticSample`: sample with `objective_value` converted to `Float64`. -Training sample for a scalar critic. Samples are produced from already-solved -TS-DDR scenarios and do not require additional optimization solves. +# Notes +Creating a `CriticSample` does not run any optimization solve; samples are +intended to be built from existing training or rollout results. """ struct CriticSample{I,W,X,L,M} initial_state::I @@ -161,6 +273,21 @@ function CriticSample( ) end +""" + CriticReplayBuffer(max_size::Integer) -> CriticReplayBuffer + +Construct a fixed-capacity FIFO replay buffer for [`CriticSample`](@ref)s. + +# Arguments +- `max_size::Integer`: maximum number of samples retained; negative values are + clamped to zero. + +# Returns +- `CriticReplayBuffer`: empty replay buffer with capacity `max(0, max_size)`. + +# Notes +A capacity of zero disables buffering, so push operations become no-ops. +""" mutable struct CriticReplayBuffer{S} samples::Vector{S} max_size::Int @@ -169,6 +296,22 @@ end CriticReplayBuffer(max_size::Integer) = CriticReplayBuffer{Any}(Any[], max(0, Int(max_size))) +""" + push_critic_sample!(buffer, sample) -> buffer + +Append one [`CriticSample`](@ref) to the buffer, evicting the oldest sample if +the buffer is at capacity. + +# Arguments +- `buffer::CriticReplayBuffer`: replay buffer to mutate. +- `sample::CriticSample`: sample to append. + +# Returns +- `buffer`: the same buffer object, after optional insertion and eviction. + +# Notes +If `buffer.max_size == 0`, the function returns without storing `sample`. +""" function push_critic_sample!(buffer::CriticReplayBuffer, sample::CriticSample) buffer.max_size == 0 && return buffer push!(buffer.samples, sample) @@ -177,6 +320,18 @@ function push_critic_sample!(buffer::CriticReplayBuffer, sample::CriticSample) return buffer end +""" + push_critic_samples!(buffer, samples) -> buffer + +Append multiple [`CriticSample`](@ref)s to the buffer in order. + +# Arguments +- `buffer::CriticReplayBuffer`: replay buffer to mutate. +- `samples`: iterable of [`CriticSample`](@ref) values. + +# Returns +- `buffer`: the same buffer after appending the supplied samples. +""" function push_critic_samples!(buffer::CriticReplayBuffer, samples) for sample in samples push_critic_sample!(buffer, sample) @@ -185,10 +340,17 @@ function push_critic_samples!(buffer::CriticReplayBuffer, samples) end """ - default_critic_featurizer(initial_state, uncertainty, xhat) + default_critic_featurizer(initial_state, uncertainty, xhat) -> AbstractVector + +Concatenate flattened critic inputs into one feature vector. + +# Arguments +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. -Default critic featurizer: concatenate flattened initial state, uncertainty, and -policy target trajectory. +# Returns +- `AbstractVector`: `vcat(vec(initial_state), vec(uncertainty), vec(xhat))`. """ default_critic_featurizer(initial_state, uncertainty, xhat) = vcat(vec(initial_state), vec(uncertainty), vec(xhat)) @@ -205,9 +367,21 @@ function _critic_value(critic, featurizer, initial_state, uncertainty, xhat) end """ - critic_value(control_variate, initial_state, uncertainty, xhat) + critic_value(control_variate, initial_state, uncertainty, xhat) -> Number Evaluate the scalar critic on one scenario. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. + +# Returns +- `Number`: scalar critic prediction. + +# Throws +- Throws an error if the critic returns a non-scalar array. """ critic_value( cv::ScalarCriticControlVariate, @@ -219,8 +393,23 @@ critic_value( """ critic_xhat_gradient(control_variate, initial_state, uncertainty, xhat) -Return `gradient(xhat -> C(initial_state, uncertainty, xhat), xhat)` and check -that it has the same shape as `xhat`. +Differentiate the scalar critic with respect to the policy target trajectory. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. + +# Returns +- An array with the same shape as `xhat`, equal to + `gradient(x -> critic_value(control_variate, initial_state, uncertainty, x), xhat)`. + +# Throws +- Throws an error if the critic gradient shape differs from `xhat`. + +# Notes +If Zygote reports `nothing`, the gradient is replaced by `zero(xhat)`. """ function critic_xhat_gradient( cv::ScalarCriticControlVariate, @@ -285,9 +474,32 @@ function _critic_loss_with( end """ - critic_loss(control_variate, samples; value_loss_weight, gradient_loss_weight) + critic_loss( + control_variate, + samples; + value_loss_weight = control_variate.value_loss_weight, + gradient_loss_weight = control_variate.gradient_loss_weight, + ) -> Real Compute the scalar critic loss on a collection of `CriticSample`s. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `samples`: collection of [`CriticSample`](@ref) values. + +# Keywords +- `value_loss_weight`: nonnegative override for objective-value loss weight. +- `gradient_loss_weight`: nonnegative override for target-gradient loss + weight. + +# Returns +- `Real`: average critic loss over `samples`, or `0.0` for an empty + collection. + +# Throws +- Throws an error if either loss weight is negative. +- Throws an error if a sample's target multipliers or critic gradient do not + match the shape of `xhat`. """ critic_loss(cv::ScalarCriticControlVariate, samples; kwargs...) = _critic_loss_with(cv.critic, cv, samples; kwargs...) @@ -303,10 +515,32 @@ function _critic_minibatch(samples, batch_size) end """ - update_critic!(opt_state, control_variate, samples; batch_size=nothing) - -Run one critic optimizer step and return the numeric loss. Only critic -parameters are updated. + update_critic!( + opt_state, + control_variate, + samples; + batch_size = nothing, + ) -> Float64 + +Run one optimizer step for the scalar critic. + +# Arguments +- `opt_state`: Flux optimizer state for `control_variate.critic`. +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `samples`: replay samples available for training. + +# Keywords +- `batch_size`: optional minibatch size; `nothing` or a value greater than the + sample count uses all samples. + +# Returns +- `Float64`: critic loss on the selected batch, or `NaN` when the selected + batch is empty. + +# Notes +Only critic parameters are updated. If the materialized gradient is `nothing` +or contains non-finite values, the optimizer update is skipped while the loss +is still reported. """ function update_critic!( opt_state, diff --git a/src/deterministic_equivalent.jl b/src/deterministic_equivalent.jl index bf672b5..6c422da 100644 --- a/src/deterministic_equivalent.jl +++ b/src/deterministic_equivalent.jl @@ -9,15 +9,32 @@ """ DeterministicEquivalentProblem -Holds an ExaModels parametric NLP for the deterministic equivalent subproblem +Container for an ExaModels parametric NLP representing a deterministic +equivalent subproblem. - Q(w, x̂) = min_{x,u,δ} Σ_t stage_cost(t, x_t, u_t, w_t) + (ρ/2)‖δ‖² - s.t. x₁ = x₀ - dynamics(t, x_t, u_t, w_t, x_{t+1}) = 0 - x̂_t - x_t - δ_t = 0 (target constraints, **added last**) +# Fields -We add the target constraints last so their multipliers are a contiguous slice of -`result.multipliers`. +- `core`: ExaModels core used to build variables, parameters, objectives, and + constraints. +- `model`: ExaModels model passed to MadNLP. +- `x`, `u`, `δ`: Flat state, control, and target-slack variables. +- `p_x0`, `p_w`, `p_target`: Initial-state, uncertainty, and target parameters. +- `nx`, `nu`, `nw`, `horizon`: State, control, uncertainty, and time dimensions. +- `target_con_range`: Range of target-constraint multipliers in solver results. + +# Notes + +The subproblem has the form + +```math +Q(w, \\hat{x}) = + \\min_{x,u,\\delta} + \\sum_t c_t(x_t, u_t, w_t) + \\frac{\\rho}{2}\\|\\delta\\|^2 +``` + +subject to the initial condition, dynamics constraints, and target constraints +``\\hat{x}_t - x_t - \\delta_t = 0``. The target constraints are added last so +their multipliers occupy the contiguous slice `target_con_range`. """ struct DeterministicEquivalentProblem core @@ -42,8 +59,17 @@ end """ MadNLPCache -Optional cache to re-use a MadNLP solver instance across repeated solves -(warm-start + reusing symbolic factorizations). +Cache for reusing a MadNLP solver across repeated solves. + +# Fields + +- `solver`: Cached `MadNLP.MadNLPSolver`. +- `last_result`: Most recent solver result, used for optional warm starts. + +# Notes + +Reusing the solver can avoid repeated symbolic setup and can warm-start the +next primal iterate from the previous solution. """ mutable struct MadNLPCache solver @@ -53,27 +79,45 @@ end """ build_deterministic_equivalent(; kwargs...) -> DeterministicEquivalentProblem -Generic builder for a deterministic-equivalent dynamic NLP. - -Keyword arguments: -- `horizon::Int` : number of stages T (states are 1..T, controls are 1..T-1) -- `nx::Int` : state dimension -- `nu::Int` : control dimension (demo assumes `nu == nx` for default dynamics) -- `nw::Int` : disturbance dimension (default: `nx`) -- `backend` : ExaModels backend (e.g., `nothing` for CPU, `CUDABackend()` for GPU) -- `float_type` : numeric type (Float64 recommended for MadNLP) -- `x_bounds` : `(lb, ub)` applied to all state variables -- `u_bounds` : `(lb, ub)` applied to all control variables -- `slack_penalty` : ρ ≥ 0, weight for (ρ/2)‖δ‖² - -- `dynamics_eq` : function `(t, i, x, u, w, nx, nu, nw) -> expr == 0` - returns the scalar equality residual for dimension `i` at stage `t`. -- `stage_cost` : function `(t, i, x, u, w, nx, nu, nw) -> expr` returns a scalar term. - -Notes: -- All variables are stored in flat vectors to keep the interface simple and robust. -- Target constraints are added *last* and written as `x̂ - x - δ = 0` so that their - multipliers are directly the gradient w.r.t. `x̂` (envelope theorem). +Build a deterministic-equivalent dynamic nonlinear program. + +# Keywords + +- `horizon::Int`: Number of stages. States are indexed over `1:horizon`; + controls and uncertainties over `1:(horizon - 1)`. +- `nx::Int`: State dimension. +- `nu::Int = nx`: Control dimension. +- `nw::Int = nx`: Uncertainty dimension. +- `backend = nothing`: ExaModels backend, such as `nothing` for CPU or a CUDA + backend for GPU execution. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: Lower and upper bounds applied + to every state variable. +- `u_bounds = (-Inf, Inf)`: Control bounds, either a scalar `(lb, ub)` tuple or + a tuple of length-`nu` lower and upper bound vectors. +- `slack_penalty::Real = 1.0`: Nonnegative target-slack penalty weight ``ρ``. +- `dynamics_eq::Function = default_dynamics_eq`: Function + `(t, i, x, u, w, nx, nu, nw) -> residual` defining one scalar dynamics + equality. +- `stage_cost::Function = default_stage_cost`: Function + `(t, i, x, u, w, nx, nu, nw) -> term` defining one scalar objective term. + +# Returns + +- `DeterministicEquivalentProblem`: A mutable problem container with parameters + that can be updated by `set_x0!`, `set_uncertainty!`, and `set_targets!`. + +# Throws + +Throws an error if dimensions are invalid, if vector control bounds have the +wrong length, or if the default dynamics/cost are used with dimensions other +than `nu == nx` and `nw == nx`. + +# Notes + +All variables are stored in flat vectors. Target constraints are added last and +written as ``\\hat{x} - x - \\delta = 0`` so the envelope theorem identifies their +multipliers with gradients with respect to the target trajectory. """ function build_deterministic_equivalent(; horizon::Int, @@ -180,11 +224,28 @@ end """ build_linear_tracking_problem(; kwargs...) -Convenience wrapper that builds a *simple* deterministic equivalent problem with: -- dynamics: x_{t+1} = x_t + u_t + w_t -- stage cost: (1/2) x_t^2 + (1/2) u_t^2 +Build the default linear-quadratic tracking demonstration problem. + +# Keywords + +- `horizon::Int`: Number of stages. +- `nx::Int = 1`: State, control, and uncertainty dimension. +- `backend = nothing`: ExaModels backend. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: State bounds. +- `u_bounds::Tuple{<:Real,<:Real} = (-1.0, 1.0)`: Control bounds. +- `slack_penalty::Real = 10.0`: Target-slack penalty weight. -This is meant as an end-to-end demo and a template for your real model. +# Returns + +- `DeterministicEquivalentProblem`: Problem with dynamics + ``x_{t+1} = x_t + u_t + w_t`` and stage cost + ``(x_t^2 + u_t^2) / 2``. + +# Notes + +This helper is intended as a small end-to-end example and as a template for +model-specific deterministic-equivalent builders. """ function build_linear_tracking_problem(; horizon::Int, @@ -215,11 +276,28 @@ end # -------------------------- """ -Default per-dimension dynamics residual for the demo: + default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) + +Return the default scalar dynamics residual. + +# Arguments - x_{t+1,i} - x_{t,i} - u_{t,i} - w_{t,i} == 0 +- `t`: Stage index. +- `i`: State component index. +- `x`: Flat state variable vector. +- `u`: Flat control variable vector. +- `w`: Flat uncertainty parameter vector. +- `nx::Int`: State dimension. +- `nu::Int`: Control dimension. +- `nw::Int`: Uncertainty dimension. -Assumes `nu == nx` and `nw == nx`. +# Returns + +- Scalar residual ``x_{t+1,i} - x_{t,i} - u_{t,i} - w_{t,i}``. + +# Notes + +The default residual assumes `nu == nx` and `nw == nx`. """ function default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) return x[x_index(nx, t + 1, i)] - @@ -229,9 +307,24 @@ function default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) end """ -Default per-dimension stage cost term for the demo: + default_stage_cost(t, i, x, u, w, nx::Int, nu::Int, nw::Int) + +Return the default scalar stage-cost term. - (1/2) x_{t,i}^2 + (1/2) u_{t,i}^2 +# Arguments + +- `t`: Stage index. +- `i`: State/control component index. +- `x`: Flat state variable vector. +- `u`: Flat control variable vector. +- `w`: Flat uncertainty parameter vector. +- `nx::Int`: State dimension. +- `nu::Int`: Control dimension. +- `nw::Int`: Uncertainty dimension. + +# Returns + +- Scalar cost ``(x_{t,i}^2 + u_{t,i}^2) / 2``. """ function default_stage_cost(t, i, x, u, w, nx::Int, nu::Int, nw::Int) return (x[x_index(nx, t, i)]^2 + u[u_index(nu, t, i)]^2) / 2 @@ -242,9 +335,22 @@ end # -------------------------- """ - set_x0!(prob, x0) + set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) + +Update the initial-state parameter. + +# Arguments -Update initial state parameter (length nx). +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `x0::AbstractVector`: Initial state with length `prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(x0) != prob.nx`. """ function set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) length(x0) == prob.nx || error("x0 length must be nx=$(prob.nx), got $(length(x0))") @@ -253,12 +359,29 @@ function set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) end """ - set_uncertainty!(prob, w) + set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVector) + +Update the disturbance-trajectory parameter. -Update disturbance trajectory parameter. -`w` must have length `(T-1)*nw` or `T*nw`; in the latter case the first `(T-1)*nw` -elements are used (the last stage's uncertainty only enters the policy rollout, not -the NLP dynamics). +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `w::AbstractVector`: Disturbance trajectory with length `(T - 1) * nw` or + `T * nw`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `w` has any length other than `(T - 1) * nw` or `T * nw`. + +# Notes + +When `w` has length `T * nw`, only the first `(T - 1) * nw` entries are used in +the NLP dynamics. The final-stage uncertainty may be needed by a policy rollout +but does not enter these dynamics constraints. """ function set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVector) expected = (prob.horizon - 1) * prob.nw @@ -274,9 +397,22 @@ function set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVecto end """ - set_targets!(prob, xhat) + set_targets!(prob::DeterministicEquivalentProblem, xhat::AbstractVector) + +Update the target-trajectory parameter. + +# Arguments -Update target trajectory parameter (length T*nx). +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `xhat::AbstractVector`: Target trajectory with length `prob.horizon * prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(xhat) != prob.horizon * prob.nx`. """ function set_targets!(prob::DeterministicEquivalentProblem, xhat::AbstractVector) expected = prob.horizon * prob.nx @@ -292,7 +428,20 @@ end """ init_madnlp_cache(prob; solver_kwargs...) -> MadNLPCache -Create and store a `MadNLP.MadNLPSolver` for repeated solves. +Create a cached MadNLP solver for repeated solves. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem whose ExaModel will be solved. + +# Keywords + +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.MadNLPSolver`. + +# Returns + +- `MadNLPCache`: Cache containing the solver and an initially empty + `last_result`. """ function init_madnlp_cache(prob::DeterministicEquivalentProblem; solver_kwargs...) solver = MadNLP.MadNLPSolver(prob.model; solver_kwargs...) @@ -300,19 +449,62 @@ function init_madnlp_cache(prob::DeterministicEquivalentProblem; solver_kwargs.. end """ - solve!(prob; solver_kwargs...) -> result + solve!(prob::DeterministicEquivalentProblem; solver_kwargs...) -> result -Solve once by instantiating a fresh MadNLP solver (simplest, but allocates). +Solve a deterministic-equivalent problem with a fresh MadNLP solver. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to solve. + +# Keywords + +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.madnlp`. + +# Returns + +- `result`: MadNLP result object. + +# Notes + +This path is simple but allocates a new solver for every call. """ function solve!(prob::DeterministicEquivalentProblem; solver_kwargs...) return MadNLP.madnlp(prob.model; solver_kwargs...) end """ - solve!(prob, cache; warmstart=true, solver_kwargs...) -> result + solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; warmstart=true, solver_kwargs...) -> result + +Solve a deterministic-equivalent problem with a cached MadNLP solver. -Solve using a cached solver instance. Optionally warm-start the primal -variables from the previous solution. +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to solve. +- `cache::MadNLPCache`: Cached solver and previous result. + +# Keywords + +- `warmstart::Bool = true`: Whether to copy the previous primal solution into + the model initial point before solving. +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.solve!`. + +# Returns + +- `result`: MadNLP result object, also stored in `cache.last_result`. + +# Notes + +Warm starts are used only when `warmstart` is true and `cache.last_result` is +available. + +MadNLP's iteration counter `cnt.k` is cumulative across `solve!` calls on the +same solver instance and is never reset by MadNLP itself. Without resetting it +(together with `cnt.acceptable_cnt` and `cnt.start_time`), repeated solves on a +cached solver eventually exhaust `max_iter` spuriously — the counters are reset +here before each solve so every call gets its intended per-solve iteration and +wall-clock budget. This mirrors the reset in `_solve!` (training.jl) and does +not change any numerics of an individual solve. """ function solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; warmstart::Bool = true, @@ -322,6 +514,10 @@ function solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; # Warm-start primal from previous solution copyto!(NLPModels.get_x0(prob.model), cache.last_result.solution) end + # Reset per-solve iteration budget (cnt.k is cumulative in MadNLP). + cache.solver.cnt.k = 0 # reset iteration counter + cache.solver.cnt.acceptable_cnt = 0 # reset acceptable-step counter + cache.solver.cnt.start_time = time() # reset wall-clock timer res = MadNLP.solve!(cache.solver; solver_kwargs...) cache.last_result = res return res @@ -332,17 +528,43 @@ end # -------------------------- """ - target_multipliers(prob, result) -> λ + target_multipliers(prob::DeterministicEquivalentProblem, result) -> λ + +Return the dual multipliers associated with target constraints. -Return the dual multipliers associated with the target constraints (∇_{x̂} Q). +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem that defines the multiplier + slice. +- `result`: MadNLP result containing `multipliers`. + +# Returns + +- `λ`: Multipliers in `result.multipliers[prob.target_con_range]`. + +# Notes + +With target constraints written as ``\\hat{x} - x - \\delta = 0``, these +multipliers are the envelope-theorem derivatives with respect to the target +trajectory. """ target_multipliers(prob::DeterministicEquivalentProblem, result) = result.multipliers[prob.target_con_range] """ - solution_components(prob, result) -> (x, u, δ) + solution_components(prob::DeterministicEquivalentProblem, result) -> (x, u, δ) Split the flat solution vector into state, control, and slack components. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem that defines component sizes. +- `result`: MadNLP result containing `solution`. + +# Returns + +- `(x, u, δ)`: Flat slices of the primal solution for states, controls, and + target slacks. """ function solution_components(prob::DeterministicEquivalentProblem, result) n_x = prob.horizon * prob.nx diff --git a/src/embedded_deterministic_equivalent.jl b/src/embedded_deterministic_equivalent.jl new file mode 100644 index 0000000..4975955 --- /dev/null +++ b/src/embedded_deterministic_equivalent.jl @@ -0,0 +1,580 @@ +# embedded_deterministic_equivalent.jl +# +# Embedded-NN deterministic equivalent: the policy π_θ is inside the NLP via +# VectorNonlinearOracle. At convergence the duals λ_t are closed-loop (joint +# NLP) and the gradient ∇_θ Q = Σ_t λ_t · ∇_θ π_θ(w_t, x*_{t-1}) follows +# from the envelope theorem — structurally identical to the open-loop formula +# but with realized states from the coupled solve. +# +# The oracle constraint is: +# π_θ(w_t, x_{t-1}) − x_t − δ_t = 0 ∀t = 1…T +# +# One oracle for all T stages guarantees sequential LSTM evaluation with +# Flux.reset! at the top of each callback invocation. Flux.reset! on the +# threaded policies is a REAL reset (it restores Flux.initialstates), and each +# stage's policy forward advances the recurrent state exactly once, so within +# every callback the policy sees the stage sequence t = 1…T from a fresh +# initial state. Audit of per-stage advancement: oracle_f! calls the policy +# once per stage; oracle_jac!/oracle_vjp! call it inside Zygote.pullback for +# t > 1 (the pullback CONSTRUCTION runs the forward exactly once; calling the +# returned back(·) does not re-run it) and as a bare call for t = 1 — one +# forward, hence one state advance, per stage in all three callbacks. +# +# Jacobian exactness caveat: the oracle callbacks (oracle_jac!, oracle_vjp!) +# compute only the DIRECT partial ∂π_t/∂x_{t-1} via a per-stage Zygote pullback. +# For a recurrent policy whose recurrent layers see x_{t-1}, stage t's output +# also depends on x_{1..t-2} through the hidden state; those cross-stage entries +# are absent from both the sparsity pattern and the pullbacks, so the oracle +# Jacobian is exact for feedforward (stateless-in-x) policies and a structural +# approximation for such recurrent ones. When the recurrent encoder reads only +# the uncertainty w_t (as in StateConditionedPolicy and HydroReachablePolicy, +# where the combiner over [h_t; x_{t-1}] is feedforward in x), the hidden state +# does not depend on x and the direct partial IS the full derivative. This +# exactness claim SURVIVES recurrent-state threading: with the encoder reading +# only w_t, the threaded hidden state depends only on the inflow history +# w_{1..t}, never on x, so threading changes the VALUE of h_t but adds no +# ∂h_t/∂x dependence — ∂π_t/∂x_{t-1} remains the full derivative. + +""" + EmbeddedDeterministicEquivalentProblem + +Container for a deterministic-equivalent NLP whose target constraints are +computed by an embedded Flux policy. + +# Fields + +- `core`: ExaModels core used to build variables, parameters, objectives, and + constraints. +- `model`: ExaModels model passed to MadNLP. +- `x`, `u`, `δ`: Flat state, control, and target-slack variables. +- `p_x0`, `p_w`: Initial-state and uncertainty parameters. +- `policy`: Flux policy evaluated by the nonlinear oracle. +- `nx`, `nu`, `nw`, `horizon`: State, control, uncertainty, and time dimensions. +- `target_con_range`: Range of oracle-constraint multipliers in solver results. +- `_w_buf`, `_x0_buf`: Mutable host buffers captured by oracle callbacks. + +# Notes + +This problem is analogous to `DeterministicEquivalentProblem`, but the explicit +target parameter is replaced by a `VectorNonlinearOracle` enforcing +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0``. The oracle closures capture +`policy` by reference, so changing Flux parameters between solves changes the +NLP callbacks without rebuilding the ExaModel. +""" +struct EmbeddedDeterministicEquivalentProblem{P} + core + model + x + u + δ + p_x0 + p_w + policy::P + nx::Int + nu::Int + nw::Int + horizon::Int + target_con_range::UnitRange{Int} + # mutable buffers captured by oracle closures + _w_buf::Vector{Float64} + _x0_buf::Vector{Float64} +end + +""" + set_x0!(prob::EmbeddedDeterministicEquivalentProblem, x0::AbstractVector) + +Update the initial state used by the embedded problem and oracle callbacks. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Embedded problem to update. +- `x0::AbstractVector`: Initial state with length `prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(x0) != prob.nx`. + +# Notes + +The value is written both to the ExaModels parameter and to the oracle closure +buffer, ensuring the policy callback sees the same initial state as the NLP. +""" +function set_x0!(prob::EmbeddedDeterministicEquivalentProblem, x0::AbstractVector) + length(x0) == prob.nx || error("x0 length must be nx=$(prob.nx), got $(length(x0))") + ExaModels.set_parameter!(prob.core, prob.p_x0, x0) + copyto!(prob._x0_buf, Float64.(x0)) + return prob +end + +""" + set_uncertainty!(prob::EmbeddedDeterministicEquivalentProblem, w::AbstractVector) + +Update the disturbance trajectory used by the embedded problem and oracle +callbacks. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Embedded problem to update. +- `w::AbstractVector`: Disturbance trajectory with length `(T - 1) * nw` or + `T * nw`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `w` has any length other than `(T - 1) * nw` or `T * nw`. + +# Notes + +The first `(T - 1) * nw` entries are passed to the ExaModels dynamics +parameter. The full length-`T * nw` trajectory is retained in the oracle buffer +when provided, because the policy callback may evaluate a final-stage +disturbance. +""" +function set_uncertainty!(prob::EmbeddedDeterministicEquivalentProblem, w::AbstractVector) + expected = (prob.horizon - 1) * prob.nw + full_len = prob.horizon * prob.nw + n = length(w) + if n == expected + ExaModels.set_parameter!(prob.core, prob.p_w, w) + copyto!(view(prob._w_buf, 1:expected), Float64.(w)) + elseif n == full_len + ExaModels.set_parameter!(prob.core, prob.p_w, view(w, 1:expected)) + copyto!(prob._w_buf, Float64.(w)) + else + error("w length must be (T-1)*nw=$expected or T*nw=$full_len, got $n") + end + return prob +end + +""" + set_targets!(::EmbeddedDeterministicEquivalentProblem, ::AbstractVector) -> nothing + +Ignore explicit targets for embedded deterministic-equivalent problems. + +# Arguments + +- `::EmbeddedDeterministicEquivalentProblem`: Embedded problem whose targets are + generated by the policy oracle. +- `::AbstractVector`: Ignored target vector. + +# Returns + +- `nothing`. + +# Notes + +Embedded problems compute targets inline through the nonlinear oracle instead of +storing them in an NLP parameter. +""" +function set_targets!(::EmbeddedDeterministicEquivalentProblem, ::AbstractVector) + return nothing +end + +""" + invalidate_policy_cache!(embedded_de) + +Invalidate policy-dependent caches for an embedded deterministic-equivalent +problem. + +# Arguments + +- `embedded_de`: Embedded deterministic-equivalent problem. + +# Returns + +- `embedded_de`. + +# Notes + +The generic embedded problem evaluates the policy directly in each oracle +callback and has no cache to invalidate. Specialized embedded problem types may +extend this hook when their oracle stores policy-dependent intermediates across +solver calls. +""" +function invalidate_policy_cache!(embedded_de) + return embedded_de +end + +""" + build_embedded_deterministic_equivalent(policy; kwargs...) + +Build a deterministic-equivalent NLP with a Flux policy embedded as a nonlinear +oracle. + +# Arguments + +- `policy`: Flux model mapping each stage input to an `nx`-vector target. + +# Keywords + +- `horizon::Int`: Number of stages. States are indexed over `1:horizon`; + controls and uncertainties over `1:(horizon - 1)`. +- `nx::Int`: State dimension and policy output dimension. +- `nu::Int = nx`: Control dimension. +- `nw::Int = nx`: Uncertainty dimension and first part of the policy input. +- `backend = nothing`: ExaModels backend, such as `nothing` for CPU or a CUDA + backend for GPU execution. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: Lower and upper bounds applied + to every state variable. +- `u_bounds = (-Inf, Inf)`: Control bounds, either a scalar `(lb, ub)` tuple or + a tuple of length-`nu` lower and upper bound vectors. +- `slack_penalty::Real = 1.0`: Nonnegative target-slack penalty weight ``ρ``. +- `dynamics_eq::Function = default_dynamics_eq`: Function + `(t, i, x, u, w, nx, nu, nw) -> residual` defining one scalar dynamics + equality. +- `stage_cost::Function = default_stage_cost`: Function + `(t, i, x, u, w, nx, nu, nw) -> term` defining one scalar objective term. + +# Returns + +- `EmbeddedDeterministicEquivalentProblem`: Problem supporting `set_x0!`, + `set_uncertainty!`, `target_multipliers`, and `solution_components`. + +# Throws + +Throws an error if dimensions are invalid, if vector control bounds have the +wrong length, or if the default dynamics/cost are used with dimensions other +than `nu == nx` and `nw == nx`. + +# Notes + +The oracle enforces +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0`` and is added last so its +multipliers form a contiguous trailing slice of `result.multipliers`. This +matches the open-loop deterministic-equivalent convention while allowing the +policy to depend on realized previous states. + +The oracle Jacobian and vector-Jacobian callbacks differentiate each stage +independently, providing only the direct partial ``\\partial \\pi_t / +\\partial x_{t-1}``. If the policy's recurrent layers consume ``x_{t-1}``, +stage ``t``'s output also depends on ``x_{1..t-2}`` through the hidden state, +and those cross-stage Jacobian entries are omitted — the reported Jacobian is +then a structural approximation. The Jacobian is exact whenever the recurrent +part of the policy reads only ``w_t`` and the state enters through a +feedforward head (the [`StateConditionedPolicy`](@ref) architecture), because +then the hidden state carries no dependence on ``x``. This holds unchanged +with recurrent-state threading: the threaded hidden state is a function of +``w_{1..t}`` only, so it changes the value of ``h_t`` but introduces no +``\\partial h_t / \\partial x`` term. + +Every callback resets the policy's recurrent state at its top and evaluates +the stages in order, advancing the state exactly once per stage (for +``t > 1`` the forward runs during `Zygote.pullback` construction; the +returned pullback does not re-run it). +""" +function build_embedded_deterministic_equivalent( + policy; + horizon::Int, + nx::Int, + nu::Int = nx, + nw::Int = nx, + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf), + u_bounds = (-Inf, Inf), + slack_penalty::Real = 1.0, + dynamics_eq::Function = default_dynamics_eq, + stage_cost::Function = default_stage_cost, +) + horizon ≥ 2 || error("horizon must be ≥ 2 (got $horizon)") + nx ≥ 1 || error("nx must be ≥ 1 (got $nx)") + nu ≥ 1 || error("nu must be ≥ 1 (got $nu)") + nw ≥ 1 || error("nw must be ≥ 1 (got $nw)") + + if (dynamics_eq === default_dynamics_eq || stage_cost === default_stage_cost) && (nu != nx || nw != nx) + error("Default dynamics/cost assume nu == nx and nw == nx.") + end + + T = horizon + n_x = T * nx + n_u = (T - 1) * nu + + core = ExaModels.ExaCore(float_type; backend = backend) + + function _u_bound(b, side) + v = b[side] + v isa AbstractVector || return float_type(v) + length(v) == nu || error("u_bounds[$side] length must be nu=$nu") + return float_type.(repeat(v, T - 1)) + end + lvar_u = _u_bound(u_bounds, 1) + uvar_u = _u_bound(u_bounds, 2) + + x = ExaModels.variable(core, n_x; + lvar = float_type(x_bounds[1]), + uvar = float_type(x_bounds[2]), + ) + u = ExaModels.variable(core, n_u; + lvar = lvar_u, + uvar = uvar_u, + ) + δ = ExaModels.variable(core, n_x) + + p_x0 = ExaModels.parameter(core, zeros(float_type, nx)) + p_w = ExaModels.parameter(core, zeros(float_type, (T - 1) * nw)) + + ExaModels.objective(core, + stage_cost(t, i, x, u, p_w, nx, nu, nw) + for t in 1:(T - 1), i in 1:nx + ) + ρ = float_type(slack_penalty) + ExaModels.objective(core, + (ρ / 2) * δ[x_index(nx, t, i)]^2 + for t in 1:T, i in 1:nx + ) + + ExaModels.constraint(core, + x[x_index(nx, 1, i)] - p_x0[i] + for i in 1:nx + ) + ExaModels.constraint(core, + dynamics_eq(t, i, x, u, p_w, nx, nu, nw) + for t in 1:(T - 1), i in 1:nx + ) + + n_con_before_oracle = nx + (T - 1) * nx + + # ── Oracle buffers (mutated by set_x0! / set_uncertainty!) ─────────── + w_buf = zeros(Float64, T * nw) + x0_buf = zeros(Float64, nx) + + nvar_total = n_x + n_u + n_x # x, u, δ + x_start = 1 + δ_start = n_x + n_u + 1 + + # ── Pre-allocated oracle buffers ──────────────────────────────────── + _x_prev = zeros(Float32, nx) + _w_t = zeros(Float32, nw) + _input = zeros(Float32, nw + nx) + _J = zeros(Float32, nx, nx) + _e = zeros(Float32, nx) + _λ_t = zeros(Float32, nx) + + function _fill_x_prev!(t, xv) + for i in 1:nx + _x_prev[i] = (t == 1) ? + Float32(x0_buf[i]) : + Float32(xv[x_start + (t-2)*nx + i - 1]) + end + return _x_prev + end + + function _fill_w_t!(t) + for j in 1:nw + _w_t[j] = Float32(w_buf[(t-1)*nw + j]) + end + return _w_t + end + + function _fill_input!(t, xv) + _fill_w_t!(t) + _fill_x_prev!(t, xv) + copyto!(view(_input, 1:nw), _w_t) + copyto!(view(_input, nw+1:nw+nx), _x_prev) + return _input + end + + # ── Oracle callbacks ───────────────────────────────────────────────── + + function oracle_f!(c, xv) + # Real reset: the stage loop below starts from the initial recurrent + # state and each policy call advances it exactly once. + Flux.reset!(policy) + for t in 1:T + _fill_input!(t, xv) + nn_out = policy(_input) + for i in 1:nx + row = (t - 1) * nx + i + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + c[row] = Float64(nn_out[i]) - xv[xi] - xv[di] + end + end + return nothing + end + + # NOTE: this Jacobian holds only the per-stage direct partial ∂π_t/∂x_{t-1}. + # Cross-stage terms through a recurrent hidden state that depends on x are + # not represented (see file-top comment); exact when the recurrent encoder + # reads only w_t, as in StateConditionedPolicy / HydroReachablePolicy. + function oracle_jac!(vals, xv) + # Real reset; each stage advances the recurrent state exactly once: + # for t > 1 the forward runs during Zygote.pullback construction, and + # calling back(·) repeatedly does NOT re-run it; t = 1 is a bare call. + Flux.reset!(policy) + k = 0 + for t in 1:T + _fill_x_prev!(t, xv) + _fill_w_t!(t) + + nn_jac_xprev = if t > 1 + _, back = Zygote.pullback(xp -> policy(vcat(_w_t, xp)), _x_prev) + fill!(_J, 0f0) + for row in 1:nx + fill!(_e, 0f0) + _e[row] = 1.0f0 + col_grad = back(_e)[1] + if col_grad !== nothing + _J[row, :] .= col_grad + end + end + _J + else + # t = 1: no x-Jacobian block, but the forward must still run + # once so the recurrent state advances to stage 2. + policy(vcat(_w_t, _x_prev)) + nothing + end + + for i in 1:nx + k += 1; vals[k] = -1.0 + k += 1; vals[k] = -1.0 + if t > 1 + for j in 1:nx + k += 1; vals[k] = Float64(nn_jac_xprev[i, j]) + end + end + end + end + return nothing + end + + function oracle_vjp!(Jtv, xv, λ) + fill!(Jtv, 0.0) + # Real reset; same one-advance-per-stage discipline as oracle_jac!. + Flux.reset!(policy) + for t in 1:T + _fill_x_prev!(t, xv) + _fill_w_t!(t) + + for i in 1:nx + _λ_t[i] = Float32(λ[(t-1)*nx + i]) + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + Jtv[xi] -= λ[(t-1)*nx + i] + Jtv[di] -= λ[(t-1)*nx + i] + end + + if t > 1 + _, back = Zygote.pullback(xp -> policy(vcat(_w_t, xp)), _x_prev) + dinput = back(_λ_t)[1] + if dinput !== nothing + for j in 1:nx + xj = x_start + (t-2)*nx + j - 1 + Jtv[xj] += Float64(dinput[j]) + end + end + else + # t = 1: no x_prev contribution, but run the forward once so + # the recurrent state advances to stage 2. + policy(vcat(_w_t, _x_prev)) + end + end + return nothing + end + + # ── Sparsity pattern ───────────────────────────────────────────────── + jac_r = Int[] + jac_c = Int[] + for t in 1:T + for i in 1:nx + row = (t - 1) * nx + i + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + push!(jac_r, row); push!(jac_c, xi) # ∂g/∂x_{t,i} + push!(jac_r, row); push!(jac_c, di) # ∂g/∂δ_{t,i} + if t > 1 + for j in 1:nx + xj = x_start + (t-2)*nx + j - 1 + push!(jac_r, row); push!(jac_c, xj) # ∂g/∂x_{t-1,j} + end + end + end + end + + oracle = ExaModels.VectorNonlinearOracle( + nvar = nvar_total, + ncon = n_x, + nnzj = length(jac_r), + jac_rows = jac_r, + jac_cols = jac_c, + lcon = zeros(n_x), + ucon = zeros(n_x), + f! = oracle_f!, + jac! = oracle_jac!, + vjp! = oracle_vjp!, + adapt = Val(true), + ) + ExaModels.constraint(core, oracle) + + model = ExaModels.ExaModel(core) + + target_start = n_con_before_oracle + 1 + target_range = target_start:(target_start + n_x - 1) + + return EmbeddedDeterministicEquivalentProblem( + core, model, x, u, δ, + p_x0, p_w, + policy, + nx, nu, nw, T, + target_range, + w_buf, x0_buf, + ) +end + +""" + target_multipliers(prob::EmbeddedDeterministicEquivalentProblem, result) -> λ + +Return the dual multipliers associated with the embedded oracle constraints. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Problem that defines the + multiplier slice. +- `result`: MadNLP result containing `multipliers`. + +# Returns + +- `λ`: Multipliers in `result.multipliers[prob.target_con_range]`. + +# Notes + +The multipliers correspond to the constraints +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0``. +""" +target_multipliers(prob::EmbeddedDeterministicEquivalentProblem, result) = + result.multipliers[prob.target_con_range] + +""" + solution_components(prob::EmbeddedDeterministicEquivalentProblem, result) -> (x, u, δ) + +Split the flat solution vector into state, control, and slack components. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Problem that defines component + sizes. +- `result`: MadNLP result containing `solution`. + +# Returns + +- `(x, u, δ)`: Flat slices of the primal solution for states, controls, and + target slacks. +""" +function solution_components(prob::EmbeddedDeterministicEquivalentProblem, result) + n_x = prob.horizon * prob.nx + n_u = (prob.horizon - 1) * prob.nu + sol = result.solution + x_sol = sol[1:n_x] + u_sol = sol[(n_x + 1):(n_x + n_u)] + δ_sol = sol[(n_x + n_u + 1):(n_x + n_u + n_x)] + return (x_sol, u_sol, δ_sol) +end diff --git a/src/policy.jl b/src/policy.jl index 8a3a6e8..132b18a 100644 --- a/src/policy.jl +++ b/src/policy.jl @@ -11,8 +11,17 @@ """ MLPPolicy(model, output_dim) -Stateless MLP policy: one call with `vcat(x0, w_flat)` returns the full target -trajectory `x̂` as a flat vector of length `T*nx`. +Wrap a stateless Flux model as a full-horizon target policy. + +The wrapped model is called once per scenario with `vcat(x0, w_flat)` and +returns the full target trajectory as a flat vector of length `output_dim`. + +# Arguments +- `model`: Flux-compatible callable. +- `output_dim::Int`: number of target components retained from `vec(model(input))`. + +# Returns +- `MLPPolicy`: a Flux layer wrapper around `model`. """ struct MLPPolicy{M} model::M @@ -21,6 +30,17 @@ end Flux.@layer MLPPolicy +""" + (policy::MLPPolicy)(input) -> AbstractVector + +Evaluate a stateless full-horizon policy. + +# Arguments +- `input`: concatenated initial state and flat uncertainty trajectory. + +# Returns +- The first `policy.output_dim` entries of `vec(policy.model(input))`. +""" function (π::MLPPolicy)(input) y = π.model(input) return vec(y)[1:π.output_dim] @@ -28,6 +48,19 @@ end """ MLPPolicy(input_dim, output_dim; hidden=(64,64), act=tanh) + +Construct a feed-forward [`MLPPolicy`](@ref). + +# Arguments +- `input_dim::Int`: length of the concatenated scenario input. +- `output_dim::Int`: length of the flattened target trajectory. + +# Keywords +- `hidden`: hidden-layer widths. +- `act`: hidden-layer activation. + +# Returns +- `MLPPolicy`: stateless full-horizon policy. """ function MLPPolicy(input_dim::Int, output_dim::Int; hidden = (64, 64), @@ -43,50 +76,531 @@ function MLPPolicy(input_dim::Int, output_dim::Int; return MLPPolicy(Flux.Chain(layers...), output_dim) end -# ── StateConditionedPolicy ──────────────────────────────────────────────────── +# ── ContextualPolicy ────────────────────────────────────────────────────────── + +""" + 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))`, with `t` advanced once per call +and reset by `Flux.reset!`. + +This keeps training and rollout loops unchanged: they still pass `[w_t; x_prev]` +to the policy, while the wrapper prepends known stage features such as seasonal +phase or forecast covariates. Only the inner policy is trainable. +""" +mutable struct ContextualPolicy{P,C} + policy::P + context::C + t::Int +end + +ContextualPolicy(policy, context) = ContextualPolicy(policy, context, 0) + +Flux.@layer ContextualPolicy trainable=(policy,) + +""" + context_at(context, t) +Return the context vector for one-based stage `t`. """ - StateConditionedPolicy{E,C} +function context_at(context::AbstractMatrix, t::Integer) + 1 <= t <= size(context, 2) || + throw(BoundsError(context, (:, t))) + return view(context, :, t) +end -Stateful LSTM policy for sequential rollout: +context_at(context::Function, t::Integer) = context(t) - x̂_t = policy(vcat(w_t, x̂_{t-1})) +function (m::ContextualPolicy)(input) + m.t += 1 + return m.policy(vcat(context_at(m.context, m.t), input)) +end -- `encoder`: LSTM chain operating on the uncertainty slice `w_t` -- `combiner`: Dense layer combining encoder output with previous state +function Flux.reset!(m::ContextualPolicy) + m.t = 0 + Flux.reset!(m.policy) + return nothing +end -Call `Flux.reset!(policy)` before each episode. +""" + stage_phase_context(T; period, include_progress=true) -# Flux 0.16 note -LSTM requires ≥2D input. The forward pass reshapes the 1D `w_t` slice to -`(n_uncertainty, 1)` before encoding and squeezes back with `vec`. +Build a `d x T` context matrix with `sin(2*pi*t/period)`, +`cos(2*pi*t/period)`, and optionally normalized horizon progress `t/T`. +The sine/cosine pair preserves cyclic adjacency between the last and first +seasonal positions while using only two bounded input features. """ -struct StateConditionedPolicy{E,C} - encoder::E - combiner::C - n_uncertainty::Int - n_state::Int +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 + +raw""" + _dense_policy_head(input_dim, output_dim, hidden; activation=tanh) + +Build the nonrecurrent target head used by state-conditioned policies. + +The head maps the concatenated features `[encoded_uncertainty; previous_state]` +to the normalized or unbounded target vector. When `hidden` is empty, this is the +historical single `Dense(input_dim => output_dim, activation)` head. When +`hidden = [h₁, …, h_L]`, the returned chain is + +```math +f(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). +``` + +The output layer intentionally uses the same `activation`. This is different +from a generic MLP helper with a linear output: bounded policies need the final +head output to stay in `[0, 1]` when `activation=sigmoid`, before affine scaling +to physical target bounds. + +# Arguments +- `input_dim::Int`: dimension of `[encoded_uncertainty; previous_state]`. +- `output_dim::Int`: number of target components produced by the policy. +- `hidden::AbstractVector{Int}`: hidden widths of the feed-forward target head. +- `activation`: activation applied at every head layer, including the output. + +# Returns +- `Flux.Dense` if `hidden` is empty, otherwise `Flux.Chain` of dense layers. + +# Examples +```julia +head = DecisionRulesExa._dense_policy_head(135, 11, [128, 128]; activation=sigmoid) +``` + +See also: [`StateConditionedPolicy`](@ref), [`bounded_state_policy`](@ref) +""" +function _dense_policy_head( + input_dim::Int, + output_dim::Int, + hidden::AbstractVector{Int}; + activation = tanh, +) + isempty(hidden) && return Flux.Dense(input_dim => output_dim, activation) + layers = Any[Flux.Dense(input_dim => hidden[1], activation)] + for i in 1:(length(hidden) - 1) + push!(layers, Flux.Dense(hidden[i] => hidden[i + 1], activation)) + end + push!(layers, Flux.Dense(hidden[end] => output_dim, activation)) + return Flux.Chain(layers...) +end + +# ── Recurrent-state threading helpers (Flux ≥ 0.16 stateless cells) ────────── +# +# In Flux 0.16 recurrent layers are stateless: `(l::LSTM)(x)` restarts from +# `Flux.initialstates(l)` on EVERY call and `Flux.reset!` on Flux layers is a +# deprecated no-op. A policy that needs cross-stage memory must therefore carry +# the recurrent state itself and advance it with one stateful cell call per +# stage. The helpers below mirror DecisionRules.jl's `_as_cell`, +# `_init_recurrent_state`, `_step_encoder`, and `_state_eltype` +# (src/dense_multilayer_nn.jl) so both packages thread recurrent encoders with +# identical semantics. Unlike DecisionRules.jl, which stores BARE cells, EXA +# encoders keep the `Flux.LSTM` wrapper layers (preserving the weight +# structure of existing checkpoints); `_as_cell` unwraps to the underlying +# cell for the stateful `(x, state) -> (output, new_state)` call. + +""" + _as_cell(layer) + +Return the underlying recurrent cell of `layer`. `Flux.LSTM`/`GRU`/`RNN` wrap a +cell (`LSTMCell`/`GRUCell`/`RNNCell`) in a `.cell` field; if `layer` has no such +field it is already a cell and is returned unchanged. +""" +# Unwrap the .cell field if present (LSTM → LSTMCell); return unchanged otherwise. +_as_cell(layer) = hasfield(typeof(layer), :cell) ? layer.cell : layer + +""" + _init_recurrent_state(encoder) + +Return the initial recurrent state for `encoder`: `Flux.initialstates` of the +underlying cell for a single layer, or a tuple of per-layer initial states for +a `Chain`. + +# Notes +`Flux.initialstates` builds zero states with `zeros_like` on the cell weights, +so the returned state inherits the encoder's device and element type — call +[`Flux.reset!`](@ref) after moving a policy between devices to re-derive the +state on the new device. +""" +# Single layer: initial state of the underlying cell (zeros for LSTM h/c). +_init_recurrent_state(layer) = Flux.initialstates(_as_cell(layer)) +# Chain: one initial state per layer, returned as a tuple. +_init_recurrent_state(chain::Flux.Chain) = map(_init_recurrent_state, chain.layers) + +""" + _step_encoder(encoder, x, state) -> (output, new_state) + +Advance `encoder` by one step on input `x` from recurrent `state`, returning +the output and the updated state. For a `Chain`, each layer's output feeds the +next layer and each layer's state is threaded independently. +""" +# Single layer: one stateful cell call returns (output, new_state). +_step_encoder(layer, x, state) = _as_cell(layer)(x, state) +function _step_encoder(chain::Flux.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 layers 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 -Flux.@layer StateConditionedPolicy +""" + _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. +""" +_state_eltype(state::Tuple) = _state_eltype(first(state)) +_state_eltype(v::AbstractVector) = eltype(v) + +""" + StateConditionedPolicy{E,C,S} + +Flux-compatible state-conditioned policy for sequential target rollout. + +At each stage the policy is called as + +```julia +xhat_t = policy(vcat(w_t, x_prev)) +``` + +where the recurrent encoder reads only `w_t` and the combiner reads +`[encoded_uncertainty; x_prev]`. +Flux's recurrent cells are stateless (Flux ≥ 0.16): each call returns +`(output, new_state)` instead of mutating internal state, and calling the +`LSTM` wrapper directly would restart from `initialstates` on every call. +`StateConditionedPolicy` therefore carries the encoder's recurrent state itself +in `state`, threading it through one cell call per stage — the same semantics +as DecisionRules.jl's `StateConditionedPolicy`. Call `Flux.reset!(policy)` to +restore it to `Flux.initialstates` at the start of a scenario. + +# Fields +- `encoder`: recurrent uncertainty encoder (`Chain` of `Flux.LSTM`-style layers). +- `combiner`: nonrecurrent target head. +- `state`: current recurrent state ``s_t``, carried across calls (not trainable). +- `n_uncertainty::Int`: number of uncertainty features at each stage. +- `n_state::Int`: number of previous-state features. +- `output_lower`: optional lower bounds for affine output scaling. +- `output_scale`: optional `upper - lower` scale for affine output scaling. + +# Notes +Call `Flux.reset!(policy)` before each scenario — the reset is REAL (it +restores the initial recurrent state), unlike the deprecated Flux-layer +`reset!` no-op. Within a differentiated rollout the threaded state is treated +as data: gradients flow into encoder/combiner parameters through each stage's +forward pass, and the state stored between calls is refreshed by mutation +(mirroring DecisionRules.jl's training semantics). +""" +mutable struct StateConditionedPolicy{E,C,S,L,U} + encoder::E # Recurrent uncertainty encoder (Chain of LSTM-style layers) + combiner::C # Nonrecurrent target head + state::S # Encoder recurrent state, carried across calls + n_uncertainty::Int # Number of uncertainty features per stage + n_state::Int # Number of previous-state features + output_lower::L # Optional affine output lower bounds + output_scale::U # Optional affine output scale (upper - lower) +end + +Flux.@layer StateConditionedPolicy trainable=(encoder, combiner) + +""" + (policy::StateConditionedPolicy)(input) -> AbstractVector + +Evaluate one stage of a state-conditioned policy, threading recurrent state. + +The input is split into the uncertainty portion ``w_t`` and the previous state +``x_{t-1}``, and the forward pass computes + +```math +h_t, s_t = \\text{encoder}(w_t, s_{t-1}), \\qquad +\\hat{x}_t = f_{\\text{combine}}([h_t;\\; x_{t-1}]), +``` + +where ``s_t`` is the updated recurrent state (stored in `policy.state` for the +next call). + +# Arguments +- `input`: concatenated vector `[w_t; x_prev]`. + +# Returns +- Raw combiner output, or affine-scaled output when `output_bounds` were + supplied at construction. +""" function (m::StateConditionedPolicy)(input) - w = reshape(input[1:m.n_uncertainty], :, 1) # (n_unc, 1) for LSTM + # Split the concatenated input into uncertainty w_t and previous state x_{t-1}. + w = input[1:m.n_uncertainty] s = input[m.n_uncertainty+1:end] - h = vec(m.encoder(w)) # (hidden,) - return m.combiner(vcat(h, s)) + + # Cast the uncertainty to the encoder precision (taken from the recurrent + # state), matching DecisionRules.jl's forward pass exactly. + T = _state_eltype(m.state) + + # Advance the recurrent encoder by one step: h_t, s_t = encoder(w_t, s_{t-1}). + h, new_state = _step_encoder(m.encoder, T.(w), m.state) + + # Persist the new recurrent state so the next call starts from s_t. + m.state = new_state + + y = m.combiner(vcat(h, s)) + if m.output_lower === nothing + return y + end + lower = _adapt_policy_bound(m.output_lower, y) + scale = _adapt_policy_bound(m.output_scale, y) + return lower .+ scale .* y end -Flux.reset!(m::StateConditionedPolicy) = Flux.reset!(m.encoder) +""" + Flux.reset!(policy::StateConditionedPolicy) -> Nothing + +Reset the encoder's recurrent state to `Flux.initialstates`, e.g. at the start +of a scenario rollout. +# Notes +The state is re-derived from the (possibly device-moved) encoder weights on +every reset, so calling `Flux.reset!` after `gpu(policy)`/`cpu(policy)` places +the state on the correct device with the correct element type. """ +function Flux.reset!(m::StateConditionedPolicy) + # Reinitialize s_0 to the cell defaults (zeros for LSTM h/c). + m.state = _init_recurrent_state(m.encoder) + return nothing +end + +""" + _adapt_policy_bound(x, ref) + +Move a policy bound vector to the same storage family and element type as +`ref`. + +# Arguments +- `x::AbstractVector`: bound vector stored on the policy. +- `ref::AbstractVector`: output vector whose element type and device should be + matched. + +# Returns +- `x` itself when the concrete vector type already matches `ref`; otherwise a + copied vector compatible with `ref`. +""" +function _adapt_policy_bound(x::AbstractVector, ref::AbstractVector) + typeof(x) === typeof(ref) && return x + y = similar(ref, length(x)) + copyto!(y, convert.(eltype(ref), x)) + return y +end + +""" + _load_encoder_state!(encoder, enc_state) -> encoder + +Load a checkpoint `encoder` state into a recurrent encoder, accepting BOTH +weight structures in use across the two packages: + +1. EXA structure — `Chain` of `Flux.LSTM` wrapper layers, so each layer state + is `(cell = (Wi, Wh, bias),)`. Loaded by the stock `Flux.loadmodel!`. +2. DecisionRules.jl (MAIN) structure — `Chain` of BARE `LSTMCell`s (MAIN's + `_as_cell` strips the wrapper at construction), so each layer state is + `(Wi, Wh, bias)` directly. `Flux.loadmodel!` matches children by key, so + loading a bare-cell layer state into an `LSTM` wrapper (children `(cell,)`) + throws; this helper falls back to cell-by-cell injection, loading each + layer state into `_as_cell(layer)` — the mathematically identical weight + assignment (the wrapper's `cell` has exactly the fields `(Wi, Wh, bias)` + MAIN saved). + +# Arguments +- `encoder`: destination encoder (typically a `Chain` of `Flux.LSTM` layers). +- `enc_state`: the checkpoint's encoder state (from `Flux.state`). + +# Returns +- `encoder`, mutated in place. + +# Throws +- Rethrows the stock `Flux.loadmodel!` error when the fallback does not apply + (non-`Chain` encoder, missing `layers`, or depth mismatch). +""" +function _load_encoder_state!(encoder, enc_state) + try + # Stock path: same weight structure (EXA wrapper layers). + Flux.loadmodel!(encoder, enc_state) + return encoder + catch err + # Fallback applies only to Chain encoders with a per-layer state list. + (encoder isa Flux.Chain && hasproperty(enc_state, :layers)) || rethrow(err) + layer_states = getproperty(enc_state, :layers) + length(layer_states) == length(encoder.layers) || rethrow(err) + for (layer, lstate) in zip(encoder.layers, layer_states) + # Wrapper-style layer state loads into the wrapper; bare-cell + # (MAIN) layer state loads into the unwrapped cell. + dest = hasproperty(lstate, :cell) ? layer : _as_cell(layer) + Flux.loadmodel!(dest, lstate) + end + return encoder + end +end + +""" + load_stateconditioned_policy!(policy, state) + +Load a Flux checkpoint into a [`StateConditionedPolicy`](@ref). + +# Arguments +- `policy::StateConditionedPolicy`: policy to update in place. +- `state`: checkpoint object accepted by `Flux.loadmodel!`. + +# Returns +- `policy`. + +# Notes +- Checkpoints saved before output bounds were added contain only the trainable + encoder and combiner state. In that case, this method restores those + trainable components and keeps the current policy's case-defined output + bounds. +- Checkpoints trained BEFORE recurrent-state threading (memoryless-encoder + era) have the SAME weight structure and load unchanged — only the runtime + semantics differ (the encoder now carries memory across stages). +- DecisionRules.jl (MAIN) checkpoints, whose encoders are `Chain`s of bare + `LSTMCell`s, load through the documented cell-by-cell fallback in + [`_load_encoder_state!`](@ref). +- The loaded policy's recurrent state is reset afterwards so the next rollout + starts from `Flux.initialstates` of the loaded weights. +""" +function load_stateconditioned_policy!(policy::StateConditionedPolicy, state) + try + Flux.loadmodel!(policy, state) + Flux.reset!(policy) + return policy + catch err + hasproperty(state, :encoder) && hasproperty(state, :combiner) || rethrow(err) + @warn "Full StateConditionedPolicy checkpoint load failed; loading encoder/combiner only and keeping current output bounds" exception=(err, catch_backtrace()) + _load_encoder_state!(policy.encoder, getproperty(state, :encoder)) + Flux.loadmodel!(policy.combiner, getproperty(state, :combiner)) + Flux.reset!(policy) + return policy + end +end + +function load_stateconditioned_policy!(policy::ContextualPolicy, state) + inner_state = hasproperty(state, :policy) ? getproperty(state, :policy) : state + load_stateconditioned_policy!(policy.policy, inner_state) + Flux.reset!(policy) + return policy +end + +raw""" StateConditionedPolicy(n_uncertainty, n_state, n_out, layers; - activation=tanh, encoder_type=Flux.LSTM) + activation=tanh, encoder_type=Flux.LSTM, + output_bounds=nothing, combiner_layers=Int[]) + +Construct a state-conditioned sequential target policy. + +The policy separates memory from state conditioning: + +```math +h_t = E_\theta(w_t, h_{t-1}), \qquad +\hat{x}_t = H_\theta([h_t;\, x_{t-1}]). +``` + +The recurrent encoder `Eθ` sees only the stage uncertainty. The previous state +enters only through the feed-forward head `Hθ`, so increasing +`combiner_layers` makes the state-to-target map nonlinear without introducing +recurrence over the state input. -Construct a `StateConditionedPolicy`. +# Arguments +- `n_uncertainty::Int`: number of uncertainty features in each stage input. +- `n_state::Int`: number of previous-state features appended after uncertainty. +- `n_out::Int`: output dimension, usually the state-target dimension. +- `layers::AbstractVector{Int}`: recurrent encoder hidden sizes. -- `layers` : hidden sizes for the LSTM encoder, e.g. `[64, 64]` -- `n_out` : output dimension (= nx = state dimension) +# Keywords +- `activation`: head activation. Use `sigmoid` with `output_bounds` when targets + must remain inside bounds. +- `encoder_type`: recurrent layer constructor, typically `Flux.LSTM`. +- `output_bounds`: optional `(lower, upper)` vectors. If provided, the raw head + output `y` is interpreted as normalized and mapped to + `lower + (upper - lower) .* y`. +- `combiner_layers`: hidden widths for the nonrecurrent target head. `Int[]` + preserves the original single Dense head. + +# Returns +- `StateConditionedPolicy` with trainable `encoder` and `combiner`. + +# Notes +The recurrent encoder receives only uncertainty. The previous state enters +through the feed-forward combiner, so `combiner_layers` increases +state-to-target expressiveness without making the state input recurrent. + +# Examples +```julia +policy = StateConditionedPolicy( + 11, 11, 11, [128, 128]; + activation = sigmoid, + output_bounds = (zeros(Float32, 11), ones(Float32, 11)), + combiner_layers = [128, 128], +) +``` + +See also: [`bounded_state_policy`](@ref) """ function StateConditionedPolicy( n_uncertainty::Int, @@ -95,11 +609,264 @@ function StateConditionedPolicy( layers::AbstractVector{Int}; activation = tanh, encoder_type = Flux.LSTM, + output_bounds = nothing, + combiner_layers = Int[], ) enc_sizes = vcat(n_uncertainty, layers) enc_layers = [encoder_type(enc_sizes[i] => enc_sizes[i+1]) for i in 1:length(layers)] encoder = Flux.Chain(enc_layers...) - combiner = Flux.Dense(layers[end] + n_state => n_out, activation) - return StateConditionedPolicy(encoder, combiner, n_uncertainty, n_state) + combiner = _dense_policy_head( + layers[end] + n_state, + n_out, + collect(Int, combiner_layers); + activation = activation, + ) + if output_bounds === nothing + # Initialize the recurrent state to Flux.initialstates for the encoder. + return StateConditionedPolicy( + encoder, combiner, _init_recurrent_state(encoder), + n_uncertainty, n_state, nothing, nothing, + ) + end + lower, upper = output_bounds + length(lower) == n_out || throw(ArgumentError("output lower bound length must be n_out=$n_out")) + length(upper) == n_out || throw(ArgumentError("output upper bound length must be n_out=$n_out")) + scale = upper .- lower + any(<(zero(eltype(scale))), scale) && + throw(ArgumentError("output upper bounds must be >= lower bounds")) + return StateConditionedPolicy( + encoder, combiner, _init_recurrent_state(encoder), + n_uncertainty, n_state, + collect(lower), collect(scale), + ) +end + +# ── Bounded state-target helpers ────────────────────────────────────────────── + +""" + ConstantStatePolicy(output_template, n_uncertainty, n_state) + +Represent a policy with no trainable target dimensions. + +The policy always returns `output_template`, adapted to the input device. + +# Arguments +- `output_template`: fixed full target vector. +- `n_uncertainty::Int`: number of uncertainty features expected in the input. +- `n_state::Int`: number of state features expected in the input. + +# Returns +- `ConstantStatePolicy`: Flux layer with no trainable parameters. +""" +struct ConstantStatePolicy{O} + output_template::O + n_uncertainty::Int + n_state::Int +end + +Flux.@layer ConstantStatePolicy trainable=() + +""" + (policy::ConstantStatePolicy)(input) -> AbstractVector + +Return the fixed target template, adapted to the input storage family. + +# Arguments +- `input`: vector used only as an adaptation reference. + +# Returns +- The fixed target vector. +""" +(m::ConstantStatePolicy)(input) = _adapt_policy_bound(m.output_template, input) + +""" + Flux.reset!(policy::ConstantStatePolicy) -> Nothing + +No-op reset method for constant policies. +""" +Flux.reset!(::ConstantStatePolicy) = nothing + +""" + FixedOutputPolicy(policy, output_template, output_expansion) + +Expand active target dimensions into a full target vector. + +The wrapped policy predicts only active dimensions. `output_template` stores +constants for inactive dimensions and zeros for active dimensions; +`output_expansion` maps active outputs into the full state vector without +mutation. + +# Arguments +- `policy`: Flux-compatible policy for active dimensions. +- `output_template`: full target vector with inactive constants. +- `output_expansion`: matrix mapping active outputs into full target space. + +# Returns +- `FixedOutputPolicy`: Flux layer wrapper around `policy`. +""" +struct FixedOutputPolicy{P,O,E} + policy::P + output_template::O + output_expansion::E +end + +Flux.@layer FixedOutputPolicy trainable=(policy,) + +""" + (policy::FixedOutputPolicy)(input) -> AbstractVector + +Evaluate the active-dimension policy and expand it into the full target vector. + +# Arguments +- `input`: stage input forwarded to the wrapped policy. + +# Returns +- Full target vector with active outputs inserted and inactive dimensions fixed. +""" +function (m::FixedOutputPolicy)(input) + y = m.policy(input) + return m.output_template .+ m.output_expansion * y +end + +""" + Flux.reset!(policy::FixedOutputPolicy) + +Reset the wrapped active-dimension policy. + +# Returns +- The result of `Flux.reset!(policy.policy)`. +""" +Flux.reset!(m::FixedOutputPolicy) = Flux.reset!(m.policy) + +""" + load_stateconditioned_policy!(policy::FixedOutputPolicy, state) + +Load checkpoint state into the wrapped active-dimension policy. + +# Arguments +- `policy::FixedOutputPolicy`: wrapper whose inner policy is updated. +- `state`: checkpoint object accepted by the inner policy loader. + +# Returns +- The result of `load_stateconditioned_policy!(policy.policy, state)`. +""" +load_stateconditioned_policy!(policy::FixedOutputPolicy, state) = + load_stateconditioned_policy!(policy.policy, state) + +raw""" + bounded_state_policy(n_uncertainty, lower, upper, layers; kwargs...) + +Build a state-conditioned policy whose full output is guaranteed to lie in +`[lower, upper]`, while avoiding trainable outputs for inactive dimensions. + +The returned policy has the same rollout semantics as +[`StateConditionedPolicy`](@ref): + +```math +\hat{x}_t = +\ell + (u - \ell) \odot H_\theta([E_\theta(w_t);\, x_{t-1}]), +``` + +where `Hθ` is a sigmoid head by default. Dimensions with `upper == lower` are +treated as constants unless `active_mask` overrides that choice. + +By default, a dimension is active when `upper > lower`. Fixed dimensions are +returned as constants, so pure pass-through or no-storage state components do +not create meaningless target parameters. Pass `active_mask` to override this +selection for case-specific target relevance. + +Set `combiner_layers` to add hidden layers after `[encoded_uncertainty; state]` +and before the bounded target output. This keeps recurrence confined to the +uncertainty encoder while making the state-to-target map nonlinear. + +# Arguments +- `n_uncertainty::Int`: number of uncertainty features in each stage input. +- `lower::AbstractVector`: lower bound for each full target dimension. +- `upper::AbstractVector`: upper bound for each full target dimension. +- `layers::AbstractVector{Int}`: recurrent uncertainty-encoder hidden sizes. + +# Keywords +- `activation`: head activation, defaulting to `sigmoid`. +- `encoder_type`: recurrent layer constructor. +- `active_mask`: optional Boolean mask selecting trainable output dimensions. +- `fixed_values`: values used for inactive target dimensions. +- `combiner_layers`: hidden widths for the nonrecurrent target head. + +# Returns +- `StateConditionedPolicy` when all dimensions are active. +- `ConstantStatePolicy` when no dimensions are active. +- `FixedOutputPolicy` when only a subset is active. + +# Throws +- `ArgumentError` if bound lengths differ, `active_mask` has the wrong length, + or any upper bound is smaller than its lower bound. + +# Examples +```julia +policy = bounded_state_policy( + 11, + min_volume, + max_volume, + [128, 128]; + combiner_layers = [256, 256], +) +``` +""" +function bounded_state_policy( + n_uncertainty::Int, + lower::AbstractVector, + upper::AbstractVector, + layers::AbstractVector{Int}; + activation = sigmoid, + encoder_type = Flux.LSTM, + active_mask = nothing, + fixed_values = lower, + combiner_layers = Int[], +) + length(lower) == length(upper) || + throw(ArgumentError("lower and upper bound vectors must have the same length")) + n_state = length(lower) + length(fixed_values) == n_state || + throw(ArgumentError("fixed_values length must match state dimension $n_state")) + + scale = upper .- lower + any(<(zero(eltype(scale))), scale) && + throw(ArgumentError("upper bounds must be >= lower bounds")) + + active = if active_mask === nothing + collect(scale .> zero(eltype(scale))) + else + length(active_mask) == n_state || + throw(ArgumentError("active_mask length must match state dimension $n_state")) + collect(Bool.(active_mask)) + end + + if all(active) + return StateConditionedPolicy( + n_uncertainty, n_state, n_state, layers; + activation = activation, + encoder_type = encoder_type, + output_bounds = (lower, upper), + combiner_layers = combiner_layers, + ) + elseif !any(active) + return ConstantStatePolicy(collect(fixed_values), n_uncertainty, n_state) + end + + idx = findall(active) + active_policy = StateConditionedPolicy( + n_uncertainty, n_state, length(idx), layers; + activation = activation, + encoder_type = encoder_type, + output_bounds = (lower[idx], upper[idx]), + combiner_layers = combiner_layers, + ) + template = collect(fixed_values) + template[idx] .= zero(eltype(template)) + expansion = zeros(eltype(template), n_state, length(idx)) + for (j, i) in enumerate(idx) + expansion[i, j] = one(eltype(template)) + end + return FixedOutputPolicy(active_policy, template, expansion) end diff --git a/src/rollout.jl b/src/rollout.jl index f998d5b..3c8dfcd 100644 --- a/src/rollout.jl +++ b/src/rollout.jl @@ -7,35 +7,287 @@ # semantics instead: at each stage, solve one stage, extract the realized next # state, and feed that state into the next policy call. +""" + _target_violation_share(objective::Real, objective_no_target_penalty::Real) -> Float64 + +Compute the fraction of a stage objective attributable to the target-tracking +penalty. + +# Arguments +- `objective::Real`: total stage objective (operational cost + target penalty). +- `objective_no_target_penalty::Real`: stage objective with target penalty + stripped out. + +# Returns +- `Float64`: fraction in ``[0, 1]``, or `NaN` if the ratio is ill-defined. + +# Notes +If the total objective includes both operational cost and a quadratic penalty +``\\lambda \\|x_t - \\hat{x}_t\\|^2``, this function returns + +```math +\\frac{\\text{objective} - \\text{objective\\_no\\_target\\_penalty}}{\\text{objective}}. +``` + +The return value is `NaN` when the objective is non-finite, the penalty is +non-finite, or the objective magnitude is below ``10^{-12}``. +""" function _target_violation_share(objective::Real, objective_no_target_penalty::Real) + # The penalty is the difference between the full and penalty-free objectives. penalty = objective - objective_no_target_penalty + # Guard against non-finite or near-zero denominators that would give NaN/Inf. (isfinite(objective) && isfinite(penalty) && abs(objective) > 1e-12) || return NaN + # Return the penalty's share of the total objective. return penalty / objective end -_cpu_vec(x) = vec(collect(Array(x))) +""" + _to_vec(x) -> AbstractVector + +Flatten `x` into a contiguous one-dimensional vector. +# Arguments +- `x`: any array-like object (matrix, vector, or view). + +# Returns +- `AbstractVector`: a one-dimensional vector with the same elements as `x`. + +# Notes +`SubArray` inputs are materialized with `collect` because downstream solvers +and array operations such as `copyto!` and `vcat` require contiguous storage. +All other array types are reshaped via `vec`. """ - rollout_tsddr(model, initial_state, stage_problem, w_flat; kwargs...) - -Evaluate `model` by solving `stage_problem` sequentially over a materialized -scenario `w_flat`. Unlike deterministic-equivalent evaluation, the optimizer -only receives one uncertainty slice at a time. - -Required callbacks: -- `set_stage_parameters!(stage_problem, state_in, w_t, target, stage)` updates the - one-stage problem before each solve. -- `realized_state(stage_problem, result)` returns the next state realized by the - solved stage problem. - -Optional callbacks: -- `objective_no_target_penalty(stage_problem, result)` returns the stage objective - with target-slack penalty removed. The default is `result.objective`. - -`policy_state = :realized` is the closed-loop deployment semantics. `:target` -keeps the policy recurrence on its own previous target, matching the target -generation used by deterministic-equivalent training while still solving stages -one by one. +_to_vec(x) = vec(x) # reshape in-place for contiguous arrays +_to_vec(x::SubArray) = collect(x) # materialize views to contiguous storage + +""" + _state_bound_vector(bound, ref::AbstractVector) -> Union{Nothing, AbstractVector} + +Convert a user-supplied state bound into a concrete vector whose element type +and device (CPU or GPU) match `ref`. + +# Arguments +- `bound`: the bound specification (`nothing`, a `Real` scalar, or an + `AbstractVector` / `SubArray`). +- `ref::AbstractVector`: reference vector whose length, element type, and + device placement determine the output format. + +# Returns +- `Nothing` if `bound === nothing`. +- `AbstractVector` matching `ref` in length, element type, and device. + +# Throws +- `ArgumentError` if a vector bound has the wrong length or `bound` is an + unsupported type. + +# Notes +Accepted bound forms are `nothing`, a scalar broadcast to all state entries, +or a vector of the same length as `ref`. `SubArray` bounds are materialized via +[`_to_vec`](@ref) before device adaptation so GPU kernels receive contiguous +storage. +""" +function _state_bound_vector(bound, ref::AbstractVector) + # Nothing means "no bound on this side" — propagate the sentinel. + bound === nothing && return nothing + if bound isa AbstractVector || bound isa SubArray + # Vector bounds must be element-wise, so lengths must agree. + length(bound) == length(ref) || + throw(ArgumentError("state bound length must match state length=$(length(ref)), got $(length(bound))")) + # Materialize, cast to ref's element type, and adapt to ref's device. + return _adapt_array(eltype(ref).(_to_vec(bound)), ref) + end + # Scalar bounds are broadcast into a constant vector matching ref's shape. + bound isa Real || + throw(ArgumentError("state bounds must be vectors, scalars, or nothing")) + # Allocate on the same device as ref using `similar`. + out = similar(ref, length(ref)) + # Fill every element with the scalar bound cast to ref's element type. + fill!(out, eltype(ref)(bound)) + return out +end + +""" + _project_state_to_bounds(state::AbstractVector, state_bounds) -> AbstractVector + +Clamp every element of `state` to lie within `[lower, upper]`. + +# Arguments +- `state::AbstractVector`: realized state vector to project. +- `state_bounds`: `nothing` (no projection), or a 2-element tuple/pair + `(lower, upper)` where each side is `nothing`, a scalar, or a vector. + +# Returns +- `AbstractVector`: projected state with the same element type as `state`. + +# Throws +- `ArgumentError` if `state_bounds` is not `nothing` and does not have + exactly two elements. + +# Notes +Given bounds ``l`` and ``u``, the projection is the element-wise clamp + +```math +\\operatorname{proj}(x)_i = \\min\\bigl(\\max(x_i,\\, l_i),\\, u_i\\bigr). +``` + +This repair is intended for solver-tolerance drift at stage interfaces, not as +a substitute for a recourse-feasible model. +""" +function _project_state_to_bounds(state::AbstractVector, state_bounds) + # No bounds means the state passes through unchanged. + state_bounds === nothing && return state + # Bounds must be a (lower, upper) pair. + length(state_bounds) == 2 || + throw(ArgumentError("state_bounds must be a pair/tuple (lower, upper)")) + # Convert each side into a concrete vector (or nothing) matching state. + lower = _state_bound_vector(state_bounds[1], state) + upper = _state_bound_vector(state_bounds[2], state) + # Apply element-wise clamping: first lower, then upper. + projected = state + lower !== nothing && (projected = max.(projected, lower)) # enforce lower bound + upper !== nothing && (projected = min.(projected, upper)) # enforce upper bound + return projected +end + +""" + _project_realized_state(state::AbstractVector, state_bounds, project_state) + -> AbstractVector + +Repair a realized state before it is fed into the next rollout stage. + +# Arguments +- `state::AbstractVector`: raw realized state from the stage solver. +- `state_bounds`: `nothing` or `(lower, upper)` passed to + [`_project_state_to_bounds`](@ref). +- `project_state`: `nothing`, or a callable `f(state) -> projected_state` + that enforces non-box feasibility constraints. + +# Returns +- `AbstractVector`: the doubly-projected state, cast to the element type + of the input `state`. + +# Notes +The function first applies box projection via +[`_project_state_to_bounds`](@ref), then applies `project_state` when a custom +projector is supplied. If `project_state === nothing`, only the box projection +is applied. +""" +function _project_realized_state(state::AbstractVector, state_bounds, project_state) + # First pass: box-clamp the raw realized state. + projected = _project_state_to_bounds(state, state_bounds) + # If no custom projector is provided, the box projection is sufficient. + project_state === nothing && return projected + # Second pass: apply the user's non-box projection and cast back to state's eltype. + return eltype(state).(_to_vec(project_state(projected))) +end + +""" + rollout_tsddr( + model, + initial_state::AbstractVector, + stage_problem, + w_flat::AbstractVector; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + policy_state::Symbol = :realized, + solver_state = nothing, + reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + ) -> Union{Nothing, NamedTuple} + +Evaluate a target-setting decision rule `model` by solving `stage_problem` +sequentially over a materialized uncertainty scenario `w_flat`. + +Unlike the deterministic-equivalent training solve (which sees the full +horizon simultaneously), this rollout mirrors deployment semantics: at each +stage the solver receives only one uncertainty slice, solves, extracts the +realized next state, and feeds that state into the next policy call. + +The stage-wise recursion is + +```math +x_0 = x_{\\text{init}}, \\quad +\\hat{x}_t = \\pi_\\theta(w_t,\\, x_{t-1}), \\quad +x_t = \\operatorname{solve}_t(x_{t-1},\\, w_t,\\, \\hat{x}_t), +``` + +where ``\\pi_\\theta`` is the learned policy (`model`), +``\\operatorname{solve}_t`` solves the single-stage optimization problem, and +``x_t`` is the realized state forwarded to stage ``t+1``. + +# Arguments +- `model`: the target-setting policy ``\\pi_\\theta``. Called as + `model(vcat(w_t, x_{t-1}))` to produce ``\\hat{x}_t``. +- `initial_state::AbstractVector`: initial state ``x_0``. +- `stage_problem`: stage optimization problem (must expose `.model`). +- `w_flat::AbstractVector`: flat uncertainty vector of length + `horizon * n_uncertainty`, sliced into per-stage windows. + +# Keywords +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: dimension of each per-stage uncertainty slice. +- `set_stage_parameters!::Function`: callback + `(stage_problem, state, w_t, target, stage) -> nothing` that writes the + current state, uncertainty, and target into the stage problem before each + solve. +- `realized_state::Function`: callback `(stage_problem, result) -> x_t` + that reads the realized next state from the solver result. +- `objective_no_target_penalty::Function`: callback + `(stage_problem, result) -> Float64` returning the stage objective with + the target-tracking penalty removed. Defaults to `result.objective`. +- `madnlp_kwargs`: keyword arguments forwarded to the MadNLP solver + constructor. +- `warmstart::Bool`: whether to warm-start the solver from the previous + stage's dual solution. +- `policy_state::Symbol`: `:realized` feeds the closed-loop realized state + ``x_t`` back to the policy; `:target` feeds the policy's own previous + target ``\\hat{x}_t`` instead, matching deterministic-equivalent training + semantics. +- `solver_state`: optional pre-built solver object to reuse across stages. +- `reuse_solver::Bool`: if `true`, a single solver object is reused (and + warm-started) across all stages. +- `state_bounds`: `nothing` or `(lower, upper)` pair to clamp realized + states via [`_project_state_to_bounds`](@ref). +- `project_state`: `nothing` or a callable `f(state) -> state` for non-box + feasibility repairs via [`_project_realized_state`](@ref). +- `retry_on_failure::Bool`: if `true`, failed or non-finite solves are + retried once with a cold-start solver. + +# Returns +- `nothing` if any stage solve fails after retry. +- A `NamedTuple` with fields: + - `objective::Float64`: cumulative objective over the horizon. + - `objective_no_target_penalty::Float64`: cumulative objective without + target penalties. + - `target_violation_share::Float64`: fraction of objective from target + penalties (see [`_target_violation_share`](@ref)). + - `final_state::AbstractVector`: realized state after the last stage. + - `state_trajectory::Vector`: realized states ``[x_0, x_1, \\ldots, x_T]``. + - `target_trajectory::Vector`: policy targets + ``[\\hat{x}_1, \\ldots, \\hat{x}_T]``. + +# Throws +- `ArgumentError` if `horizon < 1`, `n_uncertainty < 1`, `w_flat` has the + wrong length, or `policy_state` is not `:realized` or `:target`. + +# Examples +```julia +result = rollout_tsddr( + model, x0, stage_problem, w_flat; + horizon = 96, + n_uncertainty = 5, + set_stage_parameters! = my_set_params!, + realized_state = my_realized_state, +) +result !== nothing && @show result.objective +``` """ function rollout_tsddr( model, @@ -52,45 +304,76 @@ function rollout_tsddr( policy_state::Symbol = :realized, solver_state = nothing, reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, ) + # --- Input validation --------------------------------------------------- horizon >= 1 || throw(ArgumentError("horizon must be >= 1")) n_uncertainty >= 1 || throw(ArgumentError("n_uncertainty must be >= 1")) + # w_flat must contain exactly horizon slices of n_uncertainty each. length(w_flat) == horizon * n_uncertainty || throw(ArgumentError("w_flat length must be horizon*n_uncertainty=$(horizon * n_uncertainty), got $(length(w_flat))")) + # Only two feedback modes are supported. policy_state in (:realized, :target) || throw(ArgumentError("policy_state must be :realized or :target, got :$policy_state")) - F = eltype(initial_state) - state = solver_state + # --- Infer numeric types and allocate trajectory storage --------------- + F = eltype(initial_state) # element type (e.g. Float32) + nx = length(initial_state) # state dimension + state = solver_state # optional pre-built solver + w_flat = _adapt_array(F.(w_flat), initial_state) # move w_flat to same device as initial_state + # Reset any recurrent state in the policy network (e.g. LSTM hidden state). Flux.reset!(model) - realized_prev = F.(_cpu_vec(initial_state)) + # x_0: the realized state entering stage 1. + realized_prev = F.(_to_vec(initial_state)) + # Target recurrence state (used when policy_state == :target). target_prev = copy(realized_prev) - state_trajectory = Vector{Vector{F}}(undef, horizon + 1) - target_trajectory = Vector{Vector{F}}(undef, horizon) - state_trajectory[1] = copy(realized_prev) + # Pre-allocate trajectory arrays: states have T+1 entries, targets have T. + state_trajectory = Vector{AbstractVector{F}}(undef, horizon + 1) + target_trajectory = Vector{AbstractVector{F}}(undef, horizon) + state_trajectory[1] = copy(realized_prev) # store x_0 - objective = 0.0 - objective_no_penalty = 0.0 + # Pre-allocate Float64 buffers for the solver interface (solvers use Float64). + state_f64 = similar(initial_state, Float64, nx) # x_{t-1} in Float64 + w_f64 = similar(initial_state, Float64, n_uncertainty) # w_t in Float64 + target_f64 = similar(initial_state, Float64, nx) # xhat_t in Float64 + # Running sums of cumulative cost over the horizon. + objective = 0.0 # total objective (operational + target penalty) + objective_no_penalty = 0.0 # total objective without target penalty + + # --- Stage-wise forward pass ------------------------------------------- for stage in 1:horizon - lo = (stage - 1) * n_uncertainty + 1 - hi = stage * n_uncertainty - wt = F.(_cpu_vec(view(w_flat, lo:hi))) + # Slice out this stage's uncertainty window from the flat vector. + wt = view(w_flat, (stage-1)*n_uncertainty+1 : stage*n_uncertainty) + # Choose the state fed to the policy: realized (closed-loop) or target. policy_input_state = policy_state === :realized ? realized_prev : target_prev + # Evaluate the policy: pi_theta(w_t, x_{t-1}) -> xhat_t. target = model(vcat(wt, policy_input_state)) - target_trajectory[stage] = F.(_cpu_vec(target)) + # Flatten the target to a 1-D vector for downstream use. + target_vec = _to_vec(target) + # Store the target in the trajectory (cast to the state element type). + target_trajectory[stage] = F.(target_vec) + # Copy inputs into the Float64 solver buffers. + copyto!(state_f64, realized_prev) # x_{t-1} + copyto!(w_f64, wt) # w_t + copyto!(target_f64, target_vec) # xhat_t + # Write the current state, uncertainty, and target into the stage problem. set_stage_parameters!( stage_problem, - Float64.(realized_prev), - Float64.(wt), - Float64.(_cpu_vec(target)), + state_f64, + w_f64, + target_f64, stage, ) + # --- Solve the single-stage optimization problem ------------------- if reuse_solver || state !== nothing + # Reuse an existing solver; create one lazily on first call. state === nothing && (state = _make_solver(stage_problem.model, madnlp_kwargs)) result = _solve!( state, @@ -99,6 +382,7 @@ function rollout_tsddr( madnlp_kwargs = madnlp_kwargs, ) else + # Create a fresh solver for this stage (no cross-stage warm-start). stage_state = _make_solver(stage_problem.model, madnlp_kwargs) result = _solve!( stage_state, @@ -108,53 +392,226 @@ function rollout_tsddr( ) end + # --- Retry logic: cold-start fallback on solver failure ------------ + if retry_on_failure && (!solve_succeeded(result) || !isfinite(result.objective)) + # Build a fresh solver and retry without warm-start. + retry_state = _make_solver(stage_problem.model, madnlp_kwargs) + result = _solve!( + retry_state, + stage_problem.model; + warmstart = false, + madnlp_kwargs = madnlp_kwargs, + ) + end + + # Abort the rollout if the solve still failed after retry. solve_succeeded(result) || return nothing + # Abort on non-finite objective (e.g. MadNLP returning 0.0 for infeasible). isfinite(result.objective) || return nothing + # --- Extract penalty-free objective, with retry -------------------- no_penalty = objective_no_target_penalty(stage_problem, result) + if retry_on_failure && !isfinite(no_penalty) + # Non-finite penalty-free cost can indicate a solver glitch; retry. + retry_state = _make_solver(stage_problem.model, madnlp_kwargs) + result = _solve!( + retry_state, + stage_problem.model; + warmstart = false, + madnlp_kwargs = madnlp_kwargs, + ) + solve_succeeded(result) || return nothing + isfinite(result.objective) || return nothing + no_penalty = objective_no_target_penalty(stage_problem, result) + end + # Final guard: abort if the penalty-free cost is still non-finite. isfinite(no_penalty) || return nothing - objective += result.objective - objective_no_penalty += no_penalty - realized_prev = F.(_cpu_vec(realized_state(stage_problem, result))) - target_prev = F.(_cpu_vec(target)) + # --- Accumulate costs and advance the state ------------------------ + objective += result.objective # add stage cost to cumulative total + objective_no_penalty += no_penalty # add penalty-free stage cost + # Read the realized next state x_t from the solver solution. + raw_realized = F.(_to_vec(realized_state(stage_problem, result))) + # Project x_t to feasibility (box bounds + custom projector). + realized_prev = _project_realized_state(raw_realized, state_bounds, project_state) + # Update the target recurrence for :target mode. + target_prev = target_trajectory[stage] + # Record x_t in the state trajectory. state_trajectory[stage + 1] = copy(realized_prev) end + # --- Assemble and return the rollout summary --------------------------- return ( - objective = objective, - objective_no_target_penalty = objective_no_penalty, - target_violation_share = _target_violation_share(objective, objective_no_penalty), - final_state = realized_prev, - state_trajectory = state_trajectory, - target_trajectory = target_trajectory, + objective = objective, # sum of stage objectives + objective_no_target_penalty = objective_no_penalty, # sum minus target penalties + target_violation_share = _target_violation_share(objective, objective_no_penalty), # penalty fraction + final_state = realized_prev, # x_T + state_trajectory = state_trajectory, # [x_0, ..., x_T] + target_trajectory = target_trajectory, # [xhat_1, ..., xhat_T] ) end +""" + RolloutEvaluation + +Store configuration and mutable summaries for periodic rollout evaluation. + +# Fields +- `stage_problem`: the single-stage optimization problem template. +- `initial_state`: initial state ``x_0`` for every rollout. +- `scenarios::Vector`: pre-sampled uncertainty vectors, each of length + `horizon * n_uncertainty`. +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: per-stage uncertainty dimension. +- `set_stage_parameters!::Function`: callback to write stage data into the + problem (see [`rollout_tsddr`](@ref)). +- `realized_state::Function`: callback to extract ``x_t`` from a solve + result. +- `objective_no_target_penalty::Function`: callback to extract the + penalty-free stage cost. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: whether to warm-start across stages. +- `stride::Int`: evaluate every `stride` training iterations. +- `policy_state::Symbol`: `:realized` or `:target` (see + [`rollout_tsddr`](@ref)). +- `solver_state`: pre-built solver object (or `nothing`). +- `reuse_solver::Bool`: whether to reuse a single solver across stages. +- `state_bounds`: optional `(lower, upper)` feasibility bounds. +- `project_state`: optional non-box projection callback. +- `retry_on_failure::Bool`: retry failed solves with a cold start. +- `stage_problem_pool::Vector`: pool of stage problems for parallel + evaluation. +- `active_scenarios::Int`: number of scenarios to evaluate (at most + `length(scenarios)`). +- `last_objective::Float64`: mean objective across successful scenarios. +- `last_objective_no_target_penalty::Float64`: mean penalty-free objective. +- `last_violation_share::Float64`: mean target-violation share. +- `last_n_ok::Int`: number of scenarios that solved successfully. +- `last_scenario_data::Vector{Any}`: per-scenario `(index, result)` pairs + from the most recent evaluation. + +# Notes +The struct is callable as `evaluation(iter, model)`. At every `stride`-th +iteration, it runs [`rollout_tsddr`](@ref) on up to `active_scenarios` +scenarios and updates the mutable summary fields. When `stage_problem_pool` +contains more than one entry, scenarios are distributed across the pool with +`Threads.@spawn`. + +Thread-safety: the single `set_stage_parameters!`, `realized_state`, and +`objective_no_target_penalty` callbacks are shared across all spawned tasks +while each task receives its own stage problem from the pool. When +`stage_problem_pool` has more than one entry, these callbacks must therefore be +thread-safe and must write only into the stage problem they are handed — any +shared mutable buffer (e.g. a captured scratch array reused across calls) +races across tasks and silently corrupts results. +""" mutable struct RolloutEvaluation <: Function - stage_problem - initial_state - scenarios::Vector - horizon::Int - n_uncertainty::Int - set_stage_parameters!::Function - realized_state::Function - objective_no_target_penalty::Function - madnlp_kwargs - warmstart::Bool - stride::Int - policy_state::Symbol - solver_state - reuse_solver::Bool - stage_problem_pool::Vector # pool of stage problems for parallel evaluation - active_scenarios::Int # how many scenarios to evaluate (≤ length(scenarios)) - last_objective::Float64 - last_objective_no_target_penalty::Float64 - last_violation_share::Float64 - last_n_ok::Int - last_scenario_data::Vector{Any} + stage_problem # single-stage optimization problem template + initial_state # initial state x_0 for all rollouts + scenarios::Vector # pre-sampled uncertainty vectors + horizon::Int # number of decision stages T + n_uncertainty::Int # per-stage uncertainty dimension + set_stage_parameters!::Function # callback: write stage data into the problem + realized_state::Function # callback: extract realized x_t from result + objective_no_target_penalty::Function # callback: penalty-free stage cost + madnlp_kwargs # solver keyword arguments + warmstart::Bool # warm-start across stages + stride::Int # evaluate every stride-th iteration + policy_state::Symbol # :realized or :target feedback mode + solver_state # pre-built solver (or nothing) + reuse_solver::Bool # reuse single solver across stages + state_bounds # (lower, upper) feasibility bounds or nothing + project_state # non-box projection callback or nothing + retry_on_failure::Bool # retry failed solves with cold start + stage_problem_pool::Vector # pool of stage problems for parallel evaluation + active_scenarios::Int # how many scenarios to evaluate (leq length(scenarios)) + last_objective::Float64 # mean objective from last evaluation + last_objective_no_target_penalty::Float64 # mean penalty-free objective from last evaluation + last_violation_share::Float64 # mean target-violation share from last evaluation + last_n_ok::Int # number of successful scenarios in last evaluation + last_scenario_data::Vector{Any} # per-scenario (index, result) pairs from last eval end +""" + RolloutEvaluation( + stage_problem, + initial_state, + scenarios; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + stride::Int = 1, + policy_state::Symbol = :realized, + reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + stage_problem_pool::Vector = [], + active_scenarios::Int = length(scenarios), + ) -> RolloutEvaluation + +Construct a [`RolloutEvaluation`](@ref) callback for periodic out-of-sample +policy evaluation during training. + +All mutable result fields (`last_objective`, `last_n_ok`, etc.) are +initialized to `NaN` / `0` / empty and are populated on the first call. + +# Arguments +- `stage_problem`: the single-stage optimization problem (must expose + `.model`). +- `initial_state`: initial state ``x_0``. +- `scenarios`: iterable of pre-sampled uncertainty vectors, each of length + `horizon * n_uncertainty`. + +# Keywords +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: per-stage uncertainty dimension. +- `set_stage_parameters!::Function`: stage-parameter callback. +- `realized_state::Function`: realized-state extraction callback. +- `objective_no_target_penalty::Function`: penalty-free cost callback. +- `madnlp_kwargs`: solver keyword arguments. +- `warmstart::Bool`: warm-start across stages within each rollout. +- `stride::Int`: evaluate every `stride`-th training iteration. +- `policy_state::Symbol`: `:realized` (closed-loop) or `:target`. +- `reuse_solver::Bool`: reuse a single solver across stages. +- `state_bounds`: `nothing` or `(lower, upper)` for box projection. +- `project_state`: `nothing` or custom non-box projector. +- `retry_on_failure::Bool`: retry failed solves with cold start. +- `stage_problem_pool::Vector`: pool of independent stage problems for + multi-threaded evaluation. An empty pool uses sequential evaluation. + With more than one pool entry, the shared `set_stage_parameters!`, + `realized_state`, and `objective_no_target_penalty` callbacks run + concurrently on different tasks: they must be thread-safe and write only + into the stage problem passed to them (no shared mutable buffers), + otherwise results race. +- `active_scenarios::Int`: cap on how many scenarios to evaluate (defaults + to all). + +# Returns +- `RolloutEvaluation`: callable training callback with empty result summaries. + +# Throws +- `ArgumentError` if `scenarios` is empty, `stride < 1`, or `policy_state` is + not `:realized` or `:target`. + +# Examples +```julia +eval_cb = RolloutEvaluation( + stage_problem, x0, test_scenarios; + horizon = 96, + n_uncertainty = 5, + set_stage_parameters! = my_set_params!, + realized_state = my_realized_state, + stride = 10, +) +# Use as a training callback: +eval_cb(iter, model) +``` +""" function RolloutEvaluation( stage_problem, initial_state, @@ -169,18 +626,24 @@ function RolloutEvaluation( stride::Int = 1, policy_state::Symbol = :realized, reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, stage_problem_pool::Vector = [], active_scenarios::Int = length(scenarios), ) + # At least one scenario is required for meaningful evaluation. isempty(scenarios) && throw(ArgumentError("scenarios must be nonempty")) + # Stride must be positive; stride=1 evaluates every iteration. stride >= 1 || throw(ArgumentError("stride must be >= 1")) + # Validate the feedback mode. policy_state in (:realized, :target) || throw(ArgumentError("policy_state must be :realized or :target, got :$policy_state")) return RolloutEvaluation( stage_problem, initial_state, - collect(scenarios), + collect(scenarios), # materialize to a concrete Vector horizon, n_uncertainty, set_stage_parameters!, @@ -190,18 +653,42 @@ function RolloutEvaluation( warmstart, stride, policy_state, + # Pre-build the solver if reuse is requested; otherwise leave as nothing. reuse_solver ? _make_solver(stage_problem.model, madnlp_kwargs) : nothing, reuse_solver, - collect(stage_problem_pool), + state_bounds, + project_state, + retry_on_failure, + collect(stage_problem_pool), # materialize the pool to a concrete Vector active_scenarios, - NaN, - NaN, - NaN, - 0, - Any[], + NaN, # last_objective: not yet evaluated + NaN, # last_objective_no_target_penalty + NaN, # last_violation_share + 0, # last_n_ok: no successful scenarios yet + Any[], # last_scenario_data: empty until first eval ) end +""" + (evaluation::RolloutEvaluation)(iter, model) -> Nothing + +Evaluate `model` on rollout scenarios when `iter` is aligned with +`evaluation.stride`. + +# Arguments +- `evaluation::RolloutEvaluation`: callback state and rollout configuration. +- `iter`: current training iteration. +- `model`: policy model passed to [`rollout_tsddr`](@ref). + +# Returns +- `nothing`: summary fields on `evaluation` are updated in place. + +# Notes +If `iter % evaluation.stride != 0`, the method only clears stale per-scenario +data and returns. Otherwise it records mean objective, mean penalty-free +objective, mean target-violation share, the number of successful scenarios, and +per-scenario results. +""" function (evaluation::RolloutEvaluation)(iter, model) empty!(evaluation.last_scenario_data) iter % evaluation.stride == 0 || return nothing @@ -232,6 +719,9 @@ function (evaluation::RolloutEvaluation)(iter, model) policy_state = evaluation.policy_state, solver_state = evaluation.solver_state, reuse_solver = evaluation.reuse_solver, + state_bounds = evaluation.state_bounds, + project_state = evaluation.project_state, + retry_on_failure = evaluation.retry_on_failure, ) result === nothing && continue total += result.objective @@ -264,6 +754,9 @@ function (evaluation::RolloutEvaluation)(iter, model) warmstart = evaluation.warmstart, policy_state = evaluation.policy_state, reuse_solver = false, + state_bounds = evaluation.state_bounds, + project_state = evaluation.project_state, + retry_on_failure = evaluation.retry_on_failure, ) push!(tasks, t) end @@ -298,13 +791,30 @@ function (evaluation::RolloutEvaluation)(iter, model) end """ - critic_samples_from_evaluation(eval; objective_key) -> Vector{CriticSample} + critic_samples_from_evaluation( + eval_obj::RolloutEvaluation; + objective_key::Symbol = :objective, + ) -> Vector{CriticSample} Convert the last rollout evaluation results into `CriticSample`s for critic -training. Target multipliers are zero (rollout evaluation does not produce -duals), so these samples contribute only to the value loss term. By default the -critic target uses the full rollout objective; pass -`objective_key = :objective_no_target_penalty` to remove target-slack penalties. +training. + +# Arguments +- `eval_obj::RolloutEvaluation`: evaluation callback containing + `last_scenario_data` from a previous call. + +# Keywords +- `objective_key::Symbol`: field of each rollout result used as the scalar + critic target; commonly `:objective` or `:objective_no_target_penalty`. + +# Returns +- `Vector{CriticSample}`: one sample per successful rollout scenario from the + last evaluation. + +# Notes +Rollout evaluation does not produce dual multipliers, so generated samples use +zero target multipliers and contribute only to the value-loss term unless +combined with other samples. """ function critic_samples_from_evaluation( eval_obj::RolloutEvaluation; diff --git a/src/training.jl b/src/training.jl index 0407ef6..f3129ee 100644 --- a/src/training.jl +++ b/src/training.jl @@ -6,7 +6,7 @@ # 1. uncertainty_sampler() → flat w (length T × nw_per_stage) # 2. Policy rollout: x̂_t = policy(vcat(w_t, x̂_{t-1})) for t = 1..T # 3. ExaModels.set_parameter! for x0, uncertainty, targets → MadNLP.solve! -# 4. λ = result.multipliers[target_con_range] (∇_{x̂} Q, envelope theorem) +# 4. λ = target_multipliers(de, result) (∇_{x̂} Q, envelope theorem) # 5. Zygote: ∇_θ (1/n) Σ_s ⟨λ_s, x̂_s(θ)⟩ → Flux.update! # # The user passes parameter objects (p_x0, p_target, p_uncertainty) exactly as @@ -21,35 +21,265 @@ # ── Gradient materialization ────────────────────────────────────────────────── -_mat(x) = x -_mat(x::Zygote.OneElement) = collect(x) +""" + _mat(x) + _mat(x::Zygote.OneElement) + _mat(x::ChainRulesCore.Tangent) + _mat(x::ChainRulesCore.MutableTangent) + +Recursively materialize lazy Zygote / ChainRules tangent wrappers into plain +Julia values (arrays and named tuples). + +Zygote may return `OneElement` sparse arrays or `Tangent`/`MutableTangent` +wrappers instead of dense arrays or named tuples. `Flux.update!` expects +concrete data, so every tangent node must be materialized before the optimizer +step. + +# Arguments +- `x`: a gradient value — may be a plain array, a `Zygote.OneElement`, a + `ChainRulesCore.Tangent`, or a `ChainRulesCore.MutableTangent`. + +# Returns +- For plain values: returns `x` unchanged. +- For `OneElement`: returns `collect(x)`, a dense array. +- For `Tangent`/`MutableTangent`: returns a `NamedTuple` with recursively + materialized fields. +- For `Base.RefValue` (Zygote's wrapper for tangents of mutated mutable + structs, such as the state-threading policies): unwraps and recurses. +- For plain `NamedTuple`/`Tuple`: recurses so nested wrappers are stripped. +- For `NoTangent`/`ZeroTangent`: returns `nothing`. +""" +_mat(x) = x # plain value — no conversion needed +_mat(x::Zygote.OneElement) = collect(x) # sparse one-hot → dense array function _mat(x::ChainRulesCore.Tangent{<:Any}) - nt = ChainRulesCore.backing(x) - return NamedTuple{keys(nt)}(map(_mat, values(nt))) + nt = ChainRulesCore.backing(x) # extract underlying named tuple + return NamedTuple{keys(nt)}(map(_mat, values(nt))) # recursively materialize each field end function _mat(x::ChainRulesCore.MutableTangent{<:Any}) - nt = ChainRulesCore.backing(x) - return NamedTuple{keys(nt)}(map(_mat, values(nt))) + nt = ChainRulesCore.backing(x) # extract underlying named tuple + return NamedTuple{keys(nt)}(map(_mat, values(nt))) # recursively materialize each field +end +# Structural-zero tangents (non-differentiable fields) map to nothing so +# Flux.update! skips them. +_mat(::ChainRulesCore.NoTangent) = nothing +_mat(::ChainRulesCore.ZeroTangent) = nothing +# Zygote wraps tangents of MUTATED mutable structs (e.g. the state-threading +# policies, whose forward pass stores the new recurrent state via setfield!) in +# Base.RefValue and MutableTangent containers, potentially nested inside plain +# NamedTuples/Tuples. Recurse through those containers so every wrapper is +# stripped before the gradient reaches Flux.update!. +_mat(ref::Base.RefValue) = _mat(ref[]) # unwrap Ref and recurse +_mat(nt::NamedTuple{K}) where {K} = + NamedTuple{K}(map(_mat, values(nt))) # recurse plain named tuples +_mat(t::Tuple) = map(_mat, t) # recurse plain tuples + +""" + materialize_tangent(g) -> Union{Nothing, Any} + +Convert a Zygote gradient `g` into plain Julia arrays and named tuples that +`Flux.update!` can consume. Returns `nothing` when `g` is `nothing` (i.e., no +gradient was produced for that parameter). + +This is the public entry point; internally it delegates to [`_mat`](@ref). + +# Arguments +- `g`: raw gradient returned by `Zygote.gradient`; may be `nothing`. + +# Returns +- `nothing` if `g` is `nothing`. +- A materialized gradient (dense arrays / named tuples) otherwise. + +# Examples +```julia +gs = Zygote.gradient(model) do m + sum(m(x)) +end +grad = materialize_tangent(gs[1]) # NamedTuple or nothing +``` +""" +materialize_tangent(g) = isnothing(g) ? nothing : _mat(g) # guard against nothing gradients + +""" + _all_finite_gradient(x) -> Bool + +Recursively check that every element of a (possibly nested) gradient structure +is finite (no `NaN` or `±Inf`). + +After BPTT through physics-based dynamics, accumulated Jacobian products can +overflow `Float32` range (> 3.4e38), producing `Inf` or `NaN` values that +would corrupt the Adam optimizer state. This guard prevents +`Flux.update!` from being called with non-finite gradients. + +# Arguments +- `x`: gradient value — may be an `AbstractArray`, `Number`, `Nothing`, + `NamedTuple`, `Tuple`, or any other type. + +# Returns +- `true` if every numeric leaf is finite (or the value is `nothing` / an + unrecognized type that carries no numeric data). +- `false` if any leaf contains `NaN` or `±Inf`. + +# Examples +```julia +_all_finite_gradient([1.0, 2.0]) # true +_all_finite_gradient([1.0, NaN]) # false +_all_finite_gradient((a=[1.0], b=Inf)) # false +_all_finite_gradient(nothing) # true +``` +""" +_all_finite_gradient(x::AbstractArray) = all(isfinite, x) # check every element +_all_finite_gradient(x::Number) = isfinite(x) # scalar check +_all_finite_gradient(x::Nothing) = true # nothing → vacuously finite +_all_finite_gradient(x::NamedTuple) = all(_all_finite_gradient(v) for v in values(x)) # recurse named tuple fields +_all_finite_gradient(x::Tuple) = all(_all_finite_gradient(v) for v in x) # recurse tuple elements +_all_finite_gradient(x) = true # fallback: assume finite for unknown types + +""" + _status_key(status) -> String + +Convert a MadNLP solver status enum value into a safe string key by replacing +non-alphanumeric characters with underscores. Used to build human-readable +diagnostic dictionaries keyed by solver outcome. + +# Arguments +- `status`: a MadNLP status enum (e.g., `MadNLP.SOLVE_SUCCEEDED`). + +# Returns +- A sanitized `String` suitable for use as a dictionary key. + +# Examples +```julia +_status_key(MadNLP.SOLVE_SUCCEEDED) # "SOLVE_SUCCEEDED" +``` +""" +_status_key(status) = replace(string(status), r"[^A-Za-z0-9_]" => "_") # sanitize status to dict-safe key + +""" + _inc_status!(counts::Dict{String, Int}, status) -> Dict{String, Int} + +Increment the counter for the given MadNLP solver `status` in `counts`. +Converts `status` to a string key via [`_status_key`](@ref) before +incrementing. + +# Arguments +- `counts::Dict{String, Int}`: mutable dictionary of status counts. +- `status`: a MadNLP solver status enum. + +# Returns +- The mutated `counts` dictionary. +""" +function _inc_status!(counts::Dict{String, Int}, status) + key = _status_key(status) # convert enum to string key + counts[key] = get(counts, key, 0) + 1 # increment (initialize to 0 if absent) + return counts end -materialize_tangent(g) = isnothing(g) ? nothing : _mat(g) -_all_finite_gradient(x::AbstractArray) = all(isfinite, x) -_all_finite_gradient(x::Number) = isfinite(x) -_all_finite_gradient(x::Nothing) = true -_all_finite_gradient(x::NamedTuple) = all(_all_finite_gradient(v) for v in values(x)) -_all_finite_gradient(x::Tuple) = all(_all_finite_gradient(v) for v in x) -_all_finite_gradient(x) = true +""" + _inc_count!(counts::Dict{String, Int}, key::String) -> Dict{String, Int} + +Increment the counter for `key` in the diagnostics dictionary `counts`. +Used to track failure reasons (e.g., `"nonfinite_objective"`) and retry +outcomes during training batches. + +# Arguments +- `counts::Dict{String, Int}`: mutable dictionary of event counts. +- `key::String`: the event identifier to increment. + +# Returns +- The mutated `counts` dictionary. +""" +function _inc_count!(counts::Dict{String, Int}, key::String) + counts[key] = get(counts, key, 0) + 1 # increment (initialize to 0 if absent) + return counts +end + +""" + _adapt_array(x::AbstractVector, ref::AbstractVector) -> AbstractVector + +Move `x` onto the same device (CPU or GPU) as `ref`, allocating a new array +only when the concrete types differ. + +When training runs on GPU, solver outputs (multipliers, states) may be +`CuVector`s while sampled data may arrive as CPU `Vector`s (or vice versa). +This helper ensures type-homogeneous arithmetic by copying `x` into a +`similar` array derived from `ref`. + +# Arguments +- `x::AbstractVector`: the source data to adapt (may be CPU or GPU). +- `ref::AbstractVector`: a reference vector whose concrete type determines the + target device / storage backend. + +# Returns +- `x` itself when `typeof(x) === typeof(ref)` (zero-copy fast path). +- A new array on the same device as `ref`, with the element type of `x` and + the data of `x` copied in. + +# Examples +```julia +using CUDA +cpu_vec = [1.0f0, 2.0f0] +gpu_ref = CUDA.zeros(Float32, 3) +gpu_vec = _adapt_array(cpu_vec, gpu_ref) # CuVector{Float32} +``` +""" +function _adapt_array(x::AbstractVector, ref::AbstractVector) + typeof(x) === typeof(ref) && return x # same type → no copy needed + copyto!(similar(ref, eltype(x), length(x)), x) # allocate on ref's device, copy x into it +end # ── Solve-status check ──────────────────────────────────────────────────────── """ solve_succeeded(result) -> Bool + +Check whether a MadNLP solve result indicates a usable solution. + +MadNLP returns `0.0` objective for failed or infeasible solves, so checking +`isfinite(objective)` alone is not sufficient. This function inspects the +solver status directly. + +# Arguments +- `result`: a MadNLP result object with a `.status` field. + +# Returns +- `true` if `result.status` is `SOLVE_SUCCEEDED` or `SOLVED_TO_ACCEPTABLE_LEVEL`. +- `false` for all other statuses (e.g., `MAXIMUM_ITERATIONS_EXCEEDED`, + `INFEASIBLE_PROBLEM_DETECTED`). + +# Examples +```julia +result = MadNLP.solve!(solver) +if solve_succeeded(result) + @show result.objective +end +``` """ function solve_succeeded(result) - s = result.status - return s == MadNLP.SOLVE_SUCCEEDED || s == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL + s = result.status # extract MadNLP status enum + return s == MadNLP.SOLVE_SUCCEEDED || s == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL # accept both convergence levels end +""" + prepare_solve!(de, init_state, w_flat, xhat_flat) + +Hook called after standard parameter updates and before each NLP solve. + +# Arguments +- `de`: deterministic-equivalent problem. +- `init_state`: initial state used for the current sample. +- `w_flat`: flat uncertainty trajectory for the current sample. +- `xhat_flat`: flat target trajectory for the current sample. + +# Returns +- `nothing` by default. + +# Notes +Override this method for problem types that need additional parameter updates, +for example setting a reservoir parameter from `x0` and targets when the +reservoir is not a decision variable. +""" +prepare_solve!(de, init_state, w_flat, xhat_flat) = nothing + # ── Internal: one MadNLP solve with cascade-failure prevention ──────────────── # # After a failed solve the duals (y, zl, zu) are corrupted. Instead of cold- @@ -60,68 +290,256 @@ end # Per-solve iteration budget: MadNLP's cnt.k is CUMULATIVE across calls, so we # reset it before each solve to give each batch a fresh max_iter budget. +""" + _SolverState + +Mutable wrapper around a MadNLP solver that caches the last successful +primal-dual snapshot for warm-start cascade-failure prevention. + +After a failed solve, MadNLP's duals (`y`, `zl`, `zu`) are corrupted. +If the next call uses `reinitialize!()` (warm-start path), it keeps those +corrupted duals and the failure cascades. `_SolverState` stores the +last-good dual values so they can be restored after a failure, breaking +the cascade without paying the cost of a full cold start. + +# Fields +- `solver`: the `MadNLP.MadNLPSolver` instance. +- `last_good_x`: primal snapshot from the last successful solve (CPU or GPU + array), or `nothing` before the first success. +- `last_good_y`: equality dual snapshot, or `nothing`. +- `last_good_zl_vals`: lower-bound dual values snapshot, or `nothing`. +- `last_good_zu_vals`: upper-bound dual values snapshot, or `nothing`. +- `has_fixed_vars::Bool`: `true` when the NLP has fixed variables (via + `MakeParameter`), which requires a fresh solver per solve to avoid stale + KKT factorization state. +""" mutable struct _SolverState - solver - last_good_x # primal snapshot (CPU or GPU array), or nothing - last_good_y # dual snapshot, or nothing - last_good_zl_vals - last_good_zu_vals + solver # MadNLP.MadNLPSolver instance + last_good_x # primal snapshot (CPU or GPU array), or nothing + last_good_y # dual snapshot, or nothing + last_good_zl_vals # lower-bound dual values snapshot, or nothing + last_good_zu_vals # upper-bound dual values snapshot, or nothing + has_fixed_vars::Bool # true when fixed variables exist in the NLP end +""" + _make_solver(nlp, madnlp_kwargs) -> _SolverState + +Construct a [`_SolverState`](@ref) wrapping a fresh `MadNLP.MadNLPSolver`. + +Detects whether the NLP has fixed variables by comparing the solver's internal +variable count against the NLP's variable count. When fixed variables exist +(via `ExaModels.MakeParameter`), the solver cannot be safely reused across +parameter changes, so `has_fixed_vars` is set to `true`. + +# Arguments +- `nlp`: an NLPModels-compatible problem (e.g., `ExaModel`). +- `madnlp_kwargs`: `NamedTuple` of keyword arguments forwarded to + `MadNLP.MadNLPSolver`. + +# Returns +- A fresh [`_SolverState`](@ref) with no cached primal-dual snapshot. + +# Examples +```julia +state = _make_solver(det_equivalent.model, (print_level=0, tol=1e-6)) +``` +""" function _make_solver(nlp, madnlp_kwargs) - solver = MadNLP.MadNLPSolver(nlp; madnlp_kwargs...) - return _SolverState(solver, nothing, nothing, nothing, nothing) + solver = MadNLP.MadNLPSolver(nlp; madnlp_kwargs...) # build MadNLP solver with user options + nvar_solver = length(solver.x.x) # internal (reduced) variable count + nvar_nlp = length(NLPModels.get_x0(nlp)) # NLP-level variable count + has_fixed = nvar_solver != nvar_nlp # mismatch → fixed variables present + return _SolverState(solver, nothing, nothing, nothing, nothing, has_fixed) # no cached duals yet end +""" + _solve!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs) + +Solve `nlp` using the MadNLP solver cached in `state`, with cascade-failure +prevention via dual snapshot restore. + +The warm-start logic implements three paths: + +1. **Fixed variables** (`state.has_fixed_vars`): create a fresh solver each + call because MadNLP's KKT factorization becomes stale when `MakeParameter` + changes fixed-variable values between solves. + +2. **Warm start after success**: copy the last-good primal into `x0` and let + `reinitialize!()` keep the cached duals. + +3. **Warm start after failure**: restore the last-good dual snapshot + (`y`, `zl.values`, `zu.values`) and mark the solver as `SOLVE_SUCCEEDED` + so `reinitialize!()` keeps those restored duals rather than the corrupted + ones. If no good snapshot exists, fall back to `INITIAL` (cold start). + +MadNLP's `cnt.k` is cumulative and never reset internally, so we reset it +before each solve to give every call a fresh `max_iter` budget. + +# Arguments +- `state::_SolverState`: solver wrapper with optional cached primal-dual + snapshot. +- `nlp`: the NLPModels-compatible problem to solve. +- `warmstart::Bool`: `true` to reuse duals from prior solves, `false` to + cold-start. +- `madnlp_kwargs`: `NamedTuple` forwarded to `MadNLP.solve!`. + +# Returns +- A MadNLP result object with fields `.status`, `.objective`, + `.multipliers`, `.solution`. + +# Examples +```julia +state = _make_solver(nlp, (print_level=0,)) +result = _solve!(state, nlp; warmstart=true, madnlp_kwargs=(print_level=0,)) +``` +""" function _solve!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs) - solver = state.solver + solver = state.solver # cached MadNLP solver instance + + # MadNLP solver reuse with MakeParameter (fixed variables) causes INFEASIBLE + # on subsequent solves even with INITIAL status — stale KKT factorization state. + # Fix: create a fresh solver each time when fixed variables exist. + if state.has_fixed_vars + return MadNLP.madnlp(nlp; madnlp_kwargs...) # one-shot fresh solver + end - # Primal warm-start: seed NLPModel's x0 from last-good primal. + # Normal path (no fixed variables): full warm-start support. if warmstart && state.last_good_x !== nothing - copyto!(NLPModels.get_x0(nlp), state.last_good_x) + copyto!(NLPModels.get_x0(nlp), state.last_good_x) # seed primal with last-good solution end - prev_result = solver.status - prev_failed = (prev_result != MadNLP.INITIAL && - prev_result != MadNLP.SOLVE_SUCCEEDED && + prev_result = solver.status # check previous solve outcome + prev_failed = (prev_result != MadNLP.INITIAL && # any non-success, non-initial status + prev_result != MadNLP.SOLVE_SUCCEEDED && # means the duals may be corrupted prev_result != MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL) if !warmstart - solver.status = MadNLP.INITIAL + solver.status = MadNLP.INITIAL # cold start: reset x, y, zl, zu elseif prev_failed && state.last_good_y !== nothing - solver.y .= state.last_good_y - solver.zl.values .= state.last_good_zl_vals - solver.zu.values .= state.last_good_zu_vals - solver.status = MadNLP.SOLVE_SUCCEEDED + solver.y .= state.last_good_y # restore last-good equality duals + solver.zl.values .= state.last_good_zl_vals # restore last-good lower-bound duals + solver.zu.values .= state.last_good_zu_vals # restore last-good upper-bound duals + solver.status = MadNLP.SOLVE_SUCCEEDED # trick reinitialize!() into warm path elseif prev_failed - solver.status = MadNLP.INITIAL + solver.status = MadNLP.INITIAL # no snapshot available → cold start end - # Reset per-solve iteration budget. - solver.cnt.k = 0 - solver.cnt.acceptable_cnt = 0 - solver.cnt.start_time = time() + # Reset per-solve iteration budget (cnt.k is cumulative in MadNLP). + solver.cnt.k = 0 # reset iteration counter + solver.cnt.acceptable_cnt = 0 # reset acceptable-step counter + solver.cnt.start_time = time() # reset wall-clock timer - res = MadNLP.solve!(solver; madnlp_kwargs...) + res = MadNLP.solve!(solver; madnlp_kwargs...) # run the solver if solve_succeeded(res) - state.last_good_x = copy(solver.x.x) - state.last_good_y = copy(solver.y) - state.last_good_zl_vals = copy(solver.zl.values) - state.last_good_zu_vals = copy(solver.zu.values) + state.last_good_x = copy(solver.x.x) # snapshot primal (GPU-safe copy) + state.last_good_y = copy(solver.y) # snapshot equality duals + state.last_good_zl_vals = copy(solver.zl.values) # snapshot lower-bound duals + state.last_good_zu_vals = copy(solver.zu.values) # snapshot upper-bound duals end return res end +""" + _solve_with_retry!(state::_SolverState, nlp; + warmstart::Bool, madnlp_kwargs, retry_on_failure::Bool) + -> (result, retried::Bool) + +Solve `nlp` via [`_solve!`](@ref), optionally retrying with a fresh cold-start +solver if the first attempt fails or returns a non-finite objective. + +The retry creates a brand-new [`_SolverState`](@ref) (discarding any corrupted +internal state) and solves with `warmstart=false`. This is more expensive than +the dual-restore path in `_solve!` but guarantees a clean factorization. + +# Arguments +- `state::_SolverState`: primary solver state (may have cached duals). +- `nlp`: the NLPModels-compatible problem. +- `warmstart::Bool`: whether the first attempt should warm-start. +- `madnlp_kwargs`: `NamedTuple` forwarded to `MadNLP.solve!`. +- `retry_on_failure::Bool`: if `true`, retry with a fresh solver on failure. + +# Returns +- `result`: the MadNLP result from whichever attempt succeeded (or the retry + result if both failed). +- `retried::Bool`: `true` if the retry path was taken. + +# Examples +```julia +result, retried = _solve_with_retry!( + state, nlp; + warmstart=true, madnlp_kwargs=(print_level=0,), retry_on_failure=true, +) +retried && @warn "solve required retry" +``` +""" +function _solve_with_retry!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs, retry_on_failure::Bool) + result = _solve!(state, nlp; warmstart = warmstart, madnlp_kwargs = madnlp_kwargs) # primary attempt + retried = false # track whether retry was needed + if retry_on_failure && (!solve_succeeded(result) || !isfinite(result.objective)) + retry_state = _make_solver(nlp, madnlp_kwargs) # fresh solver, clean factorization + result = _solve!(retry_state, nlp; warmstart = false, madnlp_kwargs = madnlp_kwargs) # cold-start retry + retried = true + if solve_succeeded(result) + state.last_good_x = copy(retry_state.solver.x.x) + state.last_good_y = copy(retry_state.solver.y) + state.last_good_zl_vals = copy(retry_state.solver.zl.values) + state.last_good_zu_vals = copy(retry_state.solver.zu.values) + end + end + return result, retried +end + # ── simulate_tsddr ──────────────────────────────────────────────────────────── """ simulate_tsddr(model, initial_state, det_equivalent, p_x0, p_target, p_uncertainty, - uncertainty_sampler; madnlp_kwargs, warmstart) - -> (objective, lambda) or nothing + uncertainty_sampler; + madnlp_kwargs, warmstart) + -> NamedTuple{(:objective, :lambda)} or nothing + +Perform a single forward pass of the TS-DDR pipeline without a gradient +update: roll out the policy to produce target states, solve the +deterministic-equivalent NLP, and extract the envelope-theorem multipliers. + +The forward pass computes + +```math +\\hat{x}_t = \\pi_\\theta(w_t, \\hat{x}_{t-1}), \\quad t = 1, \\ldots, T, +``` -Single forward pass without a gradient update. +then solves ``\\min_z Q(z; \\hat{x}, w, x_0)`` and returns the objective value +and the multipliers ``\\lambda = \\nabla_{\\hat{x}} Q`` from the target +equality constraints. + +# Arguments +- `model`: Flux policy network (LSTM or MLP). +- `initial_state::AbstractVector`: initial state vector ``x_0``. +- `det_equivalent`: ExaModels NLP with `.core`, `.model`, `.horizon`, + `.target_con_range`. +- `p_x0`: ExaModels parameter handle for the initial state. +- `p_target`: ExaModels parameter handle for policy targets. +- `p_uncertainty`: ExaModels parameter handle for per-stage uncertainty. +- `uncertainty_sampler`: `() -> w_flat` returning a flat vector of length + ``T \\times n_w``. + +# Keywords +- `madnlp_kwargs`: `NamedTuple` forwarded to MadNLP (default `NamedTuple()`). +- `warmstart::Bool`: warm-start MadNLP (default `true`). + +# Returns +- A `NamedTuple` with fields `objective::Float64` and `lambda::Vector{F}`, + or `nothing` if the solve failed or the objective is non-finite. + +# Examples +```julia +result = simulate_tsddr(model, x0, de, p_x0, p_target, p_unc, sampler) +if result !== nothing + @show result.objective +end +``` """ function simulate_tsddr( model, @@ -134,60 +552,124 @@ function simulate_tsddr( madnlp_kwargs = NamedTuple(), warmstart::Bool = true, ) - T = det_equivalent.horizon - F = eltype(initial_state) - nx = length(initial_state) - core = det_equivalent.core - nlp = det_equivalent.model + T = det_equivalent.horizon # number of planning stages + F = eltype(initial_state) # element type (Float32 or Float64) + nx = length(initial_state) # state dimension + core = det_equivalent.core # ExaModels core (for set_parameter!) + nlp = det_equivalent.model # ExaModels NLP model (for MadNLP) - state = _make_solver(nlp, madnlp_kwargs) + state = _make_solver(nlp, madnlp_kwargs) # fresh solver — no warm-start cache - w_flat = uncertainty_sampler() - nw = length(w_flat) ÷ T + # Sample one uncertainty scenario and move to the correct device. + w_flat = uncertainty_sampler() # flat vector of length T * nw_per_stage + nw = length(w_flat) ÷ T # uncertainty dimension per stage + w_dev = _adapt_array(F.(w_flat), initial_state) # move to GPU if initial_state is on GPU - Flux.reset!(model) - xhat_stages = Vector{Vector{F}}(undef, T) - prev = F.(initial_state) + # Roll out the policy to produce target states (outside AD tape). + Flux.reset!(model) # reset LSTM hidden state + xhat_stages = Vector{AbstractVector{F}}(undef, T) # allocate per-stage target storage + prev = initial_state # first policy input is x0 for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) - xhat_stages[t] = model(vcat(wt, prev)) - prev = xhat_stages[t] + wt = view(w_dev, (t-1)*nw+1 : t*nw) # slice uncertainty for stage t + xhat_stages[t] = model(vcat(wt, prev)) # policy: [w_t; x̂_{t-1}] → x̂_t + prev = xhat_stages[t] # feed x̂_t to next stage end - xhat_flat = vcat(xhat_stages...) + xhat_flat = vcat(xhat_stages...) # flatten targets to a single vector + # Set NLP parameters: initial state, uncertainty, and policy targets. ExaModels.set_parameter!(core, p_x0, initial_state) ExaModels.set_parameter!(core, p_uncertainty, w_flat) - ExaModels.set_parameter!(core, p_target, Float64.(xhat_flat)) + ExaModels.set_parameter!(core, p_target, Float64.(xhat_flat)) # NLP uses Float64 + prepare_solve!(det_equivalent, initial_state, w_flat, xhat_flat) + # Solve the deterministic equivalent (cold start for one-shot simulation). result = _solve!(state, nlp; warmstart = false, madnlp_kwargs = madnlp_kwargs) + # Reject failed or non-finite solves. solve_succeeded(result) || return nothing isfinite(result.objective) || return nothing - λ = result.multipliers[det_equivalent.target_con_range] - return (objective = result.objective, lambda = F.(Array(λ))) + # Extract envelope-theorem multipliers λ = ∇_{x̂} Q. + λ = target_multipliers(det_equivalent, result) + return (objective = result.objective, lambda = F.(λ)) # cast λ to match initial_state eltype end +""" + _rollout_xhat_flat(model, initial_state, w_flat, T::Int, F) -> AbstractVector{F} + +Roll out the policy network over `T` stages and return the concatenated +target trajectory as a single flat vector. + +This is the differentiable inner loop used inside `Zygote.gradient` blocks. +Each stage evaluates + +```math +\\hat{x}_t = \\pi_\\theta([w_t; \\hat{x}_{t-1}]), \\quad t = 1, \\ldots, T, +``` + +and the returned vector is ``[\\hat{x}_1; \\hat{x}_2; \\ldots; \\hat{x}_T]``. + +# Arguments +- `model`: Flux policy (LSTM or MLP). +- `initial_state`: state vector ``x_0`` fed to the first policy call. +- `w_flat`: flat uncertainty vector of length ``T \\times n_w``. +- `T::Int`: number of planning stages. +- `F`: element type (e.g., `Float32`). + +# Returns +- A flat vector of length ``T \\times n_x`` containing all stage targets. +""" function _rollout_xhat_flat(model, initial_state, w_flat, T::Int, F) - nw = length(w_flat) ÷ T - nx = length(initial_state) - Flux.reset!(model) - buf = Zygote.Buffer(zeros(F, nx * T)) - prev = F.(initial_state) + nw = length(w_flat) ÷ T # uncertainty dimension per stage + Flux.reset!(model) # reset LSTM hidden state + prev = F.(initial_state) # cast initial state to element type F + stages = Vector{typeof(prev)}(undef, T) # pre-allocate per-stage output vector for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) - xt = model(vcat(wt, prev)) - for i in 1:nx - buf[(t-1)*nx + i] = xt[i] - end - prev = xt + wt = view(w_flat, (t-1)*nw+1 : t*nw) # slice uncertainty for stage t + stages[t] = model(vcat(wt, prev)) # policy forward pass + prev = stages[t] # feed target to next stage end - return copy(buf) + return vcat(stages...) # flatten to single vector end -_has_critic(::NoCriticControlVariate) = false -_has_critic(::AbstractCriticControlVariate) = true +""" + _has_critic(control_variate::AbstractCriticControlVariate) -> Bool + +Return `true` if the control variate wraps an actual critic network, `false` +for the no-op [`NoCriticControlVariate`](@ref). + +# Arguments +- `control_variate`: an [`AbstractCriticControlVariate`](@ref) instance. + +# Returns +- `false` for `NoCriticControlVariate` (recovers the original dual-only update). +- `true` for any concrete critic (e.g., [`ScalarCriticControlVariate`](@ref)). +""" +_has_critic(::NoCriticControlVariate) = false # no-op sentinel → no critic +_has_critic(::AbstractCriticControlVariate) = true # any concrete critic → active +""" + _validate_critic_training_args(; kwargs...) -> Bool + +Validate critic/control-variate keyword arguments passed to [`train_tsddr`](@ref). + +# Keywords +- `actor_gradient_mode`: must be `:control_variate` or `:surrogate`. +- `critic_cv_weight`: nonnegative control-variate weight. +- `dual_actor_weight`: nonnegative dual-gradient actor weight. +- `critic_actor_weight`: nonnegative critic actor weight. +- `critic_updates_per_batch`: nonnegative number of critic updates. +- `critic_buffer_size`: nonnegative replay-buffer capacity. +- `critic_rollout_samples_per_batch`: nonnegative integer or `nothing`. +- `num_cheap_critic_samples_per_batch`: nonnegative number of extra policy + rollouts. + +# Returns +- `true` when all arguments are valid. + +# Throws +- `ErrorException` if any argument is outside its admissible set. +""" function _validate_critic_training_args(; actor_gradient_mode, critic_cv_weight, @@ -214,6 +696,25 @@ function _validate_critic_training_args(; return true end +""" + _resolve_critic_training_target(target, has_critic::Bool) + +Resolve a user-facing critic target configuration to a concrete +`AbstractCriticTrainingTarget`. + +# Arguments +- `target`: critic target object or symbolic alias. +- `has_critic::Bool`: whether critic training is active. + +# Returns +- `DeterministicEquivalentCriticTarget()` when no critic is active or the user + selected deterministic-equivalent critic targets. +- `target` unchanged when it is already an `AbstractCriticTrainingTarget`. + +# Throws +- `ErrorException` when rollout critic training is requested without a + concrete [`RolloutCriticTarget`](@ref) configuration. +""" function _resolve_critic_training_target(target, has_critic::Bool) has_critic || return DeterministicEquivalentCriticTarget() target isa AbstractCriticTrainingTarget && return target @@ -226,6 +727,30 @@ function _resolve_critic_training_target(target, has_critic::Bool) end end +""" + _critic_sample_from_rollout(model, initial_state, target, w_flat, lambda, F, solver_state) + +Build one [`CriticSample`](@ref) by rerunning a solved scenario through +stage-wise rollout. + +# Arguments +- `model`: Flux policy being trained. +- `initial_state`: initial state vector. +- `target::RolloutCriticTarget`: rollout critic-target configuration. +- `w_flat`: uncertainty trajectory from a deterministic-equivalent sample. +- `lambda`: target multipliers from the deterministic-equivalent solve. +- `F`: element type used for critic sample arrays. +- `solver_state`: optional reusable rollout solver state. + +# Returns +- `CriticSample` when rollout succeeds. +- `nothing` when rollout fails. + +# Notes +The rollout objective supplies the critic value target. The deterministic +equivalent multipliers are sliced to the rollout target length and used as the +critic gradient target. +""" function _critic_sample_from_rollout( model, initial_state, @@ -237,11 +762,15 @@ function _critic_sample_from_rollout( ) # Keep both rollout objective variants available; target.objective_value # below selects which one is used as the critic value target. + rollout_len = target.horizon * target.n_uncertainty + length(w_flat) >= rollout_len || + error("rollout critic uncertainty has length $(length(w_flat)); expected at least $rollout_len") + w_rollout = view(w_flat, 1:rollout_len) result = rollout_tsddr( model, initial_state, target.stage_problem, - w_flat; + w_rollout; horizon = target.horizon, n_uncertainty = target.n_uncertainty, set_stage_parameters! = target.set_stage_parameters!, @@ -252,15 +781,42 @@ function _critic_sample_from_rollout( policy_state = target.policy_state, solver_state = solver_state, reuse_solver = target.reuse_solver, + state_bounds = target.state_bounds, + project_state = target.project_state, + retry_on_failure = target.retry_on_failure, ) result === nothing && return nothing objective = target.objective_value === :objective ? result.objective : result.objective_no_target_penalty xhat_flat = F.(vcat(result.target_trajectory...)) - return CriticSample(F.(initial_state), F.(w_flat), xhat_flat, objective, F.(lambda)) + λ_rollout = view(lambda, 1:length(xhat_flat)) + return CriticSample(F.(initial_state), F.(w_rollout), xhat_flat, objective, F.(λ_rollout)) end +""" + _rollout_critic_samples(model, initial_state, target, de_samples, F, max_samples, solver_state) + -> Vector{CriticSample} + +Convert deterministic-equivalent training samples into rollout critic samples. + +# Arguments +- `model`: Flux policy being trained. +- `initial_state`: initial state vector. +- `target::RolloutCriticTarget`: rollout critic-target configuration. +- `de_samples`: deterministic-equivalent samples containing uncertainty and + target multipliers. +- `F`: element type used for critic sample arrays. +- `max_samples`: maximum number of solved scenarios to rerun, or `nothing`. +- `solver_state`: optional reusable rollout solver state. + +# Returns +- `Vector{CriticSample}` containing only successful rollout conversions. + +# Notes +When `max_samples` is smaller than `length(de_samples)`, samples are selected +without replacement using `randperm`. +""" function _rollout_critic_samples( model, initial_state, @@ -298,51 +854,69 @@ end train_tsddr(model, initial_state, det_equivalent, p_x0, p_target, p_uncertainty, uncertainty_sampler; - num_batches, num_train_per_batch, optimizer, - adjust_hyperparameters, record_loss, - madnlp_kwargs, warmstart, - problem_pool) -> model - -TS-DDR policy gradient training. Mirrors `train_multistage` from DecisionRules.jl. - -Arguments: -- `model` : Flux policy (LSTM or MLP) -- `initial_state` : initial state vector -- `det_equivalent` : any ExaModels NLP with fields `.core`, `.model`, - `.horizon`, `.target_con_range` -- `p_x0` : ExaModels parameter for the initial state -- `p_target` : ExaModels parameter for policy targets -- `p_uncertainty` : ExaModels parameter for per-stage uncertainty -- `uncertainty_sampler`: `() -> w_flat` — flat vector of length `T * nw_per_stage`. - For multi-unit problems (e.g., hydro reservoirs) the sampler - should draw one joint scenario index per stage to preserve - spatial correlation; see `sample_scenario` in examples. - -Keyword arguments (mirror `train_multistage`): -- `num_batches` : total gradient steps (default 100) -- `num_train_per_batch` : scenarios averaged per step (default 1) -- `optimizer` : Flux.Optimisers optimizer -- `adjust_hyperparameters` : `(iter, opt_state, n) -> n` -- `record_loss` : `(iter, model, loss, tag) -> Bool`; return `true` to stop -- `madnlp_kwargs` : NamedTuple forwarded to MadNLP -- `warmstart` : warm-start MadNLP between solves (default `true`) -- `problem_pool` : vector of `(de, p_x0, p_target, p_uncertainty)` tuples - for parallel GPU solves; each entry gets its own MadNLP solver - and samples are distributed round-robin across the pool -- `control_variate` : optional `ScalarCriticControlVariate`; default - `NoCriticControlVariate()` recovers the original update -- `critic_training_target` : `RolloutCriticTarget(...)` for rollout-objective - critic fitting, or `DeterministicEquivalentCriticTarget()` - / `:deterministic_equivalent` for DE ablations -- `critic_rollout_samples_per_batch`: number of solved batch scenarios to rerun - through stage-wise rollout for critic targets; - `nothing` uses all successful solved scenarios -- `actor_gradient_mode` : `:control_variate` or `:surrogate` -- `num_cheap_critic_samples_per_batch`: extra policy rollouts used only for - critic actor terms; these do not trigger NLP solves -- `external_critic_samples` : mutable vector; `record_loss` can push - `CriticSample`s (e.g. from `critic_samples_from_evaluation`) - to feed the critic replay buffer without extra solves + num_batches=100, + num_train_per_batch=1, + optimizer, + adjust_hyperparameters, + record_loss, + madnlp_kwargs=NamedTuple(), + warmstart=true, + problem_pool=nothing, + kwargs...) -> model + +Train a TS-DDR policy with open-loop deterministic-equivalent solves. + +The policy rolls out a target trajectory, the ExaModels problem projects that +trajectory onto the feasible set, and target multipliers provide the actor +gradient by the envelope theorem. + +# Arguments +- `model`: Flux policy. +- `initial_state::AbstractVector`: initial state vector. +- `det_equivalent`: ExaModels deterministic-equivalent problem. +- `p_x0`: ExaModels parameter for the initial state. +- `p_target`: ExaModels parameter for policy targets. +- `p_uncertainty`: ExaModels parameter for uncertainty. +- `uncertainty_sampler`: callable returning a flat uncertainty trajectory. + +# Keywords +- `num_batches::Int`: number of gradient steps. +- `num_train_per_batch::Int`: number of scenarios averaged per step. +- `optimizer`: Flux optimizer or optimizer chain. +- `adjust_hyperparameters`: callback `(iter, opt_state, n) -> n`. +- `record_loss`: callback `(iter, model, loss, tag) -> Bool`; return `true` + to stop training. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: warm-start MadNLP between solves. +- `retry_on_failure::Bool`: retry failed solves with a fresh solver state. +- `problem_pool`: optional vector of `(de, p_x0, p_target, p_uncertainty)` + tuples for independent solves. +- `control_variate`: optional critic control variate. +- `actor_gradient_mode::Symbol`: `:control_variate` or `:surrogate`. +- `critic_cv_weight`, `dual_actor_weight`, `critic_actor_weight`: actor loss + weights. +- `critic_updates_per_batch::Int`: critic optimizer steps per batch. +- `critic_buffer_size::Int`: replay-buffer capacity. +- `critic_batch_size`: critic minibatch size, or `nothing` for all samples. +- `critic_training_target`: rollout or deterministic-equivalent critic target. +- `critic_rollout_samples_per_batch`: number of solved samples rerun through + rollout for critic targets; `nothing` uses all successful solved samples. +- `num_cheap_critic_samples_per_batch::Int`: extra policy rollouts used only + for critic actor terms. +- `critic_optimizer`: Flux optimizer for the critic. +- `external_critic_samples`: optional mutable vector of externally produced + `CriticSample`s. +- `batch_diagnostics`: callback `(iter, stats) -> nothing`. +- `reuse_solver::Bool`: force solver reuse when fixed variables have constant + bounds across solves. + +# Returns +- `model`, updated in place. + +# Notes +For multi-unit stochastic processes, `uncertainty_sampler` should preserve +within-stage spatial correlation, for example by drawing one joint scenario +index per stage. """ function train_tsddr( model, @@ -365,6 +939,7 @@ function train_tsddr( end, madnlp_kwargs = NamedTuple(), warmstart::Bool = true, + retry_on_failure::Bool = true, problem_pool = nothing, control_variate::AbstractCriticControlVariate = NoCriticControlVariate(), actor_gradient_mode::Symbol = :control_variate, @@ -379,6 +954,8 @@ function train_tsddr( num_cheap_critic_samples_per_batch::Int = 0, critic_optimizer = Flux.Adam(1f-3), external_critic_samples = nothing, + batch_diagnostics = (iter, stats) -> nothing, + reuse_solver::Bool = false, ) T = det_equivalent.horizon F = eltype(initial_state) @@ -410,6 +987,9 @@ function train_tsddr( # Single-worker: create solver on main task (no threading needed) single_state = nworkers == 1 ? _make_solver(_pool[1][1].model, madnlp_kwargs) : nothing + if reuse_solver && single_state !== nothing + single_state.has_fixed_vars = false + end # Multi-worker: persistent worker threads via channels. # Each worker creates its own MadNLP solver on its own thread so that @@ -422,8 +1002,12 @@ function train_tsddr( (de, px, pt, pu) = _pool[wi] in_ch = in_channels[wi] out_ch = out_channels[wi] + _reuse_solver = reuse_solver t = Threads.@spawn begin st = _make_solver(de.model, madnlp_kwargs) + if _reuse_solver + st.has_fixed_vars = false + end while true msg = take!(in_ch) msg === nothing && break @@ -431,15 +1015,31 @@ function train_tsddr( ExaModels.set_parameter!(de.core, px, init_state) ExaModels.set_parameter!(de.core, pu, w_flat) ExaModels.set_parameter!(de.core, pt, Float64.(xhat_flat)) - result = _solve!(st, de.model; warmstart=warmstart, madnlp_kwargs=madnlp_kwargs) - if solve_succeeded(result) && isfinite(result.objective) - λ = result.multipliers[de.target_con_range] + prepare_solve!(de, init_state, w_flat, xhat_flat) + result, retried = _solve_with_retry!( + st, + de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + failure = nothing + if !solve_succeeded(result) + failure = "status_" * _status_key(result.status) + elseif !isfinite(result.objective) + failure = "nonfinite_objective" + else + # Solve succeeded with a finite objective (both negations + # were tested above); only the multipliers remain to check. + λ = target_multipliers(de, result) if all(isfinite, λ) - put!(out_ch, (s_idx, F.(w_flat), F.(Array(λ)), result.objective)) + put!(out_ch, (s_idx, F.(w_flat), _adapt_array(F.(λ), w_flat), + result.objective, result.status, nothing, retried)) continue end + failure = "nonfinite_lambda" end - put!(out_ch, (s_idx, nothing, nothing, NaN)) + put!(out_ch, (s_idx, nothing, nothing, NaN, result.status, failure, retried)) end end push!(worker_tasks, t) @@ -461,24 +1061,36 @@ function train_tsddr( # ── Forward pass: rollout + solve (outside AD tape) ─────────────────── - # Step 1: Roll out policy for all samples (CPU, sequential) - sample_data = Vector{Tuple{Vector{F}, Vector{F}}}(undef, num_train_per_batch) + # Step 1: Roll out policy for all samples + # Precision note: uncertainties are cast to F (typically Float32, the + # policy precision) here, and this F-cast array is later written into + # the Float64 NLP via set_parameter!. The round-trip through Float32 is + # intentional: the NLP must be solved for exactly the targets the policy + # produced from these Float32 inputs, so the resulting λ multipliers + # pair with the same Float32 rollout in the gradient step below + # (train/gradient consistency). Do not "fix" this by keeping w in + # Float64 for the NLP only. + sample_data = Vector{Tuple{AbstractVector{F}, AbstractVector{F}}}(undef, num_train_per_batch) for s in 1:num_train_per_batch w_flat = uncertainty_sampler() nw = length(w_flat) ÷ T + w_dev = _adapt_array(F.(w_flat), initial_state) Flux.reset!(model) - xhat_stages = Vector{Vector{F}}(undef, T) - prev = F.(initial_state) + xhat_stages = Vector{AbstractVector{F}}(undef, T) + prev = initial_state for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) + wt = view(w_dev, (t-1)*nw+1 : t*nw) xhat_stages[t] = model(vcat(wt, prev)) - prev = xhat_stages[t] + prev = xhat_stages[t] end - sample_data[s] = (F.(w_flat), vcat(xhat_stages...)) + sample_data[s] = (w_dev, vcat(xhat_stages...)) end # Step 2: Solve — parallel across workers if pool provided - solve_ok = Vector{Union{Nothing, Tuple{Vector{F}, Vector{F}, Float64}}}(nothing, num_train_per_batch) + solve_ok = Vector{Union{Nothing, Tuple{AbstractVector{F}, AbstractVector{F}, Float64}}}(nothing, num_train_per_batch) + status_counts = Dict{String, Int}() + failure_counts = Dict{String, Int}() + retry_counts = Dict{String, Int}() if nworkers == 1 (de, px, pt, pu) = _pool[1] @@ -488,12 +1100,30 @@ function train_tsddr( ExaModels.set_parameter!(de.core, px, initial_state) ExaModels.set_parameter!(de.core, pu, w_flat) ExaModels.set_parameter!(de.core, pt, Float64.(xhat_flat)) - result = _solve!(st, de.model; warmstart=warmstart, madnlp_kwargs=madnlp_kwargs) - solve_succeeded(result) || continue - isfinite(result.objective) || continue - λ = result.multipliers[de.target_con_range] - all(isfinite, λ) || continue - solve_ok[s] = (F.(w_flat), F.(Array(λ)), result.objective) + prepare_solve!(de, initial_state, w_flat, xhat_flat) + result, retried = _solve_with_retry!( + st, + de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + retried && _inc_count!(retry_counts, solve_succeeded(result) && isfinite(result.objective) ? "retry_success" : "retry_failure") + _inc_status!(status_counts, result.status) + if !solve_succeeded(result) + _inc_count!(failure_counts, "status_" * _status_key(result.status)) + continue + end + if !isfinite(result.objective) + _inc_count!(failure_counts, "nonfinite_objective") + continue + end + λ = target_multipliers(de, result) + if !all(isfinite, λ) + _inc_count!(failure_counts, "nonfinite_lambda") + continue + end + solve_ok[s] = (F.(w_flat), _adapt_array(F.(λ), initial_state), result.objective) end else for round_start in 1:nworkers:num_train_per_batch @@ -505,7 +1135,19 @@ function train_tsddr( put!(in_channels[wi], (s, initial_state, w_flat, xhat_flat)) end for wi in 1:round_size - (s_idx, w_out, λ_out, obj_out) = take!(out_channels[wi]) + msg = take!(out_channels[wi]) + if length(msg) == 7 + (s_idx, w_out, λ_out, obj_out, status_out, failure_out, retried_out) = msg + _inc_status!(status_counts, status_out) + failure_out !== nothing && _inc_count!(failure_counts, failure_out) + retried_out && _inc_count!(retry_counts, w_out === nothing ? "retry_failure" : "retry_success") + elseif length(msg) == 6 + (s_idx, w_out, λ_out, obj_out, status_out, failure_out) = msg + _inc_status!(status_counts, status_out) + failure_out !== nothing && _inc_count!(failure_counts, failure_out) + else + (s_idx, w_out, λ_out, obj_out) = msg + end if w_out !== nothing solve_ok[s_idx] = (w_out, λ_out, obj_out) end @@ -514,7 +1156,7 @@ function train_tsddr( end # Step 3: Collect valid results - valid = Vector{Tuple{Vector{F}, Vector{F}}}() + valid = Tuple{AbstractVector{F}, AbstractVector{F}}[] de_samples = CriticSample[] obj_sum = 0.0 for (s, r) in enumerate(solve_ok) @@ -528,6 +1170,13 @@ function train_tsddr( end n_ok = length(valid) mean_obj = n_ok > 0 ? obj_sum / n_ok : NaN + batch_diagnostics(iter, Dict{String, Any}( + "n_ok" => n_ok, + "n_total" => num_train_per_batch, + "status_counts" => copy(status_counts), + "failure_counts" => copy(failure_counts), + "retry_counts" => copy(retry_counts), + )) if has_critic && n_ok > 0 && critic_updates_per_batch > 0 valid_samples = if resolved_critic_training_target isa RolloutCriticTarget @@ -574,18 +1223,18 @@ function train_tsddr( for (w_flat_s, λf) in valid nw = length(w_flat_s) ÷ T Flux.reset!(m) - prev_ad = F.(initial_state) + prev_ad = initial_state for t in 1:T - wt = F.(w_flat_s[(t-1)*nw+1 : t*nw]) + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) xt = m(vcat(wt, prev_ad)) - total = total + sum(λf[(t-1)*nx+1 : t*nx] .* xt) + total = total + sum(view(λf, (t-1)*nx+1 : t*nx) .* xt) prev_ad = xt end end total / F(n_ok) end else - solved_weights = Vector{Tuple{Vector{F}, Vector{F}}}() + solved_weights = Tuple{AbstractVector{F}, AbstractVector{F}}[] for sample in de_samples λf = F.(sample.target_multipliers) if actor_gradient_mode === :control_variate @@ -618,12 +1267,12 @@ function train_tsddr( for (w_flat_s, actor_weight) in solved_weights nw = length(w_flat_s) ÷ T Flux.reset!(m) - prev_ad = F.(initial_state) + prev_ad = initial_state for t in 1:T - wt = F.(w_flat_s[(t-1)*nw+1 : t*nw]) + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) xt = m(vcat(wt, prev_ad)) residual_total = - residual_total + sum(actor_weight[(t-1)*nx+1 : t*nx] .* xt) + residual_total + sum(view(actor_weight, (t-1)*nx+1 : t*nx) .* xt) prev_ad = xt end end @@ -637,7 +1286,7 @@ function train_tsddr( xhat_ad = _rollout_xhat_flat(m, initial_state, w_flat_s, T, F) critic_total = critic_total + critic_value( control_variate, - F.(initial_state), + initial_state, w_flat_s, xhat_ad, ) @@ -659,13 +1308,192 @@ function train_tsddr( end finally - # Shut down worker threads - for ch in in_channels - put!(ch, nothing) + # Shut down worker threads. A worker that already died never drains its + # input channel, so an unconditional put! on a full Channel{Any}(1) + # would block forever; only signal workers that are still running, and + # guard the put! itself against the check-then-put race. + for (i, ch) in enumerate(in_channels) + if !istaskdone(worker_tasks[i]) # skip dead workers (nobody consumes) + try + put!(ch, nothing) # normal shutdown sentinel + catch err + # Worker died between the istaskdone check and the put! + # (or the channel was closed) — nothing left to signal. + @warn "train_tsddr worker $i shutdown signal failed" exception=(err, catch_backtrace()) + end + end + end + for (i, t) in enumerate(worker_tasks) + try + wait(t) # join worker task + catch err + # A failed worker rethrows on wait; log instead of masking the + # original in-flight exception during cleanup. + @warn "train_tsddr worker $i failed" exception=(err, catch_backtrace()) + end + end + end + + return model +end + +# ── train_tsddr_embedded ───────────────────────────────────────────────────── + +""" + train_tsddr_embedded(model, initial_state, embedded_de, + uncertainty_sampler; kwargs...) -> model + +Train a TS-DDR policy embedded directly in the NLP. + +Unlike [`train_tsddr`](@ref), this function does not roll out targets +externally. The NLP oracle evaluates `model` inline, the solve returns +closed-loop multipliers and realized states, and the actor gradient is computed +from those realized states. + +# Arguments +- `model`: Flux policy captured by the embedded oracle closures. +- `initial_state::AbstractVector`: initial state vector. +- `embedded_de`: embedded deterministic-equivalent problem. +- `uncertainty_sampler`: callable returning a flat uncertainty trajectory. + +# Keywords +- `num_batches::Int`: number of gradient steps. +- `num_train_per_batch::Int`: number of scenarios averaged per step. +- `optimizer`: Flux optimizer or optimizer chain. +- `adjust_hyperparameters`: callback `(iter, opt_state, n) -> n`. +- `record_loss`: callback `(iter, model, loss, tag) -> Bool`; return `true` + to stop training. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: warm-start MadNLP between solves. +- `retry_on_failure::Bool`: retry failed solves with a fresh solver state. +- `get_realized_states`: optional callback `(prob, result) -> x_flat`. +- `batch_diagnostics`: callback `(iter, stats) -> nothing`. + +# Returns +- `model`, updated in place. + +# Notes +The gradient uses +`sum_t dot(lambda_t, pi_theta(w_t, x^*_{t-1}))`, where `x^*` is the realized +state trajectory from the coupled NLP solution. +""" +function train_tsddr_embedded( + model, + initial_state::AbstractVector, + embedded_de, + uncertainty_sampler; + num_batches::Int = 100, + num_train_per_batch::Int = 1, + optimizer = Flux.Optimisers.OptimiserChain( + Flux.Optimisers.ClipGrad(1.0f0), + Flux.Adam(1f-3), + ), + adjust_hyperparameters = (iter, opt_state, n) -> n, + record_loss = (iter, model, loss, tag) -> begin + println("$tag iter=$iter loss=$(round(loss; digits=4))") + return false + end, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + retry_on_failure::Bool = true, + get_realized_states = nothing, + batch_diagnostics = (iter, stats) -> nothing, +) + T = embedded_de.horizon + F = eltype(initial_state) + nx = embedded_de.nx + + _get_states = get_realized_states === nothing ? + (prob, res) -> res.solution[1 : prob.horizon * prob.nx] : + get_realized_states + + state = _make_solver(embedded_de.model, madnlp_kwargs) + opt_state = Flux.setup(optimizer, model) + + for iter in 1:num_batches + num_train_per_batch = adjust_hyperparameters(iter, opt_state, num_train_per_batch) + + valid = Tuple{AbstractVector{F}, AbstractVector{F}, AbstractVector{F}}[] + obj_sum = 0.0 + status_counts = Dict{String, Int}() + failure_counts = Dict{String, Int}() + retry_counts = Dict{String, Int}() + + for s in 1:num_train_per_batch + w_flat = uncertainty_sampler() + + set_x0!(embedded_de, initial_state) + set_uncertainty!(embedded_de, w_flat) + + result, retried = _solve_with_retry!( + state, + embedded_de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + retried && _inc_count!(retry_counts, solve_succeeded(result) && isfinite(result.objective) ? "retry_success" : "retry_failure") + + _inc_status!(status_counts, result.status) + if !solve_succeeded(result) + _inc_count!(failure_counts, "status_" * _status_key(result.status)) + continue + end + if !isfinite(result.objective) + _inc_count!(failure_counts, "nonfinite_objective") + continue + end + + λ = target_multipliers(embedded_de, result) + if !all(isfinite, λ) + _inc_count!(failure_counts, "nonfinite_lambda") + continue + end + + x_sol = _get_states(embedded_de, result) + + λf = _adapt_array(F.(λ), initial_state) + xf = _adapt_array(F.(x_sol), initial_state) + w_dev = _adapt_array(F.(w_flat), initial_state) + push!(valid, (w_dev, λf, xf)) + obj_sum += result.objective end - for t in worker_tasks - wait(t) + + n_ok = length(valid) + mean_obj = n_ok > 0 ? obj_sum / n_ok : NaN + batch_diagnostics(iter, Dict{String, Any}( + "n_ok" => n_ok, + "n_total" => num_train_per_batch, + "status_counts" => copy(status_counts), + "failure_counts" => copy(failure_counts), + "retry_counts" => copy(retry_counts), + )) + + if n_ok > 0 + gs = Zygote.gradient(model) do m + total = zero(F) + for (w_flat_s, λf, x_realized) in valid + nw = length(w_flat_s) ÷ T + Flux.reset!(m) + for t in 1:T + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) + x_prev = (t == 1) ? + initial_state : + view(x_realized, (t-2)*nx+1 : (t-1)*nx) + xt = m(vcat(wt, x_prev)) + total = total + sum(view(λf, (t-1)*nx+1 : t*nx) .* xt) + end + end + total / F(n_ok) + end + + grad = materialize_tangent(gs[1]) + if grad !== nothing && _all_finite_gradient(grad) + Flux.update!(opt_state, model, grad) + end end + + record_loss(iter, model, mean_obj, "metrics/training_loss") && break end return model diff --git a/src/utils.jl b/src/utils.jl index 4ce90ba..208492f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -2,25 +2,47 @@ # Small helpers shared across the package. """ - x_index(nx, t, i) + x_index(nx, t, i) -> Int -Linear index for state component `i ∈ 1:nx` at stage `t ∈ 1:T` -when state trajectory is stored as a flat vector of length `T*nx`. +Return the flat-vector index for state component `i` at stage `t`. + +# Arguments +- `nx::Int`: number of state components per stage. +- `t`: one-based stage index. +- `i`: one-based state-component index. + +# Returns +- `Int`: index into a stage-major state trajectory of length `T * nx`. """ @inline x_index(nx::Int, t, i) = (t - 1) * nx + i """ - u_index(nu, t, i) + u_index(nu, t, i) -> Int + +Return the flat-vector index for control component `i` at stage `t`. -Linear index for control component `i ∈ 1:nu` at stage `t ∈ 1:(T-1)` -when controls are stored as a flat vector of length `(T-1)*nu`. +# Arguments +- `nu::Int`: number of control components per stage. +- `t`: one-based stage index. +- `i`: one-based control-component index. + +# Returns +- `Int`: index into a stage-major control trajectory of length `(T - 1) * nu`. """ @inline u_index(nu::Int, t, i) = (t - 1) * nu + i """ - w_index(nw, t, i) + w_index(nw, t, i) -> Int + +Return the flat-vector index for uncertainty component `i` at stage `t`. + +# Arguments +- `nw::Int`: number of uncertainty components per stage. +- `t`: one-based stage index. +- `i`: one-based uncertainty-component index. -Linear index for disturbance component `i ∈ 1:nw` at stage `t ∈ 1:(T-1)` -when disturbances are stored as a flat vector of length `(T-1)*nw`. +# Returns +- `Int`: index into a stage-major uncertainty trajectory of length + `(T - 1) * nw`. """ @inline w_index(nw::Int, t, i) = (t - 1) * nw + i diff --git a/test/runtests.jl b/test/runtests.jl index 857e673..833b70e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,9 +1,25 @@ using Test using DecisionRulesExa +using ExaModels using Flux +using MadNLP using Random using Zygote +include(joinpath(@__DIR__, "..", "examples", "HydroPowerModels", "hydro_training_utils.jl")) +# Hydro example files (data structs, ExaModels builders, reachable policy) used +# by the reactive-deficit and recurrent-threading testsets below. +include(joinpath(@__DIR__, "..", "examples", "HydroPowerModels", "hydro_power_data.jl")) +include(joinpath(@__DIR__, "..", "examples", "HydroPowerModels", "hydro_power_exa.jl")) +include(joinpath(@__DIR__, "..", "examples", "HydroPowerModels", "hydro_reachable_policy.jl")) + +@testset "HydroPowerModels training utilities" begin + @test parse_layers("128, 64") == [128, 64] + @test parse_layers("") == Int[] + @test parse_layers(" ") == Int[] + @test parse_layers("32,,16") == [32, 16] +end + @testset "DeterministicEquivalentProblem (CPU)" begin T = 6 nx = 1 @@ -24,7 +40,7 @@ using Zygote set_uncertainty!(prob, w) set_targets!(prob, xhat) - res = solve!(prob; tol = 1e-6, max_iter = 200) + res = DecisionRulesExa.solve!(prob; tol = 1e-6, max_iter = 200) n_x = T * nx n_u = (T - 1) * nx @@ -110,6 +126,39 @@ end @test isfinite(critic_loss(hybrid, [sample])) end +@testset "Bounded state policy helper" begin + Random.seed!(21) + lower = Float32[0, 1, -2] + upper = Float32[10, 1, 2] + policy = bounded_state_policy(2, lower, upper, [4]; activation = sigmoid) + + y = policy(Float32[0.2, -0.1, 3, 1, -1]) + @test length(y) == 3 + @test lower[1] <= y[1] <= upper[1] + @test y[2] == lower[2] + @test lower[3] <= y[3] <= upper[3] + @test length(policy.policy.combiner.bias) == 2 + + deep_policy = bounded_state_policy( + 2, + lower, + upper, + [4]; + activation = sigmoid, + combiner_layers = [5, 4], + ) + @test deep_policy.policy.combiner isa Flux.Chain + y_deep = deep_policy(Float32[0.2, -0.1, 3, 1, -1]) + @test length(y_deep) == 3 + @test lower[1] <= y_deep[1] <= upper[1] + @test y_deep[2] == lower[2] + @test lower[3] <= y_deep[3] <= upper[3] + + constant_policy = bounded_state_policy(1, Float32[3, -4], Float32[3, -4], [4]) + @test constant_policy(Float32[0, 99, 100]) == Float32[3, -4] + @test isempty(Flux.trainables(constant_policy)) +end + @testset "Critic and actor update separation" begin Random.seed!(11) critic = Chain(Dense(2 => 1, bias = false)) @@ -183,3 +232,527 @@ end @test materialize_tangent(g_cv).layers[1].weight ≈ materialize_tangent(g_dual).layers[1].weight end + +@testset "EmbeddedDeterministicEquivalentProblem (CPU)" begin + T = 5 + nx = 1 + nw = 1 + + Random.seed!(42) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + nu = nx, + nw = nw, + backend = nothing, + float_type = Float64, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + @test prob isa EmbeddedDeterministicEquivalentProblem + @test prob.horizon == T + @test prob.nx == nx + + x0 = [1.0] + w = randn(T * nw) + + set_x0!(prob, x0) + set_uncertainty!(prob, w) + + n_x = T * nx + n_u = (T - 1) * nx + n_var = n_x + n_u + n_x + n_con_initial = nx + n_con_dynamics = (T - 1) * nx + n_con_oracle = n_x + n_con = n_con_initial + n_con_dynamics + n_con_oracle + + res = MadNLP.madnlp(prob.model; tol = 1e-6, max_iter = 500, print_level = MadNLP.ERROR) + @test DecisionRulesExa.solve_succeeded(res) + @test length(res.solution) == n_var + @test length(res.multipliers) == n_con + + λ = target_multipliers(prob, res) + @test length(λ) == n_x + @test all(isfinite, λ) + + x_sol, u_sol, δ_sol = solution_components(prob, res) + @test length(x_sol) == n_x + @test length(u_sol) == n_u + @test length(δ_sol) == n_x + + # Verify oracle constraint satisfaction: x_t + δ_t ≈ π_θ(w_t, x_{t-1}) + Flux.reset!(policy) + for t in 1:T + x_prev = (t == 1) ? Float32.(x0) : Float32.([x_sol[(t-2)*nx+1:(t-1)*nx]...]) + w_t = Float32.([w[(t-1)*nw+1:t*nw]...]) + nn_out = policy(vcat(w_t, x_prev)) + for i in 1:nx + xi = x_sol[(t-1)*nx+i] + di = δ_sol[(t-1)*nx+i] + @test xi + di ≈ Float64(nn_out[i]) atol = 1e-4 + end + end +end + +@testset "Embedded NN gradient (envelope theorem)" begin + T = 4 + nx = 1 + nw = 1 + + Random.seed!(7) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + nu = nx, + nw = nw, + backend = nothing, + float_type = Float64, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + x0 = [0.5] + w = randn(T * nw) + + set_x0!(prob, x0) + set_uncertainty!(prob, w) + + res = MadNLP.madnlp(prob.model; tol = 1e-6, max_iter = 500, print_level = MadNLP.ERROR) + @test DecisionRulesExa.solve_succeeded(res) + + λ = target_multipliers(prob, res) + x_sol = res.solution[1 : T * nx] + + # Zygote gradient: ∇_θ Σ_t ⟨λ_t, π_θ(w_t, x*_{t-1})⟩ + gs = Zygote.gradient(policy) do m + total = 0.0f0 + Flux.reset!(m) + for t in 1:T + wt = Float32.([w[(t-1)*nw+1:t*nw]...]) + x_prev = (t == 1) ? + Float32.(x0) : + Float32.([x_sol[(t-2)*nx+1:(t-1)*nx]...]) + xt = m(vcat(wt, x_prev)) + for i in 1:nx + total = total + Float32(λ[(t-1)*nx+i]) * xt[i] + end + end + total + end + + g = materialize_tangent(gs[1]) + @test g !== nothing + @test DecisionRulesExa._all_finite_gradient(g) +end + +@testset "train_tsddr_embedded smoke test" begin + T = 4 + nx = 1 + nw = 1 + + Random.seed!(99) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + backend = nothing, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + x0 = Float32[1.0] + losses = Float64[] + + train_tsddr_embedded( + policy, x0, prob, + () -> randn(T * nw); + num_batches = 5, + num_train_per_batch = 2, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + push!(losses, loss) + return false + end, + ) + + @test length(losses) == 5 + @test all(isfinite, losses) +end + +@testset "rollout_tsddr (CPU)" begin + horizon = 3 + nx = 1 + + # Stage problem: a horizon-2 linear tracking deterministic equivalent used + # as a one-stage projection problem. The realized next state is x_2. + stage_problem = build_linear_tracking_problem( + horizon = 2, + nx = nx, + backend = nothing, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + # Callback: write (x_{t-1}, w_t, xhat_t) into the stage problem. The stage-1 + # target equals the incoming state (its constraint is slack-absorbed anyway); + # the stage-2 target is the policy target being projected. + set_stage_params! = (prob, state, w_t, target, stage) -> begin + set_x0!(prob, state) + set_uncertainty!(prob, w_t) + set_targets!(prob, vcat(state, target)) + return nothing + end + # Callback: read the realized next state x_2 from the stage solution. + realized = (prob, result) -> begin + x_sol, _, _ = solution_components(prob, result) + return x_sol[end - prob.nx + 1 : end] + end + + Random.seed!(31) + policy = StateConditionedPolicy(1, 1, 1, [4]; activation = tanh) + x0 = Float32[0.5] + w_flat = Float32.(0.1 .* randn(horizon * 1)) + + result = rollout_tsddr( + policy, x0, stage_problem, w_flat; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + ) + @test result isa NamedTuple + @test isfinite(result.objective) + @test length(result.state_trajectory) == horizon + 1 + @test length(result.target_trajectory) == horizon + + # The :target feedback mode (policy sees its own previous target) must also run. + result_target = rollout_tsddr( + policy, x0, stage_problem, w_flat; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + policy_state = :target, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + ) + @test result_target isa NamedTuple + @test isfinite(result_target.objective) + + # A wrong-length uncertainty vector must be rejected before any solve. + @test_throws ArgumentError rollout_tsddr( + policy, x0, stage_problem, Float32[0.1]; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + ) +end + +@testset "train_tsddr open-loop smoke test" begin + T = 4 + nx = 1 + + # train_tsddr writes the full length-T*nw uncertainty sample into + # p_uncertainty via ExaModels.set_parameter!, which enforces an exact size + # match. build_linear_tracking_problem's p_w has length (T-1)*nw (dynamics + # stages only), so this test builds the same linear tracking NLP manually + # with a full-length uncertainty parameter (mirroring the hydro p_inflow + # layout, where the final-stage entry does not enter the dynamics). + core = ExaModels.ExaCore(Float64) + x = ExaModels.variable(core, T * nx) + u = ExaModels.variable(core, (T - 1) * nx; lvar = -2.0, uvar = 2.0) + δ = ExaModels.variable(core, T * nx) + p_x0 = ExaModels.parameter(core, zeros(nx)) + p_w = ExaModels.parameter(core, zeros(T * nx)) # full length T*nw + p_target = ExaModels.parameter(core, zeros(T * nx)) + # Stage cost (x_t^2 + u_t^2)/2 plus slack penalty (rho/2)*delta^2, rho = 10. + ExaModels.objective(core, + (x[x_index(nx, t, i)]^2 + u[u_index(nx, t, i)]^2) / 2 + for t in 1:(T - 1), i in 1:nx + ) + ExaModels.objective(core, + 5.0 * δ[x_index(nx, t, i)]^2 + for t in 1:T, i in 1:nx + ) + # Initial condition, dynamics x_{t+1} = x_t + u_t + w_t, then targets LAST. + ExaModels.constraint(core, + x[x_index(nx, 1, i)] - p_x0[i] + for i in 1:nx + ) + ExaModels.constraint(core, + x[x_index(nx, t + 1, i)] - x[x_index(nx, t, i)] - + u[u_index(nx, t, i)] - p_w[w_index(nx, t, i)] + for t in 1:(T - 1), i in 1:nx + ) + ExaModels.constraint(core, + p_target[x_index(nx, t, i)] - x[x_index(nx, t, i)] - δ[x_index(nx, t, i)] + for t in 1:T, i in 1:nx + ) + model = ExaModels.ExaModel(core) + target_start = nx + (T - 1) * nx + 1 + prob = DeterministicEquivalentProblem( + core, model, x, u, δ, + p_x0, p_w, p_target, + nx, nx, nx, T, + target_start:(target_start + T * nx - 1), + ) + + Random.seed!(123) + policy = StateConditionedPolicy(nx, nx, nx, [8]; activation = tanh) + x0 = Float32[1.0] + losses = Float64[] + params_before = _state_vector(policy) + + train_tsddr( + policy, x0, prob, + prob.p_x0, prob.p_target, prob.p_w, + () -> randn(T * nx); + num_batches = 2, + num_train_per_batch = 1, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + push!(losses, loss) + return false + end, + ) + + @test length(losses) == 2 + @test all(isfinite, losses) + @test _state_vector(policy) != params_before +end + +@testset "StateConditionedPolicy recurrent-state threading" begin + Random.seed!(42) + policy = StateConditionedPolicy(2, 1, 1, [4, 3]; activation = tanh) + x = Float32[0.3, -0.2, 0.5] + + # Memory: the same input twice WITHOUT reset must give different outputs + # (the recurrent state advanced between calls). A memoryless (per-call + # restart) encoder would reproduce the first output exactly. + Flux.reset!(policy) + y1 = policy(x) + y2 = policy(x) + @test y1 != y2 + + # Reset restores the exact initial recurrent state: identical output. + Flux.reset!(policy) + @test policy(x) == y1 + + # Manual LSTMCell recursion with the same weights reproduces the policy's + # stage outputs EXACTLY over a 3-stage open-loop sequence. + ws = [Float32[0.3, -0.2], Float32[0.1, 0.4], Float32[-0.5, 0.2]] + Flux.reset!(policy) + prev = Float32[0.5] + outs = Vector{Vector{Float32}}() + for t in 1:3 + y = policy(vcat(ws[t], prev)) + push!(outs, Float32.(y)) + prev = Float32.(y) + end + + cells = [l.cell for l in policy.encoder.layers] + states = Any[Flux.initialstates(c) for c in cells] + prev = Float32[0.5] + for t in 1:3 + h = ws[t] + for (i, c) in enumerate(cells) + h, states[i] = c(h, states[i]) + end + y = policy.combiner(vcat(h, prev)) + @test Float32.(y) == outs[t] + prev = Float32.(y) + end +end + +@testset "Zygote gradient through threaded 3-stage rollout" begin + Random.seed!(43) + policy = StateConditionedPolicy(2, 1, 1, [4]; activation = tanh) + ws = [Float32[0.3, -0.2], Float32[0.1, 0.4], Float32[-0.5, 0.2]] + x0 = Float32[0.5] + + # Same pattern as the training loops: reset inside the gradient block, + # thread the recurrent state through all stages via the policy forward. + gs = Zygote.gradient(policy) do m + Flux.reset!(m) + total = 0.0f0 + prev = x0 + for t in 1:3 + y = m(vcat(ws[t], prev)) + total = total + sum(y) + prev = y + end + total + end + g = materialize_tangent(gs[1]) + @test g !== nothing + @test DecisionRulesExa._all_finite_gradient(g) + + # Encoder gradients must be finite AND nonzero (the LSTM parameters + # participate in every stage of the threaded rollout). + enc_leaves = Float64[] + function _collect_leaves(x) + if x isa AbstractArray && eltype(x) <: Number + append!(enc_leaves, vec(Float64.(x))) + elseif x isa NamedTuple + foreach(_collect_leaves, values(x)) + elseif x isa Tuple + foreach(_collect_leaves, x) + end + return nothing + end + _collect_leaves(g.encoder) + @test !isempty(enc_leaves) + @test any(!=(0.0), enc_leaves) +end + +# ── Synthetic 2-bus fixture for the hydro AC builder ────────────────────────── +# One thermal generator at the reference bus, one hydro generator at the load +# bus, a single branch, and one reservoir. Small enough for fast MadNLP solves. +function _tiny_ac_case() + buses = [ + PowerBusData(1, 3, 0.0, 0.0, 0.9, 1.1), + PowerBusData(2, 1, 0.0, 0.0, 0.9, 1.1), + ] + gens = [ + PowerGenData(1, 1, 0.0, 2.0, -1.0, 1.0, 10.0, 0.0), # thermal at bus 1 + PowerGenData(2, 2, 0.0, 1.0, -1.0, 1.0, 0.0, 0.0), # hydro at bus 2 + ] + b_dc = -0.1 / (0.01^2 + 0.1^2) + branches = [ + PowerBranchData(1, 1, 2, b_dc, 2.0, -0.5, 0.5, + 0.01, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0), + ] + loads = [PowerLoadData(1, 2)] + power_data = PowerData( + 2, 2, 1, 1, buses, gens, branches, loads, + [1], 100.0, 60.0, [0.0, 0.5], [0.0, 0.1], + ) + units = [HydroUnitData(1, 2, 100.0, 0.0, 10.0, 0.0, 10.0, 0.001, 0.0, 0.0)] + hydro_data = HydroData( + 1, units, UpstreamTurn[], UpstreamSpill[], + 2.6, [50.0], [ones(2, 1)], 1, 2, + ) + return power_data, hydro_data +end + +@testset "AC reactive-deficit modes (nothing / finite / Inf)" begin + power_data, hydro_data = _tiny_ac_case() + T = 2 + nBus = power_data.nBus + + _build(cost; strict = false) = build_hydro_de(power_data, hydro_data, T; + formulation = :ac_polar, + deficit_cost = 1e4, + strict_targets = strict, + reactive_deficit_cost = cost, + ) + + prob_free = _build(nothing) + prob_pen = _build(10.0) + prob_hard = _build(Inf) + + @test prob_free.reactive_deficit_mode === :free + @test prob_pen.reactive_deficit_mode === :penalized + @test prob_hard.reactive_deficit_mode === :hard + + # Variable counts: penalized splits the slack (+T*nBus vs free); hard + # removes it (−T*nBus vs free). Constraint counts are identical. + @test prob_pen.model.meta.nvar == prob_free.model.meta.nvar + T * nBus + @test prob_hard.model.meta.nvar == prob_free.model.meta.nvar - T * nBus + @test prob_pen.model.meta.ncon == prob_free.model.meta.ncon + @test prob_hard.model.meta.ncon == prob_free.model.meta.ncon + + # target_con_range bookkeeping is unaffected in both non-strict and strict. + @test prob_pen.target_con_range == prob_free.target_con_range + @test prob_hard.target_con_range == prob_free.target_con_range + strict_free = _build(nothing; strict = true) + strict_pen = _build(10.0; strict = true) + strict_hard = _build(Inf; strict = true) + @test strict_pen.target_con_range == strict_free.target_con_range + @test strict_hard.target_con_range == strict_free.target_con_range + + # DC path rejects the kwarg (and is otherwise unaffected by it). + @test_throws ErrorException build_hydro_de(power_data, hydro_data, T; + formulation = :dc, reactive_deficit_cost = 10.0) + + # All three AC modes must solve. + x0 = [50.0] + w = [2.0, 2.0] + xhat = [50.0, 50.0] + for prob in (prob_free, prob_pen, prob_hard) + ExaModels.set_parameter!(prob.core, prob.p_x0, x0) + ExaModels.set_parameter!(prob.core, prob.p_inflow, w) + ExaModels.set_parameter!(prob.core, prob.p_target, xhat) + prepare_solve!(prob, x0, w, xhat) + res = MadNLP.madnlp(prob.model; tol = 1e-6, max_iter = 500, + print_level = MadNLP.ERROR) + @test solve_succeeded(res) + @test isfinite(res.objective) + sol = hydro_solution(prob, res) + @test size(sol.deficit_q) == (nBus, T) + @test all(isfinite, sol.deficit_q) + if prob.reactive_deficit_mode === :hard + @test all(iszero, sol.deficit_q) + end + end +end + +@testset "HydroReachablePolicy threading matches manual cell recursion" begin + power_data, hydro_data = _tiny_ac_case() + Random.seed!(44) + policy = hydro_reachable_policy(hydro_data, [4, 3]) + + # Memory across stages: identical inputs without reset differ; reset + # restores the exact first output. + xin = vcat(Float32[2.0], Float32[50.0]) + Flux.reset!(policy) + a = policy(xin) + b = policy(xin) + @test a != b + Flux.reset!(policy) + @test policy(xin) == a + + # As-is open loop vs manual LSTMCell threading (the verified diagnostic + # pattern from eval_paired_exa_strict.jl) — must agree EXACTLY. + T = 3 + w_stages = Float32[2.0, 1.5, 2.5] + Flux.reset!(policy) + prev = Float32[50.0] + asis = Vector{Vector{Float32}}() + for t in 1:T + y = policy(vcat(Float32[w_stages[t]], prev)) + push!(asis, Float32.(y)) + prev = Float32.(y) + end + + cells = [l.cell for l in policy.encoder.layers] + states = Any[Flux.initialstates(c) for c in cells] + prev = Float32[50.0] + for t in 1:T + wt = Float32[w_stages[t]] + h = wt + for (i, c) in enumerate(cells) + h, states[i] = c(h, states[i]) + end + y = policy.combiner(vcat(h, prev)) + lower, upper = _hydro_reachable_bounds(policy, wt, prev, y) + raw = lower .+ (upper .- lower) .* y + target = isempty(policy.cascade) ? raw : + min.(raw, _cascade_upper_bounds(policy, raw, wt, prev)) + @test Float32.(target) == asis[t] + prev = Float32.(target) + end +end