API Reference

AbstractGradientFallback

Abstract type governing what happens when a solver or differentiation error occurs during training.

DecisionRules ships two concrete subtypes:

Extending

Implement your own subtype to customize recovery:

struct MyFallback <: DecisionRules.AbstractGradientFallback end

function DecisionRules.handle_gradient_error(::MyFallback, e, n_state_in, n_state_out)
    # e is the caught exception
    # Return a tuple of cotangents (same shape as the rrule pullback) or rethrow
    @error "Custom handler" exception=e
    return DecisionRules._zero_cotangents(n_state_in, n_state_out)
end

function DecisionRules.handle_training_error(::MyFallback, e, iter)
    # Return true to skip this iteration, false to rethrow
    @error "Custom training handler" exception=e
    return true  # skip
end

function DecisionRules.handle_rollout_error(::MyFallback, e, iter)
    # Return true to skip this scenario, false to rethrow
    return true
end

Then pass gradient_fallback=MyFallback() to train_multistage or train_multiple_shooting.

AbstractIntegerStrategy

Abstract supertype for strategies that prepare a JuMP model before reading duals or solver sensitivities.

Arguments

This abstract type has no fields. Concrete subtypes are passed as the integer_strategy::AbstractIntegerStrategy keyword to simulation and training functions.

Examples

simulate_multistage(
    subproblems,
    state_params_in,
    state_params_out,
    initial_state,
    uncertainties,
    policy;
    integer_strategy = FixedDiscreteIntegerStrategy(),
)

ContinuousRelaxationIntegerStrategy()

Relax all binary/integer constraints to continuous bounds (binary → [0,1]), solve the resulting LP, and read duals in that relaxed state.

Mathematically, this replaces $z \in \{0,1\}$ or integer restrictions with continuous bounds before solving. The derivative signal belongs to the relaxed problem, not to the original MIP.

Compared to FixedDiscreteIntegerStrategy:

• Faster: one LP solve instead of MIP + LP.
• Smoother gradients: no integer fixing means no zero-gradient dead zones.
• Less accurate: the LP solution may have fractional integer variables, so the gradient does not correspond to any feasible integer assignment.

A practical pattern is to train with ContinuousRelaxationIntegerStrategy during warmup (smooth landscape for initial learning) and switch to FixedDiscreteIntegerStrategy later (integer-accurate gradients for fine-tuning).

Arguments

This type has no fields.

Examples

strategy = ContinuousRelaxationIntegerStrategy()

ErrorGradientFallback()

Strict fallback: re-throw any solver or differentiation error. Use this in tests to ensure that controlled problems never silently produce zero gradients.

FixedDiscreteIntegerStrategy()

Solve a mixed-integer model, fix discrete variables to their incumbent values, relax integrality, re-solve, and read duals or sensitivities from the fixed continuous model.

If $z^*$ is the incumbent binary/integer solution, this strategy reads derivative-like information from the continuous problem

\[\min_x f(x, z^*) \quad \text{subject to} \quad g(x, z^*) \le 0.\]

The result is local to the incumbent integer assignment. It is not a differentiable MIP method.

Arguments

This type has no fields.

Examples

strategy = FixedDiscreteIntegerStrategy()

NoIntegerStrategy()

Solve the model exactly as written before reading duals or sensitivities.

Use this for continuous LP, conic, or nonlinear models whose derivative information is available directly from the solved model.

Arguments

This type has no fields.

Examples

strategy = NoIntegerStrategy()

RolloutEvaluation(subproblems, state_params_in, state_params_out, initial_state,
                  scenarios; stride=1, policy_state=:realized)

Evaluation helper that assesses the policy with a stage-wise rollout (the deployment semantics of a target-trajectory policy) on a fixed held-out scenario set. Deterministic-equivalent evaluation re-optimizes all stages jointly and can absorb stage-wise-unfollowable targets through the slack penalty, silently overstating policy quality; the rollout metric is the guard that detects this.

policy_state controls what state is passed back into the policy:

• :realized pipes the previous realized state into the policy. This is the deployment/closed-loop rollout semantics.
• :target pipes the previous target state into the policy, matching the deterministic-equivalent target-generation semantics from simulate_states while still solving the stage subproblems sequentially.

scenarios must be a vector of materialized scenarios, sampled once before training (e.g. [DecisionRules.sample(uncertainty_samples) for _ in 1:n]), so every evaluation uses the same fixed set. subproblems may be the training subproblems (all stage parameters are rewritten on every solve) or a separately built copy; when training on a deterministic equivalent, pass the stage-wise subproblems here.

Call evaluation(iter, model), e.g. from within a record callback. Every stride calls it rolls the policy out over the fixed set and reports:

• metrics/rollout_objective_no_deficit: the rollout objective excluding the target-slack penalty term (the operational cost), and
• metrics/rollout_target_violation_share: the realized slack penalty divided by the full rollout objective (NaN when undefined).

Policy comparisons should only be trusted 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. The latest values are kept in last_objective_no_deficit / last_violation_share for custom logging. Calls on batches that are not a multiple of stride are a no-op and leave the cached values unchanged.

SampleLog(; on_sample=(s, models, sample_log) -> nothing,
          objective_no_deficit_fn=get_objective_no_target_deficit)

Per-sample logger with local cache state for the training loops, patterned after SaveBest. During each batch the training loop calls sample_log(s, det_equivalent_or_subproblems) right after sample s has been simulated (successful solves only; a failed solve throws exactly as before). The default behavior caches, per sample, the full objective (objective_value) and the objective excluding the target-slack penalty term (objective_no_deficit_fn). The cache is cleared at the start of every batch and handed to the per-batch record(sample_log, iter, model) callback.

on_sample is an optional hook called as on_sample(s, models, sample_log) after the default caching. It receives the live JuMP model(s), so it can inspect termination statuses or dump the details of a suspicious sample for debugging — without paying any per-sample logging cost in the default configuration.

SaveBest(best_loss::Float64, model_path::String)

Callback that saves the best policy state seen during training.

SaveBest is a small callable object used as a training callback. When called as callback(iter, model, loss), it compares loss with the best loss stored so far. If the new loss is smaller, it copies model to CPU, normalizes any recurrent layer state, and writes the Flux state to model_path with JLD2. It returns false, so it records checkpoints without stopping training.

Arguments

• best_loss::Float64: incumbent loss. Use Inf to save the first observed model.
• model_path::String: path of the JLD2 file that receives the best model state.

Examples

callback = SaveBest(Inf, "best_policy.jld2")
train_multistage(policy, x0, subproblems, state_in, state_out, sampler;
    record = (log, iter, model) -> callback(iter, model, mean(log.losses)))

ScoreFunctionConfig(
    subproblems::AbstractVector{<:JuMP.Model},
    state_params_in::AbstractVector,
    state_params_out::AbstractVector;
    dual_weight::Real = 0.5,
    perturbation_std::Real = 1.0,
    num_rollouts::Integer = 8,
    baseline::Symbol = :mean,
)

