Merge pull request #21 from findmyway/minor_enhancements

minor updates when adapting this package in RL.jl

Author: findmyway

Project.toml

@@ -4,15 +4,13 @@ version = "0.1.0"
 
 [deps]
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
-MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StackViews = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
 
 [compat]
 CircularArrayBuffers = "0.1"
 MacroTools = "0.5"
-MLUtils = "0.2"
 StackViews = "0.1"
 julia = "1.6"
 

src/patch.jl

@@ -1,7 +1,3 @@
-import MLUtils
-
-MLUtils.batch(x::AbstractArray{<:Number}) = x
-
 #####
 
 import StackViews: StackView

src/samplers.jl

@@ -1,29 +1,33 @@
 export BatchSampler, MetaSampler, MultiBatchSampler
 
-using MLUtils: batch
-
 using Random
 
 abstract type AbstractSampler end
 
-struct BatchSampler <: AbstractSampler
+struct BatchSampler{names} <: AbstractSampler
     batch_size::Int
     rng::Random.AbstractRNG
     transformer::Any
 end
 
 """
-    BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=identity)
+    BatchSampler{names}(;batch_size, rng=Random.GLOBAL_RNG, transformer=identity)
 
-Uniformly sample a batch of examples for each trace.
+Uniformly sample a batch of examples for each trace specified in `names`. By default, all the traces will be sampled.
 
 See also [`sample`](@ref).
 """
-BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=batch) = BatchSampler(batch_size, rng, transformer)
+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))
+sample(s::BatchSampler{names}, t::AbstractTraces) where {names} = sample(s, t, names)
 
-function sample(s::BatchSampler, t::AbstractTraces)
+function sample(s::BatchSampler, t::AbstractTraces, names)
     inds = rand(s.rng, 1:length(t), s.batch_size)
-    map(s.transformer, t[inds])
+    NamedTuple{names}(s.transformer(t[x][inds]) for x in names)
 end
 
 """

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
@@ -65,11 +68,11 @@ function Base.bind(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}
     bind(t.controler.ch_out, task)
 end
 
+# !!! by default we assume `x`  is a complete example which contains all the traces
+# When doing partial inserting, the result of undefined
 function Base.push!(t::Trajectory, x)
-    n_pre = length(t.container)
     push!(t.container, x)
-    n_post = length(t.container)
-    on_insert!(t.controller, n_post - n_pre)
+    on_insert!(t.controller, 1)
 end
 
 struct CallMsg
@@ -78,16 +81,17 @@ 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)
-    n_pre = length(t.container)
     append!(t.container, x)
-    n_post = length(t.container)
-    on_insert!(t.controller, n_post - n_pre)
+    on_insert!(t.controller, length(x))
 end
 
+# !!! bypass the controller
+sample(t::Trajectory) = sample(t.sampler, t.container)
+
 function Base.take!(t::Trajectory)
     res = on_sample!(t.controller)
     if isnothing(res)

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