From 6516cc64f5928e509953870366b5a009017eacd5 Mon Sep 17 00:00:00 2001
From: sdelannoypavy <solene.delannoypavy_externe@rte-france.com>
Date: Tue, 30 Dec 2025 21:15:54 +0100
Subject: [PATCH 1/6] fixed buggs

---
 src/DecisionFocusedLearningBenchmarks.jl |   3 +
 src/DynamicAssortment/environment.jl     |  11 +-
 src/Maintenance/Maintenance.jl           | 150 ++++++++++++++++++++
 src/Maintenance/environment.jl           | 167 +++++++++++++++++++++++
 src/Maintenance/instance.jl              |  36 +++++
 src/Maintenance/maximizer.jl             |  26 ++++
 src/Maintenance/policies.jl              |  10 ++
 7 files changed, 400 insertions(+), 3 deletions(-)
 create mode 100644 src/Maintenance/Maintenance.jl
 create mode 100644 src/Maintenance/environment.jl
 create mode 100644 src/Maintenance/instance.jl
 create mode 100644 src/Maintenance/maximizer.jl
 create mode 100644 src/Maintenance/policies.jl

diff --git a/src/DecisionFocusedLearningBenchmarks.jl b/src/DecisionFocusedLearningBenchmarks.jl
index 4515ec1..2f1c320 100644
--- a/src/DecisionFocusedLearningBenchmarks.jl
+++ b/src/DecisionFocusedLearningBenchmarks.jl
@@ -57,6 +57,7 @@ include("PortfolioOptimization/PortfolioOptimization.jl")
 include("StochasticVehicleScheduling/StochasticVehicleScheduling.jl")
 include("DynamicVehicleScheduling/DynamicVehicleScheduling.jl")
 include("DynamicAssortment/DynamicAssortment.jl")
+include("Maintenance/Maintenance.jl")
 
 using .Utils
 
@@ -89,6 +90,7 @@ using .PortfolioOptimization
 using .StochasticVehicleScheduling
 using .DynamicVehicleScheduling
 using .DynamicAssortment
+using .Maintenance
 
 export Argmax2DBenchmark
 export ArgmaxBenchmark
@@ -100,5 +102,6 @@ export RankingBenchmark
 export StochasticVehicleSchedulingBenchmark
 export SubsetSelectionBenchmark
 export WarcraftBenchmark
+export MaintenanceBenchmark
 
 end # module DecisionFocusedLearningBenchmarks
diff --git a/src/DynamicAssortment/environment.jl b/src/DynamicAssortment/environment.jl
index 603ca84..37a5086 100644
--- a/src/DynamicAssortment/environment.jl
+++ b/src/DynamicAssortment/environment.jl
@@ -197,16 +197,21 @@ Features observed by the agent at current step, as a concatenation of:
 - change in hype and saturation features from the starting state
 - normalized current step (divided by max steps and multiplied by 10)
 All features are normalized by dividing by 10.
+
+State
 """
 function Utils.observe(env::Environment)
     delta_features = env.features[2:3, :] .- env.instance.starting_hype_and_saturation
-    return vcat(
+    features = vcat(
         env.features,
         env.d_features,
         delta_features,
         ones(1, item_count(env)) .* (env.step / max_steps(env) * 10),
-    ) ./ 10,
-    nothing
+    ) ./ 10
+
+    state = (env.features, env.purchase_history)
+
+    return features, state
 end
 
 """
