update and cleanup

Author: BatyLeo

scripts/example_new_metrics.jl

@@ -6,27 +6,48 @@ using InferOpt
 using MLUtils
 using Plots
 
-b = ArgmaxBenchmark()
+b = ArgmaxBenchmark(; seed=42)
 initial_model = generate_statistical_model(b; seed=0)
 maximizer = generate_maximizer(b)
 dataset = generate_dataset(b, 100; seed=0);
 train_dataset, val_dataset = splitobs(dataset; at=(0.5, 0.5));
 
 algorithm = PerturbedImitationAlgorithm(;
-    nb_samples=20, ε=0.1, threaded=true, training_optimizer=Adam()
+    nb_samples=20, ε=0.1, threaded=true, training_optimizer=Adam(), seed=0
 )
 
-validation_metric = FYLLossMetric(algorithm, val_dataset, :validation_loss, maximizer);
+validation_metric = FYLLossMetric(val_dataset, :validation_loss);
+epoch_metric = FunctionMetric(ctx -> ctx.epoch, :current_epoch)
+
+dual_gap_metric = FunctionMetric(:dual_gap, (train_dataset, val_dataset)) do ctx, datasets
+    _train_dataset, _val_dataset = datasets
+    train_gap = compute_gap(b, _train_dataset, ctx.model, ctx.maximizer)
+    val_gap = compute_gap(b, _val_dataset, ctx.model, ctx.maximizer)
+    return (train_gap=train_gap, val_gap=val_gap)
+end
+
+gap_metric = FunctionMetric(:validation_gap, val_dataset) do ctx, data
+    compute_gap(b, data, ctx.model, ctx.maximizer)
+end
+periodic_gap = PeriodicMetric(gap_metric, 5)
+
+gap_metric_offset = FunctionMetric(:delayed_gap, val_dataset) do ctx, data
+    compute_gap(b, data, ctx.model, ctx.maximizer)
+end
+delayed_periodic_gap = PeriodicMetric(gap_metric_offset, 5; offset=10)
+
+# Combine metrics
+metrics = (
+    validation_metric,
+    epoch_metric,
+    dual_gap_metric,       # Outputs both train_gap and val_gap every epoch
+    periodic_gap,          # Outputs validation_gap every 5 epochs
+    delayed_periodic_gap,  # Outputs delayed_gap every 5 epochs starting at epoch 10
+);
 
 model = deepcopy(initial_model)
 history = train_policy!(
-    algorithm,
-    model,
-    maximizer,
-    train_dataset,
-    val_dataset;
-    epochs=50,
-    metrics=(validation_metric,),
+    algorithm, model, maximizer, train_dataset, val_dataset; epochs=50, metrics=metrics
 )
 X_train, Y_train = get(history, :training_loss)
 X_val, Y_val = get(history, :validation_loss)
@@ -44,3 +65,5 @@ plot!(
     label="Validation Loss",
     title="Validation Loss over Epochs",
 )
+
+plot(get(history, :validation_gap); xlabel="Epoch", title="Validation Gap over Epochs")

scripts/main.jl

@@ -0,0 +1,74 @@
+using DecisionFocusedLearningAlgorithms
+using DecisionFocusedLearningBenchmarks
+
+using Flux
+using InferOpt
+using MLUtils
+using Plots
+
+# Create Dynamic Vehicle Scheduling Problem benchmark
+b = DynamicVehicleSchedulingBenchmark(; two_dimensional_features=true)
+
+# Generate dataset and environments
+dataset = generate_dataset(b, 9)
+train_instances, val_instances, test_instances = splitobs(dataset; at=(0.5, 0.3, 0.2))
+
+train_envs = generate_environments(b, train_instances; seed=0)
+val_envs = generate_environments(b, val_instances; seed=1)
+
+# Initialize model and maximizer
+initial_model = generate_statistical_model(b; seed=0)
+maximizer = generate_maximizer(b)
+
+# Define anticipative (expert) policy
+anticipative_policy = (env; reset_env) -> generate_anticipative_solution(b, env; reset_env)
+
+# Configure training algorithm
+algorithm = PerturbedImitationAlgorithm(;
+    nb_samples=10, ε=0.1, threaded=true, training_optimizer=Adam(0.001), seed=0
+)
+
+# Define metrics to track during training
+epoch_metric = FunctionMetric(ctx -> ctx.epoch, :current_epoch)
+
+# You can add validation metrics if you have a validation function
+# For now, we'll just track epochs
+metrics = (epoch_metric,)
+
+# Train using DAgger
+println("Starting DAgger training on Dynamic Vehicle Scheduling Problem...")
+model = deepcopy(initial_model)
+
+history = DAgger_train_model!(
+    model,
+    maximizer,
+    train_envs,
+    val_envs,
+    anticipative_policy;
+    iterations=5,
+    fyl_epochs=10,
+    metrics=metrics,
+    algorithm=algorithm,
+)
+
+# Plot training progress
+X_train, Y_train = get(history, :training_loss)
+plot(
+    X_train,
+    Y_train;
+    xlabel="Epoch",
+    ylabel="Training Loss",
+    label="Training Loss",
+    title="DAgger Training on Dynamic VSP",
+    legend=:topright,
+)
+
+# Plot epoch tracking if available
+if haskey(history, :current_epoch)
+    X_epoch, Y_epoch = get(history, :current_epoch)
+    println("Tracked epochs: ", Y_epoch)
+end
+
+println("\nTraining completed!")
+println("Final training loss: ", Y_train[end])
+println("Total epochs: ", length(Y_train) - 1)  # -1 because epoch 0 is included

scripts/main_dagger.jl

@@ -12,19 +12,30 @@ using ValueHistories: MVHistory
 
 include("utils.jl")
 include("training_context.jl")
-# include("dfl_policy.jl")
-# include("callbacks.jl")
-include("metric.jl")
+
+# Metrics subsystem
+include("metrics/interface.jl")
+include("metrics/accumulators.jl")
+include("metrics/function_metric.jl")
+include("metrics/periodic.jl")
+
 include("fyl.jl")
 include("dagger.jl")
 
-export fyl_train_model!,
-    fyl_train_model, baty_train_model, DAgger_train_model!, DAgger_train_model
-export TrainingCallback, Metric, on_epoch_end, get_metric_names, run_callbacks!
-export TrainingContext, update_context
+export TrainingContext
 
 export AbstractMetric,
-    FYLLossMetric, FunctionMetric, LossAccumulator, reset!, update!, evaluate!, compute
+    FYLLossMetric,
+    FunctionMetric,
+    PeriodicMetric,
+    LossAccumulator,
+    reset!,
+    update!,
+    evaluate!,
+    compute,
+    run_metrics!
+
+export fyl_train_model, baty_train_model, DAgger_train_model!, DAgger_train_model
 export PerturbedImitationAlgorithm, train_policy!
 
 end