Merge pull request #15 from JuliaReinforcementLearning/metasampler

Metasampler, MultiBatchSampler, SamplerControler

Author: HenriDeh

src/Trajectories.jl

@@ -1,7 +1,7 @@
 module Trajectories
 
 include("samplers.jl")
-include("controlers.jl")
+include("controllers.jl")
 include("traces.jl")
 include("episodes.jl")
 include("trajectory.jl")

src/controlers.jl

@@ -0,0 +1,91 @@
+export InsertSampleRatioController, InsertSampleController, AsyncInsertSampleRatioController
+
+mutable struct InsertSampleRatioController
+    ratio::Float64
+    threshold::Int
+    n_inserted::Int
+    n_sampled::Int
+end
+
+"""
+    InsertSampleRatioController(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.
+"""
+InsertSampleRatioController(ratio, threshold) = InsertSampleRatioController(ratio, threshold, 0, 0)
+
+function on_insert!(c::InsertSampleRatioController, n::Int)
+    if n > 0
+        c.n_inserted += n
+    end
+end
+
+function on_sample!(c::InsertSampleRatioController)
+    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
+
+"""
+    InsertSampleController(n, threshold)
+
+Used in [`Trajectory`](@ref). The `threshold` means the minimal number of
+insertings before sampling. The `n` is the number of samples until stopping.
+"""
+mutable struct InsertSampleController
+    n::Int
+    threshold::Int
+    n_inserted::Int
+    n_sampled::Int
+end
+
+InsertSampleController(n, threshold) = InsertSampleController(n, threshold, 0, 0)
+
+function on_insert!(c::InsertSampleController, n::Int)
+    if n > 0
+        c.n_inserted += n
+    end
+end
+
+function on_sample!(c::InsertSampleController)
+    if c.n_inserted >= c.threshold
+        if c.n_sampled < c.n
+            c.n_sampled += 1
+            true
+        end
+    end
+end
+
+#####
+
+mutable struct AsyncInsertSampleRatioController
+    ratio::Float64
+    threshold::Int
+    n_inserted::Int
+    n_sampled::Int
+    ch_in::Channel
+    ch_out::Channel
+end
+
+function AsyncInsertSampleRatioController(
+    ratio,
+    threshold,
+    ; ch_in_sz=1,
+    ch_out_sz=1,
+    n_inserted=0,
+    n_sampled=0
+)
+    AsyncInsertSampleRatioController(
+        ratio,
+        threshold,
+        n_inserted,
+        n_sampled,
+        Channel(ch_in_sz),
+        Channel(ch_out_sz)
+    )
+end

src/controllers.jl