diff --git a/src/Maintenance/Maintenance.jl b/src/Maintenance/Maintenance.jl
new file mode 100644
index 0000000..9cacd35
--- /dev/null
+++ b/src/Maintenance/Maintenance.jl
@@ -0,0 +1,150 @@
+module Maintenance
+
+using ..Utils
+
+using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES, SIGNATURES
+using Distributions: Uniform, Categorical
+using Flux: Chain, Dense
+using LinearAlgebra: dot
+using Random: Random, AbstractRNG, MersenneTwister
+using Statistics: mean
+
+using Combinatorics: combinations
+
+"""
+$TYPEDEF
+
+Benchmark for a standard maintenance problem with resource constraints.
+Components are identical and degrade idependently over time.
+A high cost is incurred for each component that reaches the final degradation level. 
+A cost is also incurred for maintaining a component. 
+The number of simultaneous maintenance operations is limited by a maintenance capacity constraint.
+
+# Fields
+$TYPEDFIELDS
+
+"""
+struct MaintenanceBenchmark <: AbstractDynamicBenchmark{true}
+    "number of components"
+    N::Int
+    "maximum number of components that can be maintained simultaneously"
+    K::Int
+    "number of degradation states per component"
+    n::Int
+    "degradation probability"
+    p::Float64
+    "failure cost"
+    c_f::Float64
+    "maintenance cost"
+    c_m::Float64
+    "number of steps per episode"
+    max_steps::Int
+end
+
+"""
+    MaintenanceBenchmark(;
+        N=2,
+        K=1,
+        n=3,
+        p=0.2
+        c_f=10.0,
+        c_m=3.0,
+        max_steps=10,
+    )
+
+Constructor for [`MaintenanceBenchmark`](@ref).
+By default, the benchmark has 2 components, maintenance capacity 1, number of degradation levels 3, 
+degradation probability 0.2, failure cost 10.0, maintenance cost 3.0, 10 steps per episode, and is exogenous.
+"""
+
+function MaintenanceBenchmark(;
+    N=2,
+    K=1,
+    n=3,
+    p=0.2,
+    c_f=10.0,
+    c_m=3.0,
+    max_steps=10,
+)
+    return MaintenanceBenchmark(
+        N, K, n, p, c_f, c_m, max_steps
+    )
+end
+
+# Accessor functions
+component_count(b::MaintenanceBenchmark) = b.N
+maintenance_capacity(b::MaintenanceBenchmark) = b.K
+degradation_levels(b::MaintenanceBenchmark) = b.n
+degradation_probability(b::MaintenanceBenchmark) = b.p
+failure_cost(b::MaintenanceBenchmark) = b.c_f
+maintenance_cost(b::MaintenanceBenchmark) = b.c_m
+max_steps(b::MaintenanceBenchmark) = b.max_steps
+
+include("instance.jl")
+include("environment.jl")
+include("policies.jl")
+include("maximizer.jl")
+
+"""
+$TYPEDSIGNATURES
+
+Outputs a data sample containing an [`Instance`](@ref).
+"""
+function Utils.generate_sample(
+    b::MaintenanceBenchmark, rng::AbstractRNG=MersenneTwister(0)
+)
+    return DataSample(; instance=Instance(b, rng))
+end
+
+"""
+$TYPEDSIGNATURES
+
+Generates a statistical model for the maintenance benchmark.
+The model is a small neural network with one hidden layer no activation function.
+"""
+function Utils.generate_statistical_model(b::MaintenanceBenchmark; seed=nothing)
+    Random.seed!(seed)
+    N = component_count(b)
+    return Chain(Dense(N  => N), Dense(N => N), vec)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Outputs a top k maximizer, with k being the maintenance capacity of the benchmark.
+"""
+function Utils.generate_maximizer(b::MaintenanceBenchmark)
+    return TopKPositiveMaximizer(maintenance_capacity(b))
+end
+
+"""
+$TYPEDSIGNATURES
+
+Creates an [`Environment`](@ref) from an [`Instance`](@ref) of the maintenance benchmark.
+The seed of the environment is randomly generated using the provided random number generator.
+"""
+function Utils.generate_environment(
+    ::MaintenanceBenchmark, instance::Instance, rng::AbstractRNG; kwargs...
+)
+    seed = rand(rng, 1:typemax(Int))
+    return Environment(instance; seed)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Returns two policies for the dynamic assortment benchmark:
+- `Greedy`: maintains components when they are in the last state before failure, up to the maintenance capacity
+"""
+function Utils.generate_policies(::MaintenanceBenchmark)
+    greedy = Policy(
+        "Greedy",
+        "policy that maintains components when they are in the last state before failure, up to the maintenance capacity",
+        greedy_policy,
+    )
+    return (greedy)
+end
+
+export MaintenanceBenchmark
+
+end
diff --git a/src/Maintenance/environment.jl b/src/Maintenance/environment.jl
new file mode 100644
index 0000000..cc169f6
--- /dev/null
+++ b/src/Maintenance/environment.jl
@@ -0,0 +1,167 @@
+"""
+$TYPEDEF
+
+Environment for the maintenance problem.
+
+# Fields
+$TYPEDFIELDS
+"""
+@kwdef mutable struct Environment{I<:Instance,R<:AbstractRNG,S<:Union{Nothing,Int}} <:
+                      Utils.AbstractEnvironment
+    "associated instance"
+    instance::I
+    "current step"
+    step::Int
+    "degradation state"
+    degradation_state::Vector{Int}
+    "rng"
+    rng::R
+    "seed for RNG"
+    seed::S
+end
+
+"""
+$TYPEDSIGNATURES
+
+Creates an [`Environment`](@ref) from an [`Instance`](@ref) of the maintenance benchmark.
+"""
+function Environment(instance::Instance; seed=0, rng::AbstractRNG=MersenneTwister(seed))
+    degradation_state = starting_state(instance)
+    env = Environment(;
+        instance,
+        step=1,
+        degradation_state,
+        rng=rng,
+        seed=seed,
+    )
+    Utils.reset!(env; reset_rng=true)
+    return env
+end
+
+component_count(env::Environment) = component_count(env.instance) 
+maintenance_capacity(env::Environment) = maintenance_capacity(env.instance)
+degradation_levels(env::Environment) = degradation_levels(env.instance)
+degradation_probability(env::Environment) = degradation_probability(env.instance)
+failure_cost(env::Environment) = failure_cost(env.instance)
+maintenance_cost(env::Environment) = maintenance_cost(env.instance)
+max_steps(env::Environment) = max_steps(env.instance) 
+starting_state(env::Environment) = starting_state(env.instance)
+
+
+"""
+$TYPEDSIGNATURES
+Draw random degradations for all components.
+"""
+
+function degrad!(env::Environment)
+    N = component_count(env) 
+    n = degradation_levels(env)
+    p = degradation_probability(env)
+
+    for i in 1:N
+        if env.degradation_state[i] < n && rand() < p
+            env.degradation_state[i] += 1
+        end
+    end
+
+    return env.degradation_state
+end
+
+"""
+$TYPEDSIGNATURES
+Maintain components.
+"""
+
+function maintain!(env::Environment, maintenance::BitVector)
+    N = component_count(env) 
+
+    for i in 1:N
+        if maintenance[i]
+            env.degradation_state[i] = 1
+        end
+    end
+
+    return env.degradation_state
+end
+
+"""
+$TYPEDSIGNATURES
+
+Compute maintenance cost.
+"""
+function maintenance_cost(env::Environment, maintenance::BitVector)
+    return maintenance_cost(env) * sum(maintenance)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Compute degradation cost.
+"""
+function degradation_cost(env::Environment)
+    N = component_count(env) 
+    n = degradation_levels(env)
+    return failure_cost(env) * count(==(n), env.degradation_state)
+end
+
+
+"""
+$TYPEDSIGNATURES
+
+Outputs the seed of the environment.
+"""
+Utils.get_seed(env::Environment) = env.seed
+
+"""
+$TYPEDSIGNATURES
+
+Resets the environment to the initial state:
+- reset the rng if `reset_rng` is true
+- reset the step to 1
+- reset the degradation state to the starting state
+"""
+function Utils.reset!(env::Environment; reset_rng=false, seed=env.seed)
+    reset_rng && Random.seed!(env.rng, seed)
+    env.step = 1
+    env.degradation_state .= starting_state(env)
+    return nothing
+end
+
+"""
+$TYPEDSIGNATURES
+
+Checks if the environment has reached the maximum number of steps.
+"""
+function Utils.is_terminated(env::Environment)
+    return env.step > max_steps(env)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Returns features, state tuple.
+The features observed by the agent at current step are the degradation states of all components.
+It is also the internal state, so we return the same thing twice.
+
+"""
+function Utils.observe(env::Environment)
+    state = env.degradation_state
+    return state, state
+end
+
+"""
+$TYPEDSIGNATURES
+
+Performs one step in the environment given a maintenance.
+Draw random degradations for components that are not maintained.
+"""
+function Utils.step!(env::Environment, maintenance::BitVector)
+    @assert !Utils.is_terminated(env) "Environment is terminated, cannot act!"
+    reward = maintenance_cost(env, maintenance) + degradation_cost(env)
+    degrad!(env)
+    maintain!(env, maintenance)
+    env.step += 1
+    return reward
+end
+
+
diff --git a/src/Maintenance/instance.jl b/src/Maintenance/instance.jl
new file mode 100644
index 0000000..aef280e
--- /dev/null
+++ b/src/Maintenance/instance.jl
@@ -0,0 +1,36 @@
+"""
+$TYPEDEF
+
+Instance of the maintenance problem.
+
+# Fields
+$TYPEDFIELDS
+"""
+@kwdef struct Instance{B<:MaintenanceBenchmark}
+    "associated benchmark"
+    config::B
+    "starting degradation states"
+    starting_state::Vector{Int}
+end
+
+"""
+$TYPEDSIGNATURES
+
+Generates an instance with random starting degradation states uniformly in [1, n]
+"""
+function Instance(b::MaintenanceBenchmark, rng::AbstractRNG)
+    N = component_count(b)
+    n = degradation_levels(b)
+    starting_state = rand(rng, 1:n, N) 
+    return Instance(; config=b, starting_state)
+end
+
+# Accessor functions
+component_count(b::Instance) = component_count(b.config) 
+maintenance_capacity(b::Instance) = maintenance_capacity(b.config)
+degradation_levels(b::Instance) = degradation_levels(b.config)
+degradation_probability(b::Instance) = degradation_probability(b.config)
+failure_cost(b::Instance) = failure_cost(b.config)
+maintenance_cost(b::Instance) = maintenance_cost(b.config)
+max_steps(b::Instance) = max_steps(b.config) 
+starting_state(b::Instance) = b.starting_state
diff --git a/src/Maintenance/maximizer.jl b/src/Maintenance/maximizer.jl
new file mode 100644
index 0000000..a271234
--- /dev/null
+++ b/src/Maintenance/maximizer.jl
@@ -0,0 +1,26 @@
+"""
+$TYPEDEF
+
+Top k maximizer.
+"""
+struct TopKPositiveMaximizer
+    k::Int
+end
+
+"""
+$TYPEDSIGNATURES
+
+Return the top k indices of `θ`.
+"""
+
+function (m::TopKPositiveMaximizer)(θ; kwargs...)
+    N = length(θ)
+    
+    sorted_indices = sortperm(θ; rev=true)
+    positive_indices = filter(i -> θ[i] > 0, sorted_indices)
+    solution = positive_indices[1:min(m.k, length(positive_indices))]
+    
+    res = falses(N)
+    res[solution] .= 1
+    return res
+end
\ No newline at end of file
diff --git a/src/Maintenance/policies.jl b/src/Maintenance/policies.jl
new file mode 100644
index 0000000..203a4b1
--- /dev/null
+++ b/src/Maintenance/policies.jl
@@ -0,0 +1,10 @@
+
+"""
+$TYPEDSIGNATURES
+
+Greedy policy that maintains components when they are in the last state before failure, up to the maintenance capacity.
+"""
+function greedy_policy(env::Environment)
+    maximizer = generate_maximizer(env.instance.config)
+    return maximizer(prices(env)[1:item_count(env)])
+end

