Merge pull request #8 from JuliaDecisionFocusedLearning/new-benchmark-interface

Adapt to new benchmark interface

Author: BatyLeo

Project.toml

@@ -19,14 +19,14 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 ValueHistories = "98cad3c8-aec3-5f06-8e41-884608649ab7"
 
 [compat]
-DecisionFocusedLearningBenchmarks = "0.4"
+DecisionFocusedLearningBenchmarks = "0.5.0"
 DocStringExtensions = "0.9.5"
-Flux = "0.16.5"
+Flux = "0.16.9"
 InferOpt = "0.7.1"
 MLUtils = "0.4.8"
 ProgressMeter = "1.11.0"
 Random = "1.11.0"
 Statistics = "1.11.1"
-UnicodePlots = "3.8.1"
-ValueHistories = "0.5.4"
+UnicodePlots = "3.8.2"
+ValueHistories = "0.5.6"
 julia = "1.11"

src/DecisionFocusedLearningAlgorithms.jl

@@ -4,13 +4,13 @@ using DecisionFocusedLearningBenchmarks
 using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Flux, Adam
 using InferOpt: InferOpt, FenchelYoungLoss, PerturbedAdditive, PerturbedMultiplicative
-using MLUtils: splitobs, DataLoader
+using MLUtils: DataLoader
 using ProgressMeter: @showprogress
+using Random: Random, MersenneTwister
 using Statistics: mean
 using UnicodePlots: lineplot
 using ValueHistories: MVHistory
 
-include("utils.jl")
 include("training_context.jl")
 
 include("metrics/interface.jl")
@@ -39,6 +39,7 @@ export AbstractMetric,
     compute!,
     evaluate_metrics!
 
+export AbstractAlgorithm, AbstractImitationAlgorithm
 export PerturbedFenchelYoungLossImitation,
     DAgger, AnticipativeImitation, train_policy!, train_policy
 export AbstractPolicy, DFLPolicy

src/algorithms/abstract_algorithm.jl

@@ -11,6 +11,46 @@ $TYPEDEF
 An abstract type for imitation learning algorithms.
 
 All subtypes must implement:
-- `train_policy!(algorithm::AbstractImitationAlgorithm, model, maximizer, train_data; epochs, metrics)`
+- `train_policy!(algorithm::AbstractImitationAlgorithm, policy::DFLPolicy, train_data; epochs, metrics)`
 """
 abstract type AbstractImitationAlgorithm <: AbstractAlgorithm end
+
+"""
+$TYPEDSIGNATURES
+
+Train a new DFLPolicy on a benchmark using any imitation learning algorithm.
+
+Convenience wrapper that handles dataset generation, model initialization, and policy
+creation. Returns the training history and the trained policy.
+
+For dynamic benchmarks, use the algorithm-specific `train_policy` overload that accepts
+environments and an anticipative policy.
+"""
+function train_policy(
+    algorithm::AbstractImitationAlgorithm,
+    benchmark::AbstractBenchmark;
+    target_policy=nothing,
+    dataset_size=30,
+    epochs=100,
+    metrics::Tuple=(),
+    seed=nothing,
+)
+    dataset = generate_dataset(benchmark, dataset_size; target_policy)
+
+    if any(s -> isnothing(s.y), dataset)
+        error(
+            "Training dataset contains unlabeled samples (y=nothing). " *
+            "Provide a `target_policy` kwarg to label samples during dataset generation.",
+        )
+    end
+
+    model = generate_statistical_model(benchmark; seed)
+    maximizer = generate_maximizer(benchmark)
+    policy = DFLPolicy(model, maximizer)
+
+    history = train_policy!(
+        algorithm, policy, dataset; epochs, metrics, maximizer_kwargs=s -> s.context
+    )
+
+    return history, policy
+end

src/algorithms/supervised/anticipative_imitation.jl

@@ -31,12 +31,11 @@ function train_policy!(
     anticipative_policy,
     epochs=10,
     metrics::Tuple=(),
-    maximizer_kwargs=get_state,
+    maximizer_kwargs=sample -> sample.context,
 )
     # Generate anticipative solutions as training data
     train_dataset = vcat(map(train_environments) do env
-        v, y = anticipative_policy(env; reset_env=true)
-        return y
+        return anticipative_policy(env; reset_env=true)
     end...)
 
     # Delegate to inner algorithm
@@ -62,26 +61,22 @@ Uses anticipative solutions as expert demonstrations.
 """
 function train_policy(
     algorithm::AnticipativeImitation,
-    benchmark::AbstractStochasticBenchmark{true};
+    benchmark::ExogenousDynamicBenchmark;
     dataset_size=30,
-    split_ratio=(0.3, 0.3),
     epochs=10,
     metrics::Tuple=(),
     seed=nothing,
 )
