minor updates when adapting it 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,31 @@
 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(; kw...) = BatchSampler{nothing}(; 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/trajectory.jl

@@ -65,11 +65,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
@@ -81,13 +81,14 @@ 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))
 
-function Base.append!(t::Trajectory, x)
-    n_pre = length(t.container)
+function Base.append!(t::Trajectory, x::AbstractVector)
     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)