init

Author: HenriDeh

src/samplers.jl

@@ -1,8 +1,10 @@
-export BatchSampler
+export BatchSampler, MetaSampler, MultiBatchSampler
 
 using Random
 
-struct BatchSampler
+abstract type AbstractSampler end
+
+struct BatchSampler <: AbstractSampler
     batch_size::Int
     rng::Random.AbstractRNG
     transformer::Any
@@ -16,3 +18,40 @@ Uniformly sample a batch of examples for each trace.
 See also [`sample`](@ref).
 """
 BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=identity) = BatchSampler(batch_size, rng, identity)
+
+"""
+    MetaSampler(::NamedTuple)
+
+Wraps a NamedTuple containing multiple samplers. When sampled, returns a named tuple with a batch from each sampler.
+Used internally for algorithms that sample multiple times per epoch.
+
+# Example
+
+MetaSampler(policy = BatchSampler(10), critic = BatchSampler(100))
+"""
+struct MetaSampler{names, T} <: AbstractSampler
+    samplers::NamedTuple{names, T}
+end
+
+MetaSampler(; kw...) = MetaSampler(NamedTuple(kw))
+
+function sample(s::MetaSampler, t)
+   (;[(k, sample(v, t)) for (k,v) in pairs(s.samplers)]...)
+end
+
+
+"""
+    MultiBatchSampler(sampler, n)
+
+Wraps a sampler. When sampled, will sample n batches using sampler. Useful in combination with MetaSampler to allow different sampling rates between samplers.
+
+# Example
+
+MetaSampler(policy = MultiBatchSampler(BatchSampler(10), 3), critic = MultiBatchSampler(BatchSampler(100), 5))
+"""
+struct MultiBatchSampler{S <: AbstractSampler} <: AbstractSampler
+    sampler::S
+    n::Int
+end
+
+sample(m::MultiBatchSampler, t) = [sample(m.sampler, t) for _ in 1:m.n]

src/trajectory.jl

@@ -83,7 +83,7 @@ end
 
 function Base.take!(t::Trajectory)
     res = on_sample!(t.controler)
-    if isnothing(res)
+    if isnothing(res) && !isnothing(t.controler)
         nothing
     else
         sample(t.sampler, t.container)

test/samplers.jl

@@ -0,0 +1,72 @@
+using Trajectories, Test
+
+@testset "MetaSampler" begin
+    t = Trajectory(
+        container=Traces(
+            a=Int[],
+            b=Bool[]
+        ),
+        sampler = MetaSampler(policy = BatchSampler(3), critic = BatchSampler(5)),
+        controler = InsertSampleControler(10, 0)
+    )
+
+    append!(t; a=[1, 2, 3, 4], b=[false, true, false, true])
+
+    batches = []
+
+    for batch in t
+        push!(batches, batch)
+    end
+
+    @test length(batches) == 10
+    @test length(batches[1][:policy][:a]) == 3 && length(batches[1][:critic][:b]) == 5    
+end
+
+@testset "MultiBatchSampler" begin
+    t = Trajectory(
+        container=Traces(
+            a=Int[],
+            b=Bool[]
+        ),
+        sampler = MetaSampler(policy = BatchSampler(3), critic = MultiBatchSampler(BatchSampler(5), 2)),
+        controler = InsertSampleControler(10, 0)
+    )
+
+    append!(t; a=[1, 2, 3, 4], b=[false, true, false, true])
+
+    batches = []
+
+    for batch in t
+        push!(batches, batch)
+    end
+
+    @test length(batches) == 10
+    @test length(batches[1][:policy][:a]) == 3 
+    @test length(batches[1][:critic]) == 2 # we sampled 2 batches for critic
+    @test length(batches[1][:critic][1][:b]) == 5 #each batch is 5 samples 
+end
+
+@testset "async trajectories" begin
+    threshould = 100
+    ratio = 1 / 4
+    t = Trajectory(
+        container=Traces(
+            a=Int[],
+            b=Bool[]
+        ),
+        sampler=BatchSampler(3),
+        controler=AsyncInsertSampleRatioControler(ratio, threshould)
+    )
+
+    n = 100
+    insert_task = @async for i in 1:n
+        append!(t; a=[i, i, i, i], b=[false, true, false, true])
+    end
+
+    s = 0
+    sample_task = @async for _ in t
+        s += 1
+    end
+    sleep(1)
+    @test s == (n - threshould * ratio) + 1
+end