Support for contextual stochastic optimization problems

Author: BatyLeo

docs/src/tutorials/warcraft_tutorial.jl

@@ -31,7 +31,7 @@ x = sample.x
 # `y` correspond to the optimal shortest path, encoded as a binary matrix:
 y_true = sample.y
 # `context` is not used in this benchmark (no solver kwargs needed), so it is empty:
-isempty(sample.context)
+isempty(sample.instance_kwargs)
 
 # For some benchmarks, we provide the following plotting method [`plot_solution`](@ref) to visualize the data:
 plot_solution(b, sample)

ext/DFLBenchmarksPlotsExt.jl

@@ -21,7 +21,9 @@ Reconstruct a new sample with `y` overridden and delegate to the 2-arg
 function plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs...)
     return plot_solution(
         bench,
-        DataSample(; sample.context..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra);
+        DataSample(;
+            sample.instance_kwargs..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra
+        );
         kwargs...,
     )
 end

src/ContextualStochasticArgmax/ContextualStochasticArgmax.jl

@@ -0,0 +1,95 @@
+module ContextualStochasticArgmax
+
+using ..Utils
+using DocStringExtensions: TYPEDEF, TYPEDFIELDS
+using Flux: Dense
+using Random: Random, AbstractRNG, MersenneTwister
+
+function one_hot_argmax(z::AbstractVector{R}; kwargs...) where {R<:Real}
+    e = zeros(R, length(z))
+    e[argmax(z)] = one(R)
+    return e
+end
+
+"""
+$TYPEDEF
+
+Minimal contextual stochastic argmax benchmark.
+
+Per instance: `c_base ~ U[0,1]^n` (base utility, part of instance kwargs and base features).
+Per context draw: `x_raw ~ N(0, I_d)` (observable context). Features: `x = [c_base; x_raw]`.
+Per scenario: `ξ = c_base + W * x_raw + noise`, `noise ~ N(0, noise_std² I)`.
+The learner sees `x` and must predict `θ̂` so that `argmax(θ̂)` ≈ `argmax(ξ)`.
+
+A linear model `Dense(n+d → n; bias=false)` can exactly recover `[I | W]`.
+
+# Fields
+$TYPEDFIELDS
+"""
+struct ContextualStochasticArgmaxBenchmark{M<:AbstractMatrix} <:
+       AbstractStochasticBenchmark{true}
+    "number of items (argmax dimension)"
+    n::Int
+    "number of context features"
+    d::Int
+    "fixed perturbation matrix W ∈ R^{n×d}, unknown to the learner"
+    W::M
+    "noise std for scenario draws"
+    noise_std::Float32
+end
+
+function ContextualStochasticArgmaxBenchmark(;
+    n::Int=10, d::Int=5, noise_std::Float32=0.1f0, seed=nothing
+)
+    rng = MersenneTwister(seed)
+    W = randn(rng, Float32, n, d)
+    return ContextualStochasticArgmaxBenchmark(n, d, W, noise_std)
+end
+
+Utils.is_minimization_problem(::ContextualStochasticArgmaxBenchmark) = false
+Utils.generate_maximizer(::ContextualStochasticArgmaxBenchmark) = one_hot_argmax
+
+# c_base: base features (in x) and solver kwarg (in instance_kwargs for generate_scenario)
+function Utils.generate_instance(
+    bench::ContextualStochasticArgmaxBenchmark, rng::AbstractRNG; kwargs...
+)
+    c_base = rand(rng, Float32, bench.n)
+    return DataSample(; x=c_base, c_base=c_base)
+end
+
+# Enriches instance_sample: x = [c_base; x_raw], x_raw in extra for generate_scenario
+function Utils.generate_context(
+    bench::ContextualStochasticArgmaxBenchmark,
+    rng::AbstractRNG,
+    instance_sample::DataSample,
+)
+    x_raw = randn(rng, Float32, bench.d)
+    return DataSample(;
+        x=vcat(instance_sample.x, x_raw),
+        instance_sample.instance_kwargs...,
+        extra=(; x_raw),
+    )
+end
+
+# ξ = c_base + W * x_raw + noise  (c_base from instance_kwargs, x_raw from ctx.extra)
+function Utils.generate_scenario(
+    bench::ContextualStochasticArgmaxBenchmark,
+    rng::AbstractRNG;
+    c_base::AbstractVector,
+    x_raw::AbstractVector,
+    kwargs...,
+)
+    θ_true = c_base + bench.W * x_raw
+    return θ_true + bench.noise_std * randn(rng, Float32, bench.n)
+end
+
+function Utils.generate_statistical_model(
+    bench::ContextualStochasticArgmaxBenchmark; seed=nothing
+)
+    Random.seed!(seed)
+    return Dense(bench.n + bench.d => bench.n; bias=false)
+end
+
+export ContextualStochasticArgmaxBenchmark
+
+end