Configure the score-function correction used by train_multistage.

The deterministic-equivalent training path differentiates the target policy through dual information. For mixed-integer subproblems, those duals are local to a fixed integer assignment. This configuration adds a REINFORCE-style correction estimated from stage-wise rollouts with perturbed targets.

The rollout models are solved exactly as they are built. If subproblems contain binary variables, the score-function rollouts solve MIPs. If they contain relaxed variables, the score-function rollouts solve the relaxation. This is intentionally separate from the integer_strategy keyword of train_multistage, which controls only how the differentiable dual-gradient path reads local sensitivity information from the deterministic equivalent.

If $\hat{x}_{t+1}(\theta)$ is the target emitted by the policy and $\delta_t \sim \mathcal{N}(0, \sigma^2 I)$, the perturbed rollout solves with target $\hat{x}_{t+1}(\theta) + \delta_t$. The score-function surrogate loss is

\[L_{sf}(\theta) = \frac{1}{M} \sum_{m=1}^{M} (R_m - b) \sum_{t=1}^{T} \left\langle \frac{\delta_{m,t}}{\sigma^2}, \hat{x}_{t+1}(\theta) \right\rangle ,\]

and the mixed gradient is

\[\nabla L = \alpha \nabla L_{dual} + (1 - \alpha) \nabla L_{sf}.\]

Arguments

• subproblems::AbstractVector{<:JuMP.Model}: stage-wise rollout models used to estimate realized costs under perturbed targets.
• state_params_in::AbstractVector: stage input-state parameters.
• state_params_out::AbstractVector: pairs (target_parameter, state_variable) for every stage output state.

Keywords

• dual_weight::Real: mixing weight $\alpha$ on the dual-gradient term.
• perturbation_std::Real: Gaussian standard deviation $\sigma$.
• num_rollouts::Integer: number of perturbed rollouts $M$ per sample.
• baseline::Symbol: either :mean for mean-centering costs or :none.

Examples

score_function = ScoreFunctionConfig(
    rollout_subproblems,
    state_params_in,
    state_params_out;
    dual_weight = 0.5,
    perturbation_std = 1.0,
    num_rollouts = 8,
)

train_multistage(
    policy,
    initial_state,
    det_equivalent,
    state_params_in,
    state_params_out,
    uncertainty_sampler;
    score_function,
)

ScoreFunctionSchedule(config::ScoreFunctionConfig; <keyword arguments>)

Ramp a ScoreFunctionConfig into training after a pure-dual warmup.

The schedule delays score-function rollouts until sf_start, then linearly increases the score-function weight, perturbation scale, and rollout count until the final values stored in config are reached.

Let $k$ be the current iteration and $\rho_k = \operatorname{clip}((k - k_0) / r, 0, 1)$. The effective score-function weight is $\rho_k (1 - \alpha)$. The effective dual weight is one minus that value.

Arguments

• config::ScoreFunctionConfig: final score-function configuration.

Keywords

• sf_start::Integer: first iteration at which score-function rollouts are considered.
• ramp_batches::Integer: number of iterations in the linear ramp.
• perturbation_std_initial::Real: initial $\sigma$ at ramp start.
• num_rollouts_initial::Integer: initial rollout count at ramp start.

Examples

schedule = ScoreFunctionSchedule(
    score_function;
    sf_start = 200,
    ramp_batches = 300,
    perturbation_std_initial = 0.1,
    num_rollouts_initial = 2,
)

StateConditionedPolicy

A policy architecture that separates temporal encoding from state conditioning:

• encoder: a recurrent cell (LSTMCell/GRUCell/RNNCell, or a Chain of them) that encodes only the uncertainty sequence (temporal dependencies)
• combiner: a Dense layer that combines the encoder output with the previous state to produce the next state

Flux's recurrent cells are stateless (Flux >= 0.16): each call returns (output, new_state) instead of mutating an internal Recur. StateConditionedPolicy therefore carries the encoder's recurrent state itself in state, threading it through one call per stage. Call Flux.reset! to clear it (back to Flux.initialstates) at the start of a rollout.

Input format: [uncertainty..., previous_state...]

ZeroGradientFallback()

Default fallback: log a warning and return zero gradients when the solver or DiffOpt differentiation fails. Training continues with a skipped update for that iteration.

compute_parameter_dual(model::JuMP.Model, param::JuMP.VariableRef)

Compute the dual value (sensitivity) of a parameter in a solved JuMP model.

The parameter dual represents ∂(objective)/∂(parameter_value) and is computed by:

1. Finding all constraints where the parameter appears
2. For each constraint, computing: -coefficient * constraint_dual
3. For the objective, adding the coefficient (or negative for maximization)
4. Summing all contributions

This works for any solved model, not just convex ones, as long as dual values are available.

Arguments

• model: A solved JuMP model
• param: A parameter variable (created with @variable(model, p in MOI.Parameter(value)))

Returns

• The dual value (sensitivity) of the parameter

Example

model = Model(HiGHS.Optimizer)
@variable(model, x >= 0)
@variable(model, p in MOI.Parameter(1.0))
@constraint(model, con, x >= 2 * p)
@objective(model, Min, 3 * x + p)
optimize!(model)
dual_p = compute_parameter_dual(model, p)  # Should be -2 * dual(con) + 1

create_deficit!(model::JuMP.Model, len::Int; penalty_l1=nothing, penalty_l2=nothing, penalty=nothing)

Create deficit variables to penalize state deviations in a JuMP model.

Supports three modes:

• L1 norm only: Uses MOI.NormOneCone (default if no penalty specified)
• L2 squared norm only: Uses sum of squared deviations (solver-compatible alternative to SecondOrderCone)
• Both norms: Creates both constraints with separate penalties

Arguments

• model: The JuMP model to add deficit variables to
• len: Number of deficit variables (typically dimension of state)
• penalty_l1: Penalty coefficient for L1 norm (NormOneCone). If nothing and L1 is used, defaults to max objective coefficient.
• penalty_l2: Penalty coefficient for L2 squared norm (sum of squares). If nothing and L2 is used, defaults to max objective coefficient.
• penalty: Legacy argument. If provided and penaltyl1/penaltyl2 are both nothing, uses this for L1 norm only.

Returns

• norm_deficit: Single variable representing total penalized deviation (for logging compatibility)
• _deficit: Vector of deficit variables for each state dimension

Examples

# L1 norm only (default behavior, backwards compatible)
norm_deficit, _deficit = create_deficit!(model, 3; penalty=1000.0)

# L2 norm only
norm_deficit, _deficit = create_deficit!(model, 3; penalty_l2=1000.0)

# Both L1 and L2 norms
norm_deficit, _deficit = create_deficit!(model, 3; penalty_l1=1000.0, penalty_l2=500.0)

default_annealed_schedule(num_batches::Int)

Build the default annealed target-penalty schedule over num_batches training batches: multipliers 0.1 -> 1.0 -> 10.0 -> 30.0 with phase lengths proportional to 2/2/4/16 of the horizon (the last phase takes the remainder; every phase keeps at least one batch). For num_batches < 4 the last num_batches multipliers are used, one batch each, so the run always ends at the strong-penalty phase.

