fix tests

Author: findmyway

src/Trajectories.jl

@@ -1,10 +1,10 @@
 module Trajectories
 
 include("samplers.jl")
+include("controlers.jl")
 include("traces.jl")
 include("episodes.jl")
 include("trajectory.jl")
-include("async_trajectory.jl")
 include("rendering.jl")
 include("common/common.jl")
 

src/common/CircularArraySARTTraces.jl

@@ -41,7 +41,7 @@ function Random.rand(s::BatchSampler, t::CircularArraySARTTraces)
         terminal=t[:terminal][inds],
         next_state=t[:state][inds′],
         next_action=t[:state][inds′]
-    )
+    ) |> s.transformer
 end
 
 function Base.push!(t::CircularArraySARTTraces, x::NamedTuple{SA})

src/common/CircularArraySLARTTraces.jl

@@ -47,7 +47,7 @@ function sample(s::BatchSampler, t::CircularArraySLARTTraces)
         next_state=t[:state][inds′],
         next_legal_actions_mask=t[:legal_actions_mask][inds′],
         next_action=t[:state][inds′]
-    )
+    ) |> s.transformer
 end
 
 function Base.push!(t::CircularArraySLARTTraces, x::NamedTuple{SLA})

src/controlers.jl

@@ -0,0 +1,61 @@
+export InsertSampleRatioControler, AsyncInsertSampleRatioControler
+
+mutable struct InsertSampleRatioControler
+    ratio::Float64
+    threshold::Int
+    n_inserted::Int
+    n_sampled::Int
+end
+
+"""
+    InsertSampleRatioControler(ratio, threshold)
+
+Used in [`Trajectory`](@ref). The `threshold` means the minimal number of
+insertings before sampling. The `ratio` balances the number of insertings and
+the number of samplings.
+"""
+InsertSampleRatioControler(ratio, threshold) = InsertSampleRatioControler(ratio, threshold, 0, 0)
+
+function on_insert!(c::InsertSampleRatioControler, n::Int)
+    if n > 0
+        c.n_inserted += n
+    end
+end
+
+function on_sample!(c::InsertSampleRatioControler)
+    if c.n_inserted >= c.threshold
+        if c.n_sampled <= (c.n_inserted - c.threshold) * c.ratio
+            c.n_sampled += 1
+            true
+        end
+    end
+end
+
+#####
+
+mutable struct AsyncInsertSampleRatioControler
+    ratio::Float64
+    threshold::Int
+    n_inserted::Int
+    n_sampled::Int
+    ch_in::Channel
+    ch_out::Channel
+end
+
+function AsyncInsertSampleRatioControler(
+    ratio,
+    threshold,
+    ; ch_in_sz=1,
+    ch_out_sz=1,
+    n_inserted=0,
+    n_sampled=0
+)
+    AsyncInsertSampleRatioControler(
+        ratio,
+        threshold,
+        n_inserted,
+        n_sampled,
+        Channel(ch_in_sz),
+        Channel(ch_out_sz)
+    )
+end

src/episodes.jl

@@ -103,5 +103,5 @@ end
 
 function sample(s::BatchSampler, e::Episodes)
     inds = rand(s.rng, 1:length(t), s.batch_size)
-    batch([@view(s.episodes[e.inds[i][1]][e.inds[i][2]]) for i in inds])
+    batch([@view(s.episodes[e.inds[i][1]][e.inds[i][2]]) for i in inds]) |> s.transformer
 end

src/samplers.jl

@@ -5,13 +5,14 @@ using Random
 struct BatchSampler
     batch_size::Int
     rng::Random.AbstractRNG
+    transformer::Any
 end
 
 """
-    BatchSampler(batch_size; rng=Random.GLOBAL_RNG)
+    BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=identity)
 
 Uniformly sample a batch of examples for each trace.
 
 See also [`sample`](@ref).
 """
-BatchSampler(batch_size; rng=Random.GLOBAL_RNG) = BatchSampler(batch_size, rng)
+BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=identity) = BatchSampler(batch_size, rng, identity)

src/traces.jl

@@ -36,7 +36,7 @@ Base.empty!(t::Trace) = empty!(t.x)
 
 function sample(s::BatchSampler, t::Trace)
     inds = rand(s.rng, 1:length(t), s.batch_size)
-    t[inds]
+    t[inds] |> s.transformer
 end
 
 #####
@@ -84,5 +84,5 @@ function sample(s::BatchSampler, t::Traces)
     inds = rand(s.rng, 1:length(t), s.batch_size)
     map(t.traces) do x
         x[inds]
-    end
+    end |> s.transformer
 end

src/trajectory.jl

