diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 26de24d..40a3a83 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -38,3 +38,9 @@ jobs:
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v5
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false
diff --git a/Project.toml b/Project.toml
index e843595..a724657 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,12 +23,13 @@ ChainRules = "1.72.6"
 ChainRulesCore = "1.25.2"
 DiffOpt = "0.5.4"
 FillArrays = "1.14.0"
-Flux = "0.14.21"
+Flux = "0.16"
+Functors = "0.5.2"
 JLD2 = "0.6.3"
 JuMP = "1.29.4"
 MathOptInterface = "1.48.0"
-ParametricOptInterface = "0.14.1"
-Zygote = "0.6.77"
+ParametricOptInterface = "0.14.1, 0.15"
+Zygote = "0.6.77, 0.7"
 julia = "~1.9, 1.10"
 
 [extras]
diff --git a/README.md b/README.md
index 9ec8857..3e92710 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,5 @@
 [![Build Status](https://github.com/LearningToOptimize/DecisionRules.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/LearningToOptimize/DecisionRules.jl/actions/workflows/CI.yml?query=branch%3Amain)
+[![codecov](https://codecov.io/gh/LearningToOptimize/DecisionRules.jl/graph/badge.svg)](https://app.codecov.io/gh/LearningToOptimize/DecisionRules.jl)
 
 # DecisionRules.jl
 
diff --git a/examples/Atlas/train_dr_atlas.jl b/examples/Atlas/train_dr_atlas.jl
index bf7297d..6181c25 100644
--- a/examples/Atlas/train_dr_atlas.jl
+++ b/examples/Atlas/train_dr_atlas.jl
@@ -82,6 +82,7 @@ println("Number of stages: $(N-1)")
 lg = WandbLogger(
     project = "DecisionRules-Atlas",
     name = save_file,
+    save_code = false,
     config = Dict(
         "N" => N,
         "h" => h,
diff --git a/examples/Atlas/train_dr_atlas_det_eq.jl b/examples/Atlas/train_dr_atlas_det_eq.jl
index 16c97af..07993a1 100644
--- a/examples/Atlas/train_dr_atlas_det_eq.jl
+++ b/examples/Atlas/train_dr_atlas_det_eq.jl
@@ -131,6 +131,7 @@ println("Number of stages: $(N-1)")
 lg = WandbLogger(
     project = "DecisionRules-Atlas",
     name = save_file,
+    save_code = false,
     config = Dict(
         "N" => N,
         "h" => h,
diff --git a/examples/Atlas/train_dr_atlas_multipleshooting.jl b/examples/Atlas/train_dr_atlas_multipleshooting.jl
index 8086bc2..650c9cf 100644
--- a/examples/Atlas/train_dr_atlas_multipleshooting.jl
+++ b/examples/Atlas/train_dr_atlas_multipleshooting.jl
@@ -87,6 +87,7 @@ println("Window size: $window_size")
 lg = WandbLogger(
     project = "DecisionRules-Atlas",
     name = save_file,
+    save_code = false,
     config = Dict(
         "N" => N,
         "h" => h,
diff --git a/examples/HydroPowerModels/run_sddp.jl b/examples/HydroPowerModels/run_sddp.jl
index 6e765a1..c9fc7ed 100644
--- a/examples/HydroPowerModels/run_sddp.jl
+++ b/examples/HydroPowerModels/run_sddp.jl
@@ -74,6 +74,7 @@ end
 lg = WandbLogger(
     project = "HydroPowerModels",
     name = save_file,
+    save_code = false,
 )
 
 # ## Train
diff --git a/examples/HydroPowerModels/run_sddp_inconsistent.jl b/examples/HydroPowerModels/run_sddp_inconsistent.jl
index bb43fea..bb9029d 100644
--- a/examples/HydroPowerModels/run_sddp_inconsistent.jl
+++ b/examples/HydroPowerModels/run_sddp_inconsistent.jl
@@ -84,6 +84,7 @@ end
 lg = WandbLogger(
     project = "HydroPowerModels",
     name = save_file,
+    save_code = false,
 )
 
 # global num_trys = 0
diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels.jl b/examples/HydroPowerModels/train_dr_hydropowermodels.jl
index 66a5ef6..865275b 100644
--- a/examples/HydroPowerModels/train_dr_hydropowermodels.jl
+++ b/examples/HydroPowerModels/train_dr_hydropowermodels.jl
@@ -40,6 +40,9 @@ optimizers = [Flux.Adam()]
 pre_trained_model = nothing
 penalty_l2 = :auto
 penalty_l1 = :auto
+# Annealed target-penalty multipliers (relative to the :auto base above); set to `nothing`
+# to train with the constant penalties the models were built with.
+penalty_schedule = :default_annealed
 
 # Build MSP with deterministic equivalent formulation
 subproblems, state_params_in, state_params_out, uncertainty_samples, initial_state, max_volume = build_hydropowermodels(    
@@ -72,6 +75,7 @@ num_hydro = length(initial_state)
 lg = WandbLogger(
     project = "RL",
     name = save_file,
+    save_code = false,
     config = Dict(
         "layers" => layers,
         "activation" => string(activation),
@@ -81,6 +85,7 @@ lg = WandbLogger(
         "training_method" => "deterministic_equivalent",
         "penalty_l1" => string(penalty_l1),
         "penalty_l2" => string(penalty_l2),
+        "penalty_schedule" => string(penalty_schedule),
         "num_epochs" => string(num_epochs),
         "num_batches" => string(num_batches),
         "num_train_per_batch" => string(_num_train_per_batch)
@@ -118,31 +123,24 @@ model_path = joinpath(model_dir, save_file * ".jld2")
 save_control = SaveBest(best_obj, model_path)
 convergence_criterium = StallingCriterium(100, best_obj, 0)
 
-adjust_hyperparameters = (iter, opt_state, num_train_per_batch) -> begin
-    if iter % 2100 == 0
-        num_train_per_batch = num_train_per_batch * 2
-    end
-    return num_train_per_batch
-end
-
-# Train Model using deterministic equivalent
-for iter in 1:num_epochs
-    num_train_per_batch = _num_train_per_batch
-    train_multistage(models, initial_state, det_equivalent, state_params_in, state_params_out, uncertainty_samples; 
-        num_batches=num_batches,
-        num_train_per_batch=num_train_per_batch,
-        optimizer=optimizers[floor(Int, min(iter, length(optimizers)))],
-        record_loss= (iter, model, loss, tag) -> begin
-            if tag == "metrics/training_loss"
-                save_control(iter, model, loss)
-                record_loss(iter, model, loss, tag)
-                return convergence_criterium(iter, model, loss)
-            end
-            return record_loss(iter, model, loss, tag)
-        end,
-        adjust_hyperparameters=adjust_hyperparameters
-    )
-end
+# Train Model using deterministic equivalent.
+# A single call over num_epochs*num_batches batches so the penalty schedule spans the whole
+# run (this also keeps one optimizer state throughout, and a `true` return from the record
+# callback now stops the whole run).
+train_multistage(models, initial_state, det_equivalent, state_params_in, state_params_out, uncertainty_samples;
+    num_batches=num_epochs * num_batches,
+    num_train_per_batch=_num_train_per_batch,
+    optimizer=first(optimizers),
+    record_loss= (iter, model, loss, tag) -> begin
+        if tag == "metrics/training_loss"
+            save_control(iter, model, loss)
+            record_loss(iter, model, loss, tag)
+            return convergence_criterium(iter, model, loss)
+        end
+        return record_loss(iter, model, loss, tag)
+    end,
+    penalty_schedule=penalty_schedule
+)
 
 # Finish the run
 close(lg)
diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl b/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl
index 58a59c4..429e0e6 100644
--- a/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl
+++ b/examples/HydroPowerModels/train_dr_hydropowermodels_multipleshooting.jl
@@ -39,6 +39,9 @@ optimizers = [Flux.Adam()]
 pre_trained_model = nothing
 penalty_l2 = :auto
 penalty_l1 = nothing
+# Annealed target-penalty multipliers (relative to the :auto base above); set to `nothing`
+# to train with the constant penalties the models were built with.
+penalty_schedule = :default_annealed
 
 # Build MSP using subproblems (not deterministic equivalent)
 
@@ -69,6 +72,7 @@ num_hydro = length(initial_state)
 lg = WandbLogger(
     project = "RL",
     name = save_file,
+    save_code = false,
     config = Dict(
         "layers" => layers,
         "activation" => string(activation),
@@ -78,6 +82,7 @@ lg = WandbLogger(
         "window_size" => string(window_size),
         "penalty_l1" => string(penalty_l1),
         "penalty_l2" => string(penalty_l2),
+        "penalty_schedule" => string(penalty_schedule),
         "num_epochs" => string(num_epochs),
         "num_batches" => string(num_batches),
         "num_train_per_batch" => string(_num_train_per_batch),
@@ -132,35 +137,28 @@ model_path = joinpath(model_dir, save_file * ".jld2")
 save_control = SaveBest(best_obj, model_path)
 convergence_criterium = StallingCriterium(200, best_obj, 0)
 
-adjust_hyperparameters = (iter, opt_state, num_train_per_batch) -> begin
-    if iter % 2100 == 0
-        num_train_per_batch = num_train_per_batch * 2
-    end
-    return num_train_per_batch
-end
-
-# Train Model using multiple shooting
-for iter in 1:num_epochs
-    num_train_per_batch = _num_train_per_batch
-    train_multiple_shooting(
-        models,
-        initial_state,
-        windows,
-        () -> uncertainty_samples;
-        num_batches=num_batches,
-        num_train_per_batch=num_train_per_batch,
-        optimizer=optimizers[floor(Int, min(iter, length(optimizers)))],
-        record_loss=(iter, model, loss, tag) -> begin
-            if tag == "metrics/training_loss"
-                save_control(iter, model, loss)
-                record_loss(iter, model, loss, tag)
-                return convergence_criterium(iter, model, loss)
-            end
-            return record_loss(iter, model, loss, tag)
-        end,
-        adjust_hyperparameters=adjust_hyperparameters
-    )
-end
+# Train Model using multiple shooting.
+# A single call over num_epochs*num_batches batches so the penalty schedule spans the whole
+# run (this also keeps one optimizer state throughout, and a `true` return from the record
+# callback now stops the whole run).
+train_multiple_shooting(
+    models,
+    initial_state,
+    windows,
+    () -> uncertainty_samples;
+    num_batches=num_epochs * num_batches,
+    num_train_per_batch=_num_train_per_batch,
+    optimizer=first(optimizers),
+    record_loss=(iter, model, loss, tag) -> begin
+        if tag == "metrics/training_loss"
+            save_control(iter, model, loss)
+            record_loss(iter, model, loss, tag)
+            return convergence_criterium(iter, model, loss)
+        end
+        return record_loss(iter, model, loss, tag)
+    end,
+    penalty_schedule=penalty_schedule
+)
 
 # Finish the run
 close(lg)
diff --git a/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl b/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl
index eec771b..cf37814 100644
--- a/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl
+++ b/examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl
@@ -45,6 +45,9 @@ optimizers= [Flux.Adam()] # Flux.Adam(0.01), Flux.Descent(0.1), Flux.RMSProp(0.0
 pre_trained_model = nothing #joinpath(HydroPowerModels_dir, case_name, formulation, "models", "case3-ACPPowerModel-h48-2024-05-18T10:16:25.117.jld2")
 penalty_l2 = :auto
 penalty_l1 = :auto
+# Annealed target-penalty multipliers (relative to the :auto base above); set to `nothing`
+# to train with the constant penalties the models were built with.
+penalty_schedule = :default_annealed
 num_eval_scenarios = 4                   # fixed held-out scenarios for the rollout evaluation
 eval_every = 25                          # rollout-evaluate every eval_every batches
 
@@ -72,6 +75,7 @@ num_hydro = length(initial_state)
 lg = WandbLogger(
     project = "RL",
     name = save_file,
+    save_code = false,
     config = Dict(
         "layers" => layers,
         "activation" => string(activation),
@@ -81,6 +85,7 @@ lg = WandbLogger(
         "training_method" => "subproblems",
         "penalty_l1" => string(penalty_l1),
         "penalty_l2" => string(penalty_l2),
+        "penalty_schedule" => string(penalty_schedule),
         "num_epochs" => string(num_epochs),
         "num_batches" => string(num_batches),
         "num_train_per_batch" => string(_num_train_per_batch),
@@ -122,40 +127,33 @@ eval_scenarios = [DecisionRules.sample(uncertainty_samples) for _ in 1:num_eval_
 rollout_evaluation = RolloutEvaluation(subproblems, state_params_in, state_params_out,
     initial_state, eval_scenarios; stride=eval_every)
 
-adjust_hyperparameters = (iter, opt_state, num_train_per_batch) -> begin
-    if iter % 2100 == 0
-        num_train_per_batch = num_train_per_batch * 2
-    end
-    return num_train_per_batch
-end
-
-# Train Model using subproblems (not deterministic equivalent)
-for iter in 1:num_epochs
-    num_train_per_batch = _num_train_per_batch
-    train_multistage(models, initial_state, subproblems, state_params_in, state_params_out, uncertainty_samples; 
-        num_batches=num_batches,
-        num_train_per_batch=num_train_per_batch,
-        optimizer=optimizers[floor(min(iter, length(optimizers)))],
-        record=(sample_log, iter, model) -> begin
-            training_loss = mean(sample_log.objectives)
+# Train Model using subproblems (not deterministic equivalent).
+# A single call over num_epochs*num_batches batches so the penalty schedule spans the whole
+# run (this also keeps one optimizer state throughout, and a `true` return from the record
+# callback now stops the whole run).
+train_multistage(models, initial_state, subproblems, state_params_in, state_params_out, uncertainty_samples;
+    num_batches=num_epochs * num_batches,
+    num_train_per_batch=_num_train_per_batch,
+    optimizer=first(optimizers),
+    record=(sample_log, iter, model) -> begin
+        training_loss = mean(sample_log.objectives)
+        Wandb.log(lg, Dict(
+            "metrics/loss" => mean(sample_log.objectives_no_deficit),
+            "metrics/training_loss" => training_loss,
+        ))
+        rollout_evaluation(iter, model)
+        if iter % eval_every == 0
             Wandb.log(lg, Dict(
-                "metrics/loss" => mean(sample_log.objectives_no_deficit),
-                "metrics/training_loss" => training_loss,
+                "metrics/rollout_objective_no_deficit" => rollout_evaluation.last_objective_no_deficit,
+                "metrics/rollout_target_violation_share" => rollout_evaluation.last_violation_share,
             ))
-            rollout_evaluation(iter, model)
-            if iter % eval_every == 0
-                Wandb.log(lg, Dict(
-                    "metrics/rollout_objective_no_deficit" => rollout_evaluation.last_objective_no_deficit,
-                    "metrics/rollout_target_violation_share" => rollout_evaluation.last_violation_share,
-                ))
-            end
-            save_control(iter, model, training_loss)
-            return convergence_criterium(iter, model, training_loss)
-        end,
-        # ensure_feasibility=(x_out, x_in, uncertainty) -> ensure_feasibility(x_out, x_in, uncertainty, max_volume),
-        adjust_hyperparameters=adjust_hyperparameters
-    )
-end
+        end
+        save_control(iter, model, training_loss)
+        return convergence_criterium(iter, model, training_loss)
+    end,
+    # ensure_feasibility=(x_out, x_in, uncertainty) -> ensure_feasibility(x_out, x_in, uncertainty, max_volume),
+    penalty_schedule=penalty_schedule
+)
 
 # Finish the run
 close(lg)
diff --git a/examples/HydroPowerModels/train_dr_l2O_supervised.jl b/examples/HydroPowerModels/train_dr_l2O_supervised.jl
index 9af48e2..ea202e3 100644
--- a/examples/HydroPowerModels/train_dr_l2O_supervised.jl
+++ b/examples/HydroPowerModels/train_dr_l2O_supervised.jl
@@ -108,6 +108,7 @@ save_control = SaveBest(100, model_path)
 lg = WandbLogger(
     project = "HydroPowerModels",
     name = save_file,
+    save_code = false,
     config = Dict(
         "Supervised" => "Yes",
         "optimizer" => "Adam"
diff --git a/examples/Quadruped/train_policy.jl b/examples/Quadruped/train_policy.jl
index 9137c4a..5c1b44c 100644
--- a/examples/Quadruped/train_policy.jl
+++ b/examples/Quadruped/train_policy.jl
@@ -519,6 +519,7 @@ function main()
     lg = WandbLogger(
         project = "QuadrupedPolicy",
         name = run_name,
+        save_code = false,
         config = Dict(
             "state_dim" => config.state_dim,
             "action_dim" => config.action_dim,
diff --git a/src/DecisionRules.jl b/src/DecisionRules.jl
index 8247139..16f12db 100644
--- a/src/DecisionRules.jl
+++ b/src/DecisionRules.jl
@@ -11,6 +11,7 @@ using DiffOpt
 using Logging
 
 export simulate_multistage, sample, train_multistage, simulate_states, simulate_stage, dense_multilayer_nn, variable_to_parameter, create_deficit!,
+    default_annealed_schedule,
     SampleLog, default_record, RolloutEvaluation,
     SaveBest, find_variables, compute_parameter_dual, StallingCriterium, policy_input_dim, normalize_recur_state,
     StateConditionedPolicy, state_conditioned_policy, materialize_tangent,
diff --git a/src/dense_multilayer_nn.jl b/src/dense_multilayer_nn.jl
index b7b2eb3..fef5079 100644
--- a/src/dense_multilayer_nn.jl
+++ b/src/dense_multilayer_nn.jl
@@ -13,7 +13,13 @@ Create a multi-layer neural network with the specified architecture.
 - `activation`: Activation function (default: Flux.relu)
 - `dense`: Layer type (Dense, LSTM, etc.)
 """
-function dense_multilayer_nn(num_inputs::Int, num_outputs::Int, layers::Vector{Int}; activation=Flux.relu, dense=Dense)
+function dense_multilayer_nn(
+    num_inputs::Int,
+    num_outputs::Int,
+    layers::Vector{Int};
+    activation = Flux.relu,
+    dense = Dense,
+)
     if length(layers) == 0
         if dense == LSTM
             return dense(num_inputs, num_outputs)
@@ -21,11 +27,11 @@ function dense_multilayer_nn(num_inputs::Int, num_outputs::Int, layers::Vector{I
         return dense(num_inputs, num_outputs, activation)
     end
     midlayers = []
-    for i in 1:length(layers) - 1
+    for i = 1:(length(layers)-1)
         if dense == LSTM
-            push!(midlayers, dense(layers[i], layers[i + 1]))
+            push!(midlayers, dense(layers[i], layers[i+1]))
         else
-            push!(midlayers, dense(layers[i], layers[i + 1], activation))
+            push!(midlayers, dense(layers[i], layers[i+1], activation))
         end
     end
     first_layer = if dense == LSTM
@@ -82,24 +88,30 @@ end
     StateConditionedPolicy
 
 A policy architecture that separates temporal encoding from state conditioning:
-- LSTM/RNN: Encodes only the uncertainty sequence (temporal dependencies)
-- Dense: Combines LSTM output with previous state to produce next state
+- `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
 
-This design allows the LSTM to work independently of the state recurrence,
-making it more memory-efficient and compatible with the original training loop.
+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...]
 """
-struct StateConditionedPolicy{E, C}
-    encoder::E          # LSTM/RNN that processes uncertainty only
+mutable struct StateConditionedPolicy{E,C,S}
+    encoder::E          # Recurrent cell, or Chain of cells, that processes uncertainty only
     combiner::C         # Dense that combines encoder output with previous state
+    state::S            # Encoder recurrent state, carried across calls
     n_uncertainty::Int  # Number of uncertainty dimensions
     n_state::Int        # Number of state dimensions
 end
 
-# Use Functors.@functor with explicit trainable fields (excludes n_uncertainty, n_state)
+# Use Functors.@functor with explicit trainable fields (excludes state, n_uncertainty, n_state)
 # We use Functors.@functor instead of Flux.@layer to avoid MutableTangent issues with LSTM
-Functors.@functor StateConditionedPolicy (encoder, combiner,)
+Functors.@functor StateConditionedPolicy (encoder, combiner)
 
 """
     materialize_tangent(x)
@@ -140,26 +152,71 @@ function materialize_tangent(ref::Base.RefValue)
     return materialize_tangent(ref[])
 end
 
+"""
+    _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.
+"""
+_as_cell(layer) = hasfield(typeof(layer), :cell) ? layer.cell : layer
+
+"""
+    _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.
+"""
+_init_recurrent_state(cell) = Flux.initialstates(cell)
+_init_recurrent_state(chain::Chain) = map(_init_recurrent_state, chain.layers)
+
+"""
+    _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.
+"""
+_step_encoder(cell, x, state) = cell(x, state)
+_step_encoder(chain::Chain, x, states::Tuple) =
+    _step_encoder_layers(chain.layers, x, states)
+
+_step_encoder_layers(::Tuple{}, x, ::Tuple{}) = x, ()
+function _step_encoder_layers(layers::Tuple, x, states::Tuple)
+    out, new_state = _step_encoder(first(layers), x, first(states))
+    rest_out, rest_states = _step_encoder_layers(Base.tail(layers), out, Base.tail(states))
+    return rest_out, (new_state, rest_states...)
+end
+
 function (m::StateConditionedPolicy)(x)
     # Split input into uncertainty and previous state
     uncertainty = x[1:m.n_uncertainty]
-    prev_state = x[m.n_uncertainty+1:end]
-    
-    # Encode uncertainty through LSTM (temporal encoding)
-    encoded = m.encoder(uncertainty)
-    
+    prev_state = x[(m.n_uncertainty+1):end]
+
+    # Encode uncertainty through the recurrent encoder, carrying state across calls
+    encoded, new_state = _step_encoder(m.encoder, uncertainty, m.state)
+    m.state = new_state
+
     # Combine encoded uncertainty with previous state
     combined = vcat(encoded, prev_state)
-    
+
     # Output next state prediction
     return m.combiner(combined)
 end
 
-# Reset hidden state of the encoder (for LSTM/RNN)
-Flux.reset!(m::StateConditionedPolicy) = Flux.reset!(m.encoder)
+"""
+    Flux.reset!(m::StateConditionedPolicy)
 
+Reset the encoder's recurrent state to `Flux.initialstates`, e.g. before starting a
+new rollout.
 """
-    state_conditioned_policy(n_uncertainty, n_state, n_output, layers; 
+function Flux.reset!(m::StateConditionedPolicy)
+    m.state = _init_recurrent_state(m.encoder)
+    return nothing
+end
+
+"""
+    state_conditioned_policy(n_uncertainty, n_state, n_output, layers;
                              activation=Flux.relu, encoder_type=Flux.LSTM)
 
 Create a StateConditionedPolicy with the specified architecture.
@@ -170,35 +227,48 @@ Create a StateConditionedPolicy with the specified architecture.
 - `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`: Type of encoder (LSTM, RNN, GRU) (default: LSTM)
+- `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: encoder_type(n_uncertainty => layers[1]) -> ... -> layers[end]
 - Combiner: Dense(layers[end] + n_state => n_output)
 """
-function state_conditioned_policy(n_uncertainty::Int, n_state::Int, n_output::Int, layers::Vector{Int}; 
-                                  activation=Flux.relu, encoder_type=Flux.LSTM)
-    # Build encoder (LSTM stack that processes uncertainty)
+function state_conditioned_policy(
+    n_uncertainty::Int,
+    n_state::Int,
+    n_output::Int,
+    layers::Vector{Int};
+    activation = Flux.relu,
+    encoder_type = Flux.LSTM,
+)
+    # Build encoder (stack of recurrent cells that process uncertainty)
     if length(layers) == 0
-        encoder = encoder_type(n_uncertainty => n_state)
+        encoder = _as_cell(encoder_type(n_uncertainty => n_state))
         encoder_output_dim = n_state
     elseif length(layers) == 1
-        encoder = encoder_type(n_uncertainty => layers[1])
+        encoder = _as_cell(encoder_type(n_uncertainty => layers[1]))
         encoder_output_dim = layers[1]
     else
-        encoder_layers = []
-        push!(encoder_layers, encoder_type(n_uncertainty => layers[1]))
-        for i in 1:length(layers)-1
-            push!(encoder_layers, encoder_type(layers[i] => layers[i+1]))
+        encoder_layers = [_as_cell(encoder_type(n_uncertainty => layers[1]))]
+        for i = 1:(length(layers)-1)
+            push!(encoder_layers, _as_cell(encoder_type(layers[i] => layers[i+1])))
         end
         encoder = Chain(encoder_layers...)
         encoder_output_dim = layers[end]
     end
-    
+
     # Build combiner (Dense that combines encoder output with previous state)
     # Input: [encoded_uncertainty, previous_state]
     # Output: next_state
     combiner = Dense(encoder_output_dim + n_state => n_output, activation)
-    
-    return StateConditionedPolicy(encoder, combiner, n_uncertainty, n_state)
-end
\ No newline at end of file
+
+    return StateConditionedPolicy(
+        encoder,
+        combiner,
+        _init_recurrent_state(encoder),
+        n_uncertainty,
+        n_state,
+    )
+end
diff --git a/src/multiple_shooting.jl b/src/multiple_shooting.jl
index 23d4bd4..18b6e0f 100644
--- a/src/multiple_shooting.jl
+++ b/src/multiple_shooting.jl
@@ -655,9 +655,15 @@ function train_multiple_shooting(
         return false
     end,
     get_objective_no_target_deficit=get_objective_no_target_deficit,
+    penalty_schedule=nothing,
 )
     opt_state = Flux.setup(optimizer, model)
 
+    schedule = _resolve_penalty_schedule(penalty_schedule, num_batches)
+    window_models = [win.model for win in windows]
+    penalty_bases = isnothing(schedule) ? nothing : _check_deficit_penalty_bases(_deficit_penalty_bases(window_models))
+    current_multiplier = NaN
+
     # We only need the uncertainty *structure* here.
     base_uncertainty = uncertainty_sampler()
     # If uncertainty values are vectors (sample sets), draw realized values per iteration.
@@ -669,6 +675,13 @@ function train_multiple_shooting(
     initial_state_f32 = Float32.(initial_state)
 
     for iter in 1:num_batches
+        if !isnothing(schedule)
+            multiplier = _penalty_multiplier_for(schedule, iter)
+            if multiplier != current_multiplier
+                _apply_deficit_penalty_multiplier!(window_models, penalty_bases, multiplier)
+                current_multiplier = multiplier
+            end
+        end
         num_train_per_batch = adjust_hyperparameters(iter, opt_state, num_train_per_batch)
 
         objective = 0.0
diff --git a/src/simulate_multistage.jl b/src/simulate_multistage.jl
index ca596e4..632f274 100644
--- a/src/simulate_multistage.jl
+++ b/src/simulate_multistage.jl
@@ -346,13 +346,25 @@ function train_multistage(model, initial_state, subproblems::Vector{JuMP.Model},
     record_loss=nothing,
     get_objective_no_target_deficit=get_objective_no_target_deficit,
     sample_log=SampleLog(objective_no_deficit_fn=get_objective_no_target_deficit),
-    record=default_record
+    record=default_record,
+    penalty_schedule=nothing
 )
     record = _resolve_record(record, record_loss)
     # Initialise the optimiser for this model:
     opt_state = Flux.setup(optimizer, model)
 
+    schedule = _resolve_penalty_schedule(penalty_schedule, num_batches)
+    penalty_bases = isnothing(schedule) ? nothing : _check_deficit_penalty_bases(_deficit_penalty_bases(subproblems))
+    current_multiplier = NaN
+
     for iter in 1:num_batches
+        if !isnothing(schedule)
+            multiplier = _penalty_multiplier_for(schedule, iter)
+            if multiplier != current_multiplier
+                _apply_deficit_penalty_multiplier!(subproblems, penalty_bases, multiplier)
+                current_multiplier = multiplier
+            end
+        end
         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]
@@ -399,14 +411,26 @@ function train_multistage(model, initial_state, det_equivalent::JuMP.Model,
     record_loss=nothing,
     get_objective_no_target_deficit=get_objective_no_target_deficit,
     sample_log=SampleLog(objective_no_deficit_fn=get_objective_no_target_deficit),
-    record=default_record
+    record=default_record,
+    penalty_schedule=nothing
 )
     record = _resolve_record(record, record_loss)
     # Initialise the optimiser for this model:
     opt_state = Flux.setup(optimizer, model)
     num_stages = length(state_params_in)
 
+    schedule = _resolve_penalty_schedule(penalty_schedule, num_batches)
+    penalty_bases = isnothing(schedule) ? nothing : _check_deficit_penalty_bases(_deficit_penalty_bases(det_equivalent))
+    current_multiplier = NaN
+
     for iter in 1:num_batches
+        if !isnothing(schedule)
+            multiplier = _penalty_multiplier_for(schedule, iter)
+            if multiplier != current_multiplier
+                _apply_deficit_penalty_multiplier!(det_equivalent, penalty_bases, multiplier)
+                current_multiplier = multiplier
+            end
+        end
         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]
diff --git a/src/utils.jl b/src/utils.jl
index d1df096..72bbca8 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -99,6 +99,184 @@ function create_deficit!(model::JuMP.Model, len::Int; penalty_l1=nothing, penalt
     return norm_deficit, _deficit
 end
 
+"""
+    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`](@ref) 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!`](@ref)), so with `penalty=:auto` the effective penalty is
+`multiplier * max |objective coefficient|`.
+"""
+function default_annealed_schedule(num_batches::Int)
+    num_batches >= 1 || throw(ArgumentError("num_batches must be >= 1"))
+    multipliers = [0.1, 1.0, 10.0, 30.0]
+    n = length(multipliers)
+    if num_batches < n
+        mults = multipliers[(n - num_batches + 1):n]
+        return [(i, i, mults[i]) for i in 1:num_batches]
+    end
+    # Phase lengths proportional to 2/2/4/16 over 24; remainder goes to the last phase.
+    lengths = [max(1, round(Int, num_batches * f)) for f in (2 / 24, 2 / 24, 4 / 24)]
+    excess = sum(lengths) + 1 - num_batches
+    i = length(lengths)
+    while excess > 0 && i >= 1
+        take = min(lengths[i] - 1, excess)
+        lengths[i] -= take
+        excess -= take
+        i -= 1
+    end
+    push!(lengths, num_batches - sum(lengths))
+    schedule = Vector{Tuple{Int,Int,Float64}}(undef, n)
+    lo = 1
+    for k in 1:n
+        hi = lo + lengths[k] - 1
+        schedule[k] = (lo, hi, multipliers[k])
+        lo = hi + 1
+    end
+    return schedule
+end
+
+"""
+    _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.
+"""
+function _validate_penalty_schedule(schedule)
+    isempty(schedule) && throw(ArgumentError("penalty_schedule must not be empty"))
+    expected_lo = 1
+    for (lo, hi, mult) in schedule
+        lo == expected_lo || throw(ArgumentError(
+            "penalty_schedule phases must be contiguous starting at batch 1; got phase ($lo, $hi, $mult) where first_batch $expected_lo was expected"))
+        lo <= hi || throw(ArgumentError("penalty_schedule phase ($lo, $hi, $mult) has first_batch > last_batch"))
+        (isfinite(mult) && mult > 0) || throw(ArgumentError("penalty_schedule multipliers must be finite and positive; got $mult"))
+        expected_lo = hi + 1
+    end
+    return schedule
+end
+
+"""
+    _resolve_penalty_schedule(penalty_schedule, num_batches::Int)
+
+Resolve the `penalty_schedule` keyword of [`train_multistage`](@ref) 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`](@ref)`(num_batches)`;
+- any other value is checked with [`_validate_penalty_schedule`](@ref) and returned
+  as-is.
+"""
+_resolve_penalty_schedule(::Nothing, num_batches::Int) = nothing
+function _resolve_penalty_schedule(penalty_schedule::Symbol, num_batches::Int)
+    penalty_schedule === :default_annealed && return default_annealed_schedule(num_batches)
+    throw(ArgumentError("unknown penalty_schedule symbol :$penalty_schedule; use :default_annealed, an explicit vector of (first_batch, last_batch, multiplier) tuples, or nothing"))
+end
+_resolve_penalty_schedule(penalty_schedule, num_batches::Int) = _validate_penalty_schedule(penalty_schedule)
+
+"""
+    _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).
+"""
+function _penalty_multiplier_for(schedule, iter::Int)
+    for (lo, hi, mult) in schedule
+        lo <= iter <= hi && return mult
+    end
+    return schedule[end][3]
+end
+
+"""
+    _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.
+"""
+function _linear_objective_coefficient(model::JuMP.Model, variable::VariableRef)
+    obj = objective_function(model)
+    if obj isa GenericAffExpr
+        return get(obj.terms, variable, 0.0)
+    elseif obj isa GenericQuadExpr
+        return get(obj.aff.terms, variable, 0.0)
+    end
+    throw(ArgumentError("penalty_schedule requires an affine or quadratic objective; got $(typeof(obj))"))
+end
+
+"""
+    _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!`](@ref). A variable counts as a deficit
+variable if `deficit_name` occurs in its name and its linear objective coefficient
+(see [`_linear_objective_coefficient`](@ref)) is nonzero.
+
+Must be called **before** any `penalty_schedule` multiplier is applied, so the
+captured coefficients reflect the as-built penalties.
+"""
+function _deficit_penalty_bases(model::JuMP.Model; deficit_name::AbstractString="norm_deficit")
+    bases = Dict{VariableRef,Float64}()
+    for variable in all_variables(model)
+        if occursin(deficit_name, JuMP.name(variable))
+            coef = _linear_objective_coefficient(model, variable)
+            iszero(coef) || (bases[variable] = coef)
+        end
+    end
+    return bases
+end
+
+function _deficit_penalty_bases(models::Vector{JuMP.Model}; deficit_name::AbstractString="norm_deficit")
+    return [_deficit_penalty_bases(model; deficit_name=deficit_name) for model in models]
+end
+
+"""
+    _check_deficit_penalty_bases(bases) -> typeof(bases)
+
+Return `bases` (from [`_deficit_penalty_bases`](@ref)) unchanged if it has at least one
+entry; otherwise throw `ArgumentError`, since a `penalty_schedule` would then have
+nothing to scale.
+"""
+function _check_deficit_penalty_bases(bases)
+    total = bases isa Dict ? length(bases) : sum(length, bases)
+    total > 0 || throw(ArgumentError(
+        "penalty_schedule was given but no variable matching \"norm_deficit\" with a nonzero objective coefficient was found; build the model(s) with create_deficit! (or equivalent) to use a penalty schedule"))
+    return bases
+end
+
+"""
+    _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`](@ref). Return the mutated model(s).
+"""
+function _apply_deficit_penalty_multiplier!(model::JuMP.Model, bases::Dict{VariableRef,Float64}, multiplier::Real)
+    for (variable, base) in bases
+        set_objective_coefficient(model, variable, multiplier * base)
+    end
+    return model
+end
+
+function _apply_deficit_penalty_multiplier!(models::Vector{JuMP.Model}, bases::Vector, multiplier::Real)
+    for (model, base) in zip(models, bases)
+        _apply_deficit_penalty_multiplier!(model, base, multiplier)
+    end
+    return models
+end
+
 mutable struct SaveBest <: Function
     best_loss::Float64
     model_path::String
@@ -237,10 +415,15 @@ function default_record(sample_log::SampleLog, iter, model)
     return false
 end
 
-# Adapter for the deprecated 4-argument `record_loss(iter, model, loss, tag)` interface.
-# Reproduces the historical two-call contract exactly, including the short-circuit:
-# the "metrics/loss" call runs first and, if it requests a stop, the
-# "metrics/training_loss" call never happens.
+"""
+    _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.
+"""
 function _record_loss_adapter(record_loss)
     return function (sample_log::SampleLog, iter, model)
         record_loss(iter, model, _sequential_mean(sample_log.objectives_no_deficit), "metrics/loss") && return true
@@ -248,6 +431,15 @@ function _record_loss_adapter(record_loss)
     end
 end
 
+"""
+    _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`](@ref)`(record_loss)`. Throw `ArgumentError` if both a custom
+`record` and a `record_loss` are given.
+"""
 function _resolve_record(record, record_loss)
     isnothing(record_loss) && return record
     record === default_record || throw(ArgumentError("pass either `record` or the deprecated `record_loss`, not both"))
@@ -255,11 +447,11 @@ function _resolve_record(record, record_loss)
 end
 
 """
-    _target_violation_share(objective, objective_no_deficit)
+    _target_violation_share(objective::Real, objective_no_deficit::Real) -> Float64
 
-Target-violation share: the realized target-slack penalty divided by the full
-objective. Returns `NaN` when the share is undefined (nonfinite inputs or a full
-objective of magnitude ~0).
+Return the target-violation share, `(objective - objective_no_deficit) / objective`.
+Return `NaN` if either input is nonfinite or `abs(objective) <= 1e-12` (share
+undefined).
 """
 function _target_violation_share(objective::Real, objective_no_deficit::Real)
     slack = objective - objective_no_deficit
diff --git a/test/runtests.jl b/test/runtests.jl
index 4448aed..3e28a55 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -339,6 +339,152 @@ end
         @test objective_value(model6) ≈ 2510.0 rtol=1.0e-2
     end
 
+    @testset "penalty schedule (annealing)" begin
+        @testset "default_annealed_schedule" begin
+            @test default_annealed_schedule(24) == [(1, 2, 0.1), (3, 4, 1.0), (5, 8, 10.0), (9, 24, 30.0)]
+            @test default_annealed_schedule(20) == [(1, 2, 0.1), (3, 4, 1.0), (5, 7, 10.0), (8, 20, 30.0)]
+            @test default_annealed_schedule(4) == [(1, 1, 0.1), (2, 2, 1.0), (3, 3, 10.0), (4, 4, 30.0)]
+            # Fewer batches than multipliers: keep the last multipliers, always end at 30x
+            @test default_annealed_schedule(2) == [(1, 1, 10.0), (2, 2, 30.0)]
+            @test default_annealed_schedule(1) == [(1, 1, 30.0)]
+            @test_throws ArgumentError default_annealed_schedule(0)
+            # Every phase covered, contiguous from batch 1
+            for n in (4, 5, 6, 24, 100, 3000)
+                sched = default_annealed_schedule(n)
+                @test sched[1][1] == 1
+                @test sched[end][2] == n
+                @test all(sched[k+1][1] == sched[k][2] + 1 for k in 1:length(sched)-1)
+            end
+        end
+
+        @testset "schedule resolution and validation" begin
+            sched = default_annealed_schedule(4)
+            @test DecisionRules._penalty_multiplier_for(sched, 1) == 0.1
+            @test DecisionRules._penalty_multiplier_for(sched, 4) == 30.0
+            @test DecisionRules._penalty_multiplier_for(sched, 99) == 30.0  # past-end hold
+            @test DecisionRules._resolve_penalty_schedule(nothing, 10) === nothing
+            @test DecisionRules._resolve_penalty_schedule(:default_annealed, 24) == default_annealed_schedule(24)
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule(:not_a_schedule, 10)
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([], 10)                              # empty
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([(2, 3, 1.0)], 10)                   # starts after 1
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([(1, 2, 1.0), (4, 5, 2.0)], 10)      # gap
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([(1, 3, 1.0), (3, 5, 2.0)], 10)      # overlap
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([(1, 2, -1.0)], 10)                  # nonpositive
+            @test_throws ArgumentError DecisionRules._resolve_penalty_schedule([(2, 1, 1.0)], 10)                   # lo > hi
+        end
+
+        @testset "deficit penalty scaling across create_deficit! modes" begin
+            # L1-only: penalty is the objective coefficient of norm_deficit
+            ml1 = Model()
+            @variable(ml1, yl1 >= 0)
+            @objective(ml1, Min, 10 * yl1)
+            ndl1, _ = create_deficit!(ml1, 2; penalty_l1=100.0)
+            bases = DecisionRules._deficit_penalty_bases(ml1)
+            @test bases == Dict(ndl1 => 100.0)
+            DecisionRules._apply_deficit_penalty_multiplier!(ml1, bases, 10.0)
+            @test coefficient(objective_function(ml1), ndl1) ≈ 1000.0
+            DecisionRules._apply_deficit_penalty_multiplier!(ml1, bases, 1.0)
+            @test coefficient(objective_function(ml1), ndl1) ≈ 100.0
+
+            # L2-only: same entry point
+            ml2 = Model()
+            @variable(ml2, yl2 >= 0)
+            @objective(ml2, Min, 10 * yl2)
+            ndl2, _ = create_deficit!(ml2, 2; penalty_l2=100.0)
+            DecisionRules._apply_deficit_penalty_multiplier!(ml2, DecisionRules._deficit_penalty_bases(ml2), 0.5)
+            @test coefficient(objective_function(ml2), ndl2) ≈ 50.0
+
+            # Both-mode: penalties live in the linking constraint; the objective coefficient
+            # of norm_deficit is 1.0 and scaling it scales L1 and L2 uniformly
+            mb = Model()
+            @variable(mb, yb >= 0)
+            @objective(mb, Min, 10 * yb)
+            ndb, _ = create_deficit!(mb, 2; penalty_l1=100.0, penalty_l2=50.0)
+            basesb = DecisionRules._deficit_penalty_bases(mb)
+            @test basesb == Dict(ndb => 1.0)
+            DecisionRules._apply_deficit_penalty_multiplier!(mb, basesb, 30.0)
+            @test coefficient(objective_function(mb), ndb) ≈ 30.0
+
+            # No matching variable with a schedule active is an error, not a silent no-op
+            mnone = Model()
+            @variable(mnone, ynone >= 0)
+            @objective(mnone, Min, ynone)
+            @test_throws ArgumentError DecisionRules._check_deficit_penalty_bases(
+                DecisionRules._deficit_penalty_bases(mnone))
+        end
+
+        @testset "train_multistage penalty_schedule end-to-end" begin
+            # nothing (default) leaves coefficients untouched
+            sp1, si1, so1, sov1, u1 = build_subproblem(10; subproblem=DiffOpt.conic_diff_model(Ipopt.Optimizer))
+            sp2, si2, so2, sov2, u2 = build_subproblem(10; state_i_val=1.0, state_out_val=9.0, subproblem=DiffOpt.conic_diff_model(Ipopt.Optimizer))
+            sps = [sp1, sp2]
+            spi = Vector{Vector{Any}}(undef, 2); spi .= [[si1], [si2]]
+            spo = Vector{Vector{Tuple{Any, VariableRef}}}(undef, 2)
+            spo .= [[(so1, sov1)], [(so2, sov2)]]
+            usamples = [[(u1, [2.0])], [(u2, [1.0])]]
+            nd1 = variable_by_name(sp1, "norm_deficit")
+            nd2 = variable_by_name(sp2, "norm_deficit")
+            Random.seed!(222)
+            m = Chain(Dense(2, 4), Dense(4, 1))
+            train_multistage(m, [5.0], sps, spi, spo, usamples; num_batches=2, num_train_per_batch=1)
+            @test coefficient(objective_function(sp1), nd1) ≈ 1.0e4
+            @test coefficient(objective_function(sp2), nd2) ≈ 1.0e4
+
+            # Explicit two-phase schedule crosses a boundary on DiffOpt models and the final
+            # coefficients hold the last multiplier times the built base
+            train_multistage(m, [5.0], sps, spi, spo, usamples;
+                num_batches=4, num_train_per_batch=1,
+                penalty_schedule=[(1, 2, 1.0), (3, 4, 3.0)])
+            @test coefficient(objective_function(sp1), nd1) ≈ 3.0e4
+            @test coefficient(objective_function(sp2), nd2) ≈ 3.0e4
+        end
+
+        @testset "deterministic-equivalent overload applies the schedule to the copies" begin
+            sp1, si1, so1, sov1, u1 = build_subproblem(10)
+            sp2, si2, so2, sov2, u2 = build_subproblem(10; state_i_val=4.0, state_out_val=3.0, uncertainty_val=1.0)
+            sps = [sp1, sp2]
+            spi = Vector{Vector{Any}}(undef, 2); spi .= [[si1], [si2]]
+            spo = Vector{Vector{Tuple{Any, VariableRef}}}(undef, 2)
+            spo .= [[(so1, sov1)], [(so2, sov2)]]
+            usamples = [[(u1, [2.0])], [(u2, [1.0])]]
+            det_equivalent, usamples = DecisionRules.deterministic_equivalent!(
+                DiffOpt.nonlinear_diff_model(Ipopt.Optimizer),
+                sps, spi, spo, [5.0], usamples)
+            nd_copies = [v for v in all_variables(det_equivalent) if occursin("norm_deficit", JuMP.name(v))]
+            @test length(nd_copies) == 2  # one renamed copy per stage
+            @test all(coefficient(objective_function(det_equivalent), v) ≈ 1.0e4 for v in nd_copies)
+            Random.seed!(222)
+            m = Chain(Dense(2, 10), Dense(10, 1))
+            train_multistage(m, [5.0], det_equivalent, spi, spo, usamples;
+                num_batches=4, num_train_per_batch=1,
+                penalty_schedule=[(1, 2, 1.0), (3, 4, 3.0)])
+            @test all(coefficient(objective_function(det_equivalent), v) ≈ 3.0e4 for v in nd_copies)
+        end
+
+        @testset "train_multiple_shooting applies the schedule to the window models" begin
+            sp, si, so, sov, u = build_subproblem(10)
+            spi = Vector{Vector{Any}}(undef, 1); spi .= [[si]]
+            spo = Vector{Vector{Tuple{Any, VariableRef}}}(undef, 1)
+            spo .= [[(so, sov)]]
+            usamples = [[(u, [2.0])]]
+            windows = DecisionRules.setup_shooting_windows(
+                [sp], spi, spo, [5.0], usamples;
+                window_size=1,
+                model_factory=() -> DiffOpt.conic_diff_model(Ipopt.Optimizer))
+            wm = windows[1].model
+            nd_window = [v for v in all_variables(wm) if occursin("norm_deficit", JuMP.name(v))]
+            @test length(nd_window) == 1
+            @test coefficient(objective_function(wm), nd_window[1]) ≈ 1.0e4
+            model = Dense(2, 1; bias=false)
+            model.weight .= 0.5f0
+            DecisionRules.train_multiple_shooting(model, [5.0], windows, () -> usamples;
+                num_batches=4, num_train_per_batch=1, optimizer=Flux.Descent(0.0),
+                record_loss=(iter, m, loss, tag) -> false,
+                penalty_schedule=[(1, 2, 1.0), (3, 4, 3.0)])
+            @test coefficient(objective_function(wm), nd_window[1]) ≈ 3.0e4
+        end
+    end
+
     @testset "sample logger and per-batch record" begin
         function build_two_stage()
             sp1, si1, so1, sov1, u1 = build_subproblem(10; subproblem=DiffOpt.conic_diff_model(Ipopt.Optimizer))
@@ -511,7 +657,18 @@ end
             @test length(next_state) == n_output
             prev_state = next_state
         end
-        
+
+        # Test the encoder's recurrent state is carried across calls: repeating the
+        # same input twice in a row must give different outputs (state changed after
+        # the first call), and Flux.reset! must restore the initial-state output.
+        Flux.reset!(policy)
+        repeated_input = rand(Float32, n_uncertainty + n_state)
+        out1 = policy(repeated_input)
+        out2 = policy(repeated_input)
+        @test out1 != out2
+        Flux.reset!(policy)
+        @test policy(repeated_input) ≈ out1
+
         # Test gradient computation with Flux.gradient
         function test_loss(m, n_uncertainty, n_state)
             Flux.reset!(m)
@@ -756,9 +913,11 @@ end
             @variable(subproblems[1], uncertainty in MOI.Parameter(0.1))
             @variable(subproblems[1], state_out)
             @variable(subproblems[1], state_out_var)
+            # Extra free variable so Ipopt has >=1 degree of freedom (n_vars > n_eq_constraints);
+            # otherwise it throws TOO_FEW_DOF before attempting to solve.
+            @variable(subproblems[1], _dof_slack)
             @constraint(subproblems[1], state_out_var == state_in + uncertainty)
             @constraint(subproblems[1], x == state_out_var)
-            @constraint(subproblems[1], state_out_var == state_out)
             @objective(subproblems[1], Min, x)
 
             state_params_in[1] = [state_in]
@@ -1073,9 +1232,11 @@ end
             @variable(subproblems[1], uncertainty in MOI.Parameter(0.0))
             @variable(subproblems[1], state_out in MOI.Parameter(0.0))
             @variable(subproblems[1], state_out_var)
+            # Extra free variable so Ipopt has >=1 degree of freedom (n_vars > n_eq_constraints);
+            # otherwise it throws TOO_FEW_DOF before attempting to solve.
+            @variable(subproblems[1], _dof_slack)
             @constraint(subproblems[1], x == state_in + uncertainty)
             @constraint(subproblems[1], state_out_var == x)
-            @constraint(subproblems[1], state_out_var == state_out)
             @objective(subproblems[1], Min, x)
 
             state_params_in[1] = [state_in]