Returns a Vector{Tuple{Int,Int,Float64}} of (first_batch, last_batch, multiplier) entries suitable for the penalty_schedule keyword of train_multistage and train_multiple_shooting. Multipliers are applied relative to the penalty the model was built with (the objective coefficient of the norm_deficit variables created by create_deficit!), so with penalty=:auto the effective penalty is multiplier * max |objective coefficient|.

default_record(sample_log, iter, model)

Default per-batch recording callback: prints the same two per-batch lines as the historical record_loss default (metrics/loss = mean objective excluding the target-slack penalty, then metrics/training_loss = mean full objective) and returns false (training continues). Return true from a custom record to stop training.

dense_multilayer_nn(num_inputs, num_outputs, layers; activation=Flux.relu, dense=Dense)

Create a multi-layer neural network with the specified architecture.

Arguments

• num_inputs::Int: Number of input features
• num_outputs::Int: Number of output features
• layers::Vector{Int}: Hidden layer sizes
• activation: Activation function (default: Flux.relu)
• dense: Layer type (Dense, LSTM, etc.)

materialize_tangent(x)

Recursively convert ChainRulesCore tangent types (MutableTangent, Tangent) to plain NamedTuples/Arrays that Flux.update! can handle.

This is needed because Zygote produces MutableTangent for mutable structs (like Flux.Recur), but Flux.update!/Optimisers.jl expects plain NamedTuples.

normalize_recur_state(state)

Return a copy of a Flux.state object where any Recur-like nodes have their state field set to cell.state0. This avoids Flux.loadmodel! tie errors when loading into freshly constructed recurrent layers.

policy_input_dim(num_uncertainties, num_states)

Compute the input dimension for a policy network.

Policy networks receive [uncertainty..., previous_state...] as input, so the input dimension is num_uncertainties + num_states.

This format is consistent between subproblems and deterministic equivalent formulations, enabling warmstarting policies trained with det_eq for use with subproblems.

Arguments

• num_uncertainties::Int: Number of uncertainty parameters per stage
• num_states::Int: Number of state variables

Returns

• Int: Total input dimension for the policy network

policy_input_dim(uncertainty_samples, initial_state)

Compute the input dimension for a policy network from problem data.

Arguments

• uncertainty_samples: Uncertainty samples from problem construction
• initial_state: Initial state vector

Returns

• Int: Total input dimension for the policy network

predict_window_targets(decision_rule, s_in, uncertainties_vec)

Predict one target per stage in a window. This is an AD-friendly scan: target1 = π([u1; sin]) target2 = π([u2; target1]) ...

sample(sampler::Function, T::Int)

Draw a full trajectory using a callable trajectory sampler with temporal dependence.

sampler(t, past) receives the current stage t and a vector of all previously realized samples past[1:t-1], and returns the realized sample for stage t.

This enables autoregressive, Markovian, or any custom temporal correlation between stages — something the data-based pool formats cannot express.

See sample for the full API and examples.

sample(sampler::Function)

Call a zero-argument trajectory sampler that returns a complete trajectory.

This is the dispatch used by train_multistage and train_multiple_shooting when uncertainty_sampler is a callable. Wrap a trajectory sampler as:

uncertainty_sampler = () -> sample(my_stage_sampler, T)

sample(uncertainty_pool) -> Vector{Vector{Tuple{VariableRef, T}}}

Draw one full uncertainty trajectory from a DecisionRules uncertainty pool.

The returned trajectory is a length-$T$ vector where each element is Vector{Tuple{VariableRef, Float64}} — one realized value per uncertain parameter for that stage. This is the format consumed by simulate_multistage, train_multistage, and all other training/evaluation functions.

Three pool formats are supported, offering increasing levels of correlation:

1. Independent sampling (per-unit pools)

Each uncertain parameter has its own finite support; sampling draws independently from each support at each stage.

sample(multistage_pool::Vector{Vector{Tuple{VariableRef, Vector{T}}}})

multistage_pool[t] is [(param₁, [v₁₁, v₁₂, …]), (param₂, [v₂₁, v₂₂, …]), …]. Each parameter picks one value uniformly at random from its own support. No spatial or temporal correlation is preserved.

2. Joint-scenario sampling (spatial correlation)

Scenarios are pre-defined joint realizations across all parameters at each stage. Sampling picks one complete scenario per stage uniformly, preserving cross-parameter correlations (e.g., spatially correlated inflows across hydro reservoirs). Stages are still drawn independently.

sample(multistage_joint::Vector{Vector{Vector{Tuple{VariableRef, T}}}})

multistage_joint[t] is [scenario₁, scenario₂, …] where each scenario is [(param₁, val₁), (param₂, val₂), …].

3. Trajectory sampler (spatial + temporal correlation)

A callable sampler(t, past) -> Vector{Tuple{VariableRef, T}} that generates stage t's realization given the realized values from stages 1:t-1. This enables autoregressive, Markovian, or any custom temporal dependence.

sample(sampler::Function, T::Int)

The callable receives:

• t::Int — the current stage (1-indexed)
• past::Vector{Vector{Tuple{VariableRef, T}}} — realized samples from stages 1:t-1 (empty vector for t=1)

and must return Vector{Tuple{VariableRef, T}} — the realized sample for stage t.

Output format

All three methods return Vector{Vector{Tuple{VariableRef, T}}} — a length-$T$ vector of per-stage realized samples. This is the universal input to simulate_multistage, train_multistage, simulate_multiple_shooting, and all evaluation functions.

Examples

# 1. Independent sampling (each unit draws independently):
independent_pool = [
    [(inflow_1, [10.0, 15.0, 12.0]), (inflow_2, [8.0, 12.0, 9.0])],
    [(inflow_1, [11.0, 14.0, 13.0]), (inflow_2, [7.0, 11.0, 10.0])],
]
path = sample(independent_pool)

# 2. Joint-scenario sampling (preserves spatial correlation):
joint_pool = [
    [[(inflow_1, 10.0), (inflow_2, 8.0)],   # scenario 1
     [(inflow_1, 15.0), (inflow_2, 12.0)]],  # scenario 2 — stage 1
    [[(inflow_1, 11.0), (inflow_2, 7.0)],
     [(inflow_1, 14.0), (inflow_2, 11.0)]],  # stage 2
]
path = sample(joint_pool)

# 3. Trajectory sampler (preserves temporal + spatial correlation):
function my_sampler(t, past)
    if t == 1
        ω = rand(1:nScenarios)
        return [(inflow_params[t][r], data[r][t, ω]) for r in 1:nHyd]
    else
        # AR(1): next inflow depends on previous realized inflow
        prev_values = [pair[2] for pair in past[end]]
        noise = randn(nHyd) .* σ
        return [(inflow_params[t][r], ρ * prev_values[r] + noise[r]) for r in 1:nHyd]
    end