@@ -110,7 +110,7 @@ function Base.convert(r::Type{Term.AbstractRenderable}, t::Trajectory; width=88)
     Panel(
         convert(r, t.container; width=width - 8) /
         Panel(convert(Term.Tree, t.sampler); title="sampler", style="yellow3", fit=true, width=width - 8) /
-        Panel(convert(Term.Tree, t.controler); title="controler", style="yellow3", fit=true, width=width - 8);
+        Panel(convert(Term.Tree, t.controller); title="controller", style="yellow3", fit=true, width=width - 8);
         title="Trajectory",
         style="yellow3",
         width=width,

src/rendering.jl

@@ -1,8 +1,10 @@
-export BatchSampler
+export BatchSampler, MetaSampler, MultiBatchSampler
 
 using Random
 
-struct BatchSampler
+abstract type AbstractSampler end
+
+struct BatchSampler <: AbstractSampler
     batch_size::Int
     rng::Random.AbstractRNG
     transformer::Any
@@ -16,3 +18,40 @@ Uniformly sample a batch of examples for each trace.
 See also [`sample`](@ref).
 """
 BatchSampler(batch_size; rng=Random.GLOBAL_RNG, transformer=identity) = BatchSampler(batch_size, rng, identity)
+
+"""
+    MetaSampler(::NamedTuple)
+
+Wraps a NamedTuple containing multiple samplers. When sampled, returns a named tuple with a batch from each sampler.
+Used internally for algorithms that sample multiple times per epoch.
+
+# Example
+
+MetaSampler(policy = BatchSampler(10), critic = BatchSampler(100))
+"""
+struct MetaSampler{names, T} <: AbstractSampler
+    samplers::NamedTuple{names, T}
+end
+
+MetaSampler(; kw...) = MetaSampler(NamedTuple(kw))
+
+function sample(s::MetaSampler, t)
+   (;[(k, sample(v, t)) for (k,v) in pairs(s.samplers)]...)
+end
+
+
+"""
+    MultiBatchSampler(sampler, n)
+
+Wraps a sampler. When sampled, will sample n batches using sampler. Useful in combination with MetaSampler to allow different sampling rates between samplers.
+
+# Example
+
+MetaSampler(policy = MultiBatchSampler(BatchSampler(10), 3), critic = MultiBatchSampler(BatchSampler(100), 5))
+"""
+struct MultiBatchSampler{S <: AbstractSampler} <: AbstractSampler
+    sampler::S
+    n::Int
+end
+
+sample(m::MultiBatchSampler, t) = [sample(m.sampler, t) for _ in 1:m.n]

src/samplers.jl

@@ -4,11 +4,11 @@ using Base.Threads
 
 
 """
-    Trajectory(container, sampler, controler)
+    Trajectory(container, sampler, controller)
 
 The `container` is used to store experiences. Common ones are [`Traces`](@ref)
 or [`Episodes`](@ref). The `sampler` is used to sample experience batches from
-the `container`. The `controler` controls whether it is time to sample a batch
+the `container`. The `controller` controls whether it is time to sample a batch
 or not.
 
 Supported methoes are:
@@ -21,35 +21,35 @@ Supported methoes are:
 Base.@kwdef struct Trajectory{C,S,T}
     container::C
     sampler::S
-    controler::T
+    controller::T
 
     Trajectory(c::C, s::S, t::T) where {C,S,T} = new{C,S,T}(c, s, t)
 
-    function Trajectory(container::C, sampler::S, controler::T) where {C,S,T<:AsyncInsertSampleRatioControler}
+    function Trajectory(container::C, sampler::S, controller::T) where {C,S,T<:AsyncInsertSampleRatioController}
         t = Threads.@spawn while true
-            for msg in controler.ch_in
+            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)
-                    controler.n_inserted += n_post - n_pre
+                    controller.n_inserted += n_post - n_pre
                 else
                     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 controller.n_inserted >= controller.threshold
+                    if controller.n_sampled <= (controller.n_inserted - controller.threshold) * controller.ratio
                         batch = sample(sampler, container)
-                        put!(controler.ch_out, batch)
-                        controler.n_sampled += 1
+                        put!(controller.ch_out, batch)
+                        controller.n_sampled += 1
                     end
                 end
             end
         end
 
-        bind(controler.ch_in, t)
-        bind(controler.ch_out, t)
-        new{C,S,T}(container, sampler, controler)
+        bind(controller.ch_in, t)
+        bind(controller.ch_out, t)
+        new{C,S,T}(container, sampler, controller)
     end
 end
 
@@ -60,7 +60,7 @@ function Base.push!(t::Trajectory, x)
     n_pre = length(t.container)
     push!(t.container, x)
     n_post = length(t.container)
-    on_insert!(t.controler, n_post - n_pre)
+    on_insert!(t.controller, n_post - n_pre)
 end
 
 struct CallMsg
@@ -69,21 +69,21 @@ struct CallMsg
     kw::Any
 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...) = put!(t.controler.ch_in, CallMsg(Base.append!, args, kw))
+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.append!(t::Trajectory; kw...) = append!(t, values(kw))
 
 function Base.append!(t::Trajectory, x)
     n_pre = length(t.container)
     append!(t.container, x)
     n_post = length(t.container)
-    on_insert!(t.controler, n_post - n_pre)
+    on_insert!(t.controller, n_post - n_pre)
 end
 
 function Base.take!(t::Trajectory)
-    res = on_sample!(t.controler)
-    if isnothing(res)
+    res = on_sample!(t.controller)
+    if isnothing(res) && !isnothing(t.controller)
         nothing
     else
         sample(t.sampler, t.container)
@@ -101,5 +101,5 @@ end
 
 Base.iterate(t::Trajectory, state) = iterate(t)
 
-Base.iterate(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}, args...) = iterate(t.controler.ch_out, args...)
-Base.take!(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}) = take!(t.controler.ch_out)
+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)

src/trajectory.jl

@@ -4,4 +4,5 @@ using Test
 @testset "Trajectories.jl" begin
     include("traces.jl")
     include("trajectories.jl")
+    include("samplers.jl")
 end