Fix tests

Author: BatyLeo

src/Argmax2D/Argmax2D.jl

@@ -7,7 +7,7 @@ using Flux: Chain, Dense
 using LaTeXStrings: @L_str
 using LinearAlgebra: dot, norm
 using Plots: Plots
-using Random: Random, MersenneTwister
+using Random: Random, MersenneTwister, AbstractRNG
 
 include("polytope.jl")
 
@@ -53,20 +53,16 @@ maximizer(θ; instance, kwargs...) = instance[argmax(dot(θ, v) for v in instanc
 """
 $TYPEDSIGNATURES
 
-Generate a dataset for the [`Argmax2DBenchmark`](@ref).
+Generate a sample for the [`Argmax2DBenchmark`](@ref).
 """
-function Utils.generate_dataset(
-    bench::Argmax2DBenchmark, dataset_size=10; seed=nothing, rng=MersenneTwister(seed)
-)
+function Utils.generate_sample(bench::Argmax2DBenchmark, rng::AbstractRNG)
     (; nb_features, encoder, polytope_vertex_range) = bench
-    return map(1:dataset_size) do _
-        x = randn(rng, Float32, nb_features)
-        θ_true = encoder(x)
-        θ_true ./= 2 * norm(θ_true)
-        instance = build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng))
-        y_true = maximizer(θ_true; instance)
-        return DataSample(; x=x, θ_true=θ_true, y_true=y_true, instance=instance)
-    end
+    x = randn(rng, Float32, nb_features)
+    θ_true = encoder(x)
+    θ_true ./= 2 * norm(θ_true)
+    instance = build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng))
+    y_true = maximizer(θ_true; instance)
+    return DataSample(; x=x, θ_true=θ_true, y_true=y_true, instance=instance)
 end
 
 """

src/DynamicAssortment/DynamicAssortment.jl

@@ -48,6 +48,7 @@ end
 
 include("instance.jl")
 include("environment.jl")
+include("policies.jl")
 
 customer_choice_model(b::DynamicAssortmentBenchmark) = b.customer_choice_model
 item_count(b::DynamicAssortmentBenchmark) = b.N
@@ -56,9 +57,8 @@ assortment_size(b::DynamicAssortmentBenchmark) = b.K
 max_steps(b::DynamicAssortmentBenchmark) = b.max_steps
 
 function Utils.generate_sample(
-    b::DynamicAssortmentBenchmark, rng::AbstractRNG=MersenneTwister(0); seed=nothing
+    b::DynamicAssortmentBenchmark, rng::AbstractRNG=MersenneTwister(0)
 )
-    Random.seed!(rng, seed)
     return DataSample(; instance=Instance(b, rng))
 end
 

src/DynamicAssortment/environment.jl

@@ -166,45 +166,3 @@ function compute_expected_revenue(env::Environment, S)
     expected_revenue = dot(probs, r)
     return expected_revenue
 end
-
-function expert_solution(env::Environment)
-    N = item_count(env)
-    K = assortment_size(env)
-    best_S = falses(N)
-    best_revenue = -1.0
-    S_vec = falses(N)
-    for S in combinations(1:N, K)
-        S_vec .= false
-        S_vec[S] .= true
-        expected_revenue = compute_expected_revenue(env, S_vec)
-        if expected_revenue > best_revenue
-            best_S, best_revenue = copy(S_vec), expected_revenue
-        end
-    end
-    return best_S
-end
-
-function greedy_solution(env::Environment)
-    maximizer = generate_maximizer(env.instance.config)
-    return maximizer(prices(env))
-end
-
-function run_policy(env::Environment, episodes::Int; first_seed=1, policy=expert_solution)
-    dataset = []
-    rev_global = Float64[]
-    for i in 1:episodes
-        rev_episode = 0.0
-        CommonRLInterface.reset!(env; seed=first_seed - 1 + i, reset_seed=true)
-        training_instances = []
-        while !CommonRLInterface.terminated(env)
-            S = policy(env)
-            features = CommonRLInterface.observe(env)
-            push!(training_instances, DataSample(; x=features, y_true=S))
-            reward = CommonRLInterface.act!(env, S)
-            rev_episode += reward
-        end
-        push!(rev_global, rev_episode)
-        push!(dataset, training_instances)
-    end
-    return mean(rev_global), rev_global, dataset
-end

src/DynamicAssortment/policies.jl

@@ -0,0 +1,41 @@
+function expert_policy(env::Environment)
+    N = item_count(env)
+    K = assortment_size(env)
+    best_S = falses(N)
+    best_revenue = -1.0
+    S_vec = falses(N)
+    for S in combinations(1:N, K)
+        S_vec .= false
+        S_vec[S] .= true
+        expected_revenue = compute_expected_revenue(env, S_vec)
+        if expected_revenue > best_revenue
+            best_S, best_revenue = copy(S_vec), expected_revenue
+        end
+    end
+    return best_S
+end
+
+function greedy_policy(env::Environment)
+    maximizer = generate_maximizer(env.instance.config)
+    return maximizer(prices(env))
+end
+
+function run_policy(env::Environment, episodes::Int; first_seed=1, policy=expert_policy)
+    dataset = []
+    rev_global = Float64[]
+    for i in 1:episodes
+        rev_episode = 0.0
+        CommonRLInterface.reset!(env; seed=first_seed - 1 + i, reset_seed=true)
+        training_instances = []
+        while !CommonRLInterface.terminated(env)
+            S = policy(env)
+            features = CommonRLInterface.observe(env)
+            push!(training_instances, DataSample(; x=features, y_true=S))
+            reward = CommonRLInterface.act!(env, S)
+            rev_episode += reward
+        end
+        push!(rev_global, rev_episode)
+        push!(dataset, training_instances)
+    end
+    return mean(rev_global), rev_global, dataset
+end

src/DynamicVehicleScheduling/algorithms/anticipative_solver.jl

@@ -201,9 +201,11 @@ end
 end
 
 function (solver::AnticipativeSolver)(env::DVSPEnv, scenario=env.scenario; reset_env=false)
-    if solver.is_2D
-        return anticipative_solver(env, scenario; model_builder=highs_model_2d, reset_env)
-    else
-        return anticipative_solver(env, scenario; model_builder=highs_model, reset_env)
-    end
+    return anticipative_solver(
+        env,
+        scenario;
+        model_builder=highs_model,
+        reset_env,
+        two_dimensional_features=solver.is_2D,
+    )
 end