DecisionRulesExa.jl

GPU-accelerated training of parametric decision rules through multi-stage optimization, using ExaModels.jl and solvers like MadNLP.jl.

This package replicates the core functionality of DecisionRules.jl but replaces JuMP with ExaModels for the optimization backend. JuMP relies on MathOptInterface operations that are not differentiable on GPU and cannot exploit GPU-native sparse linear solvers. ExaModels provides a fully SIMD-compatible algebraic modeling layer whose parameters can be updated at runtime via ExaModels.set_parameter!, enabling efficient repeated solves on GPU with MadNLPGPU (CUDSS-backed interior point).

Motivation

In the Two-Stage Deep Decision Rule (TS-DDR) framework, a neural-network policy predicts target state trajectories that are projected onto the feasible set by solving a parametric NLP (the deterministic equivalent). Training uses the envelope theorem: dual multipliers on the target constraints give the policy gradient without differentiating through the solver.

The inner NLP solve is the bottleneck. By formulating it in ExaModels and solving with MadNLP + CUDSS on GPU, DecisionRulesExa.jl achieves significant speedups over the CPU JuMP-based workflow in DecisionRules.jl, especially for large-scale problems (AC-OPF, multi-stage hydro scheduling).

Installation

using Pkg
Pkg.add(url="https://github.com/LearningToOptimize/DecisionRulesExa.jl.git")

Quick start

using DecisionRulesExa
using ExaModels, MadNLP, Flux, Random

Random.seed!(1)

T  = 8   # horizon
nx = 1   # state dimension

# Build a simple linear-tracking NLP on CPU
prob = build_linear_tracking_problem(
    horizon       = T,
    nx            = nx,
    backend       = nothing,       # CPU (use CUDABackend() for GPU)
    slack_penalty = 10.0,
    u_bounds      = (-2.0, 2.0),
)

# LSTM policy: input = [w_t ; x_{t-1}], called once per stage
policy = StateConditionedPolicy(nx, nx, nx, [64, 64])

# Uncertainty sampler (returns flat vector of length T*nw)
sampler() = Float32.(0.1 .* randn(T * nx))

# Train with TS-DDR policy gradient
train_tsddr(
    policy,
    Float32.([1.0]),        # initial state
    prob,
    prob.p_x0,
    prob.p_target,
    prob.p_w,
    sampler;
    num_batches         = 100,
    num_train_per_batch = 4,
    optimizer           = Flux.Adam(1f-3),
    madnlp_kwargs       = (print_level = MadNLP.ERROR, tol = 1e-6),
)

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

DecisionRulesExa.jl is model-first: you describe your multi-stage NLP in ExaModels, then the package handles simulation and training.

For a custom problem you need:

• An ExaModels deterministic-equivalent NLP with parametric initial state, uncertainty trajectory, and target trajectory. Target constraints must be added last so their multipliers form a contiguous slice of result.multipliers.
• An uncertainty sampler () -> w_flat returning a flat Float32/Float64 vector of length T * nw.
• A Flux policy (LSTM or MLP) mapping (w_t, x_{t-1}) to target x_t at each stage.

The package provides build_deterministic_equivalent for generic problems and build_linear_tracking_problem as a ready-made demo. For domain-specific models, build the ExaModels NLP directly; examples/BatteryStorageOPF/ contains an AC-OPF battery example.

Strict reachable target equality

The usual TS-DDR deterministic equivalent uses slack-penalized target constraints,

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:

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.

The package supports two ways to make strict mode safe:

• Embedded strict DE (train_tsddr_embedded, build_embedded_deterministic_equivalent) evaluates the policy inside the NLP against realized states, so the policy always 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.

The published hydro case study uses the second path; see examples/HydroPowerModels.

Do not use strict equality for a generic open-loop target policy. For unreachable targets, the slack-penalty formulation is the robust fallback.

The reachable map depends on the incoming state, and that dependence must be differentiated. Declaring the interval endpoints non-differentiable still produces a gradient and still lowers the loss, while descending a different direction — measured on the hydro case, 6% of the true magnitude and 48 degrees off. Check the complete actor gradient against finite differences before drawing hyperparameter conclusions from it.

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:

pool = [(prob, prob.p_x0, prob.p_target, prob.p_w)]
for _ in 2:num_workers
    p = build_my_problem(...)
    push!(pool, (p, p.p_x0, p.p_target, p.p_w))
end

train_tsddr(policy, x0, prob, ..., sampler;
    problem_pool = pool,
    num_train_per_batch = num_workers,
)

Each pool entry gets its own MadNLP solver on a dedicated thread, with CUDA handles properly bound.

Optional Critic Control Variates

train_tsddr can optionally train a scalar critic C(w, xhat). The critic does not replace the deterministic-equivalent solve in the default :control_variate mode: solved target-constraint multipliers remain the actor's primary local sensitivity signal. Instead, the critic supplies a learned rollout-value guide and optional control variate.

For critic fitting, the preferred target is the stage-wise rollout objective via RolloutCriticTarget, with policy_state = :target by default. This matches the differentiable target recurrence used by the actor while evaluating the true stage-by-stage objective. Set policy_state = :realized to train the critic on closed-loop realized-state rollout labels. For ablations, use DeterministicEquivalentCriticTarget() or critic_training_target = :deterministic_equivalent to fit the deterministic-equivalent objective instead.

