DecisionRules.jl

DecisionRules.jl is a Julia package for training parametric decision rules for multistage stochastic programs in which every stage is an optimization model, such as an optimal power flow or a model-predictive control problem. It implements two-stage generalized decision rules (TS-GDR) and their deep instantiation, two-stage deep decision rules (TS-DDR), introduced in arXiv:2405.14973. A policy, from an affine rule to a deep recurrent network, maps the observed uncertainty to a target for the state at the end of every stage; an optimization problem enforces the targets against the stage models; and the multipliers of the target constraints provide the gradient with which the policy is trained. The stage models are written in JuMP, and the companion package DecisionRulesExa.jl trains the same policies on GPUs.

Workflow

In TS-GDR the policy does not output a control action. It outputs targets, typically for the state at the end of every stage, and an optimization model enforces them through target constraints: with a penalized slack in the general formulation, or as hard equalities in the strict formulation when the policy can be restricted to targets reachable in one stage. For a sampled uncertainty trajectory $w_{1:T}$, a training step is:

1. sample $w_{1:T}$
2. predict targets $\hat x_{1:T} = \pi(\cdot;\theta)$
3. solve an optimization problem that projects targets onto the feasible set (dynamics + constraints)
4. differentiate to update $\theta$ using dual information and/or implicit sensitivities (via DiffOpt)

DecisionRules.jl implements this workflow in three flavors:

• Deterministic equivalent (direct transcription): one coupled optimization over the full horizon.
• Stage-wise decomposition (single shooting): one optimization per stage in a sequential rollout.
• Windowed decomposition (multiple shooting): one coupled optimization per window, chained by the realized end-state.

Installation

using Pkg
Pkg.add("DecisionRules")

What you need to provide

DecisionRules.jl is intentionally “model-first”: you describe your problem in JuMP (DiffOpt-enabled for all but the deterministic equivalent approach), then the package handles simulation and training.

For any multi-stage model you will need:

• subproblems::Vector{JuMP.Model}: one JuMP model per stage (DiffOpt-enabled).
• state_params_in[t]: a vector of parameter variables for the incoming state at stage t.
• state_params_out[t]: a vector of (target_param, realized_state_var) tuples at stage t. The target_param is the parameter variable that the policy sets; the realized_state_var is the JuMP decision variable whose value becomes the realized state.
• an uncertainty_sampler that returns per-stage samples in the format used by DecisionRules.sample(...).
• Differentiable policies built with Flux.jl or similar compatible libraries. Input size is the number of uncertainty components per stage plus the size of the initial state; output size is the size of the target state at each stage.

Working patterns are provided in examples/.

Deterministic equivalent (TS-DDR / direct transcription)

This corresponds to solving one coupled optimization over all stages (the deterministic equivalent of a sampled trajectory). You build the deterministic-equivalent JuMP model with DecisionRules.deterministic_equivalent! and then train with the deterministic-equivalent overload of train_multistage.

using DecisionRules, JuMP, DiffOpt, Flux
using SCS

# 1) Build per-stage subproblems (DiffOpt-enabled) and collect:
#    subproblems, state_params_in, state_params_out, uncertainty_samples

# 2) Build the deterministic equivalent over the full horizon
det = DiffOpt.diff_model(() -> DiffOpt.diff_optimizer(SCS.Optimizer))

det, uncertainty_samples_det = DecisionRules.deterministic_equivalent!(
    det,
    subproblems,
    state_params_in,
    state_params_out,
    Float64.(initial_state),
    uncertainty_samples,
)

# deterministic_equivalent! remaps state_params_in/state_params_out in place.
# Copy those arrays first if you also need the original stage-wise refs later.

# 3) Train a TS-DDR policy end-to-end
num_uncertainties = length(uncertainty_samples[1])  # number of uncertainty components per stage
policy = Chain(
    Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu),
    Dense(64, length(initial_state)),
)

DecisionRules.train_multistage(
    policy,
    initial_state,
    det,
    state_params_in,
    state_params_out,
    uncertainty_samples_det;
    num_batches=100,
    num_train_per_batch=32,
    optimizer=Flux.Adam(1e-3),
)

This mode typically gives the most faithful gradient signal (full coupling across the horizon), but it requires solving the largest inner problem per sample.

Stage-wise decomposition (single shooting)