From 0cf9022e853271ea15411be077a8fb37861da7ae Mon Sep 17 00:00:00 2001
From: sdelannoypavy <solene.delannoypavy_externe@rte-france.com>
Date: Thu, 8 Jan 2026 17:36:25 +0100
Subject: [PATCH 2/6] Added maintenance benchmark

---
 src/Maintenance/Maintenance.jl |   4 +-
 src/Maintenance/environment.jl |   4 +-
 src/Maintenance/policies.jl    |  21 ++-
 test/maintenance.jl            | 231 +++++++++++++++++++++++++++++++++
 test/runtests.jl               |   1 +
 5 files changed, 255 insertions(+), 6 deletions(-)
 create mode 100644 test/maintenance.jl

diff --git a/src/Maintenance/Maintenance.jl b/src/Maintenance/Maintenance.jl
index 9cacd35..fa99cfe 100644
--- a/src/Maintenance/Maintenance.jl
+++ b/src/Maintenance/Maintenance.jl
@@ -64,7 +64,7 @@ function MaintenanceBenchmark(;
     p=0.2,
     c_f=10.0,
     c_m=3.0,
-    max_steps=10,
+    max_steps=80,
 )
     return MaintenanceBenchmark(
         N, K, n, p, c_f, c_m, max_steps
@@ -142,7 +142,7 @@ function Utils.generate_policies(::MaintenanceBenchmark)
         "policy that maintains components when they are in the last state before failure, up to the maintenance capacity",
         greedy_policy,
     )
-    return (greedy)
+    return (greedy,)
 end
 
 export MaintenanceBenchmark
diff --git a/src/Maintenance/environment.jl b/src/Maintenance/environment.jl
index cc169f6..b17d3d3 100644
--- a/src/Maintenance/environment.jl
+++ b/src/Maintenance/environment.jl
@@ -157,11 +157,11 @@ Draw random degradations for components that are not maintained.
 """
 function Utils.step!(env::Environment, maintenance::BitVector)
     @assert !Utils.is_terminated(env) "Environment is terminated, cannot act!"
-    reward = maintenance_cost(env, maintenance) + degradation_cost(env)
+    cost = maintenance_cost(env, maintenance) + degradation_cost(env)
     degrad!(env)
     maintain!(env, maintenance)
     env.step += 1
-    return reward
+    return cost
 end
 
 
diff --git a/src/Maintenance/policies.jl b/src/Maintenance/policies.jl
index 203a4b1..100b2ed 100644
--- a/src/Maintenance/policies.jl
+++ b/src/Maintenance/policies.jl
@@ -5,6 +5,23 @@ $TYPEDSIGNATURES
 Greedy policy that maintains components when they are in the last state before failure, up to the maintenance capacity.
 """
 function greedy_policy(env::Environment)
-    maximizer = generate_maximizer(env.instance.config)
-    return maximizer(prices(env)[1:item_count(env)])
+    state = env.degradation_state
+    N = component_count(env)
+    K = maintenance_capacity(env)
+    res = falses(N)   
+    n = degradation_levels(env)
+    
+    
+    idx_max = findall(==(n), state)
+    take = first(idx_max, min(K, length(idx_max)))
+    res[take] .= true
+    remaining = K - length(take)
+
+    if remaining > 0
+        idx_second = findall(==(n-1), state)
+        take2 = first(idx_second, min(remaining, length(idx_second)))
+        res[take2] .= true
+    end
+
+    return res 
 end
