Add target-penalty annealing support (penalty_schedule) to the training loops

.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

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]

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
 

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,

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,

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,

examples/HydroPowerModels/run_sddp.jl

@@ -74,6 +74,7 @@ end
 lg = WandbLogger(
     project = "HydroPowerModels",
     name = save_file,
+    save_code = false,
 )
 
 # ## Train

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

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)

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)

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)

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"

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,

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,