Merge pull request #28 from findmyway/add_controller_and_sampler

Add ElasticArraySARTTraces

Author: findmyway

Project.toml

@@ -1,9 +1,10 @@
 name = "ReinforcementLearningTrajectories"
 uuid = "6486599b-a3cd-4e92-a99a-2cea90cc8c3c"
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
+ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/common/ElasticArraySARTTraces.jl

@@ -0,0 +1,40 @@
+export ElasticArraySARTTraces
+
+using ElasticArrays: ElasticArray, resize_lastdim!
+
+const ElasticArraySARTTraces = Traces{
+    SS′AA′RT,
+    <:Tuple{
+        <:MultiplexTraces{SS′,<:Trace{<:ElasticArray}},
+        <:MultiplexTraces{AA′,<:Trace{<:ElasticArray}},
+        <:Trace{<:ElasticArray},
+        <:Trace{<:ElasticArray},
+    }
+}
+
+function ElasticArraySARTTraces(;
+    state=Int => (),
+    action=Int => (),
+    reward=Float32 => (),
+    terminal=Bool => ()
+)
+    state_eltype, state_size = state
+    action_eltype, action_size = action
+    reward_eltype, reward_size = reward
+    terminal_eltype, terminal_size = terminal
+
+    MultiplexTraces{SS′}(ElasticArray{state_eltype}(undef, state_size..., 0)) +
+    MultiplexTraces{AA′}(ElasticArray{action_eltype}(undef, action_size..., 0)) +
+    Traces(
+        reward=ElasticArray{reward_eltype}(undef, reward_size..., 0),
+        terminal=ElasticArray{terminal_eltype}(undef, terminal_size..., 0),
+    )
+end
+
+#####
+# extensions for ElasticArrays
+#####
+
+Base.push!(a::ElasticArray, x) = append!(a, x)
+Base.push!(a::ElasticArray{T,1}, x) where {T} = append!(a, [x])
+Base.empty!(a::ElasticArray) = resize_lastdim!(a, 0)

src/common/common.jl

@@ -15,3 +15,4 @@ include("sum_tree.jl")
 include("CircularArraySARTTraces.jl")
 include("CircularArraySLARTTraces.jl")
 include("CircularPrioritizedTraces.jl")
+include("ElasticArraySARTTraces.jl")

src/controllers.jl

@@ -24,9 +24,10 @@ 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
+            return true
         end
     end
+    return false
 end
 
 #####
@@ -56,4 +57,4 @@ function AsyncInsertSampleRatioController(
         Channel(ch_in_sz),
         Channel(ch_out_sz)
     )
-end
+end

src/samplers.jl

@@ -2,6 +2,16 @@ using Random
 
 abstract type AbstractSampler end
 
+#####
+# DummySampler
+#####
+
+export DummySampler
+
+struct DummySampler end
+
+sample(s::DummySampler, t::AbstractTraces) = t
+
 #####
 # BatchSampler
 #####
@@ -160,4 +170,4 @@ function sample(s::NStepBatchSampler{names}, t::CircularPrioritizedTraces) where
         (key=t.keys[inds], priority=priorities),
         sample(s, t.traces, Val(names), inds)
     )
-end
+end

src/trajectory.jl

@@ -22,7 +22,7 @@ Supported methoes are:
 """
 Base.@kwdef struct Trajectory{C,S,T}
     container::C
-    sampler::S
+    sampler::S = DummySampler()
     controller::T = InsertSampleRatioController()
     transformer::Any = identity
 
@@ -96,12 +96,13 @@ end
 # !!! bypass the controller
 sample(t::Trajectory) = sample(t.sampler, t.container)
 
+on_sample!(t::Trajectory) = on_sample!(t.controller)
+
 function Base.take!(t::Trajectory)
-    res = on_sample!(t.controller)
-    if isnothing(res)
-        nothing
-    else
+    if on_sample!(t)
         sample(t.sampler, t.container) |> t.transformer
+    else
+        nothing
     end
 end
 
@@ -120,4 +121,4 @@ Base.iterate(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioController}, 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(::Trajectory) = Base.SizeUnknown()

test/common.jl

@@ -91,6 +91,26 @@ end
     @test batch.terminal == Bool[0, 0, 0]
 end
 
+@testset "ElasticArraySARTTraces" begin
+    t = ElasticArraySARTTraces(;
+        state=Float32 => (2, 3),
+        action=Int => (),
+        reward=Float32 => (),
+        terminal=Bool => ()
+    )
+
+    @test t isa ElasticArraySARTTraces
+
+    push!(t, (state=ones(Float32, 2, 3), action=1))
+    push!(t, (reward=1.0f0, terminal=false, state=ones(Float32, 2, 3) * 2, action=2))
+
+    @test length(t) == 1
+
+    empty!(t)
+
+    @test length(t) == 0
+end
+
 @testset "CircularArraySLARTTraces" begin
     t = CircularArraySLARTTraces(;
         capacity=3,