Single shooting solves one optimization per stage and rolls forward using the realized state returned by the solver. The policy can be closed-loop because it receives the realized state $x_{t-1}$ when predicting the next target $\hat x_t$.

using DecisionRules, Flux

num_uncertainties = length(uncertainty_samples[1])
policy = Chain(
    Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu),
    Dense(64, length(initial_state)),
)

DecisionRules.train_multistage(
    policy,
    initial_state,
    subproblems,
    state_params_in,
    state_params_out,
    uncertainty_samples;
    num_batches=100,
    num_train_per_batch=32,
    optimizer=Flux.Adam(1e-3),
)

Internally, gradients are obtained by combining (i) dual information for target parameters and (ii) solution sensitivities computed through DiffOpt along the rollout.

Windowed decomposition (multiple shooting)

Multiple shooting partitions the horizon into windows of length window_size. Each window solves a deterministic equivalent over its stages, then passes the realized end state to the next window. This can stabilize learning over long horizons compared to pure single shooting, while remaining cheaper than a full-horizon deterministic equivalent.

using DecisionRules, Flux, DiffOpt
using SCS

num_uncertainties = length(uncertainty_samples[1])
policy = Chain(
    Dense(DecisionRules.policy_input_dim(num_uncertainties, length(initial_state)), 64, relu),
    Dense(64, length(initial_state)),
)

windows = DecisionRules.setup_shooting_windows(
    subproblems,
    state_params_in,
    state_params_out,
    Float64.(initial_state),
    uncertainty_samples;
    window_size=24,
    model_factory=() -> DiffOpt.nonlinear_diff_model(optimizer_with_attributes(
        SCS.Optimizer,
        "verbose" => 0,
    )),
)

DecisionRules.train_multiple_shooting(
    policy,
    initial_state,
    windows,
    uncertainty_samples;
    num_batches=100,
    num_train_per_batch=32,
    optimizer=Flux.Adam(1e-3),
)

Evaluation: stage-wise rollout and target-violation share

Evaluating a trained policy only through the deterministic equivalent can overstate its quality: the coupled solve re-optimizes all stages jointly and absorbs targets that are not followable stage by stage through the slack penalty — exactly what deployment cannot do. The stage-wise rollout is the deployment semantics of a target-trajectory policy, so report it as the headline metric, together with a target-violation measure.

The training loops record metrics through a per-sample SampleLog cache and a per-batch record(sample_log, iter, model) callback. RolloutEvaluation is a ready-made helper that evaluates the policy stage-wise on a fixed held-out scenario set; call it from within record:

using DecisionRules, Random

# Materialize a FIXED held-out evaluation set once, before training.
# Use stage-wise subproblems and parameter refs, not DE-remapped refs.
Random.seed!(1234)
eval_scenarios = [DecisionRules.sample(uncertainty_samples) for _ in 1:8]

rollout_eval = RolloutEvaluation(
    subproblems, state_params_in, state_params_out, initial_state, eval_scenarios;
    stride=25,  # evaluate every 25 batches
    policy_state=:realized,
)

train_multistage(policy, initial_state, det, state_params_in, state_params_out, uncertainty_samples_det;
    num_batches=100,
    record=(sample_log, iter, model) -> begin
        rollout_eval(iter, model)
        return false
    end)

policy_state selects which state is fed back to the policy between stage solves:

• :realized feeds the previous realized optimizer state to the policy. This is the closed-loop/deployment rollout and is the default.
• :target feeds the previous target/predicted state to the policy, matching the deterministic-equivalent target-generation path while still solving the stage subproblems sequentially.

Use policy_state=:target when comparing against a deterministic-equivalent training loss, because both metrics then call the policy on the same target-state history. Log :realized separately as the harder deployment diagnostic.

Each evaluation reports (a) the rollout objective excluding the target-slack penalty term (the operational cost) and (b) the target-violation share — the realized slack penalty divided by the full objective. Policy comparisons are only trustworthy when the violation share is small (≤ ~0.05): a larger share means the policy's targets are not followable stage by stage and the reported cost is not what deployment would realize. When training drives the violation share to ~0, the deterministic-equivalent and rollout views are expected to coincide; the rollout metric is the guard that detects when they don't.

Per-sample debugging hooks can be attached with SampleLog(on_sample=(s, models, log) -> ...); the training loop calls the hook after each sample's solve with the live JuMP model(s). The previous record_loss=(iter, model, loss, tag) -> ... keyword keeps working as a deprecated adapter.