The default control_variate = NoCriticControlVariate() recovers the original dual-only behavior. A scalar critic can be attached with:

input_dim = length(x0) + 2 * T * nx
critic = Chain(
    Dense(input_dim => 128, tanh),
    Dense(128 => 128, tanh),
    Dense(128 => 1),
)

cv = ScalarCriticControlVariate(
    critic;
    featurizer = default_critic_featurizer,
    value_loss_weight = 1.0,
    gradient_loss_weight = 0.0,
)

critic_target = RolloutCriticTarget(
    stage_problem;
    horizon = T,
    n_uncertainty = nx,
    set_stage_parameters! = set_stage_parameters!,
    realized_state = realized_state,
    objective_no_target_penalty = objective_no_target_penalty,
    policy_state = :target,
    objective_value = :objective_no_target_penalty,
)

train_tsddr(
    policy, x0, prob, prob.p_x0, prob.p_target, prob.p_w, sampler;
    control_variate = cv,
    critic_training_target = critic_target,
    critic_rollout_samples_per_batch = 1,
    actor_gradient_mode = :control_variate,
    critic_cv_weight = 1.0,
    num_cheap_critic_samples_per_batch = 32,
    critic_updates_per_batch = 1,
    critic_optimizer = Flux.Adam(1f-3),
)

The critic loss combines value matching against the selected target objective and optional gradient matching against target multipliers:

value_loss_weight * mse(C(w, xhat), objective) +
gradient_loss_weight * mse(gradient(xhat -> C(w, xhat), xhat), lambda)

Set either weight to zero for objective-only or gradient-only critic training. For rollout targets, objective-only critic training is usually the clean default, because DE target multipliers are not exact gradients of the realized rollout objective. If objectives and multipliers have very different magnitudes, prefer a custom featurizer and tuned loss weights; the Hydro example normalizes volumes and inflows before critic evaluation.

Actor modes:

• :control_variate: subtracts critic_cv_weight * gradient_xhat(C) from the solved dual signal and adds the critic actor gradient back on solved or cheap rollout samples. This is the recommended mode when dual multipliers are reliable. critic_cv_weight = 0.0 recovers dual-only updates.
• :surrogate: uses a practical hybrid of solved dual gradients and critic actor gradients, controlled by dual_actor_weight and critic_actor_weight. This is useful when raw dual/subgradient signals are empirically noisy or unstable, but it is no longer a pure unbiased control-variate estimator.

num_cheap_critic_samples_per_batch draws additional uncertainty samples, rolls out the current policy, and evaluates the critic actor term without any extra MadNLP or ExaModels solve.

Rollout evaluation

RolloutEvaluation evaluates the policy in deployment semantics (stage-by-stage sequential solves) on held-out scenarios:

eval = RolloutEvaluation(
    stage_problem, x0, eval_scenarios;
    horizon              = T,
    n_uncertainty        = nw,
    set_stage_parameters! = my_stage_setter!,
    realized_state       = my_realized_state,
    stride               = 25,
    policy_state         = :realized,
)

Supports parallel evaluation across a pool of stage problems via stage_problem_pool, and dynamic scenario count via active_scenarios.

Relationship to DecisionRules.jl

DecisionRules.jl implements the same TS-DDR algorithm using JuMP + DiffOpt for CPU-based training. It supports stage-wise decomposition, multiple shooting, and integer strategies that are not yet ported here.

DecisionRulesExa.jl focuses on the deterministic-equivalent training mode with GPU acceleration. Choose this package when:

• Your NLP is large enough that GPU acceleration matters (e.g., AC-OPF with hundreds of buses)
• You need to run many training samples per gradient step
• You want to leverage CUDA-native sparse solvers (CUDSS)

Choose DecisionRules.jl when:

• You need stage-wise or multiple-shooting decomposition
• Your problem is naturally expressed in JuMP
• You need DiffOpt-based sensitivity computation
• You want CPU-only deployment

Examples

examples/end_to_end_cpu.jl and examples/end_to_end_gpu.jl train a policy on a small linear tracking problem on the CPU and on a GPU. The two case studies of the paper, examples/HydroPowerModels/ and examples/BatteryStorageOPF/, hold the GPU trainers and the training schedules; the cases, the SDDP baselines and the JuMP engines are in the same directories of DecisionRules.jl. To reproduce a case study, start from the examples README of DecisionRules.jl. How the code of both packages is organized is described in Reading the code.

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/	Bolivian hydrothermal case study: GPU trainer, lineage driver, training schedules
examples/BatteryStorageOPF/	Battery-storage AC-OPF case study: GPU trainer
test/runtests.jl	Unit and smoke tests

Citation

If you use this package in academic work, please cite:

@article{rosemberg2024efficiently,
  title={Efficiently Training Deep-Learning Parametric Policies using Lagrangian Duality},
  author={Rosemberg, Andrew and Street, Alexandre and Vallad{\~a}o, Davi M and Van Hentenryck, Pascal},
  journal={arXiv preprint arXiv:2405.14973},
  year={2024}
}

License

MIT. See LICENSE.