feat!: Better reproducibility management of dynamic problems with SeededEnvironment (#73)

* refactor: replace MersenneTwister with Xoshiro

* refactor: replace MersenneTwister with Xoshiro

* feat: implement the SeededEnv wrapper, and adapt and cleanup policy evaluation and various dynamic environments

* feat: user now needs to implement build_environment insteadof generate_environment

* test: update tests + unit tests for SeededEnvironment

* doc: update documentation

* cleanup

Author: BatyLeo

docs/src/custom_benchmarks.md

@@ -121,26 +121,39 @@ Dynamic benchmarks extend stochastic ones with an environment-based rollout inte
 ### Environment generation
 
 ```julia
-# Strategy A: generate one environment at a time (default implementation of
-#              generate_environments calls this repeatedly)
-generate_environment(bench::MyDynamicBenchmark, rng::AbstractRNG; kwargs...) -> AbstractEnvironment
+# Option A: build one bare environment (the framework wraps it in a SeededEnvironment)
+build_environment(bench::MyDynamicBenchmark, rng::AbstractRNG; kwargs...) -> AbstractEnvironment
 
-# Strategy B: override when environments are not independent (e.g. loaded from files)
-generate_environments(bench::MyDynamicBenchmark, n::Int; rng, kwargs...) -> Vector{<:AbstractEnvironment}
+# Option B: override when environments are not independent (e.g. loaded from files)
+generate_environments(bench::MyDynamicBenchmark, n::Int; seed=nothing, kwargs...) -> Vector{SeededEnvironment}
 ```
 
+You implement **`build_environment`**, which returns a *bare* environment. The package then automatically wraps it in a [`SeededEnvironment`](@ref) (attaching a seed and RNG) through two public entry points:
+
+```julia
+generate_environment(bench; seed=nothing, kwargs...)      -> SeededEnvironment       # one env
+generate_environments(bench, n; seed=nothing, kwargs...)  -> Vector{SeededEnvironment} # n envs
+```
+
+For environments that cannot be drawn independently (e.g. loaded from files), override `generate_environments` instead of implementing `build_environment`.
+An override must return already-wrapped `SeededEnvironment`s (use `SeededEnvironment(env; seed=...)`).
+Do not override `generate_environment`: it just delegates to `generate_environments`.
+
 ### `AbstractEnvironment` interface
 
-Your environment type must implement:
+Your environment must be **stateless about randomness**: it does *not* store its own RNG or seed.
+All randomness is owned by the wrapping [`SeededEnvironment`](@ref), which passes its `rng` into every `reset!` and `step!`.
+Draw all stochasticity from that `rng` so that re-seeding the wrapper (via [`reset_to_initial!`](@ref)) replays an episode exactly.
 
 ```julia
-get_seed(env::MyEnv)                             # Return the RNG seed used at creation
-reset!(env::MyEnv; reset_rng::Bool, seed=get_seed(env))  # Reset to initial state
-observe(env::MyEnv) -> (observation, info)       # Current observation
-step!(env::MyEnv, action) -> reward              # Apply action, advance state
-is_terminated(env::MyEnv) -> Bool                # True when episode has ended
+reset!(env::MyEnv, rng::AbstractRNG)                    # Reset to a starting state, sampling from rng
+observe(env::MyEnv) -> (observation, state)             # Current observation and internal state
+step!(env::MyEnv, action, rng::AbstractRNG) -> reward   # Apply action, advance state (draw from rng)
+is_terminated(env::MyEnv) -> Bool                       # True when the episode has ended
 ```
 
+Environments may ignore the `rng` argument.
+
 ### Baseline policies (required for `generate_dataset`)
 
 ```julia

docs/src/using_benchmarks.md

@@ -132,15 +132,25 @@ rollout and returns the resulting trajectory.
 
 ## Seed / RNG control
 
-All `generate_dataset` and `generate_environments` calls accept either `seed`
-(creates an internal `MersenneTwister`) or `rng` for full control:
+`generate_dataset` accepts either `seed` (creates an internal `Xoshiro`) or `rng` (any `AbstractRNG`)::
 
 ```julia
 using Random
-rng = MersenneTwister(42)
+rng = Xoshiro(42)
 dataset = generate_dataset(bench, 50; rng=rng)
 ```
 
+`generate_environments` takes a `seed`:
+
+```julia
+envs = generate_environments(bench, 10; seed=0)
+```
+
+### Reproducibility in dynamic benchmarks
+
+Each environment returned by `generate_environments` or `generate_environment` is a `SeededEnvironment`: a wrapper that owns the random number generator and a stored seed.
+It is the single source of randomness for the episode, so re-running a policy on the same environment replays the exact same trajectory. `evaluate_policy!` resets to the stored seed before each run, which is what makes evaluation reproducible (pass `seed=...` to override the stored seed for a given run).
+
 ---
 
 ## Evaluation

src/Argmax2D/Argmax2D.jl

@@ -4,7 +4,7 @@ using ..Utils
 using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Chain, Dense
 using LinearAlgebra: dot, norm
-using Random: Random, MersenneTwister, AbstractRNG
+using Random: Random, Xoshiro, AbstractRNG
 
 include("polytope.jl")
 
@@ -79,7 +79,7 @@ $TYPEDSIGNATURES
 Generate a statistical model for the [`Argmax2DBenchmark`](@ref).
 """
 function Utils.generate_statistical_model(
-    bench::Argmax2DBenchmark; seed=nothing, rng=MersenneTwister(seed)
+    bench::Argmax2DBenchmark; seed=nothing, rng=Xoshiro(seed)
 )
     Random.seed!(rng, seed)
     (; nb_features) = bench

src/ContextualStochasticArgmax/ContextualStochasticArgmax.jl

@@ -4,7 +4,7 @@ using ..Utils
 using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Dense
 using LinearAlgebra: dot
-using Random: Random, AbstractRNG, MersenneTwister
+using Random: Random, AbstractRNG, Xoshiro
 using Statistics: mean
 
 """
@@ -37,7 +37,7 @@ end
 function ContextualStochasticArgmaxBenchmark(;
     n::Int=10, d::Int=5, noise_std::Float32=0.1f0, seed=nothing
 )
-    rng = MersenneTwister(seed)
+    rng = Xoshiro(seed)
     W = randn(rng, Float32, n, d)
     return ContextualStochasticArgmaxBenchmark(n, d, W, noise_std)
 end

src/DecisionFocusedLearningBenchmarks.jl

@@ -66,12 +66,17 @@ using .Utils
 export AbstractBenchmark, AbstractStochasticBenchmark, AbstractDynamicBenchmark, DataSample
 export ExogenousStochasticBenchmark,
     EndogenousStochasticBenchmark, ExogenousDynamicBenchmark, EndogenousDynamicBenchmark
-export AbstractEnvironment, get_seed, is_terminated, observe, reset!, step!
+export AbstractEnvironment, SeededEnvironment
+export get_seed, is_terminated, observe, reset!, reset_to_initial!, step!
 
 export Policy, evaluate_policy!
 
 export generate_instance,
-    generate_sample, generate_dataset, generate_environments, generate_environment
+    generate_sample,
+    generate_dataset,
+    build_environment,
+    generate_environment,
+    generate_environments
 export generate_scenario, generate_context
 export generate_baseline_policies
 export SampleAverageApproximation

src/DynamicAssortment/DynamicAssortment.jl

@@ -6,7 +6,7 @@ using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES, SIGNATURES
 using Distributions: Uniform, Categorical
 using Flux: Chain, Dense
 using LinearAlgebra: dot
-using Random: Random, AbstractRNG, MersenneTwister
+using Random: Random, AbstractRNG, Xoshiro
 using Statistics: mean
 
 using Combinatorics: combinations
@@ -100,14 +100,11 @@ end
 $TYPEDSIGNATURES
 
 Creates an [`Environment`](@ref) for the dynamic assortment benchmark.
-The instance and seed are randomly generated using the provided random number generator.
+The instance is randomly generated using the provided random number generator.
 """
-function Utils.generate_environment(
-    b::DynamicAssortmentBenchmark, rng::AbstractRNG; kwargs...
-)
+function Utils.build_environment(b::DynamicAssortmentBenchmark, rng::AbstractRNG; kwargs...)
     instance = Instance(b, rng)
-    seed = rand(rng, 1:typemax(Int))
-    return Environment(instance; seed)
+    return Environment(instance)
 end
 
 """

src/DynamicAssortment/environment.jl

@@ -6,18 +6,13 @@ Environment for the dynamic assortment problem.
 # Fields
 $TYPEDFIELDS
 """
-@kwdef mutable struct Environment{I<:Instance,R<:AbstractRNG,S<:Union{Nothing,Int}} <:
-                      AbstractEnvironment
+@kwdef mutable struct Environment{I<:Instance} <: AbstractEnvironment
     "associated instance"
     instance::I
     "current step"
     step::Int
     "purchase history (used to update hype feature)"
     purchase_history::Vector{Int}
-    "rng"
-    rng::R
-    "seed for RNG"
-    seed::S
     "customer utility for each item"
     utility::Vector{Float64}
     "current full features"
@@ -31,25 +26,21 @@ $TYPEDSIGNATURES
 
 Creates an [`Environment`](@ref) from an [`Instance`](@ref) of the dynamic assortment benchmark.
 """
-function Environment(instance::Instance; seed=0, rng::AbstractRNG=MersenneTwister(seed))
+function Environment(instance::Instance)
     N = item_count(instance)
     (; prices, features, starting_hype_and_saturation) = instance
     full_features = vcat(
         reshape(prices[1:(end - 1)], 1, :), starting_hype_and_saturation, features
     )
     model = customer_choice_model(instance)
-    env = Environment(;
+    return Environment(;
         instance,
         step=1,
         purchase_history=Int[],
-        rng=rng,
-        seed=seed,
         utility=model(full_features),
         features=full_features,
         d_features=zeros(2, N),
     )
-    Utils.reset!(env; reset_rng=true)
-    return env
 end
 
 customer_choice_model(env::Environment) = customer_choice_model(env.instance)
@@ -144,36 +135,26 @@ 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 features to the initial features
 - reset the change in features to zero
 - reset the utility to the initial utility
 - clear the purchase history
 """
-function Utils.reset!(env::Environment; reset_rng=false, seed=env.seed)
-    reset_rng && Random.seed!(env.rng, seed)
-
+function Utils.reset!(env::Environment, ::AbstractRNG)
     env.step = 1
 
     (; prices, starting_hype_and_saturation, features) = env.instance
-    features = vcat(
+    full_features = vcat(
         reshape(prices[1:(end - 1)], 1, :), starting_hype_and_saturation, features
     )
-    env.features .= features
+    env.features .= full_features
 
     env.d_features .= 0.0
 
     model = customer_choice_model(env)
-    env.utility .= model(features)
+    env.utility .= model(full_features)
 
     empty!(env.purchase_history)
     return nothing
@@ -224,12 +205,11 @@ $TYPEDSIGNATURES
 Performs one step in the environment given an assortment.
 Draw an item according to the customer choice model and updates the environment state.
 """
-function Utils.step!(env::Environment, assortment::BitVector)
+function Utils.step!(env::Environment, assortment::BitVector, rng::AbstractRNG)
     @assert !Utils.is_terminated(env) "Environment is terminated, cannot act!"
-    r = prices(env)
     probs = choice_probabilities(env, assortment)
-    item = rand(env.rng, Categorical(probs))
-    reward = r[item]
+    item = rand(rng, Categorical(probs))
+    reward = prices(env)[item]
     buy_item!(env, item)
     return reward
 end

src/DynamicVehicleScheduling/DynamicVehicleScheduling.jl

@@ -13,7 +13,7 @@ using IterTools: partition
 using JSON
 using JuMP
 using Printf: @printf, @sprintf
-using Random: Random, AbstractRNG, MersenneTwister, seed!, randperm
+using Random: Random, AbstractRNG, Xoshiro, seed!, randperm
 using Requires: @require
 using Statistics: mean, quantile
 
@@ -59,50 +59,37 @@ include("policy.jl")
 $TYPEDSIGNATURES
 
 Generate environments for the dynamic vehicle scheduling benchmark.
-Reads from pre-existing DVRPTW files and creates [`DVSPEnv`](@ref) environments.
+Reads from pre-existing DVRPTW files and creates [`DVSPEnv`](@ref) environments
+wrapped in [`SeededEnvironment`](@ref) for reproducible episode replay.
 """
 function Utils.generate_environments(
-    b::DynamicVehicleSchedulingBenchmark,
-    n::Int;
-    seed=nothing,
-    rng=MersenneTwister(seed),
-    kwargs...,
+    b::DynamicVehicleSchedulingBenchmark, n::Int; seed=nothing, rng=Xoshiro(seed), kwargs...
 )
     (; max_requests_per_epoch, Δ_dispatch, epoch_duration, two_dimensional_features) = b
     files = readdir(datadep"dvrptw"; join=true)
     n = min(n, length(files))
+    gen_rng = Xoshiro(rand(rng, UInt))
+    seed_rng = Xoshiro(rand(rng, UInt))
     return [
-        generate_environment(
-            b,
-            Instance(
-                read_vsp_instance(files[i]);
-                max_requests_per_epoch,
-                Δ_dispatch,
-                epoch_duration,
-                two_dimensional_features,
-            ),
-            rng;
-            kwargs...,
+        Utils.SeededEnvironment(
+            DVSPEnv(
+                Instance(
+                    read_vsp_instance(files[i]);
+                    max_requests_per_epoch,
+                    Δ_dispatch,
+                    epoch_duration,
+                    two_dimensional_features,
+                ),
+                gen_rng,
+            );
+            seed=rand(seed_rng, UInt),
         ) for i in 1:n
     ]
 end
 
 """
 $TYPEDSIGNATURES
 
-Creates an environment from an [`Instance`](@ref) of the dynamic vehicle scheduling benchmark.
-The seed of the environment is randomly generated using the provided random number generator.
-"""
-function generate_environment(
-    ::DynamicVehicleSchedulingBenchmark, instance::Instance, rng::AbstractRNG; kwargs...
-)
-    seed = rand(rng, 1:typemax(Int))
-    return DVSPEnv(instance; seed)
-end
-
-"""
-$TYPEDSIGNATURES
-
 Returns a linear maximizer for the dynamic vehicle scheduling benchmark, of the form:
 θ ↦ argmax_{y} θᵀg(y) + h(y)
 """
@@ -119,8 +106,11 @@ as a `Vector{DataSample}`. Set `reset_env=true` (default) to reset the environme
 before solving, or `reset_env=false` to plan from the current state.
 """
 function Utils.generate_anticipative_solver(::DynamicVehicleSchedulingBenchmark)
-    return (env; reset_env=true, kwargs...) -> begin
-        _, trajectory = anticipative_solver(env; reset_env, kwargs...)
+    return (env::Utils.SeededEnvironment; reset_env=true, kwargs...) -> begin
+        # Anchor to the wrapper's initial seed so the solved scenario is reproducible,
+        # then solve from that state (the env was already reset).
+        reset_env && Utils.reset_to_initial!(env)
+        _, trajectory = anticipative_solver(env.env, env.rng; reset_env=false, kwargs...)
         return trajectory
     end
 end

src/DynamicVehicleScheduling/anticipative_solver.jl

@@ -45,16 +45,16 @@ For this, it uses the current environment history, so make sure that the environ
 """
 function anticipative_solver(
     env::DVSPEnv,
+    rng::AbstractRNG,
     scenario=env.scenario;
     model_builder=highs_model,
     two_dimensional_features=env.instance.two_dimensional_features,
     reset_env=true,
     nb_epochs=nothing,
-    seed=get_seed(env),
     verbose=false,
 )
     if reset_env
-        reset!(env; reset_rng=true, seed)
+        reset!(env, rng)
         scenario = env.scenario
     end