Strict mode and reachable policies

The standard TS-DDR target constraint uses slack:

$$x_t + \delta_t = \hat{x}_t, \qquad \text{objective} += C_\delta \|\delta_t\|.$$

Slack makes training robust to unreachable targets, but it also makes the dual signal depend on the target-penalty calibration. Strict mode removes the slack:

$$x_t = \hat{x}_t.$$

The resulting dual is the clean shadow price of imposing the target. The price of that cleaner signal is feasibility: every policy target must be reachable from the state used to condition the policy.

This is automatic in the hydro strict subproblem path because each stage is solved sequentially and the policy receives the realized previous reservoir state. It is also possible in regular deterministic equivalents when the target trajectory is rolled out from the true initial state using a reachable policy:

$$\hat{x}_0 = x_0,\qquad \hat{x}_t = \pi_\theta(w_t, \hat{x}_{t-1}),\qquad \hat{x}_t \in R(\hat{x}_{t-1}, w_t).$$

By induction, all targets are feasible, and the strict equalities force the realized trajectory to match that reachable path. See examples/HydroPowerModels and the DecisionRulesExa.jl companion for the GPU strict regular-DE implementation.

The reachable map R(\hat x_{t-1}, w_t) depends on the state, and that dependence must be differentiated. Treating the interval endpoints as constants still trains and still lowers the loss while descending a materially different direction — on the hydro case, a gradient carrying 6% of the true magnitude and pointing 48 degrees away from it. Verify the complete actor gradient against finite differences before trusting any hyperparameter conclusion drawn on top of it.

The policy helpers separate two architectural choices:

• recurrent layers / DR_ENCODER_LAYERS process uncertainty history only;
• combiner_layers / DR_HEAD_LAYERS add a nonlinear feed-forward state-to-target head without recurrence over the state input.

GPU acceleration with DecisionRulesExa.jl

For large-scale problems where the inner NLP solve is the bottleneck (e.g., AC-OPF with hundreds of buses), DecisionRulesExa.jl provides a GPU-accelerated backend that replaces JuMP with ExaModels.jl and solves with MadNLP.jl + CUDSS on GPU.

DecisionRulesExa.jl implements the same TS-DDR algorithm (deterministic-equivalent mode) with the same envelope-theorem gradient computation but formulates the NLP in ExaModels' SIMD-compatible modeling layer. This enables:

• GPU-native interior-point solves via MadNLP + CUDSS
• Parallel GPU solves for multiple training samples per gradient step
• Runtime parameter updates via ExaModels.set_parameter! (no model reconstruction)

See the GPU Acceleration page in the documentation for a tutorial on getting started with DecisionRulesExa.jl.

Examples and tests

The case studies of the paper and smaller demonstrations live in examples/. To reproduce a case study, start from examples/README.md, which says where each case lives, which engine trains which method, and what to expect from a run. The documentation explains the method, the case studies and how the code is organized. Run the tests with:

julia --project -e 'using Pkg; Pkg.test()'

Repository Map

Path	Purpose
src/DecisionRules.jl	Module entrypoint and exports
src/dense_multilayer_nn.jl	MLP helpers, state-conditioned recurrent policies, nonlinear target heads
src/simulate_multistage.jl	Stage-wise and deterministic-equivalent simulation logic
src/multiple_shooting.jl	Windowed multiple-shooting setup, simulation, and training
src/utils.jl	Target-parameter utilities, deficit construction, rollout evaluation
src/integer_strategies.jl	Strategies for extracting gradients from integer/mixed-integer models
src/score_function.jl	Score-function gradient correction for nonsmooth/integer problems
src/parameter_duals.jl	Dual/sensitivity helpers for parameterized JuMP models
docs/src/	Documenter.jl manual pages
examples/HydroPowerModels/	Bolivian hydrothermal case study: case, JuMP engine, SDDP and RL baselines
examples/BatteryStorageOPF/	Battery-storage AC-OPF case study: case construction, JuMP engine, SDDP
examples/inventory_control/	Inventory-control example and dynamic-programming/SDDP comparisons
examples/rocket_control/	Rocket MPC/control example
examples/Experimental/	Research prototypes and robotics/control explorations
test/runtests.jl	Package test suite

Citation

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

@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}
}

DiffOpt (for differentiating through optimization): https://github.com/jump-dev/DiffOpt.jl

License

MIT. See LICENSE.