test: connect and improve tests

Author: BatyLeo

test/Project.toml

@@ -7,6 +7,7 @@ InferOpt = "4846b161-c94e-4150-8dac-c7ae193c601f"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ValueHistories = "98cad3c8-aec3-5f06-8e41-884608649ab7"
 

test/mirror_descent.jl

@@ -4,6 +4,19 @@ using Test
 using ValueHistories
 using Statistics: mean
 
+function _val_obj_metric(benchmark, val_data)
+    return FunctionMetric(:val_obj, val_data) do ctx, data
+        vals = map(data) do s
+            θ = ctx.policy.statistical_model(s.x)
+            y = ctx.policy.maximizer(θ; s.context...)
+            return Float64(
+                DecisionFocusedLearningBenchmarks.objective_value(benchmark, s, y)
+            )
+        end
+        return (val_obj=mean(vals),)
+    end
+end
+
 @testset "MirrorDescent Training" begin
     @testset "MirrorDescent - ContextualStochasticArgmax basic" begin
         benchmark = ContextualStochasticArgmaxBenchmark()
@@ -105,4 +118,44 @@ using Statistics: mean
 
         @test all(haskey(h, :epoch) for h in histories)
     end
+
+    @testset "MirrorDescent - trained beats untrained (ContextualStochasticArgmax)" begin
+        benchmark = ContextualStochasticArgmaxBenchmark()
+        val_data = generate_dataset(benchmark, 50; seed=99)
+
+        histories, _ = train_policy(
+            MirrorDescent(),
+            benchmark;
+            dataset_size=20,
+            epochs=5,
+            iterations=5,
+            seed=0,
+            metrics=(_val_obj_metric(benchmark, val_data),),
+        )
+
+        obj_untrained = get(histories[1], :val_obj)[2][1]
+        obj_trained = get(histories[end], :val_obj)[2][end]
+
+        @test obj_trained > obj_untrained
+    end
+
+    @testset "MirrorDescent - trained beats untrained (StochasticVehicleScheduling)" begin
+        benchmark = StochasticVehicleSchedulingBenchmark()
+        val_data = generate_dataset(benchmark, 10; seed=99)
+
+        histories, _ = train_policy(
+            MirrorDescent(),
+            benchmark;
+            dataset_size=10,
+            epochs=5,
+            iterations=2,
+            seed=0,
+            metrics=(_val_obj_metric(benchmark, val_data),),
+        )
+
+        obj_untrained = get(histories[1], :val_obj)[2][1]
+        obj_trained = get(histories[end], :val_obj)[2][end]
+
+        @test obj_trained < obj_untrained
+    end
 end

test/runtests.jl

@@ -13,4 +13,8 @@ using DecisionFocusedLearningAlgorithms
     @testset "DAgger" begin
         include("dagger.jl")
     end
+
+    @testset "MirrorDescent" begin
+        include("mirror_descent.jl")
+    end
 end