diff --git a/Project.toml b/Project.toml
index 15380ec..c014192 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,13 +5,17 @@ version = "1.0.0-DEV"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffOpt = "930fe3bc-9c6b-11ea-2d94-6184641e85e7"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
 julia = "1"
@@ -21,6 +25,7 @@ CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13"
 
 [targets]
-test = ["Test", "HiGHS", "Zygote"]
+test = ["Test", "SCS", "Zygote"]
diff --git a/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl b/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl
index 806787c..627e15c 100644
--- a/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl
+++ b/examples/HydroPowerModels/gen_inputs_l2O_hydropowermodels.jl
@@ -21,7 +21,7 @@ subproblems, state_params_in, state_params_out, uncertainty_samples, initial_sta
     joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages
 )
 
-det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent(subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
+det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent!(JuMP.Model(), subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
 
 # Remove state imposing constraints
 
diff --git a/examples/HydroPowerModels/load_hydropowermodels.jl b/examples/HydroPowerModels/load_hydropowermodels.jl
index b31d189..484bf10 100644
--- a/examples/HydroPowerModels/load_hydropowermodels.jl
+++ b/examples/HydroPowerModels/load_hydropowermodels.jl
@@ -57,7 +57,7 @@ function read_inflow(file::String, nHyd::Int; num_stages=nothing)
     return vector_inflows, nCen, num_stages
 end
 
-function build_hydropowermodels(case_folder::AbstractString, subproblem_file::AbstractString; num_stages=nothing, param_type=:Var, penalty=nothing) # :Param, :Cons, :Var
+function build_hydropowermodels(case_folder::AbstractString, subproblem_file::AbstractString; num_stages=nothing, penalty=nothing)
     hydro_file = JSON.parsefile(joinpath(case_folder, "hydro.json"))["Hydrogenerators"]
     nHyd = length(hydro_file)
     vector_inflows, nCen, num_stages = read_inflow(joinpath(case_folder, "inflows.csv"), nHyd; num_stages=num_stages)
@@ -77,9 +77,9 @@ function build_hydropowermodels(case_folder::AbstractString, subproblem_file::Ab
             delete(subproblems[t], con)
         end
         state_params_in[t], state_param_out, inflow = find_reservoirs_and_inflow(subproblems[t])
-        state_params_in[t] = variable_to_parameter.(subproblems[t], state_params_in[t], param_type=param_type)
-        state_params_out[t] = [variable_to_parameter(subproblems[t], state_param_out[i]; deficit=_deficit[i], param_type=param_type) for i in 1:nHyd]
-        inflow = [(variable_to_parameter(subproblems[t], inflow[i]; param_type=param_type),  vector_inflows[i][t, :] .+ 0.0) for i in 1:nHyd]
+        state_params_in[t] = variable_to_parameter.(subproblems[t], state_params_in[t])
+        state_params_out[t] = [variable_to_parameter(subproblems[t], state_param_out[i]; deficit=_deficit[i]) for i in 1:nHyd]
+        inflow = [(variable_to_parameter(subproblems[t], inflow[i]),  vector_inflows[i][t, :] .+ 0.0) for i in 1:nHyd]
         uncertainty_samples[t] = inflow
     end
 
diff --git a/examples/HydroPowerModels/test_dr_hydropowermodels.jl b/examples/HydroPowerModels/test_dr_hydropowermodels.jl
index 34eafa3..7ab5013 100644
--- a/examples/HydroPowerModels/test_dr_hydropowermodels.jl
+++ b/examples/HydroPowerModels/test_dr_hydropowermodels.jl
@@ -6,10 +6,9 @@ using Gurobi
 using MosekTools
 using Ipopt, HSL_jll
 import MathOptSymbolicAD
-import ParametricOptInterface as POI
 using JLD2
 using HydroPowerModels
-
+using DiffOpt
 
 HydroPowerModels_dir = dirname(@__FILE__)
 include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl"))
@@ -40,26 +39,16 @@ HydroPowerModels.gather_useful_info!(data)
 # Build MSP
 
 subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels(    
-    joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, param_type=:Var
+    joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages
 )
 
-det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent(subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
-
-set_optimizer(det_equivalent, Gurobi.Optimizer)
-
-# set_optimizer(det_equivalent, Mosek.Optimizer)
-
-# set_optimizer(det_equivalent, optimizer_with_attributes(Ipopt.Optimizer, 
-#     "print_level" => 0,
-#     "hsllib" => HSL_jll.libhsl_path,
-#     "linear_solver" => "ma27"
-# ))
+det_equivalent = DiffOpt.diff_model(optimizer_with_attributes(Ipopt.Optimizer, 
+    "print_level" => 0,
+    "hsllib" => HSL_jll.libhsl_path,
+    "linear_solver" => "ma27"
+))
 
-# set_attribute(
-#     det_equivalent,
-#     MOI.AutomaticDifferentiationBackend(),
-#     MathOptSymbolicAD.DefaultBackend(),
-# )
+det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent!(det_equivalent, subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
 
 num_hydro = length(initial_state)
 
@@ -103,13 +92,13 @@ for i in 1:num_samples
         inflow_var = inflow_var[findfirst(x -> occursin("_inflow[$j]", JuMP.name(x)), inflow_var)]
         inflows[i, j, t] = uncertainty_s[t][inflow_var]
     end
-    objective_values[i] = simulate_multistage(
+    simulate_multistage(
         det_equivalent, state_params_in, state_params_out, 
         initial_state, uncertainty_s, 
         models;
-        ensure_feasibility=(x_out, x_in, _sa) -> ensure_feasibility(x_out, x_in, _sa, max_volume),
-        _objective_value = DecisionRules.get_objective_no_target_deficit,
+        ensure_feasibility=(x_out, x_in, _sa) -> ensure_feasibility(x_out, x_in, _sa, max_volume)
     )
+    objective_values[i] = DecisionRules.get_objective_no_target_deficit(det_equivalent)
     for _var in record_variables_names
         num_vars = length(find_variables(det_equivalent, [_var; r"#1$"]))
         for j in 1:num_vars, t in 1:num_stages
diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels.jl b/examples/HydroPowerModels/train_dr_hydropowermodels.jl
index 15bc580..17ec2a6 100644
--- a/examples/HydroPowerModels/train_dr_hydropowermodels.jl
+++ b/examples/HydroPowerModels/train_dr_hydropowermodels.jl
@@ -20,6 +20,7 @@ using Ipopt, HSL_jll # Gurobi, MosekTools, Ipopt, MadNLP
 # import ParametricOptInterface as POI
 using Wandb, Dates, Logging
 using JLD2
+using DiffOpt
 
 HydroPowerModels_dir = dirname(@__FILE__)
 include(joinpath(HydroPowerModels_dir, "load_hydropowermodels.jl"))
@@ -52,41 +53,16 @@ pre_trained_model = nothing #joinpath(HydroPowerModels_dir, case_name, formulati
 # Build MSP
 
 subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels(    
-    joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, param_type=:Var
+    joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages
 )
 
-det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent(subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples) #; model = JuMP.Model(() -> POI_cached_optimizer()))
-
-set_optimizer(det_equivalent, optimizer_with_attributes(Ipopt.Optimizer, 
+det_equivalent = DiffOpt.diff_model(optimizer_with_attributes(Ipopt.Optimizer, 
     "print_level" => 0,
     "hsllib" => HSL_jll.libhsl_path,
     "linear_solver" => "ma27"
 ))
 
-# set_optimizer(det_equivalent, Gurobi.Optimizer)
-
-# set_optimizer(det_equivalent, Mosek.Optimizer)
-
-# ipopt = Ipopt.Optimizer()
-# MOI.set(ipopt, MOI.RawOptimizerAttribute("print_level"), 0)
-# MOI.set(ipopt, MOI.RawOptimizerAttribute("hsllib"), HSL_jll.libhsl_path)
-# MOI.set(ipopt, MOI.RawOptimizerAttribute("linear_solver"), "ma97")
-# cached =
-#     () -> MOI.Bridges.full_bridge_optimizer(
-#         MOI.Utilities.CachingOptimizer(
-#             MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
-#             ipopt,
-#         ),
-#         Float64,
-# )
-# POI_cached_optimizer() = POI.Optimizer(cached())
-
-# set_optimizer(det_equivalent, () -> POI.Optimizer(Ipopt.Optimizer()))
-
-# set_optimizer(det_equivalent, () -> POI_cached_optimizer())
-# set_optimizer(det_equivalent, () -> Mosek.Optimizer())
-# set_attribute(det_equivalent, "QUIET", true)
-# set_attributes(det_equivalent, "OutputFlag" => 0)
+det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent!(det_equivalent, subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
 
 num_hydro = length(initial_state)
 # for subproblem in subproblems
diff --git a/examples/RL/hydro_pre_trained.jl b/examples/RL/hydro_pre_trained.jl
index 51898b9..4a6c688 100644
--- a/examples/RL/hydro_pre_trained.jl
+++ b/examples/RL/hydro_pre_trained.jl
@@ -56,7 +56,7 @@ end
     formulation_file = formulation * ".mof.json"
     Random.seed!(1234)
     subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels(    
-        joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, param_type=:Var, penalty=penalty
+        joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, penalty=penalty
     )
 
     for subproblem in subproblems
@@ -88,7 +88,7 @@ end
             uncertainty_sample = [(uncertainty_samples[j][i][1], rain_state[i]) for i in 1:length(rain)]
             simulate_stage(subproblems[j], state_params_in[j], state_params_out[j], uncertainty_sample, state_in, state_out)
             r = _objective_value(subproblems[j])
-            next_volume = DecisionRules.get_next_state(subproblems[j], state_params_out[j], state_in, state_out)
+            next_volume = DecisionRules.get_next_state(subproblems[j], state_params_in[j], state_params_out[j], state_in, state_out)
             @info "Stage t=$j" sum(state_in) sum(rain_state) sum(state_out) sum(next_volume) r
             sp = [next_volume; rain; j+1]
             o = rain #+ next_volume - state_out
diff --git a/examples/RL/hydropowermodels_rl.jl b/examples/RL/hydropowermodels_rl.jl
index f853059..c4fd309 100644
--- a/examples/RL/hydropowermodels_rl.jl
+++ b/examples/RL/hydropowermodels_rl.jl
@@ -58,7 +58,7 @@ end
     formulation_file = formulation * ".mof.json"
     Random.seed!(1234)
     subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels(    
-        joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, param_type=:Var, penalty=penalty
+        joinpath(HydroPowerModels_dir, case_name), formulation_file; num_stages=num_stages, penalty=penalty
     )
 
     for subproblem in subproblems
@@ -95,7 +95,7 @@ end
             uncertainty_sample = [(uncertainty_samples[j][i][1], rain_state[i]) for i in 1:length(rain)]
             simulate_stage(subproblems[j], state_params_in[j], state_params_out[j], uncertainty_sample, state_in, state_out)
             r = _objective_value(subproblems[j])
-            sp = DecisionRules.get_next_state(subproblems[j], state_params_out[j], state_in, state_out)
+            sp = DecisionRules.get_next_state(subproblems[j], state_params_in[j], state_params_out[j], state_in, state_out)
             @info "Stage t=$j" sum(state_in) sum(rain_state) sum(state_out) sum(sp) r
             sp = [sp; rain; j+1]
             o = sp # no hidden state
diff --git a/examples/rocket_control/build_rocket_problem.jl b/examples/rocket_control/build_rocket_problem.jl
index 2719f31..bc8021d 100644
--- a/examples/rocket_control/build_rocket_problem.jl
+++ b/examples/rocket_control/build_rocket_problem.jl
@@ -5,6 +5,7 @@ using Statistics
 
 using JuMP
 import Ipopt, HSL_jll
+using DiffOpt
 # import Plots
 
 
@@ -32,8 +33,7 @@ function build_rocket_problem(;
     # program, we need to use a nonlinear solver like Ipopt. We cannot use a linear
     # solver like HiGHS.
 
-    det_equivalent = Model()
-    set_optimizer(det_equivalent, optimizer_with_attributes(Ipopt.Optimizer, 
+    det_equivalent = DiffOpt.diff_model(optimizer_with_attributes(Ipopt.Optimizer, 
         "print_level" => 0,
         "hsllib" => HSL_jll.libhsl_path,
         "linear_solver" => "ma27"
diff --git a/src/DecisionRules.jl b/src/DecisionRules.jl
index dcd3d28..16d3a68 100644
--- a/src/DecisionRules.jl
+++ b/src/DecisionRules.jl
@@ -4,8 +4,10 @@ using JuMP
 import ParametricOptInterface as POI
 using Flux
 using JLD2
+using ChainRules: @ignore_derivatives
 using ChainRulesCore
 import ChainRulesCore.rrule
+using DiffOpt
 
 export simulate_multistage, sample, train_multistage, simulate_states, simulate_stage, dense_multilayer_nn, variable_to_parameter, create_deficit!, 
     SaveBest, find_variables, identity
diff --git a/src/simulate_multistage.jl b/src/simulate_multistage.jl
index db6fda0..88970e2 100644
--- a/src/simulate_multistage.jl
+++ b/src/simulate_multistage.jl
@@ -1,18 +1,7 @@
-function set_parameter(subproblem, _var::VariableRef, val)
-    @assert owner_model(_var) == subproblem
-    if is_parameter(_var)
-        MOI.set(subproblem, POI.ParameterValue(), _var, val)
-    else
-        fix(_var, val)
-    end
-end
-set_parameter(subproblem, _var::ConstraintRef, val) = set_normalized_rhs(_var, val)
-
 function simulate_states(
     initial_state::Vector{Float64},
     uncertainties,
     decision_rule::F;
-    ensure_feasibility=(x_out, x_in, uncertainty) -> x_out
 ) where {F}
     num_stages = length(uncertainties)
     states = Vector{Vector{Float64}}(undef, num_stages + 1)
@@ -22,8 +11,7 @@ function simulate_states(
         if stage == 1
             uncertainties_stage = initial_state .+ uncertainties_stage
         end
-        state_out = decision_rule(uncertainties_stage)
-        states[stage + 1] = ensure_feasibility(state_out, states[stage], uncertainties_stage)
+        states[stage + 1] = decision_rule(uncertainties_stage)
     end
     return states
 end
@@ -32,7 +20,6 @@ function simulate_states(
     initial_state::Vector{Float64},
     uncertainties,
     decision_rules::Vector{F};
-    ensure_feasibility=(x_out, x_in, uncertainty) -> x_out
 ) where {F}
     num_stages = length(uncertainties)
     states = Vector{Vector{Float64}}(undef, num_stages + 1)
@@ -41,8 +28,7 @@ function simulate_states(
         uncertainties_stage_vec = vcat(initial_state, [[uncertainties[j][i][2] for i in 1:length(uncertainties[j])] for j in 1:stage]...)
         uncertainties_stage = [uncertainties[stage][i][2] for i in 1:length(uncertainties[stage])]
         decision_rule = decision_rules[stage]
-        state_out = decision_rule(uncertainties_stage_vec)
-        states[stage + 1] = ensure_feasibility(state_out, states[stage], uncertainties_stage)
+        states[stage + 1] = decision_rule(uncertainties_stage_vec)
     end
     return states
 end
@@ -51,18 +37,18 @@ function simulate_stage(subproblem::JuMP.Model, state_param_in::Vector{Any}, sta
 ) where {T <: Real, V <: Real, Z <: Real}
     # Update state parameters
     for (i, state_var) in enumerate(state_param_in)
-        set_parameter(subproblem, state_var, state_in[i])
+        set_parameter_value(state_var, state_in[i])
     end
 
     # Update uncertainty
     for (uncertainty_param, uncertainty_value) in uncertainty
-        set_parameter(subproblem, uncertainty_param, uncertainty_value)
+        set_parameter_value(uncertainty_param, uncertainty_value)
     end
 
     # Update state parameters out
     for i in 1:length(state_param_in)
         state_var = state_param_out[i][1]
-        set_parameter(subproblem, state_var, state_out_target[i])
+        set_parameter_value(state_var, state_out_target[i])
     end
 
     # Solve subproblem
@@ -74,30 +60,77 @@ function simulate_stage(subproblem::JuMP.Model, state_param_in::Vector{Any}, sta
     return obj
 end
 
-function get_next_state(subproblem::JuMP.Model, state_param_out::Vector{Tuple{Any, VariableRef}}, state_in::Vector{T}, state_out_target::Vector{Z}) where {T <: Real, Z <: Real}
+function get_next_state(subproblem::JuMP.Model, state_param_in::Vector{Any}, state_param_out::Vector{Tuple{Any, VariableRef}}, state_in::Vector{T}, state_out_target::Vector{Z}) where {T <: Real, Z <: Real}
     state_out = [value(state_param_out[i][2]) for i in 1:length(state_param_out)]
     return state_out
 end
 
-# Define rrule of get_next_state
-# This is simplified. This actual jacobian will require DiffOpt
-function rrule(::typeof(get_next_state), subproblem, state_param_out, state_in, state_out_target)
-    state_out = get_next_state(subproblem, state_param_out, state_in, state_out_target)
-    function _pullback(Δstate_out)
-        d_state_in = zeros(length(state_in))
-        d_state_out = zeros(length(state_out_target))
-        for i in 1:length(state_in)
-            s_out_target = state_out_target[i]
-            s_in = state_in[i]
-            if s_out_target < s_in && s_out_target >= 0.0
-                d_state_out[i] = Δstate_out[i]
-            elseif s_out_target > s_in && s_out_target >= 0.0
-                d_state_in[i] = Δstate_out[i]
-            end
+"""
+    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 (state_param_in, state_param_out 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
+"""
+function ChainRulesCore.rrule(::typeof(get_next_state),
+    subproblem::JuMP.Model,
+    state_param_in,
+    state_param_out,
+    state_in,
+    state_out_target
+)
+
+    # Forward pass: run the solver via the user's function
+    y = DecisionRules.get_next_state(subproblem, state_param_in, state_param_out, state_in, state_out_target)
+
+    function pullback(Δy)
+        Δy = collect(Δy)  # ensure indexable, concrete element type
+
+        # Best practice: clear previous seeds
+        DiffOpt.empty_input_sensitivities!(subproblem)
+
+        # 1) Seed reverse on the realized output variables with Δy
+        #    Each entry in state_param_out is (param_target, realized_state_var)
+        @inbounds for i in eachindex(state_param_out)
+            realized_var = state_param_out[i][2]
+            # J' * Δ: set reverse seed on variable primal
+            DiffOpt.set_reverse_variable(subproblem, realized_var, Δy[i])
+        end
+
+        # 2) Reverse differentiate
+        DiffOpt.reverse_differentiate!(subproblem)  # computes all needed products
+
+        # 3) Read sensitivities w.r.t. parameter variables
+        #    These are vector-Jacobian products dL/d(param) = (∂y/∂param)^T * Δy
+        d_state_in = similar(state_in, promote_type(eltype(state_in), eltype(Δy)))
+        @inbounds for i in eachindex(state_param_in)
+            pin = state_param_in[i]  # JuMP.Parameter variable            
+            d_state_in[i] = DiffOpt.get_reverse_parameter(subproblem, pin)
         end
-        return (NoTangent(), NoTangent(), NoTangent(), d_state_in , d_state_out)
+
+        d_state_out_target = similar(state_out_target, promote_type(eltype(state_out_target), eltype(Δy)))
+        @inbounds for i in eachindex(state_param_out)
+            pout = state_param_out[i][1]  # target Parameter variable
+            d_state_out_target[i] = DiffOpt.get_reverse_parameter(subproblem, pout)
+        end
+
+        # Optional: clear seeds so they don't accumulate between calls
+        DiffOpt.empty_input_sensitivities!(subproblem)
+
+        # Return cotangents for each primal argument, in order:
+        #  (f, subproblem, state_param_in, state_param_out, state_in, state_out_target)
+        return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), d_state_in, d_state_out_target)
     end
-    return state_out, _pullback
+
+    return y, pullback
 end
 
 function get_objective_no_target_deficit(subproblem::JuMP.Model; norm_deficit::AbstractString="norm_deficit")
@@ -111,8 +144,16 @@ function get_objective_no_target_deficit(subproblem::JuMP.Model; norm_deficit::A
     return objective_val
 end
 
-# define rrule of get_objective_no_target_deficit
-function rrule(::typeof(get_objective_no_target_deficit), subproblem; norm_deficit="norm_deficit")
+function get_objective_no_target_deficit(subproblems::Vector{JuMP.Model}; norm_deficit::AbstractString="norm_deficit")
+    total_objective = 0.0
+    for subproblem in subproblems
+        total_objective += get_objective_no_target_deficit(subproblem, norm_deficit=norm_deficit)
+    end
+    return total_objective
+end
+
+# define ChainRulesCore.rrule of get_objective_no_target_deficit
+function ChainRulesCore.rrule(::typeof(get_objective_no_target_deficit), subproblem; norm_deficit="norm_deficit")
     objective_val = get_objective_no_target_deficit(subproblem, norm_deficit=norm_deficit)
     function _pullback(Δobjective_val)
         return (NoTangent(), NoTangent())
@@ -120,27 +161,35 @@ function rrule(::typeof(get_objective_no_target_deficit), subproblem; norm_defic
     return objective_val, _pullback
 end
 
+function apply_rule(::Int, decision_rule::T, uncertainty, state_in) where {T}
+    return decision_rule(vcat([uncertainty[i][2] for i in 1:length(uncertainty)], state_in))
+end
+
+function apply_rule(stage::Int, decision_rules::Vector{T}, uncertainty, state_in) where {T}
+    return apply_rule(stage, decision_rules[stage], uncertainty, state_in)
+end
+
 function simulate_multistage(
     subproblems::Vector{JuMP.Model},
-    state_params_in::Vector{Vector{Any}},
-    state_params_out::Vector{Vector{Tuple{Any, VariableRef}}},
+    state_params_in::Vector{Vector{U}},
+    state_params_out::Vector{Vector{Tuple{U, VariableRef}}},
+    initial_state::Vector{T},
     uncertainties,
-    states::Vector{Vector{T}};
-    _objective_value = objective_value
-    ) where {T <: Real}
+    decision_rules
+) where {T <: Real, U}
+    @ignore_derivatives Flux.reset!.(decision_rules)
     
     # Loop over stages
     objective_value = 0.0
-    state_in = states[1]
+    state_in = initial_state
     for stage in 1:length(subproblems)
-        state_out = states[stage + 1]
         subproblem = subproblems[stage]
         state_param_in = state_params_in[stage]
         state_param_out = state_params_out[stage]
         uncertainty = uncertainties[stage]
-        simulate_stage(subproblem, state_param_in, state_param_out, uncertainty, state_in, state_out)
-        objective_value += _objective_value(subproblem)
-        state_in = DecisionRules.get_next_state(subproblem, state_param_out, state_in, state_out)
+        state_out = apply_rule(stage, decision_rules, uncertainty, state_in)
+        objective_value += simulate_stage(subproblem, state_param_in, state_param_out, uncertainty, state_in, state_out)
+        state_in = DecisionRules.get_next_state(subproblem, state_param_in, state_param_out, state_in, state_out)
     end
     
     # Return final objective value
@@ -152,8 +201,7 @@ function simulate_multistage(
     state_params_in::Vector{Vector{Z}},
     state_params_out::Vector{Vector{Tuple{Z, VariableRef}}},
     uncertainties,
-    states;
-    _objective_value = objective_value #get_objective_no_target_deficit
+    states
     ) where {Z}
     
     for t in  1:length(state_params_in)
@@ -161,54 +209,58 @@ function simulate_multistage(
         # Update state parameters in
         if t == 1
             for (i, state_var) in enumerate(state_params_in[t])
-                set_parameter(det_equivalent, state_var, state[i])
+                set_parameter_value(state_var, state[i])
             end
         end
 
         # Update uncertainty
         for (uncertainty_param, uncertainty_value) in uncertainties[t]
-            set_parameter(det_equivalent, uncertainty_param, uncertainty_value)
+            set_parameter_value(uncertainty_param, uncertainty_value)
         end
 
         # Update state parameters out
         for i in 1:length(state_params_out[t])
             state_var = state_params_out[t][i][1]
-            set_parameter(det_equivalent, state_var, states[t + 1][i])
+            set_parameter_value(state_var, states[t + 1][i])
         end
     end
 
     # Solve det_equivalent
     optimize!(det_equivalent)
 
-    return _objective_value(det_equivalent)
+    return objective_value(det_equivalent)
 end
 
 function simulate_multistage(
-    subproblems::Union{Vector{JuMP.Model}, JuMP.Model},
+    subproblems::JuMP.Model,
     state_params_in::Vector{Vector{U}},
     state_params_out::Vector{Vector{Tuple{U, VariableRef}}},
     initial_state::Vector{T},
     uncertainties,
-    decision_rules;
-    ensure_feasibility=(x_out, x_in, uncertainty) -> x_out,
-    _objective_value=objective_value
+    decision_rules
 ) where {T <: Real, U}
     Flux.reset!.(decision_rules)
-    states = simulate_states(initial_state, uncertainties, decision_rules, ensure_feasibility=ensure_feasibility)
-    return simulate_multistage(subproblems, state_params_in, state_params_out, uncertainties, states; _objective_value=_objective_value)
+    states = simulate_states(initial_state, uncertainties, decision_rules)
+    return simulate_multistage(subproblems, state_params_in, state_params_out, uncertainties, states)
 end
 
 function pdual(v::VariableRef)
     if is_parameter(v)
-        return MOI.get(JuMP.owner_model(v), POI.ParameterDual(), v)
+        # Prefer objective sensitivity via dual(v) when available (DiffOpt >= objective-sensitivity API)
+        try
+            return dual(v)
+        catch
+            # Fallback to POI's ParameterDual for older setups
+            return MOI.get(JuMP.owner_model(v), POI.ParameterDual(), v)
+        end
     else
-        return dual(FixRef(v))
+        error("Variable is not a parameter")
     end
 end
-pdual(v::ConstraintRef) = dual(v) # this needs to be fixed to not depend com how the constraint is created
+
 pdual(vs::Vector) = [pdual(v) for v in vs]
 
-function rrule(::typeof(simulate_stage), subproblem, state_param_in, state_param_out, uncertainty, state_in, state_out)
+function ChainRulesCore.rrule(::typeof(simulate_stage), subproblem, state_param_in, state_param_out, uncertainty, state_in, state_out)
     y = simulate_stage(subproblem, state_param_in, state_param_out, uncertainty, state_in, state_out)
     function _pullback(Δy)
         return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), pdual.(state_param_in) * Δy, pdual.([s[1] for s in state_param_out]) * Δy)
@@ -216,8 +268,8 @@ function rrule(::typeof(simulate_stage), subproblem, state_param_in, state_param
     return y, _pullback
 end
 
-# Define rrule of simulate_multistage
-function rrule(::typeof(simulate_multistage), det_equivalent, state_params_in, state_params_out, uncertainties, states)
+# Define ChainRulesCore.rrule of simulate_multistage
+function ChainRulesCore.rrule(::typeof(simulate_multistage), det_equivalent::JuMP.Model, state_params_in, state_params_out, uncertainties, states)
     y = simulate_multistage(det_equivalent, state_params_in, state_params_out, uncertainties, states)
     function _pullback(Δy)
         Δ_states = similar(states)
@@ -230,8 +282,7 @@ function rrule(::typeof(simulate_multistage), det_equivalent, state_params_in, s
     return y, _pullback
 end
 
-function sample(uncertainty_samples::Vector{Tuple{VariableRef, Vector}})
-    T = eltype(uncertainty_samples[1][2])
+function sample(uncertainty_samples::Vector{Tuple{VariableRef, Vector{T}}}) where {T<:Real}
     uncertainty_sample = Vector{Tuple{VariableRef,T}}(undef, length(uncertainty_samples))
     for i in 1:length(uncertainty_samples)
         uncertainty_sample[i] = (uncertainty_samples[i][1], rand(uncertainty_samples[i][2]))
@@ -239,57 +290,48 @@ function sample(uncertainty_samples::Vector{Tuple{VariableRef, Vector}})
     return uncertainty_sample
 end
 
-sample(uncertainty_samples::Vector{Vector{Tuple{VariableRef, Vector}}}) = [sample(uncertainty_samples[t]) for t in 1:length(uncertainty_samples)]
-
-# function train_multistage(model, initial_state, subproblems::Vector{JuMP.Model}, 
-#     state_params_in, state_params_out, uncertainty_sampler; 
-#     num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01), ensure_feasibility=(x_out, x_in, uncertainty) -> x_out,
-#     adjust_hyperparameters=(iter, opt_state, num_train_per_batch) -> num_train_per_batch,
-#     record_loss=(iter, model, loss, tag) -> begin println("tag: $tag, Iter: $iter, Loss: $loss")
-#         return false
-#     end,
-#     get_objective_no_target_deficit=get_objective_no_target_deficit
-# )
-#     # Initialise the optimiser for this model:
-#     opt_state = Flux.setup(optimizer, model)
+function sample(uncertainty_samples::Vector{Vector{Tuple{VariableRef, Vector{T}}}}) where {T<:Real}
+    [sample(uncertainty_samples[t]) for t in 1:length(uncertainty_samples)]
+end
 
-#     for iter in 1:num_batches
-#         num_train_per_batch = adjust_hyperparameters(iter, opt_state, num_train_per_batch)
-#         # Sample uncertainties
-#         uncertainty_samples = [sample(uncertainty_sampler) for _ in 1:num_train_per_batch]
-#         uncertainty_samples_vec = [[collect(values(uncertainty_sample[j])) for j in 1:length(uncertainty_sample)] for uncertainty_sample in uncertainty_samples]
+function train_multistage(model, initial_state, subproblems::Vector{JuMP.Model}, 
+    state_params_in, state_params_out, uncertainty_sampler; 
+    num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01),
+    adjust_hyperparameters=(iter, opt_state, num_train_per_batch) -> num_train_per_batch,
+    record_loss=(iter, model, loss, tag) -> begin println("tag: $tag, Iter: $iter, Loss: $loss")
+        return false
+    end,
+    get_objective_no_target_deficit=get_objective_no_target_deficit
+)
+    # Initialise the optimiser for this model:
+    opt_state = Flux.setup(optimizer, model)
 
-#         # Calculate the gradient of the objective
-#         # with respect to the parameters within the model:
-#         objective = 0.0
-#         eval_loss = 0.0
-#         grads = Flux.gradient(model) do m
-#             for s in 1:num_train_per_batch
-#                 Flux.reset!(m)
-#                 # m(initial_state) # Breaks Everything
-#                 # state_in = initial_state
-#                 for (j, subproblem) in enumerate(subproblems)
-#                     state_out = m(uncertainty_samples_vec[s][j])
-#                     state_out = ensure_feasibility(state_out, state_in, uncertainty_samples_vec[s][j])
-#                     objective += simulate_stage(subproblem, state_params_in[j], state_params_out[j], uncertainty_samples[s][j], state_in, state_out)
-#                     eval_loss += get_objective_no_target_deficit(subproblem)
-#                     state_in = get_next_state(subproblem, state_params_out[j], state_in, state_out)
-#                 end
-#             end
-#             objective /= num_train_per_batch
-#             eval_loss /= num_train_per_batch
-#             return objective
-#         end
-#         record_loss(iter, model, eval_loss, "metrics/loss") && break
-#         record_loss(iter, model, objective, "metrics/training_loss") && break
+    for iter in 1:num_batches
+        num_train_per_batch = adjust_hyperparameters(iter, opt_state, num_train_per_batch)
+        # Sample uncertainties
+        uncertainty_samples = [sample(uncertainty_sampler) for _ in 1:num_train_per_batch]
+        objective = 0.0
+        eval_loss = 0.0
+        grads = Flux.gradient(model) do m
+            for s in 1:num_train_per_batch
+                Flux.reset!(m)
+                objective += simulate_multistage(subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples[s], m)
+                eval_loss += get_objective_no_target_deficit(subproblems)
+            end
+            objective /= num_train_per_batch
+            eval_loss /= num_train_per_batch
+            return objective
+        end
+        record_loss(iter, model, eval_loss, "metrics/loss") && break
+        record_loss(iter, model, objective, "metrics/training_loss") && break
 
-#         # Update the parameters so as to reduce the objective,
-#         # according the chosen optimisation rule:
-#         Flux.update!(opt_state, model, grads[1])
-#     end
+        # Update the parameters so as to reduce the objective,
+        # according the chosen optimisation rule:
+        Flux.update!(opt_state, model, grads[1])
+    end
     
-#     return model
-# end
+    return model
+end
 
 function sim_states(t, m, initial_state, uncertainty_sample_vec)
     if t == 1
@@ -360,7 +402,7 @@ end
 
 # function train_multistage(models::Vector, initial_state, subproblems::Vector{JuMP.Model}, 
 #     state_params_in, state_params_out, uncertainty_sampler; 
-#     num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01), ensure_feasibility=(x_out, x_in, uncertainty) -> x_out,
+#     num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01),
 #     adjust_hyperparameters=(iter, opt_state, num_train_per_batch) -> num_train_per_batch,
 #     record_loss=(iter, model, loss, tag) -> begin println("tag: $tag, Iter: $iter, Loss: $loss")
 #         return false
@@ -388,7 +430,7 @@ end
 #                 states = m(uncertainty_samples_vec[s])
 #                 state_in = initial_state
 #                 for (j, subproblem) in enumerate(subproblems)
-#                     state_out = ensure_feasibility(states[j], state_in, uncertainty_samples_vecs[s][j])
+#                     state_out = states[j]
 #                     objective += simulate_stage(subproblem, state_params_in[j], state_params_out[j], uncertainty_samples[s][j], state_in, state_out)
 #                     eval_loss += get_objective_no_target_deficit(subproblem)
 #                     state_in = get_next_state(subproblem, state_params_out[j], state_in, state_out)
@@ -412,7 +454,7 @@ end
 
 function train_multistage(models::Vector, initial_state, det_equivalent::JuMP.Model, 
     state_params_in, state_params_out, uncertainty_sampler; 
-    num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01), ensure_feasibility=(x_out, x_in, uncertainty) -> x_out,
+    num_batches=100, num_train_per_batch=32, optimizer=Flux.Adam(0.01),
     adjust_hyperparameters=(iter, opt_state, num_train_per_batch) -> num_train_per_batch,
     record_loss=(iter, model, loss, tag) -> begin println("tag: $tag, Iter: $iter, Loss: $loss")
         return false
diff --git a/src/utils.jl b/src/utils.jl
index ac4e32f..34c22d2 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1,33 +1,12 @@
-function variable_to_parameter(model::JuMP.Model, variable::JuMP.VariableRef; initial_value=0.0, deficit=nothing, param_type=:Param)
-    if param_type === :Param
-        parameter = @variable(model; base_name = "_" * name(variable), set=MOI.Parameter(initial_value))
-        # bind the parameter to the variable
-        if isnothing(deficit)
-            @constraint(model, variable == parameter)
-            return parameter
-        else
-            @constraint(model, variable + deficit == parameter)
-            return parameter, variable
-        end
-    elseif param_type === :Cons
-        if isnothing(deficit)
-            c = @constraint(model, variable == 0.0)
-            return c
-        else
-            c = @constraint(model, variable + deficit == 0.0)
-            return c, variable
-        end
+function variable_to_parameter(model::JuMP.Model, variable::JuMP.VariableRef; initial_value=0.0, deficit=nothing)
+    parameter = @variable(model; base_name = "_" * name(variable), set=MOI.Parameter(initial_value))
+    # bind the parameter to the variable
+    if isnothing(deficit)
+        @constraint(model, variable == parameter)
+        return parameter
     else
-        parameter = @variable(model; base_name = "_" * name(variable))
-        # fix(parameter, initial_value)
-        # bind the parameter to the variable
-        if isnothing(deficit)
-            @constraint(model, variable == parameter)
-            return parameter
-        else
-            @constraint(model, variable + deficit == parameter)
-            return parameter, variable
-        end
+        @constraint(model, variable + deficit == parameter)
+        return parameter, variable
     end
 end
 
@@ -259,15 +238,15 @@ function add_child_model_exps!(model::JuMP.Model, subproblem::JuMP.Model, var_sr
 end
 
 "Create Single JuMP.Model from subproblems. rename variables to avoid conflicts by adding [t] at the end of the variable name where t is the subproblem index"
-function deterministic_equivalent(subproblems::Vector{JuMP.Model},
+function deterministic_equivalent!(model::JuMP.Model,
+    subproblems::Vector{JuMP.Model},
     state_params_in::Vector{Vector{Any}},
     state_params_out::Vector{Vector{Tuple{Any, VariableRef}}},
     initial_state::Vector{Float64},
-    uncertainties::Vector{Vector{Tuple{VariableRef, Vector}}};
-    model = JuMP.Model()
+    uncertainties::Vector{Vector{Tuple{VariableRef, Vector{Float64}}}},
 )
     set_objective_sense(model, objective_sense(subproblems[1]))
-    uncertainties_new = Vector{Vector{Tuple{VariableRef, Vector}}}(undef, length(uncertainties))
+    uncertainties_new = Vector{Vector{Tuple{VariableRef, Vector{Float64}}}}(undef, length(uncertainties))
     var_src_to_dest = Dict{VariableRef, VariableRef}()
     for t in 1:length(subproblems)
         DecisionRules.add_child_model_vars!(model, subproblems[t], t, state_params_in, state_params_out, initial_state, var_src_to_dest)
@@ -281,7 +260,7 @@ function deterministic_equivalent(subproblems::Vector{JuMP.Model},
     if uncertainties[1][1][1] isa VariableRef
         # use var_src_to_dest
         for t in 1:length(subproblems)
-            uncertainties_new[t] = Vector{Tuple{VariableRef, Vector}}(undef, length(uncertainties[t]))
+            uncertainties_new[t] = Vector{Tuple{VariableRef, Vector{Float64}}}(undef, length(uncertainties[t]))
             for (i, tup) in enumerate(uncertainties[t])
                 ky, val = tup
                 uncertainties_new[t][i] = (var_src_to_dest[ky],val)
@@ -290,7 +269,7 @@ function deterministic_equivalent(subproblems::Vector{JuMP.Model},
     else
         # use cons_to_cons
         for t in 1:length(subproblems)
-            uncertainties_new[t] = Vector{Tuple{VariableRef, Vector}}(undef, length(uncertainties[t]))
+            uncertainties_new[t] = Vector{Tuple{VariableRef, Vector{Float64}}}(undef, length(uncertainties[t]))
             for (i, tup) in enumerate(uncertainties[t])
                 ky, val = tup
                 uncertainties_new[t] = (cons_to_cons[t][ky],val)
diff --git a/test/runtests.jl b/test/runtests.jl
index 9bd542a..4eb2cfb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,15 +1,22 @@
 using DecisionRules
 using Test
-using Gurobi
-import ParametricOptInterface as POI
+using SCS
 using JuMP
 using Zygote
 using Flux
 using Random
 
-function build_subproblem(d; state_i_val=5.0, state_out_val=4.0, uncertainty_val=2.0)
-    subproblem = JuMP.Model(() -> POI.Optimizer(Gurobi.Optimizer()))
-    set_attributes(subproblem, "output_flag" => false)
+import Pkg
+
+Pkg.add(;
+    url = "https://github.com/jump-dev/DiffOpt.jl",
+    rev = "ar/dualparameter",
+)
+
+using DiffOpt
+
+function build_subproblem(d; state_i_val=5.0, state_out_val=4.0, uncertainty_val=2.0, subproblem=JuMP.Model())
+    # set_attributes(subproblem, "output_flag" => false)
     @variable(subproblem, x >= 0.0)
     @variable(subproblem, 0.0 <= y <= 8.0)
     @variable(subproblem, 0.0 <= state_out_var <= 8.0)
@@ -21,89 +28,120 @@ function build_subproblem(d; state_i_val=5.0, state_out_val=4.0, uncertainty_val
     @constraint(subproblem, state_out_var == state_in + uncertainty - x)
     @constraint(subproblem, x + y >= d)
     @constraint(subproblem, _deficit == state_out_var - state_out)
-    @constraint(subproblem, [norm_deficit; _deficit] in MOI.NormOneCone(2))
-    @objective(subproblem, Min, 30 * y + norm_deficit * 10^7)
+    @constraint(subproblem, [norm_deficit; _deficit] in SecondOrderCone())
+    @objective(subproblem, Min, 30 * y + norm_deficit * 10^4)
     return subproblem, state_in, state_out, state_out_var, uncertainty
 end
 
 @testset "DecisionRules.jl" begin
-    subproblem1, state_in_1, state_out_1, state_out_var_1, uncertainty_1 = build_subproblem(10)
+    @testset "pdual at infeasibility" begin
+        subproblem1, state_in_1, state_out_1, state_out_var_1, uncertainty_1 = build_subproblem(10; subproblem=DiffOpt.conic_diff_model(SCS.Optimizer), state_out_val=9.0)
+        optimize!(subproblem1)
+        @test DecisionRules.pdual(state_in_1) ≈ -1.0e4 rtol=1.0e-1
+        @test DecisionRules.pdual(state_out_1) ≈ 1.0e4 rtol=1.0e-1
+    end
+
+    subproblem1, state_in_1, state_out_1, state_out_var_1, uncertainty_1 = build_subproblem(10; subproblem=DiffOpt.conic_diff_model(SCS.Optimizer))
 
     optimize!(subproblem1)
 
     @testset "pdual" begin
-        @test DecisionRules.pdual(state_in_1) ≈ -30.0
-        @test DecisionRules.pdual(state_out_1) ≈ 30.0
+        @test DecisionRules.pdual(state_in_1) ≈ -30.0 rtol=1.0e-1
+        @test DecisionRules.pdual(state_out_1) ≈ 30.0 rtol=1.0e-1
     end
 
     @testset "simulate_stage" begin
         inflow = 2.0
-        state_param_in = [state_in_1]
-        state_param_out = [(state_out_1, state_out_var_1)]
-        uncertainty_sample = Dict(uncertainty_1 => inflow)
+        state_param_in = Vector{Any}(undef, 1)
+        state_param_out = Vector{Tuple{Any, VariableRef}}(undef, 1)
+        state_param_in .= [state_in_1]
+        state_param_out .= [(state_out_1, state_out_var_1)]
+        uncertainty_sample = [(uncertainty_1, inflow)]
         state_in_val = [5.0]
         state_out_val = [4.0]
         # Test simulate_stage
-        @test DecisionRules.simulate_stage(subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val) ≈ 210
-        grad = gradient(DecisionRules.simulate_stage, subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val)
-        @test grad[5] ≈ [-30.0]
-        @test grad[6] ≈ [30.0]
+        @test DecisionRules.simulate_stage(subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val) ≈ 210 rtol=1.0e-1
+        grad = Zygote.gradient(DecisionRules.simulate_stage, subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val)
+        @test grad[5][1] ≈ -30.0 rtol=1.0e-1
+        @test grad[6][1] ≈ 30.0 rtol=1.0e-1
+        # Test get next state
+        jac = Zygote.jacobian(DecisionRules.get_next_state, subproblem1, state_param_in, state_param_out, state_in_val, state_out_val)
+        @test jac[4][1] ≈ 0.0 atol=1.0e-4 # ∂next_state/∂state_in
+        @test jac[5][1] ≈ 1.0 rtol=1.0e-1 # ∂next_state/∂state_out
+
         # Train flux DR
         subproblem = subproblem1
         Random.seed!(222)
         m = Chain(Dense(1, 10), Dense(10, 1))
+        # 90 is what we expect after training, so we start above that for a random policy
         @test DecisionRules.simulate_stage(subproblem, state_param_in, state_param_out, uncertainty_sample, state_in_val, m([inflow])) > 90.0
         for _ in 1:2050
             _inflow = rand(1.:5)
-            uncertainty_samp = Dict(uncertainty => _inflow)
+            uncertainty_samp = [(uncertainty_1, _inflow)]
             Flux.train!((m, inflow) -> DecisionRules.simulate_stage(subproblem, state_param_in, state_param_out, uncertainty_sample, state_in_val, m(inflow)), m, [[_inflow] for _ =1:10], Flux.Adam())
         end
+        # since we trained towards 90, we should be close to it now
         @test DecisionRules.simulate_stage(subproblem, state_param_in, state_param_out, uncertainty_sample, state_in_val, m([inflow])) <= 92
     end
 
-    # @testset "get_next_state" begin # WRONG get_next_state needs DiffOpt.jl
-    #     inflow = 2.0
-    #     state_param_in = [state_in_1]
-    #     state_param_out = [(state_out_1, state_out_var_1)]
-    #     uncertainty_sample = Dict(uncertainty_1 => inflow)
-    #     state_in_val = [5.0]
-    #     state_out_val = [4.0]
-    #     DecisionRules.simulate_stage(subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val)
-    #     # Test 1st case
-    #     state_out_res1 = DecisionRules.get_next_state(subproblem1, state_param_out, max_state_out, state_out_val)
-    #     @test state_out_res1 ≈ [4.0]
-    #     jacob = jacobian(DecisionRules.get_next_state, subproblem1, state_param_out, max_state_out, state_out_val)
-    #     state_out_val = state_out_val .+ 0.0001
-    #     DecisionRules.simulate_stage(subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val)
-    #     state_out_res2 = DecisionRules.get_next_state(subproblem1, state_param_out, max_state_out, state_out_val)
-    #     @test state_out_res2 ≈ state_out_res1 .+ 0.0001* jacob[4][1,1]
-    #     # Test 2nd case
-    #     state_out_val = max_state_out .+ 1.0
-    #     DecisionRules.simulate_stage(subproblem1, state_param_in, state_param_out, uncertainty_sample, state_in_val, state_out_val)
-    #     state_out_res3 = DecisionRules.get_next_state(subproblem1, state_param_out, max_state_out, state_out_val)
-    #     jacob = jacobian(DecisionRules.get_next_state, subproblem1, state_param_out, max_state_out, state_out_val)
-    # end
-
-    @test "deterministic_equivalent" begin
+    @testset "simulate_multistage (per-stage)" begin
+        subproblem1, state_in_1, state_out_1, state_out_var_1, uncertainty_1 = build_subproblem(10; subproblem=DiffOpt.conic_diff_model(SCS.Optimizer))
+        subproblem2, state_in_2, state_out_2, state_out_var_2, uncertainty_2 = build_subproblem(10; state_i_val=1.0, state_out_val=9.0, uncertainty_val=2.0, subproblem=DiffOpt.conic_diff_model(SCS.Optimizer))
+
+        subproblems = [subproblem1, subproblem2]
+        state_params_in = Vector{Vector{Any}}(undef, 2)
+        state_params_out = Vector{Vector{Tuple{Any, VariableRef}}}(undef, 2)
+        state_params_in .= [[state_in_1], [state_in_2]]
+        state_params_out .= [[(state_out_1, state_out_var_1)], [(state_out_2, state_out_var_2)]]
+        uncertainty_samples = [[(uncertainty_1, [2.0])], [(uncertainty_2, [1.0])]]
+        initial_state = [5.0]
+
+        function simulate(initial_state, params)
+            i = 0
+            m = (ars...) -> begin i= i+1; return params[i] end
+            DecisionRules.simulate_multistage(subproblems, state_params_in, state_params_out, initial_state, sample(uncertainty_samples), m)
+        end
+        grad = Zygote.gradient(simulate, initial_state, [[4.0], [3.]])
+        @test grad[2][1][1] + grad[2][2][1] ≈ 30.0 atol=1.0e-2
+
+        rand_policy = simulate([5.0], [[4.0], [3.]])
+
+        @test rand_policy ≈ 450 rtol=1.0e-1
+
+        @test simulate([9.0], [[7.], [4.000]]) ≈ 359 rtol=1.0e-1
+
+        m = Chain(Dense(2, 10), Dense(10, 1))
+        train_multistage(m, initial_state, subproblems, state_params_in, state_params_out, uncertainty_samples)
+        obj_val = DecisionRules.simulate_multistage(
+            subproblems, state_params_in, state_params_out, 
+            initial_state, sample(uncertainty_samples), 
+            m
+        )
+
+        @test obj_val < rand_policy
+    end
+
+    @testset "deterministic_equivalent" begin
+        subproblem1, state_in_1, state_out_1, state_out_var_1, uncertainty_1 = build_subproblem(10)
         subproblem2, state_in_2, state_out_2, state_out_var_2, uncertainty_2 = build_subproblem(10; state_i_val=4.0, state_out_val=3.0, uncertainty_val=1.0)
-        optimize!(subproblem2)
-        objective_value(subproblem2)
-        value(state_out_var_2)
 
         subproblems = [subproblem1, subproblem2]
-        state_params_in = [[state_in_1], [state_in_2]]
-        state_params_out = [[(state_out_1, state_out_var_1)], [(state_out_2, state_out_var_2)]]
-        uncertainty_samples = [Dict(uncertainty_1 => [2.0]), Dict(uncertainty_2 => [1.0])]
+        state_params_in = Vector{Vector{Any}}(undef, 2)
+        state_params_out = Vector{Vector{Tuple{Any, VariableRef}}}(undef, 2)
+        state_params_in .= [[state_in_1], [state_in_2]]
+        state_params_out .= [[(state_out_1, state_out_var_1)], [(state_out_2, state_out_var_2)]]
+        uncertainty_samples = [[(uncertainty_1, [2.0])], [(uncertainty_2, [1.0])]]
         initial_state = [5.0]
 
-        det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent(subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
-        set_optimizer(det_equivalent, () -> POI.Optimizer(HiGHS.Optimizer()))
+        det_equivalent, uncertainty_samples = DecisionRules.deterministic_equivalent!(DiffOpt.diff_model(SCS.Optimizer), subproblems, state_params_in, state_params_out, initial_state, uncertainty_samples)
+        # set_optimizer(det_equivalent, DiffOpt.diff_model(SCS.Optimizer))
         JuMP.optimize!(det_equivalent)
-        obj_val = objective_value(det_equivalent)
+        objective_value(det_equivalent)
         DecisionRules.pdual.(state_params_in[1])
         DecisionRules.pdual.(state_params_out[1][1][1])
-        DecisionRules.simulate_multistage(det_equivalent, state_params_in, state_params_out, sample(uncertainty_samples), [[9.0], [7.], [4.000]])
-        grad = gradient(DecisionRules.simulate_multistage, det_equivalent, state_params_in, state_params_out, sample(uncertainty_samples), [[9.0], [7.], [4.0]])
+        obj_val = DecisionRules.simulate_multistage(det_equivalent, state_params_in, state_params_out, sample(uncertainty_samples), [[9.0], [7.], [4.000]])
+        @test obj_val ≈ 450 rtol=1.0e-1
+        grad = Zygote.gradient(DecisionRules.simulate_multistage, det_equivalent, state_params_in, state_params_out, sample(uncertainty_samples), [[9.0], [7.], [4.0]])
 
         m = Chain(Dense(1, 10), Dense(10, 1))
         obj_val = DecisionRules.simulate_multistage(