src/DecisionFocusedLearningBenchmarks.jl

@@ -55,6 +55,7 @@ include("Warcraft/Warcraft.jl")
 include("FixedSizeShortestPath/FixedSizeShortestPath.jl")
 include("PortfolioOptimization/PortfolioOptimization.jl")
 include("StochasticVehicleScheduling/StochasticVehicleScheduling.jl")
+include("ContextualStochasticArgmax/ContextualStochasticArgmax.jl")
 include("DynamicVehicleScheduling/DynamicVehicleScheduling.jl")
 include("DynamicAssortment/DynamicAssortment.jl")
 include("Maintenance/Maintenance.jl")
@@ -71,8 +72,9 @@ export Policy, evaluate_policy!
 
 export generate_instance,
     generate_sample, generate_dataset, generate_environments, generate_environment
-export generate_scenario
+export generate_scenario, generate_context
 export generate_baseline_policies
+export SAA
 export generate_statistical_model
 export generate_maximizer
 export generate_anticipative_solver, generate_parametric_anticipative_solver
@@ -91,6 +93,7 @@ using .Warcraft
 using .FixedSizeShortestPath
 using .PortfolioOptimization
 using .StochasticVehicleScheduling
+using .ContextualStochasticArgmax
 using .DynamicVehicleScheduling
 using .DynamicAssortment
 using .Maintenance
@@ -106,5 +109,6 @@ export StochasticVehicleSchedulingBenchmark
 export SubsetSelectionBenchmark
 export WarcraftBenchmark
 export MaintenanceBenchmark
+export ContextualStochasticArgmaxBenchmark
 
 end # module DecisionFocusedLearningBenchmarks

src/StochasticVehicleScheduling/StochasticVehicleScheduling.jl

@@ -116,7 +116,7 @@ Returns a [`DataSample`](@ref) with features `x` and `instance` set, but `y=noth
 To obtain labeled samples, pass a `target_policy` to [`generate_dataset`](@ref):
 
 ```julia