end
path = sample(my_sampler, T)

See the Uncertainty Sampling documentation page for a complete guide.

setup_shooting_windows(subproblems, state_params_in, state_params_out, initial_state,
                       uncertainties; window_size, model_factory=() -> JuMP.Model())

Build window models for multiple shooting.

Notes:

• We store only the uncertainty PARAMETER refs (not sample sets) in WindowData.

sf_params(config::ScoreFunctionConfig, iteration::Integer)
sf_params(schedule::ScoreFunctionSchedule, iteration::Integer)

Return the effective score-function parameters for iteration.

Arguments

• config::ScoreFunctionConfig: unscheduled score-function configuration.
• schedule::ScoreFunctionSchedule: scheduled score-function configuration.
• iteration::Integer: one-based training iteration.

Returns

A named tuple with fields:

• alpha::Float64: weight on the dual-gradient term.
• score_weight::Float64: weight on the score-function term.
• perturbation_std::Float64: Gaussian standard deviation $\sigma$.
• num_rollouts::Int: number of perturbed rollouts.
• active::Bool: whether rollout estimation should run.

Examples

params = sf_params(schedule, 250)
params.active && @show params.score_weight

simulate_multiple_shooting(windows, decision_rule, initial_state, uncertainty_sample, uncertainties_vec)

• uncertainty_sample: per-stage sampled tuples (param, value) matching your existing sampler output: Vector{Vector{Tuple{VariableRef,<:Real}}}
• uncertainties_vec: per-stage vectors (Float32) used as policy inputs

Returns total objective across windows. Gradients flow through:

• targets within each window (via solve_window rrule)
• realized end state between windows (via solve_window rrule seeding on end vars)

simulate_multistage(det_equivalent::JuMP.Model, state_params_in, state_params_out,
                    initial_state, uncertainties, decision_rules) -> Float64

Convenience overload: rolls out decision_rules to produce target states, then calls the deterministic-equivalent simulate_multistage to solve the coupled problem.

simulate_multistage(subproblems, state_params_in, state_params_out,
                    initial_state, uncertainties, decision_rules) -> Float64

Stage-wise (single shooting) forward simulation. Rolls decision_rules over uncertainties, solving one subproblem per stage. The realized state from each stage feeds the next via get_next_state. Returns the total objective across all stages (Extension §2, Eq. 2.1–2.4).

simulate_multistage(det_equivalent, state_params_in, state_params_out,
                    uncertainties, states) -> Float64

Deterministic-equivalent (direct transcription) forward pass. Sets all parameter values from states and uncertainties into the coupled det_equivalent model, solves it, and returns the objective value (Extension §1, Eq. 1.1).

simulate_stage(subproblem, state_param_in, state_param_out, uncertainty,
               state_in, state_out_target) -> Float64

Set parameter values on subproblem (incoming state, outgoing target, uncertainty), solve it, and return the objective value. Used as the inner solve in single-shooting rollouts (Extension §2, Eq. 2.1).

simulate_states(initial_state, uncertainties, decision_rule) -> Vector{Vector}

Roll out decision_rule over uncertainties to produce a target-state trajectory. At each stage the policy receives [uncertainty..., previous_state...] and outputs the next target state. Returns a length-(T+1) vector of states starting with initial_state.

solve_window(window_model, window_state_in_params, window_state_out_params,
             s_in, targets)

Solve a deterministic-equivalent window model.

Arguments

• window_model: JuMP model (DiffOpt-enabled) for the window
• window_state_in_params: Vector of MOI.Parameter vars for window initial state
• window_state_out_params: per-stage vector of tuples (targetparam, realizedvar)
• s_in: numeric initial state
• targets: vector of numeric targets, one per stage in the window

Returns

• (objective, s_out): objective value, realized end state (Float32 vector)

state_conditioned_policy(n_uncertainty, n_state, n_output, layers;
                         activation=Flux.relu, encoder_type=Flux.LSTM)

Create a StateConditionedPolicy with the specified architecture.

Arguments

• n_uncertainty::Int: Number of uncertainty input dimensions
• n_state::Int: Number of state dimensions (both input and output)
• n_output::Int: Number of output dimensions (typically same as n_state)
• layers::Vector{Int}: Hidden layer sizes for the encoder
• activation: Activation function for dense layers (default: relu)
• encoder_type: Recurrent layer/cell type (LSTM, GRU, RNN, or their *Cell variants; default: Flux.LSTM). Must support Flux.initialstates and the stateful (x, state) -> (output, new_state) call (Flux >= 0.16).

Architecture

• Encoder: encodertype(nuncertainty => layers[1]) -> ... -> layers[end]
• Combiner: Dense(layers[end] + nstate => noutput)

train_multiple_shooting(model, initial_state, windows, uncertainty_sampler; ...)

Train a target-state policy with multiple-shooting decomposition (windowed).

uncertainty_sampler controls how trajectories are drawn at each SGD step. Three formats are accepted (same API as train_multistage):

1. Per-unit pool (Vector{Vector{Tuple{VariableRef, Vector{T}}}}): independent sampling per parameter per stage.
2. Joint-scenario pool (Vector{Vector{Vector{Tuple{VariableRef, T}}}}): one scenario drawn per stage, preserving spatial correlation.
3. Callable (() -> Vector{Vector{Tuple{VariableRef, T}}}): a zero-arg function returning a realized trajectory. Use this for temporal correlation; see sample.

See the Uncertainty Sampling documentation for details.

train_multistage(model, initial_state, det_equivalent::JuMP.Model,
                 state_params_in, state_params_out, uncertainty_sampler;
                 score_function=nothing, kwargs...)

Train a target-state policy with a deterministic equivalent (direct transcription).

For one sampled trajectory $w_{1:T}$, the policy first produces the full target trajectory

\[\hat{x}_{1:T}(\theta) = \pi_\theta(w_{1:T}, x_0).\]

The coupled implementation problem is

\[\begin{aligned} Q(w; \theta) = \min_{\{x_t, y_t, \delta_t\}_{t=1}^{T}} \quad & \sum_{t=1}^{T} f_t(x_t, y_t) + C_\delta \sum_{t=1}^{T} \|\delta_t\| \\ \text{s.t.}\quad & x_t = T_t(w_t, y_t, x_{t-1}) && t=1,\ldots,T, \\ & x_t + \delta_t = \hat{x}_t(\theta) && : \lambda_t,\quad t=1,\ldots,T, \\ & h_t(x_t, y_t) \ge 0 && t=1,\ldots,T . \end{aligned}\]

The target trajectory appears as right-hand-side parameters. If $\lambda_t$ is the dual multiplier of the target constraint, the envelope gradient used by this overload is

\[\nabla_\theta \mathbb{E}[Q(w; \theta)] \approx \frac{1}{S} \sum_{s=1}^{S} \sum_{t=1}^{T} \lambda_t^s \odot \nabla_\theta \hat{x}_t^s(\theta),\]

where $S$ is num_train_per_batch and $\odot$ denotes componentwise multiplication.

