Fix #29

src/samplers.jl

@@ -1,23 +1,33 @@
 using Random
 
-abstract type AbstractSampler end
+struct SampleGenerator{S,T}
+    sampler::S
+    traces::T
+end
+
+Base.iterate(s::SampleGenerator) = sample(s.sampler, s.traces), nothing
+Base.iterate(s::SampleGenerator, ::Nothing) = nothing
 
 #####
 # DummySampler
 #####
 
 export DummySampler
 
+"""
+Just return the underlying traces.
+"""
 struct DummySampler end
 
-sample(s::DummySampler, t::AbstractTraces) = t
+sample(::DummySampler, t) = t
 
 #####
 # BatchSampler
 #####
 
 export BatchSampler
-struct BatchSampler{names} <: AbstractSampler
+
+struct BatchSampler{names}
     batch_size::Int
     rng::Random.AbstractRNG
 end
@@ -26,10 +36,8 @@ end
     BatchSampler{names}(;batch_size, rng=Random.GLOBAL_RNG)
     BatchSampler{names}(batch_size ;rng=Random.GLOBAL_RNG)
 
-Uniformly sample a batch of examples for each trace specified in `names`. 
-By default, all the traces will be sampled.
-
-See also [`sample`](@ref).
+Uniformly sample **ONE** batch of `batch_size` examples for each trace specified
+in `names`. If `names` is not set, all the traces will be sampled.
 """
 BatchSampler(batch_size; kw...) = BatchSampler(; batch_size=batch_size, kw...)
 BatchSampler(; kw...) = BatchSampler{nothing}(; kw...)
@@ -44,6 +52,15 @@ function sample(s::BatchSampler, t::AbstractTraces, names)
     NamedTuple{names}(map(x -> t[x][inds], names))
 end
 
+# !!! avoid iterating an empty trajectory
+function Base.iterate(s::SampleGenerator{<:BatchSampler})
+    if length(s.traces) > 0
+        sample(s.sampler, s.traces), nothing
+    else
+        nothing
+    end
+end
+
 #####
 
 sample(s::BatchSampler{nothing}, t::CircularPrioritizedTraces) = sample(s, t, keys(t.traces))
@@ -75,7 +92,7 @@ initializing an agent.
 MetaSampler(policy = BatchSampler(10), critic = BatchSampler(100))
 ```
 """
-struct MetaSampler{names,T} <: AbstractSampler
+struct MetaSampler{names,T}
     samplers::NamedTuple{names,T}
 end
 
@@ -104,7 +121,7 @@ MetaSampler(policy = MultiBatchSampler(BatchSampler(10), 3),
             critic = MultiBatchSampler(BatchSampler(100), 5))
 ```
 """
-struct MultiBatchSampler{S<:AbstractSampler} <: AbstractSampler
+struct MultiBatchSampler{S}
     sampler::S
     n::Int
 end

src/trajectory.jl

@@ -69,15 +69,12 @@ 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)
-    push!(t.container, x)
-    on_insert!(t.controller, 1)
-end
-
 Base.setindex!(t::Trajectory, v, I...) = setindex!(t.container, v, I...)
 
+#####
+# in
+#####
+
 struct CallMsg
     f::Any
     args::Tuple
@@ -93,32 +90,33 @@ function Base.append!(t::Trajectory, x)
     on_insert!(t.controller, length(x))
 end
 
-# !!! bypass the controller
-sample(t::Trajectory) = sample(t.sampler, t.container)
+# !!! 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)
+    push!(t.container, x)
+    on_insert!(t)
+end
 
-on_sample!(t::Trajectory) = on_sample!(t.controller)
+on_insert!(t::Trajectory) = on_insert!(t, 1)
+on_insert!(t::Trajectory, n::Int) = on_insert!(t.controller, n)
 
-function Base.take!(t::Trajectory)
-    if on_sample!(t)
-        sample(t.sampler, t.container) |> t.transformer
-    else
-        nothing
-    end
-end
+#####
+# out
+#####
 
-function Base.iterate(t::Trajectory)
-    x = take!(t)
-    if isnothing(x)
-        nothing
-    else
-        x, true
-    end
-end
+SampleGenerator(t::Trajectory) = SampleGenerator(t.sampler, t.container)
+
+on_sample!(t::Trajectory) = on_sample!(t.controller)
+sample(t::Trajectory) = sample(t.sampler, t.container)
+
+"""
+Keep sampling batches from the trajectory until the trajectory is not ready to
+be sampled yet due to the `controller`.
+"""
+iter(t::Trajectory) = Iterators.takewhile(_ -> on_sample!(t), Iterators.cycle(SampleGenerator(t)))
 
-Base.iterate(t::Trajectory, state) = iterate(t)
+Base.iterate(t::Trajectory, args...) = iterate(iter(t), args...)
+Base.IteratorSize(t::Trajectory) = Base.IteratorSize(iter(t))
 
 Base.iterate(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, args...) = iterate(t.controller.ch_out, args...)
-Base.take!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}) = take!(t.controller.ch_out)
-
-Base.IteratorSize(::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}) = Base.IsInfinite()
-Base.IteratorSize(::Trajectory) = Base.SizeUnknown()
+Base.IteratorSize(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}) = Base.IteratorSize(t.controller.ch_out)