-    # Generate instances and environments
-    dataset = generate_dataset(benchmark, dataset_size)
-    train_instances, validation_instances, _ = splitobs(dataset; at=split_ratio)
-    train_environments = generate_environments(benchmark, train_instances)
+    # Generate environments
+    train_environments = generate_environments(benchmark, dataset_size; seed)
 
     # Initialize model and create policy
     model = generate_statistical_model(benchmark; seed)
     maximizer = generate_maximizer(benchmark)
     policy = DFLPolicy(model, maximizer)
 
     # Define anticipative policy from benchmark
-    anticipative_policy =
-        (env; reset_env) -> generate_anticipative_solution(benchmark, env; reset_env)
+    anticipative_policy = generate_anticipative_solver(benchmark)
 
     # Train policy
     history = train_policy!(

src/algorithms/supervised/dagger.jl

@@ -8,7 +8,7 @@ Reference: <https://arxiv.org/abs/2402.04463>
 # Fields
 $TYPEDFIELDS
 """
-@kwdef struct DAgger{A} <: AbstractImitationAlgorithm
+@kwdef struct DAgger{A,S} <: AbstractImitationAlgorithm
     "inner imitation algorithm for supervised learning"
     inner_algorithm::A = PerturbedFenchelYoungLossImitation()
     "number of DAgger iterations"
@@ -17,6 +17,11 @@ $TYPEDFIELDS
     epochs_per_iteration::Int = 3
     "decay factor for mixing expert and learned policy"
     α_decay::Float64 = 0.9
+    "random seed for the expert/policy mixing coin-flip (nothing = non-reproducible)"
+    seed::S = nothing
+    "maximum dataset size across iterations (nothing keeps all samples,
+    an integer caps to the most recent N samples via FIFO)"
+    max_dataset_size::Union{Int,Nothing} = nothing
 end
 
 """
@@ -34,24 +39,24 @@ function train_policy!(
     train_environments;
     anticipative_policy,
     metrics::Tuple=(),
-    maximizer_kwargs=get_state,
+    maximizer_kwargs=sample -> sample.context,
 )
-    (; inner_algorithm, iterations, epochs_per_iteration, α_decay) = algorithm
+    (; inner_algorithm, iterations, epochs_per_iteration, α_decay, seed) = algorithm
     (; statistical_model, maximizer) = policy
 
+    rng = isnothing(seed) ? MersenneTwister() : MersenneTwister(seed)
     α = 1.0
 
     # Initial dataset from expert demonstrations
     train_dataset = vcat(map(train_environments) do env
-        v, y = anticipative_policy(env; reset_env=true)
-        return y
+        return anticipative_policy(env; reset_env=true)
     end...)
 
     dataset = deepcopy(train_dataset)
 
     # Initialize combined history for all DAgger iterations
     combined_history = MVHistory()
-    global_epoch = 0
+    epoch_offset = 0
 
     for iter in 1:iterations
         println("DAgger iteration $iter/$iterations (α=$(round(α, digits=3)))")
@@ -68,53 +73,26 @@ function train_policy!(
 
         # Merge iteration history into combined history
         for key in keys(iter_history)
-            epochs, values = get(iter_history, key)
-            for i in eachindex(epochs)
-                # Calculate global epoch number
-                if iter == 1
-                    # First iteration: use epochs as-is [0, 1, 2, ...]
-                    global_epoch_value = epochs[i]
-                else
-                    # Later iterations: skip epoch 0 and renumber starting from global_epoch
-                    if epochs[i] == 0
-                        continue  # Skip epoch 0 for iterations > 1
-                    end
-                    # Map epoch 1 → global_epoch, epoch 2 → global_epoch+1, etc.
-                    global_epoch_value = global_epoch + epochs[i] - 1
-                end
-
-                # For the epoch key, use global_epoch_value as both time and value
-                # For other keys, use global_epoch_value as time and original value
-                if key == :epoch
-                    push!(combined_history, key, global_epoch_value, global_epoch_value)
-                else
-                    push!(combined_history, key, global_epoch_value, values[i])
-                end
+            local_epochs, values = get(iter_history, key)
+            for i in eachindex(local_epochs)
+                # Skip epoch 0 for all iterations after the first
+                local_epochs[i] == 0 && epoch_offset > 0 && continue
+                global_e = epoch_offset + local_epochs[i]
+                push!(combined_history, key, global_e, key == :epoch ? global_e : values[i])
             end
         end
 
-        # Update global_epoch for next iteration
-        # After each iteration, advance by the number of non-zero epochs processed
-        if iter == 1
-            # First iteration processes all epochs [0, 1, ..., epochs_per_iteration]
-            # Next iteration should start at epochs_per_iteration + 1
-            global_epoch = epochs_per_iteration + 1
-        else
-            # Subsequent iterations skip epoch 0, so they process epochs_per_iteration epochs
-            # Next iteration should start epochs_per_iteration later
-            global_epoch += epochs_per_iteration
-        end
+        epoch_offset += epochs_per_iteration
 
         # Dataset update - collect new samples using mixed policy
         new_samples = eltype(dataset)[]
         for env in train_environments
             DecisionFocusedLearningBenchmarks.reset!(env; reset_rng=false)
             while !is_terminated(env)
-                x_before = copy(observe(env)[1])
-                _, anticipative_solution = anticipative_policy(env; reset_env=false)
-                p = rand()
+                anticipative_solution = anticipative_policy(env; reset_env=false)
+                p = rand(rng)
                 target = anticipative_solution[1]
-                x, state = observe(env)
+                x, _ = observe(env)
                 if size(target.x) != size(x)
                     @error "Mismatch between expert and observed state" size(target.x) size(
                         x
@@ -124,14 +102,16 @@ function train_policy!(
                 if p < α
                     action = target.y
                 else
-                    x, state = observe(env)
                     θ = statistical_model(x)
                     action = maximizer(θ; maximizer_kwargs(target)...)
                 end
                 step!(env, action)
             end
         end
-        dataset = new_samples  # TODO: replay buffer
+        dataset = vcat(dataset, new_samples)
+        if !isnothing(algorithm.max_dataset_size)
+            dataset = last(dataset, algorithm.max_dataset_size)
+        end
         α *= α_decay  # Decay factor for mixing expert and learned policy
     end
 
@@ -149,25 +129,21 @@ This high-level function handles all setup from the benchmark and returns a trai
 """
 function train_policy(
     algorithm::DAgger,
-    benchmark::AbstractStochasticBenchmark{true};
+    benchmark::ExogenousDynamicBenchmark;
     dataset_size=30,
-    split_ratio=(0.3, 0.3, 0.4),
     metrics::Tuple=(),
-    seed=0,
+    seed=nothing,
 )
-    # Generate dataset and environments
-    dataset = generate_dataset(benchmark, dataset_size)
-    train_instances, validation_instances, _ = splitobs(dataset; at=split_ratio)
-    train_environments = generate_environments(benchmark, train_instances; seed)
+    # Generate environments
+    train_environments = generate_environments(benchmark, dataset_size; seed)
 
     # Initialize model and create policy
-    model = generate_statistical_model(benchmark)
+    model = generate_statistical_model(benchmark; seed)
     maximizer = generate_maximizer(benchmark)
     policy = DFLPolicy(model, maximizer)
 
     # Define anticipative policy from benchmark
-    anticipative_policy =
-        (env; reset_env) -> generate_anticipative_solution(benchmark, env; reset_env)
+    anticipative_policy = generate_anticipative_solver(benchmark)
 
     # Train policy
     history = train_policy!(
@@ -176,7 +152,7 @@ function train_policy(
         train_environments;
         anticipative_policy=anticipative_policy,
         metrics=metrics,
-        maximizer_kwargs=get_state,
+        maximizer_kwargs=sample -> sample.context,
     )
 
     return history, policy

src/algorithms/supervised/fyl.jl

@@ -45,7 +45,7 @@ function train_policy!(
     train_dataset::DataLoader;
     epochs=100,
     metrics::Tuple=(),
-    maximizer_kwargs=get_info,
+    maximizer_kwargs=sample -> sample.context,
 )
     (; nb_samples, ε, threaded, training_optimizer, seed) = algorithm
     (; statistical_model, maximizer) = policy
@@ -106,7 +106,7 @@ function train_policy!(
     train_dataset::AbstractArray{<:DataSample};
     epochs=100,
     metrics::Tuple=(),
-    maximizer_kwargs=get_info,
+    maximizer_kwargs=sample -> sample.context,
 )
     data_loader = DataLoader(train_dataset; batchsize=1, shuffle=false)
     return train_policy!(
@@ -118,38 +118,3 @@ function train_policy!(
         maximizer_kwargs=maximizer_kwargs,
     )
 end
-
-"""
-$TYPEDSIGNATURES
-
-Train a DFLPolicy using the Perturbed Fenchel-Young Loss Imitation Algorithm on a benchmark.
-
-# Benchmark convenience wrapper
-
-This high-level function handles all setup from the benchmark and returns a trained policy.
-"""
-function train_policy(
-    algorithm::PerturbedFenchelYoungLossImitation,
-    benchmark::AbstractBenchmark;
-    dataset_size=30,
-    split_ratio=(0.3, 0.3),
-    epochs=100,
-    metrics::Tuple=(),
-    seed=nothing,
-)
-    # Generate dataset and split
-    dataset = generate_dataset(benchmark, dataset_size)
-    train_instances, _, _ = splitobs(dataset; at=split_ratio)
-
-    # Initialize model and create policy
-    model = generate_statistical_model(benchmark; seed)
-    maximizer = generate_maximizer(benchmark)
-    policy = DFLPolicy(model, maximizer)
-
-    # Train policy
-    history = train_policy!(
-        algorithm, policy, train_instances; epochs, metrics, maximizer_kwargs=get_info
-    )
-
-    return history, policy
-end