Pass a ScoreFunctionConfig or ScoreFunctionSchedule via score_function to mix the dual gradient with a REINFORCE correction estimated from rollouts under perturbed targets.

When score_function is used, there are two separate solve paths:

1. integer_strategy applies to det_equivalent and controls how local dual information is read for the differentiable dual-gradient term.
2. score_function owns separate rollout subproblems. Those models are solved exactly as they are built, and their realized costs define the Monte Carlo score-function term.

For a mixed-integer model, this usually means integer_strategy = FixedDiscreteIntegerStrategy() for the dual path and MIP rollout subproblems inside ScoreFunctionConfig for the score-function path.

Arguments

• model: differentiable Flux-compatible policy. It is rolled forward over uncertainty values to produce $\hat{x}_{1:T}$.
• initial_state::AbstractVector{<:Real}: state $x_0$.
• det_equivalent::JuMP.Model: full-horizon JuMP model for one sampled trajectory.
• state_params_in: input-state parameters in the deterministic equivalent.
• state_params_out: (target_parameter, realized_state_variable) pairs for each target state.
• uncertainty_sampler: source of uncertainty trajectories, passed to sample. Three formats are accepted:
  1. Per-unit pool (Vector{Vector{Tuple{VariableRef, Vector{T}}}}): independent sampling per parameter per stage.
  2. Joint-scenario pool (Vector{Vector{Vector{Tuple{VariableRef, T}}}}): one scenario drawn per stage, preserving spatial correlation.
  3. Callable (() -> Vector{Vector{Tuple{VariableRef, T}}}): a zero-arg function returning a full trajectory. Use this for temporal correlation; see sample.

Keywords

• num_batches::Integer: number of SGD batches.
• num_train_per_batch::Integer: sampled trajectories per batch $S$.
• optimizer: Flux optimizer used to update model.
• adjust_hyperparameters::Function: optional hook returning the batch size for the current iteration.
• record_loss: legacy logging callback.
• sample_log::SampleLog: per-batch objective cache.
• record::Function: callback called as record(sample_log, iter, model).
• penalty_schedule: optional multiplier schedule for target-penalty terms.
• integer_strategy::AbstractIntegerStrategy: strategy used to read local dual information from det_equivalent when it has discrete variables.
• score_function: optional ScoreFunctionConfig or ScoreFunctionSchedule for mixed dual/score-function gradients.

Examples

train_multistage(
    policy,
    initial_state,
    det_equivalent,
    state_params_in,
    state_params_out,
    uncertainty_sampler;
    num_batches = 200,
    num_train_per_batch = 16,
    optimizer = Flux.Adam(1.0e-3),
    integer_strategy = FixedDiscreteIntegerStrategy(),
    score_function = nothing,
)

train_multistage(model, initial_state, subproblems::Vector{JuMP.Model},
                 state_params_in, state_params_out, uncertainty_sampler;
                 kwargs...)

Train a target-state policy with stage-wise decomposition (single shooting).

For one sampled uncertainty trajectory $w_{1:T}$, this overload solves one optimization problem per stage. At stage $t$, given the realized incoming state $x_{t-1}$, the policy predicts a target $\hat{x}_t = \pi_\theta(w_t, x_{t-1})$ and the stage problem is

\[\begin{aligned} q_t(x_{t-1}, w_t; \hat{x}_t) = \min_{x_t, y_t, \delta_t} \quad & f_t(x_t, y_t) + C_\delta \|\delta_t\| \\ \text{s.t.}\quad & x_t = T_t(w_t, y_t, x_{t-1}) && : \mu_t, \\ & x_t + \delta_t = \hat{x}_t && : \lambda_t, \\ & h_t(x_t, y_t) \ge 0 . \end{aligned}\]

The rollout objective is the sum of stage values,

\[Q(\theta; w) = \sum_{t=1}^{T} q_t(x_{t-1}, w_t; \hat{x}_t),\]

where each realized $x_t$ is read from the previous stage solve. The gradient therefore contains both the target duals $\lambda_t$ and the sensitivity of later realized states with respect to earlier targets. In the notation of the extension note,

\[\nabla_\theta Q(\theta; w) = \sum_{t=1}^{T} \left[ \frac{\partial q_t}{\partial \hat{x}_t} + \sum_{k=t+1}^{T} \frac{\partial q_k}{\partial x_{k-1}} \prod_{j=t+1}^{k-1} \frac{\partial x_j}{\partial x_{j-1}} \frac{\partial x_t}{\partial \hat{x}_t} \right] \nabla_\theta \pi_\theta(w_t, x_{t-1}).\]

The dual terms come from target and transition constraints; the state sensitivities are computed through DiffOpt in the rrules for simulate_stage and get_next_state.

Arguments

• model: differentiable Flux-compatible policy. It receives vcat(stage_uncertainty, realized_state) and returns the next target state.
• initial_state::AbstractVector{<:Real}: state $x_0$ entering stage 1.
• subproblems::Vector{JuMP.Model}: one JuMP model per stage.
• state_params_in: stage input-state parameters.
• state_params_out: (target_parameter, realized_state_variable) pairs for each stage output state.
• uncertainty_sampler: source of uncertainty trajectories, passed to sample. Three formats are accepted:
  1. Per-unit pool (Vector{Vector{Tuple{VariableRef, Vector{T}}}}): independent sampling per parameter per stage.
  2. Joint-scenario pool (Vector{Vector{Vector{Tuple{VariableRef, T}}}}): one scenario drawn per stage, preserving spatial correlation.
  3. Callable (() -> Vector{Vector{Tuple{VariableRef, T}}}): a zero-arg function returning a full trajectory. Use this for temporal correlation by wrapping a trajectory sampler: () -> sample(my_stage_sampler, T) where my_stage_sampler(t, past) generates stage t conditioned on past realizations.

Keywords

• num_batches::Integer: number of SGD batches.
• num_train_per_batch::Integer: sampled trajectories per batch.
• optimizer: Flux optimizer used to update model.
• adjust_hyperparameters::Function: optional hook returning the batch size for the current iteration.
• record_loss: legacy logging callback.
• sample_log::SampleLog: per-batch objective cache.
• record::Function: callback called as record(sample_log, iter, model).
• penalty_schedule: optional multiplier schedule for target-penalty terms.
• integer_strategy::AbstractIntegerStrategy: strategy used when a stage model has discrete variables and derivative information must be read.

Examples

# With data pool (independent or joint):
train_multistage(
    policy, initial_state, subproblems,
    state_params_in, state_params_out, uncertainty_pool;
    num_batches=200, optimizer=Flux.Adam(1e-3),
)

# With trajectory sampler (temporal correlation):
ar_sampler(t, past) = my_ar1_model(t, past, inflow_params)
train_multistage(
    policy, initial_state, subproblems,
    state_params_in, state_params_out,
    () -> sample(ar_sampler, T);
    num_batches=200, optimizer=Flux.Adam(1e-3),
)

variable_to_parameter(model, variable; initial_value=0.0, deficit=nothing)