-policy = sample -> DataSample(; sample.context..., x=sample.x,
+policy = sample -> DataSample(; sample.instance_kwargs..., x=sample.x,
                                 y=column_generation_algorithm(sample.instance))
 dataset = generate_dataset(benchmark, N; target_policy=policy)
 ```

src/StochasticVehicleScheduling/policies.jl

@@ -5,10 +5,17 @@ SAA baseline policy: builds a stochastic instance from all K scenarios and solve
 via column generation.
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_saa_policy(sample, scenarios)
-    stochastic_inst = build_stochastic_instance(sample.instance, scenarios)
+function svs_saa_policy(instance_sample, ctx_sample, scenarios)
+    stochastic_inst = build_stochastic_instance(instance_sample.instance, scenarios)
     y = column_generation_algorithm(stochastic_inst)
-    return [DataSample(; sample.context..., x=sample.x, y, extra=(; scenarios))]
+    return [
+        DataSample(;
+            instance_sample.instance_kwargs...,
+            x=ctx_sample.x,
+            y,
+            extra=(; ctx_sample.extra..., scenarios),
+        ),
+    ]
 end
 
 """
@@ -17,9 +24,18 @@ $TYPEDSIGNATURES
 Deterministic baseline policy: solves the deterministic MIP (ignores scenario delays).
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_deterministic_policy(sample, scenarios; model_builder=highs_model)
-    y = deterministic_mip(sample.instance; model_builder)
-    return [DataSample(; sample.context..., x=sample.x, y, extra=(; scenarios))]
+function svs_deterministic_policy(
+    instance_sample, ctx_sample, scenarios; model_builder=highs_model
+)
+    y = deterministic_mip(instance_sample.instance; model_builder)
+    return [
+        DataSample(;
+            instance_sample.instance_kwargs...,
+            x=ctx_sample.x,
+            y,
+            extra=(; ctx_sample.extra..., scenarios),
+        ),
+    ]
 end
 
 """
@@ -29,24 +45,60 @@ Local search baseline policy: builds a stochastic instance from all K scenarios
 solves via local search heuristic.
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_local_search_policy(sample, scenarios)
-    stochastic_inst = build_stochastic_instance(sample.instance, scenarios)
+function svs_local_search_policy(instance_sample, ctx_sample, scenarios)
+    stochastic_inst = build_stochastic_instance(instance_sample.instance, scenarios)
     y = local_search(stochastic_inst)
-    return [DataSample(; sample.context..., x=sample.x, y, extra=(; scenarios))]
+    return [
+        DataSample(;
+            instance_sample.instance_kwargs...,
+            x=ctx_sample.x,
+            y,
+            extra=(; ctx_sample.extra..., scenarios),
+        ),
+    ]
+end
+
+"""
+$TYPEDSIGNATURES
+
+Exact SAA MIP policy (linearized): solves the stochastic VSP exactly for the given
+scenarios via [`compact_linearized_mip`](@ref).
+Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
+
+Prefer this over [`svs_saa_policy`](@ref) when an exact solution is needed; requires
+SCIP (default) or Gurobi.
+"""
+function svs_saa_mip_policy(
+    instance_sample, ctx_sample, scenarios; model_builder=scip_model
+)
+    y = compact_linearized_mip(instance_sample.instance, scenarios; model_builder)
+    return [
+        DataSample(;
+            instance_sample.instance_kwargs...,
+            x=ctx_sample.x,
+            y,
+            extra=(; ctx_sample.extra..., scenarios),
+        ),
+    ]
 end
 
 """
 $TYPEDSIGNATURES
 
 Return the named baseline policies for [`StochasticVehicleSchedulingBenchmark`](@ref).
-Each policy has signature `(sample, scenarios) -> Vector{DataSample}`.
+Each policy has signature `(instance_sample, ctx_sample, scenarios) -> Vector{DataSample}`.
 """
 function svs_generate_baseline_policies(::StochasticVehicleSchedulingBenchmark)
     return (;
         deterministic=Policy(
             "Deterministic MIP", "Ignores delays", svs_deterministic_policy
         ),
         saa=Policy("SAA (col gen)", "Stochastic MIP over K scenarios", svs_saa_policy),
+        saa_mip=Policy(
+            "SAA (exact MIP)",
+            "Exact stochastic MIP over K scenarios via compact linearized formulation",
+            svs_saa_mip_policy,
+        ),
         local_search=Policy(
             "Local search", "Heuristic with K scenarios", svs_local_search_policy
         ),

src/StochasticVehicleScheduling/solution/algorithms/mip.jl

@@ -84,6 +84,18 @@ end
 """
 $TYPEDSIGNATURES
 
+SAA variant: build stochastic instance from `scenarios` then solve via
+[`compact_linearized_mip`](@ref).
+"""
+function compact_linearized_mip(
+    instance::Instance, scenarios::Vector{VSPScenario}; kwargs...
+)
+    return compact_linearized_mip(build_stochastic_instance(instance, scenarios); kwargs...)
+end
+
+"""
+$TYPEDSIGNATURES
+
 Returns the optimal solution of the Stochastic VSP instance, by solving the associated compact quadratic MIP.
 Note: If you have Gurobi, use `grb_model` as `model_builder` instead of `highs_model`.
 
@@ -151,3 +163,13 @@ function compact_mip(
     sol = solution_from_JuMP_array(solution, graph)
     return sol.value
 end
+
+"""
+$TYPEDSIGNATURES
+
+SAA variant: build stochastic instance from `scenarios` then solve via
+[`compact_mip`](@ref).
+"""
+function compact_mip(instance::Instance, scenarios::Vector{VSPScenario}; kwargs...)
+    return compact_mip(build_stochastic_instance(instance, scenarios); kwargs...)
+end

src/Utils/Utils.jl

@@ -31,8 +31,9 @@ export ExogenousStochasticBenchmark,
     EndogenousStochasticBenchmark, ExogenousDynamicBenchmark, EndogenousDynamicBenchmark
 export generate_instance, generate_sample, generate_dataset
 export generate_statistical_model, generate_maximizer
-export generate_scenario
+export generate_scenario, generate_context
 export generate_environment, generate_environments
+export SAA
 export generate_baseline_policies
 export generate_anticipative_solver, generate_parametric_anticipative_solver