diff --git a/test/maintenance.jl b/test/maintenance.jl
new file mode 100644
index 0000000..53ac7b3
--- /dev/null
+++ b/test/maintenance.jl
@@ -0,0 +1,231 @@
+const maintenance = DecisionFocusedLearningBenchmarks.Maintenance
+
+@testset "Maintenance - Benchmark Construction" begin
+    # Test default constructor
+    b = MaintenanceBenchmark()
+    @test b.N == 2
+    @test b.K == 1
+    @test b.n == 3
+    @test b.p == 0.2
+    @test b.c_f == 10.0
+    @test b.c_m == 3.0
+    @test b.max_steps == 80
+    @test is_exogenous(b)
+    @test !is_endogenous(b)
+
+    # Test custom constructor
+    b_custom = MaintenanceBenchmark(; N=10, K=3, n=5, p=0.3, c_f=5.0, c_m=3.0, max_steps=50)
+    @test b_custom.N == 10
+    @test b_custom.K == 3
+    @test b_custom.n == 5
+    @test b_custom.p == 0.3
+    @test b_custom.c_f == 5.0
+    @test b_custom.c_m == 3.0
+    @test b_custom.max_steps == 50
+
+    # Test accessor functions
+    @test maintenance.component_count(b) == 2
+    @test maintenance.maintenance_capacity(b) == 1
+    @test maintenance.degradation_levels(b) == 3
+    @test maintenance.degradation_probability(b) == 0.2
+    @test maintenance.failure_cost(b) == 10.0
+    @test maintenance.maintenance_cost(b) == 3.0
+    @test maintenance.max_steps(b) == 80
+end
+
+@testset "Maintenance - Instance Generation" begin
+    b = MaintenanceBenchmark(; N=10, K=3, n=5, p=0.3, c_f=5.0, c_m=3.0, max_steps=50)
+    rng = MersenneTwister(42)
+
+    instance = maintenance.Instance(b, rng)
+
+    # test state is randomly initialized
+    state1 = maintenance.starting_state(instance)
+    rng2 = MersenneTwister(43)
+    instance2 = maintenance.Instance(b, rng2)
+    state2 = maintenance.starting_state(instance2)
+    @test state1 != state2
+
+    # Test instance structure
+    @test length(instance.starting_state) == 10 
+    @test all(1.0 ≤ s ≤ 5 for s in instance.starting_state)  
+
+    # Test accessor functions
+    @test maintenance.component_count(instance) == 10
+    @test maintenance.maintenance_capacity(instance) == 3
+    @test maintenance.degradation_levels(instance) == 5
+    @test maintenance.degradation_probability(instance) == 0.3
+    @test maintenance.failure_cost(instance) == 5.0
+    @test maintenance.maintenance_cost(instance) == 3.0
+    @test maintenance.max_steps(instance) == 50
+end
+
+@testset "Maintenance - Environment Initialization" begin
+    b = MaintenanceBenchmark()
+    instance = maintenance.Instance(b, MersenneTwister(42))
+
+    env = maintenance.Environment(instance; seed=123)
+
+    # Test initial state
+    @test env.step == 1
+    @test env.seed == 123
+    @test !is_terminated(env)
+
+    # Test accessor functions
+    @test maintenance.component_count(env) == 2
+    @test maintenance.maintenance_capacity(env) == 1
+    @test maintenance.degradation_levels(env) == 3
+    @test maintenance.degradation_probability(env) == 0.2
+    @test maintenance.failure_cost(env) == 10.0
+    @test maintenance.maintenance_cost(env) == 3.0
+    @test maintenance.max_steps(env) == 80
+end
+
+@testset "Maintenance - Environment Reset" begin
+    b = MaintenanceBenchmark()
+    instance = maintenance.Instance(b, MersenneTwister(42))
+    env = maintenance.Environment(instance; seed=123)
+
+    # Modify environment state
+    env.step = 3
+
+    # Reset environment
+    reset!(env)
+
+    # Check reset state
+    @test env.step == 1
+end
+
+@testset "Maintenance - Cost" begin
+    b = MaintenanceBenchmark()
+    instance = maintenance.Instance(b, MersenneTwister(42))
+    env = maintenance.Environment(instance; seed=123)
+
+    env.degradation_state = [1,1]
+    @test maintenance.maintenance_cost(env, BitVector([false, false])) == 0.0 
+    @test maintenance.maintenance_cost(env, BitVector([false, true])) == 3.0 
+    @test maintenance.maintenance_cost(env, BitVector([true, true])) == 6.0 
+
+    @test maintenance.degradation_cost(env) == 0.0 
+    env.degradation_state = [2,2]
+    @test maintenance.degradation_cost(env) == 0.0 
+    env.degradation_state = [3,2]
+    @test maintenance.degradation_cost(env) == 10.0 
+    env.degradation_state = [3,3]
+    @test maintenance.degradation_cost(env) == 20.0 
+end
+
+@testset "Maintenance - Environment Step" begin
+    b = MaintenanceBenchmark()
+    instance = maintenance.Instance(b, MersenneTwister(42))
+    env = maintenance.Environment(instance; seed=123)
+
+    maintenance_vect = BitVector([false, false])
+
+    initial_step = env.step
+    # Take a step
+    reward = step!(env, maintenance_vect)
+
+    # Check step progression
+    @test env.step == initial_step + 1
+    @test reward ≥ 0.0  # Reward should be non-negative 
+
+    # Test termination
+    for _ in 1:(maintenance.max_steps(env) - 1)
+        if !is_terminated(env)
+            step!(env, maintenance_vect)
+        end
+    end
+    @test is_terminated(env)
+
+    # Test error on terminated environment
+    @test_throws AssertionError step!(env, maintenance_vect)
+end
+
+@testset "Maintenance - Observation" begin
+    b = MaintenanceBenchmark()
+    instance = maintenance.Instance(b, MersenneTwister(42))
+    env = maintenance.Environment(instance; seed=123)
+    env.degradation_state = [1,1]
+
+    state, features = observe(env)
+
+    @test state == [1,1]
+    @test features === state
+
+    env.degradation_state = [2,3]
+    state2, _ = observe(env)
+
+    @test state != state2  # Observations should differ after purchase
+end
+
+
+@testset "Maintenance - Policies" begin
+    using Statistics: mean
+
+    b = MaintenanceBenchmark()
+
+    # Generate test data
+    dataset = generate_dataset(b, 10; seed=0)
+    environments = generate_environments(b, dataset)
+
+    # Get policies
+    policies = generate_policies(b)
+    greedy = policies[1]
+
+    @test greedy.name == "Greedy"
+
+    # Test policy evaluation
+    r_greedy, _ = evaluate_policy!(greedy, environments, 10)
+
+    @test length(r_greedy) == length(environments)
+    @test all(r_greedy .≥ 0.0)
+
+    # Test policy output format
+    env = environments[1]
+    reset!(env)
+
+    greedy_action = greedy(env)
+    @test greedy_action isa BitVector && length(greedy_action) == 2
+end
+
+
+@testset "Maintenance - Model and Maximizer Integration" begin
+    b = MaintenanceBenchmark()
+
+    # Test statistical model generation
+    model = generate_statistical_model(b; seed=42)
+    # Test maximizer generation
+    maximizer = generate_maximizer(b)
+
+    # Test integration with sample data
+    sample = generate_sample(b, MersenneTwister(42))
+    @test hasfield(typeof(sample), :info)
+
+    dataset = generate_dataset(b, 3; seed=42)
+    environments = generate_environments(b, dataset)
+
+    # Evaluate policy to get data samples
+    policies = generate_policies(b)
+    _, data_samples = evaluate_policy!(policies[1], environments)
+
+    # Test model-maximizer pipeline
+    sample = data_samples[1]
+    x = sample.x
+    θ = model(x)
+    y = maximizer(θ)
+
+    @test length(θ) == 2
+
+    θ = [1,2]
+    @test maximizer(θ) == BitVector([false, true])
+
+    b = MaintenanceBenchmark(; N=10, K=3, n=5, p=0.3, c_f=5.0, c_m=3.0, max_steps=50)
+    θ = [i for i in 1:10]
+    maximizer = generate_maximizer(b)
+    @test maximizer(θ) == BitVector([false, false, false, false, false, false, false, true, true, true])
+
+
+
+    #test maximizer output
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 914a6e4..00bea25 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,6 +14,7 @@ using Random
     include("ranking.jl")
     include("subset_selection.jl")
     include("fixed_size_shortest_path.jl")
+    include("maintenance.jl")
     include("warcraft.jl")
     include("vsp.jl")
     include("portfolio_optimization.jl")

From 9cfc83b30a0497a7c868d8f4bba59eae92de0c69 Mon Sep 17 00:00:00 2001
From: sdelannoypavy <solene.delannoypavy_externe@rte-france.com>
Date: Thu, 8 Jan 2026 18:06:30 +0100
Subject: [PATCH 3/6] Completed documentation

---
 docs/src/index.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/src/index.md b/docs/src/index.md
