big update

Author: BatyLeo

.gitignore

@@ -8,3 +8,4 @@ tensorboard_logs
 .vscode
 Manifest.toml
 examples
+scripts

README.md

@@ -6,3 +6,43 @@
 [![Coverage](https://codecov.io/gh/JuliaDecisionFocusedLearning/DecisionFocusedLearningAlgorithms.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/JuliaDecisionFocusedLearning/DecisionFocusedLearningAlgorithms.jl)
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 [![Aqua](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
+
+> [!WARNING]  
+>  This package is currently under active development. The API may change in future releases.
+>  Please refer to the [documentation](https://JuliaDecisionFocusedLearning.github.io/DecisionFocusedLearningAlgorithms.jl/stable/) for the latest updates.
+
+## Overview
+
+This package provides a unified interface for training decision-focused learning algorithms that combine machine learning with combinatorial optimization. It implements several state-of-the-art algorithms for learning to predict parameters of optimization problems.
+
+### Key Features
+
+- **Unified Interface**: Consistent API across all algorithms via `train_policy!`
+- **Policy-Centric Design**: `DFLPolicy` encapsulates statistical models and optimizers
+- **Flexible Metrics**: Track custom metrics during training
+- **Benchmark Integration**: Seamless integration with DecisionFocusedLearningBenchmarks.jl
+
+### Quick Start
+
+```julia
+using DecisionFocusedLearningAlgorithms
+using DecisionFocusedLearningBenchmarks
+
+# Create a policy
+benchmark = ArgmaxBenchmark()
+model = generate_statistical_model(benchmark)
+maximizer = generate_maximizer(benchmark)
+policy = DFLPolicy(model, maximizer)
+
+# Train with FYL algorithm
+algorithm = PerturbedFenchelYoungLossImitation()
+result = train_policy(algorithm, benchmark; epochs=50)
+```
+
+See the [documentation](https://JuliaDecisionFocusedLearning.github.io/DecisionFocusedLearningAlgorithms.jl/stable/) for more details.
+
+## Available Algorithms
+
+- **Perturbed Fenchel-Young Loss Imitation**: Differentiable imitation learning with perturbed optimization
+- **AnticipativeImitation**: Imitation of anticipative solutions for dynamic problems
+- **DAgger**: DAgger algorithm for dynamic problems

docs/make.jl

@@ -2,21 +2,6 @@ using DecisionFocusedLearningAlgorithms
 using Documenter
 using Literate
 
-DocMeta.setdocmeta!(
-    DecisionFocusedLearningAlgorithms,
-    :DocTestSetup,
-    :(
-        begin
-            using DecisionFocusedLearningAlgorithms
-            using DecisionFocusedLearningBenchmarks
-            using Flux
-            using MLUtils
-            using Plots
-        end
-    );
-    recursive=true,
-)
-
 # Generate markdown files from tutorial scripts
 tutorial_dir = joinpath(@__DIR__, "src", "tutorials")
 tutorial_files = filter(f -> endswith(f, ".jl"), readdir(tutorial_dir))
@@ -29,22 +14,27 @@ end
 
 # Get list of generated markdown files for the docs
 md_tutorial_files = [
-    "tutorials/" * replace(file, ".jl" => ".md") for file in tutorial_files
+    joinpath("tutorials", replace(file, ".jl" => ".md")) for file in tutorial_files
 ]
 
 makedocs(;
     modules=[DecisionFocusedLearningAlgorithms],
     authors="Members of JuliaDecisionFocusedLearning and contributors",
     sitename="DecisionFocusedLearningAlgorithms.jl",
-    format=Documenter.HTML(;
-        canonical="https://JuliaDecisionFocusedLearning.github.io/DecisionFocusedLearningAlgorithms.jl",
-        edit_link="main",
-        assets=String[],
-    ),
-    pages=["Home" => "index.md", "Tutorials" => md_tutorial_files],
+    format=Documenter.HTML(; size_threshold=typemax(Int)),
+    pages=[
+        "Home" => "index.md",
+        "Interface Guide" => "interface.md",
+        "Tutorials" => md_tutorial_files,
+        "API Reference" => "api.md",
+    ],
 )
 
 deploydocs(;
     repo="github.com/JuliaDecisionFocusedLearning/DecisionFocusedLearningAlgorithms.jl",
     devbranch="main",
 )
+
+for file in md_tutorial_files
+    rm(joinpath(@__DIR__, "src", file))
+end

docs/src/api.md

@@ -0,0 +1,6 @@
+```@index
+```
+
+```@autodocs
+Modules = [DecisionFocusedLearningAlgorithms]
+```

docs/src/index.md

@@ -2,9 +2,38 @@
 
 Documentation for [DecisionFocusedLearningAlgorithms](https://github.com/JuliaDecisionFocusedLearning/DecisionFocusedLearningAlgorithms.jl).
 
-```@index
-```
+## Overview
+
+This package provides a unified interface for training decision-focused learning algorithms that combine machine learning with combinatorial optimization. It implements several state-of-the-art algorithms for learning to predict parameters of optimization problems.
+
+### Key Features
+
+- **Unified Interface**: Consistent API across all algorithms via `train_policy!`
+- **Policy-Centric Design**: `DFLPolicy` encapsulates statistical models and optimizers
+- **Flexible Metrics**: Track custom metrics during training
+- **Benchmark Integration**: Seamless integration with DecisionFocusedLearningBenchmarks.jl
+
+### Quick Start
 
-```@autodocs
-Modules = [DecisionFocusedLearningAlgorithms]
+```julia
+using DecisionFocusedLearningAlgorithms
+using DecisionFocusedLearningBenchmarks
+
+# Create a policy
+benchmark = ArgmaxBenchmark()
+model = generate_statistical_model(benchmark)
+maximizer = generate_maximizer(benchmark)
+policy = DFLPolicy(model, maximizer)
+
+# Train with FYL algorithm
+algorithm = PerturbedFenchelYoungLossImitation()
+result = train_policy(algorithm, benchmark; epochs=50)
 ```
+
+See the [Interface Guide](interface.md) and [Tutorials](tutorials/tutorial.md) for more details.
+
+## Available Algorithms
+
+- **Perturbed Fenchel-Young Loss Imitation**: Differentiable imitation learning with perturbed optimization
+- **AnticipativeImitation**: Imitation of anticipative solutions for dynamic problems
+- **DAgger**: DAgger algorithm for dynamic problems

docs/src/interface.md

@@ -0,0 +1,100 @@
+# Algorithm Interface
+
+This page describes the unified interface for Decision-Focused Learning algorithms provided by this package.
+
+## Core Concepts
+
+### DFLPolicy
+
+The [`DFLPolicy`](@ref) is the central abstraction that encapsulates a decision-focused learning policy. It combines:
+- A **statistical model** (typically a neural network) that predicts parameters from input features
+- A **combinatorial optimizer** (maximizer) that solves optimization problems using the predicted parameters
+
+```julia
+policy = DFLPolicy(
+    Chain(Dense(input_dim => hidden_dim, relu), Dense(hidden_dim => output_dim)),
+    my_optimizer
+)
+```
+
+### Training Interface
+
+All algorithms in this package follow a unified training interface with two main functions:
+
+#### Core Training Method
+
+```julia
+history = train_policy!(algorithm, policy, training_data; epochs=100, metrics=(), maximizer_kwargs=get_info)
+```
+
+**Arguments:**
+- `algorithm`: An algorithm instance (e.g., `PerturbedFenchelYoungLossImitation`, `DAgger`, `AnticipativeImitation`)
+- `policy::DFLPolicy`: The policy to train (contains the model and maximizer)
+- `training_data`: Either a dataset of `DataSample` objects or `Environment` (depends on algorithm)
+- `epochs::Int`: Number of training epochs (default: 100)
+- `metrics::Tuple`: Metrics to evaluate during training (default: empty)
+- `maximizer_kwargs::Function`: Function that extracts keyword arguments for the maximizer from data samples (default: `get_info`)
+
+**Returns:**
+- `history::MVHistory`: Training history containing loss values and metric evaluations
+
+#### Benchmark Convenience Wrapper
+
+```julia
+result = train_policy(algorithm, benchmark; dataset_size=30, split_ratio=(0.3, 0.3), epochs=100, metrics=())
+```
+
+This high-level function handles all setup from a benchmark and returns a trained policy along with training history.
+
+**Arguments:**
+- `algorithm`: An algorithm instance
+- `benchmark::AbstractBenchmark`: A benchmark from DecisionFocusedLearningBenchmarks.jl
+- `dataset_size::Int`: Number of instances to generate
+- `split_ratio::Tuple`: Train/validation/test split ratios
+- `epochs::Int`: Number of training epochs
+- `metrics::Tuple`: Metrics to track during training
+
+**Returns:**
+- `(; policy, history)`: Named tuple with trained policy and training history
+
+## Metrics
+
+Metrics allow you to track additional quantities during training.
+
+### Built-in Metrics
+
+#### FYLLossMetric
+
+Evaluates Fenchel-Young loss on a validation dataset.
+
+```julia
+val_metric = FYLLossMetric(validation_data, :validation_loss)
+```
+
+#### FunctionMetric
+
+Custom metric defined by a function.
+
+```julia
+# Simple metric (no stored data)
+epoch_metric = FunctionMetric(ctx -> ctx.epoch, :epoch)
+
+# Metric with stored data
+gap_metric = FunctionMetric(:validation_gap, validation_data) do ctx, data
+    compute_gap(benchmark, data, ctx.policy.statistical_model, ctx.policy.maximizer)
+end
+```
+
+### TrainingContext
+
+Metrics receive a `TrainingContext` object containing:
+- `policy::DFLPolicy`: The policy being trained
+- `epoch::Int`: Current epoch number
+- `maximizer_kwargs::Function`: Maximizer kwargs extractor
+- `other_fields`: Algorithm-specific fields (e.g., `loss` for FYL)
+
+Access policy components:
+```julia
+ctx.policy.statistical_model  # Neural network
+ctx.policy.maximizer          # Combinatorial optimizer
+```

docs/src/tutorials/tutorial.jl

@@ -1,48 +1,46 @@
-# Tutorial
+# # Basic Tutorial: Training with FYL on Argmax Benchmark
+#
+# This tutorial demonstrates the basic workflow for training a policy
+# using the Perturbed Fenchel-Young Loss algorithm.
+
+# ## Setup
 using DecisionFocusedLearningAlgorithms
 using DecisionFocusedLearningBenchmarks
 using MLUtils: splitobs
 using Plots
 
+# ## Create Benchmark and Data
 b = ArgmaxBenchmark()
 dataset = generate_dataset(b, 100)
-train_instances, validation_instances, test_instances = splitobs(
-    dataset; at=(0.3, 0.3, 0.4)
-)
+train_data, val_data, test_data = splitobs(dataset; at=(0.3, 0.3, 0.4))
 
+# ## Create Policy
 model = generate_statistical_model(b; seed=0)
 maximizer = generate_maximizer(b)
+policy = DFLPolicy(model, maximizer)
 
-# Compute initial gap
-initial_gap = compute_gap(b, test_instances, model, maximizer)
-println("Initial test gap: $initial_gap")
-
-# Configure the training algorithm
-algorithm = PerturbedImitationAlgorithm(; nb_samples=10, ε=0.1, threaded=true, seed=0)
+# ## Configure Algorithm
+algorithm = PerturbedFenchelYoungLossImitation(;
+    nb_samples=10, ε=0.1, threaded=true, seed=0
+)
 
-# Define metrics to track during training
-validation_loss_metric = FYLLossMetric(validation_instances, :validation_loss)
+# ## Define Metrics to track during training
+validation_loss_metric = FYLLossMetric(val_data, :validation_loss)
 
-# Validation gap metric
-val_gap_metric = FunctionMetric(:val_gap, validation_instances) do ctx, data
-    compute_gap(b, data, ctx.model, ctx.maximizer)
+val_gap_metric = FunctionMetric(:val_gap, val_data) do ctx, data
+    compute_gap(b, data, ctx.policy.statistical_model, ctx.policy.maximizer)
 end
 
-# Test gap metric
-test_gap_metric = FunctionMetric(:test_gap, test_instances) do ctx, data
-    compute_gap(b, data, ctx.model, ctx.maximizer)
+test_gap_metric = FunctionMetric(:test_gap, test_data) do ctx, data
+    compute_gap(b, data, ctx.policy.statistical_model, ctx.policy.maximizer)
 end
 
-# Combine metrics
 metrics = (validation_loss_metric, val_gap_metric, test_gap_metric)
 
-# Train the model
-fyl_model = deepcopy(model)
-history = train_policy!(
-    algorithm, fyl_model, maximizer, train_instances; epochs=100, metrics=metrics
-)
+# ## Train the Policy
+history = train_policy!(algorithm, policy, train_data; epochs=100, metrics=metrics)
 
-# Plot validation and test gaps
+# ## Plot Results
 val_gap_epochs, val_gap_values = get(history, :val_gap)
 test_gap_epochs, test_gap_values = get(history, :test_gap)
 
@@ -55,7 +53,7 @@ plot(
     title="Gap Evolution During Training",
 )
 
-# Plot validation loss
+# Plot loss evolution
 train_loss_epochs, train_loss_values = get(history, :training_loss)
 val_loss_epochs, val_loss_values = get(history, :validation_loss)