@@ -1,72 +1,8 @@
-export Trajectory, InsertSampleRatioControler
+export Trajectory
 
 using Base.Threads
 
 
-#####
-
-mutable struct InsertSampleRatioControler
-    ratio::Float64
-    threshold::Int
-    n_inserted::Int
-    n_sampled::Int
-end
-
-"""
-    InsertSampleRatioControler(ratio, threshold)
-
-Used in [`Trajectory`](@ref). The `threshold` means the minimal number of
-insertings before sampling. The `ratio` balances the number of insertings and
-the number of samplings.
-"""
-InsertSampleRatioControler(ratio, threshold) = InsertSampleRatioControler(ratio, threshold, 0, 0)
-
-function on_insert!(c::InsertSampleRatioControler, n::Int)
-    if n > 0
-        c.n_inserted += n
-    end
-end
-
-function on_sample!(c::InsertSampleRatioControler)
-    if c.n_inserted >= c.threshold
-        if c.n_sampled < (c.n_inserted - n.threshold) * c.ratio
-            c.n_sampled += 1
-            true
-        end
-    end
-end
-
-#####
-
-mutable struct AsyncInsertSampleRatioControler
-    ratio::Float64
-    threshold::Int
-    n_inserted::Int
-    n_sampled::Int
-    ch_in::Channel
-    ch_out::Channel
-end
-
-function AsyncInsertSampleRatioControler(
-    ratio,
-    threshold,
-    ; ch_in_sz=1,
-    ch_out_sz=1,
-    n_inserted=0,
-    n_sampled=0
-)
-    AsyncInsertSampleRatioControler(
-        ratio,
-        threshold,
-        n_inserted,
-        n_sampled,
-        Channel(ch_in_sz),
-        Channel(ch_out_sz)
-    )
-end
-
-#####
-
 """
     Trajectory(container, sampler, controler)
 
@@ -91,18 +27,18 @@ Base.@kwdef struct Trajectory{C,S,T}
 
     function Trajectory(container::C, sampler::S, controler::T) where {C,S,T<:AsyncInsertSampleRatioControler}
         t = Threads.@spawn while true
-            for msg in controler.in
+            for msg in controler.ch_in
                 if msg.f === Base.push! || msg.f === Base.append!
-                    n_pre = length(trajectory)
-                    msg.f(trajectory, msg.args...; msg.kw...)
-                    n_post = length(trajectory)
+                    n_pre = length(container)
+                    msg.f(container, msg.args...; msg.kw...)
+                    n_post = length(container)
                     controler.n_inserted += n_post - n_pre
                 else
-                    msg.f(trajectory, msg.args...; msg.kw...)
+                    msg.f(container, msg.args...; msg.kw...)
                 end
 
                 if controler.n_inserted >= controler.threshold
-                    if controler.n_sampled < (controler.n_inserted - controler.threshold) * controler.ratio
+                    if controler.n_sampled <= (controler.n_inserted - controler.threshold) * controler.ratio
                         batch = sample(sampler, container)
                         put!(controler.ch_out, batch)
                         controler.n_sampled += 1
@@ -111,8 +47,8 @@ Base.@kwdef struct Trajectory{C,S,T}
             end
         end
 
-        bind(controler.in, t)
-        bind(controler.out, t)
+        bind(controler.ch_in, t)
+        bind(controler.ch_out, t)
         new{C,S,T}(container, sampler, controler)
     end
 end
@@ -134,7 +70,7 @@ struct CallMsg
 end
 
 Base.push!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}, args...; kw...) = put!(t.controler.ch_in, CallMsg(Base.push!, args, kw))
-Base.append!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}, args...; kw...) = append!(t.controler.ch_in, CallMsg(Base.push!, args, kw))
+Base.append!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}, args...; kw...) = put!(t.controler.ch_in, CallMsg(Base.append!, args, kw))
 
 Base.append!(t::Trajectory; kw...) = append!(t, values(kw))
 

test/trajectories.jl

@@ -30,15 +30,15 @@
         push!(batches, batch)
     end
 
-    @test length(batches) == 1  # 4 inserted, ratio is 0.25
+    @test length(batches) == 1  # 4 inserted, threshold is 4, ratio is 0.25
 
     append!(t; a=[5, 6, 7], b=[true, true, true])
 
     for batch in t
         push!(batches, batch)
     end
 
-    @test length(batches) == 2  # 7 inserted, ratio is 0.25
+    @test length(batches) == 1  # 7 inserted, threshold is 4, ratio is 0.25
 
     push!(t; a=8, b=true)
 
@@ -58,4 +58,29 @@
         s += 1
     end
     @test s == n
+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