Replace a decision variable with an MOI.Parameter and bind them via an equality constraint. When deficit is provided, the constraint becomes variable + deficit == parameter and both the parameter and the deficit variable are returned.

STRICT_GRADIENTS

Global flag controlling gradient fallback behavior in rrules.

When false (default), rrule pullbacks return zero gradients with a warning when the solver terminates unsuccessfully — this keeps training alive when a few samples hit numerical trouble.

When true, the same situation throws an error instead. Enable this in tests to verify that controlled test cases never silently fall through to zero gradients:

DecisionRules.STRICT_GRADIENTS[] = true

ChainRulesCore.rrule for getlastrealized_state

Computes gradients w.r.t.:

• s_in (window start numeric state)
• targets (all target vectors in the window)

Given cotangents (Δs_out), we:

• seed reverse variables (realized end state vars) with Δs_out
• reverse_differentiate!
• read reverse sensitivities w.r.t. windowstatein_params and all target params

ChainRulesCore.rrule(get_next_state, subproblem, state_param_in, state_param_out, state_in, state_out_target)

Correct reverse-mode rule using DiffOpt:

• Seeds reverse on realized output variables with Δstate_out
• Calls DiffOpt.reverse_differentiate!
• Reads sensitivities wrt parameter vars (stateparamin, stateparamout parameters)
• Returns VJP wrt the numeric inputs state_in and state_out_target

Assumptions:

• subproblem is a JuMP.Model constructed with Model(() -> DiffOpt.diff_optimizer(...))
• state_param_in::Vector{JuMP.VariableRef} are JuMP Parameter variables (incoming state parameters)
• state_param_out::Vector{Tuple{JuMP.VariableRef,JuMP.VariableRef}} holds (target-Parameter variable, realized-state Variable) per component
• get_next_state(...) updates parameter values, optimize!s, and returns a Vector matching the realized-state variables

ChainRulesCore.rrule(::typeof(simulate_multistage), det_equivalent, state_params_in,
                     state_params_out, uncertainties, states)

Reverse-mode rule for the deterministic-equivalent (full-horizon) solve.

Mathematical basis (TS-DDR, arXiv:2405.14973, Eq. 1.2; Extension §1)

For the coupled problem $Q(w;θ) = \min \sum_t f_t + C_δ\|δ_t\|$ the gradient estimator is:

∇_θ E[Q] ≈ (1/S) Σ_s  λ^s ⊙ ∇_θ π(·; θ)

where $λ_t$ is the dual of the target constraint $x_t + δ_t = \hat{x}_t$. The pullback returns $Δ_{states}$ such that $Δ_{states}[1]$ holds the parameter duals of the initial-state parameters and $Δ_{states}[t+1]$ holds the target-constraint duals $λ_t$ for each stage.

Fallback strategy

Same as simulate_stage: tries compute_parameter_dual first, falls back to DiffOpt reverse differentiation if pdual raises. Solver failure (bad termination status) returns zero gradients or throws depending on STRICT_GRADIENTS.

ChainRulesCore.rrule(::typeof(simulate_stage), subproblem, state_param_in,
                     state_param_out, uncertainty, state_in, state_out)

Reverse-mode rule for a single-stage subproblem solve.

Mathematical basis (TS-DDR, arXiv:2405.14973; Extension §2 Eq. 2.5)

For stage problem $q_t(x_{t-1}, w_t; \hat{x}_t)$, the sensitivities are:

∂q_t/∂(state_in)  = μ_t    (dual of dynamics constraint w.r.t. incoming state)
∂q_t/∂(target)     = λ_t    (dual of target constraint w.r.t. target x̂_t)

These are the Lagrange multipliers that compute_parameter_dual (pdual) extracts from the solved model. This is the preferred path: closed-form and exact whenever the solver exposes constraint duals.

Fallback strategy

1. pdual (parameter duals) — tried first.
2. DiffOpt reverse differentiation — if pdual raises (e.g. the optimizer wrapper does not expose conic duals). Computes the same sensitivities via implicit differentiation of the KKT system.
3. Zero gradients — only when the solver terminated with an unsuccessful status (not OPTIMAL / ALMOSTOPTIMAL / LOCALLYSOLVED). A warning is emitted. Set DecisionRules.STRICT_GRADIENTS[] = true to throw instead.

ChainRulesCore.rrule for solve_window

Computes gradients w.r.t.:

• s_in (window start numeric state)
• targets (all target vectors in the window)

Given cotangents (Δobj_val), we:

• seed reverse variables (objective and realized end state vars) with Δobj_val
• reverse_differentiate!
• read reverse sensitivities w.r.t. windowstatein_params and all target params

ChainRulesCore.rrule(::typeof(setwindowuncertainties!), window::WindowData, uncertainty_sample)

Declare setwindowuncertainties! as non-differentiable (mutates solver state).

_apply_deficit_penalty_multiplier!(model::JuMP.Model, bases::Dict{VariableRef,Float64}, multiplier::Real) -> JuMP.Model
_apply_deficit_penalty_multiplier!(models::Vector{JuMP.Model}, bases::Vector, multiplier::Real) -> Vector{JuMP.Model}

Mutate model (or each model in models) in place, setting every deficit variable's objective coefficient to multiplier * base using the bases from _deficit_penalty_bases. Return the mutated model(s).

_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.

_assert_successful_solve(model::JuMP.Model; context::AbstractString = "solve")

Throw an error unless model terminated with an accepted success status.

Arguments

• model::JuMP.Model: model whose termination status is checked.
• context::AbstractString: human-readable phrase included in the error.

Examples

DecisionRules._assert_successful_solve(model; context = "fixed LP solve")

_center_rollout_costs(costs::AbstractVector{<:Real}, baseline::Symbol)

Convert rollout costs into score-function advantages.

Arguments

• costs::AbstractVector{<:Real}: operational costs from perturbed rollouts.
• baseline::Symbol: either :mean or :none.

Examples

advantages = DecisionRules._center_rollout_costs([10.0, 12.0], :mean)

_check_deficit_penalty_bases(bases) -> typeof(bases)

Return bases (from _deficit_penalty_bases) unchanged if it has at least one entry; otherwise throw ArgumentError, since a penalty_schedule would then have nothing to scale.

_deficit_penalty_bases(model::JuMP.Model; deficit_name="norm_deficit") -> Dict{VariableRef,Float64}
_deficit_penalty_bases(models::Vector{JuMP.Model}; deficit_name="norm_deficit") -> Vector{Dict{VariableRef,Float64}}

Capture the current objective coefficient of every deficit variable as the multiplier base for _apply_deficit_penalty_multiplier!. A variable counts as a deficit variable if deficit_name occurs in its name and its linear objective coefficient (see _linear_objective_coefficient) is nonzero.

Must be called before any penalty_schedule multiplier is applied, so the captured coefficients reflect the as-built penalties.

_get_dual_from_affine_constraints(model::JuMP.Model, param::JuMP.VariableRef, S)