index 0e81d5a..363abee 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -61,6 +61,7 @@ Single-stage optimization problems under uncertainty:
 Multi-stage sequential decision-making problems:
 - [`DynamicVehicleSchedulingBenchmark`](@ref): multi-stage vehicle scheduling under customer uncertainty
 - [`DynamicAssortmentBenchmark`](@ref): sequential product assortment selection with endogenous uncertainty
+- [`MaintenanceBenchmark`](@ref): maintenance problem with resource constraint
 
 ## Getting Started
 

From 778e7ece6d602fb12f03819f78db50534a613039 Mon Sep 17 00:00:00 2001
From: sdelannoypavy <solene.delannoypavy_externe@rte-france.com>
Date: Fri, 9 Jan 2026 18:43:43 +0100
Subject: [PATCH 4/6] Corrected some bugs and added documentation

---
 Project.toml                         |   8 +-
 docs/src/api.md                      |  12 +++
 docs/src/benchmarks/maintenance.md   | 107 +++++++++++++++++++++++++++
 src/DynamicAssortment/environment.jl |  20 +++--
 src/Maintenance/Maintenance.jl       |  35 ++++-----
 src/Maintenance/environment.jl       |  28 ++-----
 src/Maintenance/instance.jl          |  12 +--
 src/Maintenance/maximizer.jl         |  24 +++---
 src/Maintenance/policies.jl          |   9 +--
 test/dynamic_assortment.jl           |  20 +++--
 test/maintenance.jl                  |  43 +++++------
 11 files changed, 217 insertions(+), 101 deletions(-)
 create mode 100644 docs/src/benchmarks/maintenance.md

diff --git a/Project.toml b/Project.toml
index 85e6211..b8daf1c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,9 +3,6 @@ uuid = "2fbe496a-299b-4c81-bab5-c44dfc55cf20"
 authors = ["Members of JuliaDecisionFocusedLearning"]
 version = "0.4.0"
 
-[workspace]
-projects = ["docs", "test"]
-
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
@@ -23,6 +20,7 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc"
@@ -54,6 +52,7 @@ Ipopt = "1.6"
 IterTools = "1.10.0"
 JSON = "1"
 JuMP = "1.22"
+JuliaFormatter = "1.0.62"
 LaTeXStrings = "1.4.0"
 LinearAlgebra = "1"
 Metalhead = "0.9.4"
@@ -68,3 +67,6 @@ SparseArrays = "1"
 Statistics = "1"
 StatsBase = "0.34.4"
 julia = "1.10"
+
+[workspace]
+projects = ["docs", "test"]
diff --git a/docs/src/api.md b/docs/src/api.md
index 9f45d03..16d3658 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -72,6 +72,18 @@ Modules = [DecisionFocusedLearningBenchmarks.FixedSizeShortestPath]
 Public = false
 ```
 
+## Maintenance
+
+```@autodocs
+Modules = [DecisionFocusedLearningBenchmarks.Maintenance]
+Private = false
+```
+
+```@autodocs
+Modules = [DecisionFocusedLearningBenchmarks.Maintenance]
+Public = false
+```
+
 ## Portfolio Optimization
 
 ```@autodocs
diff --git a/docs/src/benchmarks/maintenance.md b/docs/src/benchmarks/maintenance.md
new file mode 100644
index 0000000..236501c
--- /dev/null
+++ b/docs/src/benchmarks/maintenance.md
@@ -0,0 +1,107 @@
+# Maintenance problem with resource constraint
+
+The Maintenance problem with resource constraint is a sequential decision-making benchmark where an agent must repeatedly decide which components to maintain over time. The goal is to minimize total expected cost while accounting for independent degradation of components and limited maintenance capacity.
+
+
+## Problem Description
+
+### Overview
+
+In this benchmark, a system consists of $N$ identical components, each of which can degrade over $n$ discrete states. State $1$ means that the component is new, state $n$ means that the component is failed. At each time step, the agent can maintain up to $K$ components.  
+
+This forms an endogenous multistage stochastic optimization problem, where the agent must plan maintenance actions over the horizon.
+
+### Mathematical Formulation
+
+The maintenance problem can be formulated as a finite-horizon Markov Decision Process (MDP) with the following components:
+
+**State Space** $\mathcal{S}$: At time step $t$, the state $s_t \in [1:n]^N$ is the degradation state for each component.
+
+**Action Space** $\mathcal{A}$: The action at time $t$ is the set of components that are maintained at time $t$:
+```math
+a_t \subseteq \{1, 2, \ldots, N\} \text{ such that } |a_t| \leq K
+```
+### Transition Dynamics
+
+The state transitions depend on whether a component is maintained or not:
+
+For each component \(i\) at time \(t\):
+
+- **Maintained component** (\(i \in a_t\)):
+
+\[
+s_{t+1}^i = 1 \quad \text{(perfect maintenance)}
+\]
+
+- **Unmaintained component** (\(i \notin a_t\)):
+
+\[
+s_{t+1}^i =
+\begin{cases}
+\min(s_t^i + 1, n) & \text{with probability } p,\\
+s_t^i & \text{with probability } 1-p.
+\end{cases}
+\]
+
+Here, \(p\) is the degradation probability, \(s_t^i\) is the current state of component \(i\), and \(n\) is the maximum (failed) state.
+
+---
+
+### Cost Function
+
+The immediate cost at time \(t\) is:
+
+$$
+c(s_t, a_t) = \Big( c_m \cdot |a_t| + c_f \cdot \#\{ i : s_t^i = n \} \Big)
+$$
+
+Where:
+
+- $c_m$ is the maintenance cost per component.  
+- $|a_t|$ is the number of components maintained.  
+- $c_f$ is the failure cost per failed component.  
+- $\#\{ i : s_t^i = n \}$ counts the number of components in the failed state.
+
+This formulation captures the total cost for maintaining components and penalizing failures.
+
+**Objective**: Find a policy $\pi: \mathcal{S} \to \mathcal{A}$ that minimizes the expected cumulative cost:
+```math
+\min_\pi \mathbb{E}\left[\sum_{t=1}^T c(s_t, \pi(s_t)) \right]
+```
+
+**Terminal Condition**: The episode terminates after $T$ time steps, with no terminal reward.
+
+## Key Components
+
+### [`MaintenanceBenchmark`](@ref)
+
+The main benchmark configuration with the following parameters:
+
+- `N`: number of components (default: 2)
+- `K`: maximum number of components that can be maintained simultaneously (default: 1) 
+- `n`: number of degradation states per component (default: 3)
+- `p`: degradation probability (default: 0.2)
+- `c_f`: failure cost (default: 10.0)
+- `c_m`: maintenance cost (default: 3.0)
+- `max_steps`: Number of time steps per episode (default: 80)
+
+### Instance Generation
+
+Each problem instance includes:
+
+- **Starting State**: Random starting degradation state in $[1,n]$ for each components.
+
+### Environment Dynamics
+
+The environment tracks:
+- Current time step
+- Current degradation state.
+
+**State Observation**: Agents observe a normalized feature vector containing the degradation state of each component.
+
+## Benchmark Policies
+
+### Greedy Policy  
+
+Greedy policy that maintains components in the last two degradation states, up to the maintenance capacity. This provides a simple baseline.
+
diff --git a/src/DynamicAssortment/environment.jl b/src/DynamicAssortment/environment.jl
index 37a5086..665395e 100644
--- a/src/DynamicAssortment/environment.jl
+++ b/src/DynamicAssortment/environment.jl
@@ -199,17 +199,21 @@ Features observed by the agent at current step, as a concatenation of:
 All features are normalized by dividing by 10.
 
 State
+Return as a tuple:
+- `env.features`: the current feature matrix (feature vector for all items).
+- `env.purchase_history`: the purchase history over the most recent steps.
 """
 function Utils.observe(env::Environment)
     delta_features = env.features[2:3, :] .- env.instance.starting_hype_and_saturation
-    features = vcat(
-        env.features,
-        env.d_features,
-        delta_features,
-        ones(1, item_count(env)) .* (env.step / max_steps(env) * 10),
-    ) ./ 10
-
-    state = (env.features, env.purchase_history)
+    features =
+        vcat(
+            env.features,
+            env.d_features,
+            delta_features,
+            ones(1, item_count(env)) .* (env.step / max_steps(env) * 10),
+        ) ./ 10
+
+    state = (copy(env.features), copy(env.purchase_history))
 
     return features, state
 end
diff --git a/src/Maintenance/Maintenance.jl b/src/Maintenance/Maintenance.jl
index fa99cfe..bd65bfd 100644
--- a/src/Maintenance/Maintenance.jl
+++ b/src/Maintenance/Maintenance.jl
@@ -15,7 +15,7 @@ using Combinatorics: combinations
 $TYPEDEF
 
 Benchmark for a standard maintenance problem with resource constraints.
-Components are identical and degrade idependently over time.
+Components are identical and degrade independently over time.
 A high cost is incurred for each component that reaches the final degradation level. 
 A cost is also incurred for maintaining a component. 
 The number of simultaneous maintenance operations is limited by a maintenance capacity constraint.
@@ -39,6 +39,14 @@ struct MaintenanceBenchmark <: AbstractDynamicBenchmark{true}
     c_m::Float64
     "number of steps per episode"
     max_steps::Int
+
+    function MaintenanceBenchmark(N, K, n, p, c_f, c_m, max_steps)
+        @assert K <= N "number of maintained components $K > number of components $N"
+        @assert K >= 0 && N >= 0 "number of components should be positive"
+        @assert 0 <= p <= 1 "degradation probability $p is not in [0, 1]"
+        # ...
+        return new(N, K, n, p, c_f, c_m, max_steps)
+    end
 end
 
 """
@@ -49,26 +57,15 @@ end
         p=0.2
         c_f=10.0,
         c_m=3.0,
-        max_steps=10,
+        max_steps=80,
     )
 
 Constructor for [`MaintenanceBenchmark`](@ref).
 By default, the benchmark has 2 components, maintenance capacity 1, number of degradation levels 3, 
