fix tests

Author: findmyway

src/samplers.jl

@@ -17,7 +17,9 @@ Uniformly sample a batch of examples for each trace specified in `names`. By def
 
 See also [`sample`](@ref).
 """
+BatchSampler(batch_size; kw...) = BatchSampler(; batch_size=batch_size, kw...)
 BatchSampler(; kw...) = BatchSampler{nothing}(; kw...)
+BatchSampler{names}(batch_size; kw...) where {names} = BatchSampler{names}(; batch_size=batch_size, kw...)
 BatchSampler{names}(; batch_size, rng=Random.GLOBAL_RNG, transformer=identity) where {names} = BatchSampler{names}(batch_size, rng, transformer)
 
 sample(s::BatchSampler{nothing}, t::AbstractTraces) = sample(s, t, keys(t))

src/traces.jl

@@ -298,7 +298,7 @@ end
 
 Base.size(t::Traces) = (mapreduce(length, min, t.traces),)
 
-for f in (:push!, :pushfirst!, :append!, :prepend!)
+for f in (:push!, :pushfirst!)
     @eval function Base.$f(ts::Traces, xs::NamedTuple)
         for (k, v) in pairs(xs)
             t = ts.traces[ts.inds[k]]
@@ -311,6 +311,15 @@ for f in (:push!, :pushfirst!, :append!, :prepend!)
     end
 end
 
+for f in (:append!, :prepend!)
+    @eval function Base.$f(ts::Traces, xs::Traces)
+        for k in keys(xs)
+            t = ts.traces[ts.inds[k]]
+            $f(t, xs[k])
+        end
+    end
+end
+
 for f in (:pop!, :popfirst!, :empty!)
     @eval function Base.$f(ts::Traces)
         for t in ts.traces

src/trajectory.jl

@@ -30,11 +30,14 @@ Base.@kwdef struct Trajectory{C,S,T}
     function Trajectory(container::C, sampler::S, controller::T) where {C,S,T<:AsyncInsertSampleRatioController}
         t = Threads.@spawn while true
             for msg in controller.ch_in
-                if msg.f === Base.push! || msg.f === Base.append!
-                    n_pre = length(container)
-                    msg.f(container, msg.args...; msg.kw...)
-                    n_post = length(container)
-                    controller.n_inserted += n_post - n_pre
+                if msg.f === Base.push!
+                    x, = msg.args
+                    msg.f(container, x)
+                    controller.n_inserted += 1
+                elseif msg.f === Base.append!
+                    x, = msg.args
+                    msg.f(container, x)
+                    controller.n_inserted += length(x)
                 else
                     msg.f(container, msg.args...; msg.kw...)
                 end
@@ -78,10 +81,10 @@ struct CallMsg
     kw::Any
 end
 
-Base.push!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, args...; kw...) = put!(t.controller.ch_in, CallMsg(Base.push!, args, kw))
-Base.append!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, args...; kw...) = put!(t.controller.ch_in, CallMsg(Base.append!, args, kw))
+Base.push!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, x) = put!(t.controller.ch_in, CallMsg(Base.push!, (x,), NamedTuple()))
+Base.append!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, x) = put!(t.controller.ch_in, CallMsg(Base.append!, (x,), NamedTuple()))
 
-function Base.append!(t::Trajectory, x::AbstractVector)
+function Base.append!(t::Trajectory, x)
     append!(t.container, x)
     on_insert!(t.controller, length(x))
 end

test/samplers.jl

@@ -33,7 +33,7 @@ end
         sampler=MetaSampler(policy=BatchSampler(3), critic=BatchSampler(5)),
     )
 
-    append!(t, (a=rand(Int, 10), b=rand(Bool, 10)))
+    append!(t, Traces(a=rand(Int, 10), b=rand(Bool, 10)))
 
     batches = collect(t)
 
@@ -50,37 +50,12 @@ end
         sampler=MetaSampler(policy=BatchSampler(3), critic=MultiBatchSampler(BatchSampler(5), 2)),
     )
 
-    append!(t, (a=rand(Int, 10), b=rand(Bool, 10)))
+    append!(t, Traces(a=rand(Int, 10), b=rand(Bool, 10)))
 
     batches = collect(t)
 
     @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),
-        controller=AsyncInsertSampleRatioController(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

test/traces.jl

@@ -11,7 +11,7 @@
     @test t[:a][end] == 3
     @test t[:b][end] == true
 
-    append!(t, (a=[4, 5], b=[false, false]))
+    append!(t, Traces(a=[4, 5], b=[false, false]))
     @test length(t[:a]) == 5
     @test t[:b][end-1:end] == [false, false]
 
@@ -71,7 +71,7 @@ end
     @test length(t3) == 1
     @test t3[1] == (a=1, b=false)
 
-    append!(t3, (; a=[2, 3], b=[false, true]))
+    append!(t3, Traces(; a=[2, 3], b=[false, true]))
     @test length(t3) == 3
 
     @test t3[:a][1:3] == [1, 2, 3]
@@ -102,7 +102,7 @@ end
 
     empty!(t8)
     push!(t8, (a=1, b=false, aa=1, bb=false))
-    append!(t8, (a=[2, 3], b=[true, true], aa=[2, 3], bb=[true, true]))
+    append!(t8, Traces(a=[2, 3], b=[true, true], aa=[2, 3], bb=[true, true]))
 
     @test length(t8) == 3
 
@@ -124,7 +124,7 @@ end
     push!(t, (state=1, action=1.0))
     @test length(t) == 1
 
-    append!(t, (state=[2, 3], action=[2.0, 3.0]))
+    append!(t, Traces(state=[2, 3], action=[2.0, 3.0]))
     @test length(t) == 3
 
     @test t[:state] == [1, 2, 3]

test/trajectories.jl

@@ -32,7 +32,7 @@ end
 
     @test length(batches) == 0  # threshold not reached yet
 
-    append!(t, (a=[1, 2, 3], b=[false, true, false]))
+    append!(t, Traces(a=[1, 2, 3], b=[false, true, false]))
 
     for batch in t
         push!(batches, batch)
@@ -48,7 +48,7 @@ end
 
     @test length(batches) == 1  # 4 inserted, threshold is 4, ratio is 0.25
 
-    append!(t, (a=[5, 6, 7], b=[true, true, true]))
+    append!(t, Traces(a=[5, 6, 7], b=[true, true, true]))
 
     for batch in t
         push!(batches, batch)
@@ -66,7 +66,7 @@ end
 
     n = 100
     for i in 1:n
-        append!(t, (a=[i, i, i, i], b=[false, true, false, true]))
+        append!(t, Traces(a=[i, i, i, i], b=[false, true, false, true]))
     end
 
     s = 0
@@ -90,7 +90,7 @@ end
 
     n = 100
     insert_task = @async for i in 1:n
-        append!(t, (a=[i, i, i, i], b=[false, true, false, true]))
+        append!(t, Traces(a=[i, i, i, i], b=[false, true, false, true]))
     end
 
     s = 0