Get dual contribution from scalar affine constraints of set type S.

_get_dual_from_constraints(model::JuMP.Model, param::JuMP.VariableRef)

Compute the dual contribution from all constraints containing the parameter.

_get_dual_from_objective(model::JuMP.Model, param::JuMP.VariableRef)

Get the dual contribution from the objective function. If parameter appears in the objective with coefficient c, contribution is:

• +c for minimization
• +c for maximization

_get_dual_from_quadratic_constraints(model::JuMP.Model, param::JuMP.VariableRef, S)

Get dual contribution from quadratic constraints of set type S. Handles both affine terms (parameter appears linearly) and quadratic terms (parameter appears in products pv or pp).

_get_dual_from_vector_affine_constraints(model::JuMP.Model, param::JuMP.VariableRef, F, S)

Get dual contribution from vector affine constraints (like conic constraints).

_get_objective_parameter_coefficient(obj, param::JuMP.VariableRef)

Get the coefficient of a parameter in the objective function.

_get_parameter_coefficient(expr::JuMP.GenericAffExpr, param::JuMP.VariableRef)

Get the coefficient of a parameter in an affine expression.

_get_parameter_coefficient_from_affine(expr::JuMP.GenericQuadExpr, param::JuMP.VariableRef)

Get the coefficient of a parameter from the affine part of a quadratic expression.

_get_parameter_coefficient_from_quadratic(expr::JuMP.GenericQuadExpr, param::JuMP.VariableRef)

Get the effective coefficient of a parameter from quadratic terms. For terms like coef * p * v, the effective coefficient is coef * value(v). For terms like coef * p * p, the effective coefficient is 2 * coef * value(p).

_init_recurrent_state(encoder)

Return the initial recurrent state for encoder: Flux.initialstates(encoder) for a single cell, or a tuple of per-layer initial states for a Chain of cells.

_linear_objective_coefficient(model::JuMP.Model, variable::VariableRef) -> Float64

Return variable's linear coefficient in model's objective, or 0.0 if variable does not appear in it. Supports affine and quadratic objectives (quadratic terms are ignored); throw ArgumentError for any other objective type.

_penalty_multiplier_for(schedule, iter::Int) -> Float64

Return the multiplier of the phase containing batch iter. If iter is past the schedule's last phase, return that phase's multiplier (hold the final value steady).

_record_loss_adapter(record_loss)

Adapt the deprecated 4-argument record_loss(iter, model, loss, tag) callback to the record(sample_log, iter, model) interface, reproducing the historical two-call contract: record_loss is called first with tag = "metrics/loss", and only if that returns false is it called again with tag = "metrics/training_loss". Return the result of whichever call is made last.

_remap_uncertainties(uncertainties, var_src_to_dest, cons_to_cons)

Replace source-model VariableRef keys in an uncertainty pool with their destination-model counterparts (using the variable or constraint mapping built by deterministic_equivalent!).

Two methods dispatch on the pool format:

• Per-unit pools (Vector{Vector{Tuple{VariableRef, Vector{T}}}}): each stage maps [(param₁, [v₁, …]), …] independently.
• Joint-scenario pools (Vector{Vector{Vector{Tuple{VariableRef, T}}}}): each stage maps [[scenario₁…], [scenario₂…], …] preserving the grouped structure.

This is an internal helper; users interact with it indirectly through deterministic_equivalent!.

_resolve_penalty_schedule(penalty_schedule, num_batches::Int)

Resolve the penalty_schedule keyword of train_multistage and train_multiple_shooting into a Vector{Tuple{Int,Int,Float64}} of (first_batch, last_batch, multiplier) phases, or nothing if penalty scaling is disabled:

• nothing returns nothing (no scaling);
• :default_annealed returns default_annealed_schedule(num_batches);
• any other value is checked with _validate_penalty_schedule and returned as-is.

_resolve_record(record, record_loss)

Resolve the record/record_loss keywords of the training loops into a single per-batch callback. Return record unchanged if record_loss is nothing; otherwise require record === default_record and return _record_loss_adapter(record_loss). Throw ArgumentError if both a custom record and a record_loss are given.

_sample_target_perturbations(num_stages::Integer, state_dimension::Integer, sigma::Real)

Draw Gaussian target perturbations for one score-function rollout.

Arguments

• num_stages::Integer: number of stage targets to perturb.
• state_dimension::Integer: length of each target state vector.
• sigma::Real: Gaussian standard deviation $\sigma$.

Examples

perturbations = DecisionRules._sample_target_perturbations(3, 2, 0.5)

_score_function_rollouts(
    config::ScoreFunctionConfig,
    initial_state::AbstractVector,
    uncertainties,
    targets;
    perturbation_std = config.perturbation_std,
    num_rollouts = config.num_rollouts,
) -> (advantages, perturbations)

Estimate rollout advantages for the score-function term.

Arguments

• config::ScoreFunctionConfig: score-function rollout configuration.
• initial_state::AbstractVector: state entering stage 1.
• uncertainties: sampled uncertainty trajectory.
• targets: target trajectory, including the initial state.
• perturbation_std::Real: Gaussian standard deviation $\sigma$.
• num_rollouts::Integer: number of perturbed rollouts to sample.

Examples

advantages, perturbations = DecisionRules._score_function_rollouts(
    score_function,
    initial_state,
    uncertainty_sample,
    targets;
    perturbation_std = 0.5,
    num_rollouts = 4,
)

_score_function_surrogate(
    advantage::Real,
    perturbations,
    targets,
    perturbation_std::Real,
) -> Real

Build the differentiable scalar whose gradient is the Gaussian score estimate.

For fixed rollout cost advantage $A$ and perturbations $\delta_t$, the surrogate is

\[A \sum_t \left\langle \delta_t / \sigma^2, \hat{x}_{t+1}(\theta) \right\rangle .\]

Arguments

• advantage::Real: centered rollout cost $R - b$.
• perturbations: stage perturbations $\delta_t$.
• targets: differentiable target trajectory produced by the policy.
• perturbation_std::Real: Gaussian standard deviation $\sigma$.

Examples

loss = DecisionRules._score_function_surrogate(
    3.0,
    perturbations,
    targets,
    0.5,
)

_sensitivity_forward_status(model::JuMP.Model, strategy) -> MOI.TerminationStatusCode

Return the termination status that an rrule should use for gradient fallback.

Arguments

• model::JuMP.Model: model inspected after the forward pass.
• strategy::AbstractIntegerStrategy: integer strategy used for the solve.

Examples

status = DecisionRules._sensitivity_forward_status(model, strategy)

_set_score_function_stage_parameters!(
    state_params_in,
    state_params_out,
    uncertainties,
    state,
    target,
) -> Nothing

Set the JuMP parameters needed for one perturbed rollout stage.

Arguments

• state_params_in::AbstractVector: parameters receiving the current state.
• state_params_out::AbstractVector: (target_parameter, state_variable) pairs receiving the target state.
• uncertainties::AbstractVector: (parameter, value) pairs for stage uncertainty.
• state::AbstractVector{<:Real}: realized state entering this stage.
• target::AbstractVector{<:Real}: perturbed target for the output state.