-degradation probability 0.2, failure cost 10.0, maintenance cost 3.0, 10 steps per episode, and is exogenous.
+degradation probability 0.2, failure cost 10.0, maintenance cost 3.0, 80 steps per episode, and is exogenous.
 """
-
-function MaintenanceBenchmark(;
-    N=2,
-    K=1,
-    n=3,
-    p=0.2,
-    c_f=10.0,
-    c_m=3.0,
-    max_steps=80,
-)
-    return MaintenanceBenchmark(
-        N, K, n, p, c_f, c_m, max_steps
-    )
+function MaintenanceBenchmark(; N=2, K=1, n=3, p=0.2, c_f=10.0, c_m=3.0, max_steps=80)
+    return MaintenanceBenchmark(N, K, n, p, c_f, c_m, max_steps)
 end
 
 # Accessor functions
@@ -90,9 +87,7 @@ $TYPEDSIGNATURES
 
 Outputs a data sample containing an [`Instance`](@ref).
 """
-function Utils.generate_sample(
-    b::MaintenanceBenchmark, rng::AbstractRNG=MersenneTwister(0)
-)
+function Utils.generate_sample(b::MaintenanceBenchmark, rng::AbstractRNG)
     return DataSample(; instance=Instance(b, rng))
 end
 
@@ -105,7 +100,7 @@ The model is a small neural network with one hidden layer no activation function
 function Utils.generate_statistical_model(b::MaintenanceBenchmark; seed=nothing)
     Random.seed!(seed)
     N = component_count(b)
-    return Chain(Dense(N  => N), Dense(N => N), vec)
+    return Chain(Dense(N => N), Dense(N => N), vec)
 end
 
 """
diff --git a/src/Maintenance/environment.jl b/src/Maintenance/environment.jl
index b17d3d3..60193c6 100644
--- a/src/Maintenance/environment.jl
+++ b/src/Maintenance/environment.jl
@@ -26,40 +26,33 @@ $TYPEDSIGNATURES
 Creates an [`Environment`](@ref) from an [`Instance`](@ref) of the maintenance benchmark.
 """
 function Environment(instance::Instance; seed=0, rng::AbstractRNG=MersenneTwister(seed))
-    degradation_state = starting_state(instance)
-    env = Environment(;
-        instance,
-        step=1,
-        degradation_state,
-        rng=rng,
-        seed=seed,
-    )
+    degradation_state = copy(starting_state(instance))
+    env = Environment(; instance, step=1, degradation_state, rng=rng, seed=seed)
     Utils.reset!(env; reset_rng=true)
     return env
 end
 
-component_count(env::Environment) = component_count(env.instance) 
+component_count(env::Environment) = component_count(env.instance)
 maintenance_capacity(env::Environment) = maintenance_capacity(env.instance)
 degradation_levels(env::Environment) = degradation_levels(env.instance)
 degradation_probability(env::Environment) = degradation_probability(env.instance)
 failure_cost(env::Environment) = failure_cost(env.instance)
 maintenance_cost(env::Environment) = maintenance_cost(env.instance)
-max_steps(env::Environment) = max_steps(env.instance) 
+max_steps(env::Environment) = max_steps(env.instance)
 starting_state(env::Environment) = starting_state(env.instance)
 
-
 """
 $TYPEDSIGNATURES
 Draw random degradations for all components.
 """
-
 function degrad!(env::Environment)
-    N = component_count(env) 
+    N = component_count(env)
     n = degradation_levels(env)
     p = degradation_probability(env)
+    rng = env.rng
 
     for i in 1:N
-        if env.degradation_state[i] < n && rand() < p
+        if env.degradation_state[i] < n && rand(rng) < p
             env.degradation_state[i] += 1
         end
     end
@@ -71,9 +64,8 @@ end
 $TYPEDSIGNATURES
 Maintain components.
 """
-
 function maintain!(env::Environment, maintenance::BitVector)
-    N = component_count(env) 
+    N = component_count(env)
 
     for i in 1:N
         if maintenance[i]
@@ -99,12 +91,10 @@ $TYPEDSIGNATURES
 Compute degradation cost.
 """
 function degradation_cost(env::Environment)