Examples

DecisionRules._set_score_function_stage_parameters!(
    spi[t],
    spo[t],
    uncertainty_sample[t],
    state,
    target,
)

_sf_config(score_function) -> Union{Nothing,ScoreFunctionConfig}

Extract the underlying ScoreFunctionConfig, if one exists.

Arguments

• score_function::Nothing: score-function correction is disabled.
• score_function::ScoreFunctionConfig: returned as-is.
• score_function::ScoreFunctionSchedule: unwraps score_function.config.

Examples

config = DecisionRules._sf_config(score_function)

_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 of cells, each layer's output feeds the next and each layer's state is threaded independently.

_target_violation_share(objective::Real, objective_no_deficit::Real) -> Float64

Return the target-violation share, (objective - objective_no_deficit) / objective. Return NaN if either input is nonfinite or abs(objective) <= 1e-12 (share undefined).

_validate_penalty_schedule(schedule) -> typeof(schedule)

Validate an explicit penalty_schedule, a Vector of (first_batch, last_batch, multiplier) tuples: phases must be non-empty, contiguous, start at batch 1, satisfy first_batch <= last_batch, and have finite positive multipliers. Return schedule unchanged, or throw ArgumentError describing the first violation found.

_with_current_or_sensitivity_solution(f, model, integer_strategy)

Run f(model) directly for continuous models and through with_sensitivity_solution for integer strategies.

Arguments

• f::Function: callback that reads values, duals, or sensitivities.
• model::JuMP.Model: model to inspect.
• integer_strategy::AbstractIntegerStrategy: current integer strategy.

Examples

value = DecisionRules._with_current_or_sensitivity_solution(
    m -> JuMP.objective_value(m),
    model,
    strategy,
)

deterministic_equivalent!(model, subproblems, state_params_in, state_params_out,
                          initial_state, uncertainties)

Build the deterministic-equivalent (direct transcription) JuMP model by copying all stage subproblems into model. Variables are renamed with a #t suffix to avoid conflicts. Stage coupling is enforced by identifying the realized state variable of stage t with the incoming state parameter of stage t+1.

uncertainties accepts both sampling formats (see sample):

• Per-unit pools: Vector{Vector{Tuple{VariableRef, Vector{T}}}} — one pool per parameter, drawing independently per parameter.
• Joint-scenario pools: Vector{Vector{Vector{Tuple{VariableRef, T}}}} — pre-built joint scenarios preserving cross-parameter correlations.

Returns (model, uncertainties_new) where uncertainties_new has the same format as the input but with variable refs remapped to the deterministic-equivalent model.

discrete_variables(model::JuMP.Model)

Return the binary or integer variables in model.

Arguments

• model::JuMP.Model: model to inspect.

Examples

vars = DecisionRules.discrete_variables(model)

extract_uncertainty_params(window_uncertainties) -> Vector{Vector{VariableRef}}

Extract the JuMP parameter VariableRefs from each stage of an uncertainty pool.

Handles three possible input shapes (automatically detected):

• Already-extracted Vector{Vector{VariableRef}} — returned as-is.
• Per-unit pool Vector{Vector{Tuple{VariableRef, Vector}}} — extracts the first element of each tuple.
• Joint-scenario pool Vector{Vector{Vector{Tuple{VariableRef, T}}}} — extracts params from the first scenario of each stage (all scenarios share the same params).

get_last_realized_state(window_model, window_state_in_params, window_state_out_params,
                        s_in, targets)

Get the realized end state from the window model after solving.

get_next_state(subproblem, state_param_in, state_param_out, state_in,
               state_out_target) -> Vector

Return the realized state from the most recent solve of subproblem by reading the values of the realized-state variables in state_param_out.

has_discrete_variables(model::JuMP.Model) -> Bool

Return whether model contains at least one binary or integer variable.

Arguments

• model::JuMP.Model: model to inspect.

Examples

if DecisionRules.has_discrete_variables(model)
    @info "MIP model"
end

pdual(v::VariableRef) -> Float64

Compute $∂Q/∂p$ for a JuMP parameter variable $p$ in a solved model, where $Q$ is the optimal objective value. By the envelope theorem / Lagrangian duality this equals the sum of $-\text{coef} \times \text{dual}$ over all constraints where $p$ appears, plus the objective coefficient of $p$.

This is the key quantity in TS-DDR (arXiv:2405.14973): the dual $λ_t$ of the target constraint gives the sensitivity $∂Q/∂\hat{x}_t$ used in Eq. 1.2.

rollout_with_perturbation(
    config::ScoreFunctionConfig,
    initial_state::AbstractVector,
    uncertainties,
    targets,
    perturbations,
) -> Float64

Run one stage-wise rollout with fixed target perturbations.

The rollout target at stage t is targets[t + 1] + perturbations[t]. The returned cost excludes the target-deficit penalty so the score-function signal estimates operational cost rather than target-following slack.

Arguments

• config::ScoreFunctionConfig: rollout models and parameter mappings.
• initial_state::AbstractVector: state entering stage 1.
• uncertainties: sampled uncertainty trajectory.
• targets: target trajectory, including targets[1] == initial_state.
• perturbations: one perturbation vector for each stage target.

Examples

cost = DecisionRules.rollout_with_perturbation(
    score_function,
    initial_state,
    uncertainty_sample,
    targets,
    perturbations,
)

set_window_uncertainties!(window, uncertainty_sample)

Set sampled (realized) uncertainty values into the window model's JuMP parameters.

uncertainty_sample is a realized trajectory (output of sample), so each stage is Vector{Tuple{VariableRef, Float64}} regardless of whether the original pool used independent or joint-scenario format.

windows_equivalent!(model, subproblems, state_params_in, state_params_out,
                     initial_state, uncertainties)

Build a coupled JuMP model for a contiguous window of stages by copying all variables, constraints, and objectives from subproblems into model. Stage coupling is enforced by identifying each stage's realized state variable with the next stage's incoming state parameter (same approach as deterministic_equivalent!, but scoped to a window).

uncertainties accepts both per-unit and joint-scenario pool formats (see sample). The returned uncertainties_new preserves the input format with variable refs remapped to the window model.

Returns (model, state_params_in_new, state_params_out_new, uncertainties_new).

with_sensitivity_solution(f, model, integer_strategy)

Run f(model) while model is in a state suitable for reading duals or DiffOpt sensitivities.

Arguments

• f::Function: callback that reads values, duals, or sensitivities.
• model::JuMP.Model: model to solve and inspect.
• integer_strategy::AbstractIntegerStrategy: strategy used to prepare models with binary or integer variables.

Examples

objective = with_sensitivity_solution(model, FixedDiscreteIntegerStrategy()) do m
    JuMP.objective_value(m)
end

Flux.reset!(m::StateConditionedPolicy)

Reset the encoder's recurrent state to Flux.initialstates, e.g. before starting a new rollout.