-    N = component_count(env) 
     n = degradation_levels(env)
     return failure_cost(env) * count(==(n), env.degradation_state)
 end
 
-
 """
 $TYPEDSIGNATURES
 
@@ -163,5 +153,3 @@ function Utils.step!(env::Environment, maintenance::BitVector)
     env.step += 1
     return cost
 end
-
-
diff --git a/src/Maintenance/instance.jl b/src/Maintenance/instance.jl
index aef280e..4a2f4fb 100644
--- a/src/Maintenance/instance.jl
+++ b/src/Maintenance/instance.jl
@@ -6,9 +6,9 @@ Instance of the maintenance problem.
 # Fields
 $TYPEDFIELDS
 """
-@kwdef struct Instance{B<:MaintenanceBenchmark}
+@kwdef struct Instance{MaintenanceBenchmark}
     "associated benchmark"
-    config::B
+    config::MaintenanceBenchmark
     "starting degradation states"
     starting_state::Vector{Int}
 end
@@ -21,16 +21,16 @@ Generates an instance with random starting degradation states uniformly in [1, n
 function Instance(b::MaintenanceBenchmark, rng::AbstractRNG)
     N = component_count(b)
     n = degradation_levels(b)
-    starting_state = rand(rng, 1:n, N) 
-    return Instance(; config=b, starting_state)
+    starting_state = rand(rng, 1:n, N)
+    return Instance(; config=b, starting_state=starting_state)
 end
 
 # Accessor functions
-component_count(b::Instance) = component_count(b.config) 
+component_count(b::Instance) = component_count(b.config)
 maintenance_capacity(b::Instance) = maintenance_capacity(b.config)
 degradation_levels(b::Instance) = degradation_levels(b.config)
 degradation_probability(b::Instance) = degradation_probability(b.config)
 failure_cost(b::Instance) = failure_cost(b.config)
 maintenance_cost(b::Instance) = maintenance_cost(b.config)
-max_steps(b::Instance) = max_steps(b.config) 
+max_steps(b::Instance) = max_steps(b.config)
 starting_state(b::Instance) = b.starting_state
diff --git a/src/Maintenance/maximizer.jl b/src/Maintenance/maximizer.jl
index a271234..1e7a653 100644
--- a/src/Maintenance/maximizer.jl
+++ b/src/Maintenance/maximizer.jl
@@ -12,15 +12,21 @@ $TYPEDSIGNATURES
 
 Return the top k indices of `θ`.
 """
-
 function (m::TopKPositiveMaximizer)(θ; kwargs...)
     N = length(θ)
-    
-    sorted_indices = sortperm(θ; rev=true)
-    positive_indices = filter(i -> θ[i] > 0, sorted_indices)
-    solution = positive_indices[1:min(m.k, length(positive_indices))]
-    
+
+    positive_indices = findall(x -> x > 0, θ)
+    nb_positive = length(positive_indices)
     res = falses(N)
-    res[solution] .= 1
-    return res
-end
\ No newline at end of file
+
+    if nb_positive == 0
+        return res
+    elseif nb_positive <= m.k
+        res[positive_indices] .= true
+        return res
+    else
+        idx = partialsortperm(θ[positive_indices], 1:(m.k); rev=true)
+        res[positive_indices[idx]] .= true
+        return res
+    end
+end
diff --git a/src/Maintenance/policies.jl b/src/Maintenance/policies.jl
index 100b2ed..89de852 100644
--- a/src/Maintenance/policies.jl
+++ b/src/Maintenance/policies.jl
@@ -8,20 +8,19 @@ function greedy_policy(env::Environment)
     state = env.degradation_state
     N = component_count(env)
     K = maintenance_capacity(env)
-    res = falses(N)   
+    res = falses(N)
     n = degradation_levels(env)
-    
-    
+
     idx_max = findall(==(n), state)
     take = first(idx_max, min(K, length(idx_max)))
     res[take] .= true
     remaining = K - length(take)
 
     if remaining > 0
-        idx_second = findall(==(n-1), state)
+        idx_second = findall(==(n - 1), state)
         take2 = first(idx_second, min(remaining, length(idx_second)))
         res[take2] .= true
     end
 
-    return res 
+    return res
 end
diff --git a/test/dynamic_assortment.jl b/test/dynamic_assortment.jl
index a3feddb..97674bf 100644
--- a/test/dynamic_assortment.jl
+++ b/test/dynamic_assortment.jl
@@ -246,22 +246,28 @@ end
     instance = DAP.Instance(b, MersenneTwister(42))
     env = DAP.Environment(instance; seed=123)
 
-    obs, info = observe(env)
+    features, state = observe(env)
 
     # Check observation dimensions: (d + 8, N)
     # Features: prices(1) + hype_sat(2) + static(d) + d_features(2) + delta_features(2) + step(1)
     expected_rows = 2 + 8  # d + 8 where d=2
-    @test size(obs) == (expected_rows, 3)
-    @test info === nothing
+    @test size(features) == (expected_rows, 3)
 
-    @test all(-1.0 ≤ x ≤ 1.0 for x in obs)
+    v, purchase_history = state
+    @test size(v) == (5, 3)
+
+    @test all(-1.0 ≤ x ≤ 1.0 for x in features)
 
     # Test observation changes with step
-    obs1, _ = observe(env)
+    features1, state1 = observe(env)
+    v, purchase_history1 = state1
+
     DAP.buy_item!(env, 1)
-    obs2, _ = observe(env)
+    features2, state2 = observe(env)
+    _, purchase_history2 = state2
 
-    @test obs1 != obs2  # Observations should differ after purchase
+    @test purchase_history1 != purchase_history2  # Observations should differ after purchase
+    @test features1 != features2
 end
 
 @testset "DynamicAssortment - Policies" begin
diff --git a/test/maintenance.jl b/test/maintenance.jl
index 53ac7b3..6cdc51d 100644
--- a/test/maintenance.jl
+++ b/test/maintenance.jl
@@ -47,8 +47,8 @@ end
     @test state1 != state2
 
     # Test instance structure
-    @test length(instance.starting_state) == 10 
-    @test all(1.0 ≤ s ≤ 5 for s in instance.starting_state)  
+    @test length(instance.starting_state) == 10
+    @test all(1.0 ≤ s ≤ 5 for s in instance.starting_state)
 
     # Test accessor functions
     @test maintenance.component_count(instance) == 10
@@ -101,18 +101,18 @@ end
     instance = maintenance.Instance(b, MersenneTwister(42))
     env = maintenance.Environment(instance; seed=123)
 
-    env.degradation_state = [1,1]
-    @test maintenance.maintenance_cost(env, BitVector([false, false])) == 0.0 
-    @test maintenance.maintenance_cost(env, BitVector([false, true])) == 3.0 
-    @test maintenance.maintenance_cost(env, BitVector([true, true])) == 6.0 
-
-    @test maintenance.degradation_cost(env) == 0.0 
-    env.degradation_state = [2,2]
-    @test maintenance.degradation_cost(env) == 0.0 
-    env.degradation_state = [3,2]
-    @test maintenance.degradation_cost(env) == 10.0 
-    env.degradation_state = [3,3]
-    @test maintenance.degradation_cost(env) == 20.0 
+    env.degradation_state = [1, 1]
+    @test maintenance.maintenance_cost(env, BitVector([false, false])) == 0.0
+    @test maintenance.maintenance_cost(env, BitVector([false, true])) == 3.0
+    @test maintenance.maintenance_cost(env, BitVector([true, true])) == 6.0
+
+    @test maintenance.degradation_cost(env) == 0.0
+    env.degradation_state = [2, 2]
+    @test maintenance.degradation_cost(env) == 0.0
+    env.degradation_state = [3, 2]
+    @test maintenance.degradation_cost(env) == 10.0
+    env.degradation_state = [3, 3]
+    @test maintenance.degradation_cost(env) == 20.0
 end
 
 @testset "Maintenance - Environment Step" begin
@@ -146,20 +146,19 @@ end
     b = MaintenanceBenchmark()
     instance = maintenance.Instance(b, MersenneTwister(42))
     env = maintenance.Environment(instance; seed=123)
-    env.degradation_state = [1,1]
+    env.degradation_state = [1, 1]
 
     state, features = observe(env)
 
-    @test state == [1,1]
+    @test state == [1, 1]
     @test features === state
 
-    env.degradation_state = [2,3]
+    env.degradation_state = [2, 3]
     state2, _ = observe(env)
 
     @test state != state2  # Observations should differ after purchase
 end
 
-
 @testset "Maintenance - Policies" begin
     using Statistics: mean
 
@@ -189,7 +188,6 @@ end
     @test greedy_action isa BitVector && length(greedy_action) == 2
 end
 
-
 @testset "Maintenance - Model and Maximizer Integration" begin
     b = MaintenanceBenchmark()
 
@@ -217,15 +215,14 @@ end
 
     @test length(θ) == 2
 
-    θ = [1,2]
+    θ = [1, 2]
     @test maximizer(θ) == BitVector([false, true])
 
     b = MaintenanceBenchmark(; N=10, K=3, n=5, p=0.3, c_f=5.0, c_m=3.0, max_steps=50)
     θ = [i for i in 1:10]
     maximizer = generate_maximizer(b)
-    @test maximizer(θ) == BitVector([false, false, false, false, false, false, false, true, true, true])
-
-
+    @test maximizer(θ) ==
+        BitVector([false, false, false, false, false, false, false, true, true, true])
 
     #test maximizer output
 end

From 873111396ed08310083fc57508a25a6b8e3f45f7 Mon Sep 17 00:00:00 2001
From: BatyLeo <leo.baty67@gmail.com>
Date: Fri, 9 Jan 2026 19:23:44 +0100
Subject: [PATCH 5/6] Fix tests

---
 Project.toml                         | 10 ++++------
 src/DynamicAssortment/environment.jl |  2 +-
 src/Maintenance/Maintenance.jl       |  1 -
 src/Maintenance/environment.jl       |  6 +++---
 4 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/Project.toml b/Project.toml
index b8daf1c..a95bb3e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,10 @@
 name = "DecisionFocusedLearningBenchmarks"
 uuid = "2fbe496a-299b-4c81-bab5-c44dfc55cf20"
-authors = ["Members of JuliaDecisionFocusedLearning"]
 version = "0.4.0"
+authors = ["Members of JuliaDecisionFocusedLearning"]
+
+[workspace]
+projects = ["docs", "test"]
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
@@ -20,7 +23,6 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
-JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc"
@@ -52,7 +54,6 @@ Ipopt = "1.6"
 IterTools = "1.10.0"
 JSON = "1"
 JuMP = "1.22"
-JuliaFormatter = "1.0.62"
 LaTeXStrings = "1.4.0"
 LinearAlgebra = "1"
 Metalhead = "0.9.4"
@@ -67,6 +68,3 @@ SparseArrays = "1"
 Statistics = "1"
 StatsBase = "0.34.4"
 julia = "1.10"
-
-[workspace]
-projects = ["docs", "test"]
diff --git a/src/DynamicAssortment/environment.jl b/src/DynamicAssortment/environment.jl
index 665395e..5236e2c 100644
--- a/src/DynamicAssortment/environment.jl
+++ b/src/DynamicAssortment/environment.jl
@@ -7,7 +7,7 @@ Environment for the dynamic assortment problem.
 $TYPEDFIELDS
 """
 @kwdef mutable struct Environment{I<:Instance,R<:AbstractRNG,S<:Union{Nothing,Int}} <:
-                      Utils.AbstractEnvironment
+                      AbstractEnvironment
     "associated instance"
     instance::I
     "current step"
diff --git a/src/Maintenance/Maintenance.jl b/src/Maintenance/Maintenance.jl
index bd65bfd..5dc0580 100644
--- a/src/Maintenance/Maintenance.jl
+++ b/src/Maintenance/Maintenance.jl
@@ -44,7 +44,6 @@ struct MaintenanceBenchmark <: AbstractDynamicBenchmark{true}
         @assert K <= N "number of maintained components $K > number of components $N"
         @assert K >= 0 && N >= 0 "number of components should be positive"
         @assert 0 <= p <= 1 "degradation probability $p is not in [0, 1]"
-        # ...
         return new(N, K, n, p, c_f, c_m, max_steps)
     end
 end
diff --git a/src/Maintenance/environment.jl b/src/Maintenance/environment.jl
index 60193c6..6d29ec2 100644
--- a/src/Maintenance/environment.jl
+++ b/src/Maintenance/environment.jl
@@ -6,10 +6,10 @@ Environment for the maintenance problem.
 # Fields
 $TYPEDFIELDS
 """
-@kwdef mutable struct Environment{I<:Instance,R<:AbstractRNG,S<:Union{Nothing,Int}} <:
-                      Utils.AbstractEnvironment
+@kwdef mutable struct Environment{R<:AbstractRNG,S<:Union{Nothing,Int}} <:
+                      AbstractEnvironment
     "associated instance"
-    instance::I
+    instance::Instance
     "current step"
     step::Int
     "degradation state"

From 9894b43496760a2ea139d14c33d3692cc44e51d1 Mon Sep 17 00:00:00 2001
From: BatyLeo <leo.baty67@gmail.com>
Date: Mon, 12 Jan 2026 08:30:42 +0100
Subject: [PATCH 6/6] coverage

---
 test/maintenance.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/maintenance.jl b/test/maintenance.jl
index 6cdc51d..d7183f5 100644
--- a/test/maintenance.jl
+++ b/test/maintenance.jl
@@ -224,5 +224,6 @@ end
     @test maximizer(θ) ==
         BitVector([false, false, false, false, false, false, false, true, true, true])
 
-    #test maximizer output
+    θ = fill(-1.0, 10)
+    @test maximizer(θ) == falses(10)
 end