add more traces

Author: findmyway

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
+ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/common/ElasticArraySARTTraces.jl

@@ -0,0 +1,43 @@
+using ElasticArrays
+
+const ElasticSARTTraces = Traces{
+    SART,
+    <:Tuple{
+        <:Trace{<:ElasticArray},
+        <:Trace{<:ElasticArray},
+        <:Trace{<:ElasticArray},
+        <:Trace{<:ElasticArray},
+    }
+}
+
+function ElasticSARTTraces(;
+    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
+
+    Traces(
+        state=ElasticArray{state_eltype}(state_size..., 0),
+        action=ElasticArray{action_eltype}(action_size..., 0),
+        reward=ElasticArray{reward_eltype}(reward_size..., 0),
+        terminal=ElasticArray{terminal_eltype}(terminal_size..., 0),
+    )
+end
+
+function Random.rand(s::BatchSampler, t::ElasticSARTTraces)
+    inds = rand(s.rng, 1:length(t), s.batch_size)
+    inds′ = inds .+ 1
+    (
+        state=t[:state][inds],
+        action=t[:action][inds],
+        reward=t[:reward][inds],
+        terminal=t[:terminal][inds],
+        next_state=t[:state][inds′],
+        next_action=t[:state][inds′]
+    ) |> s.transformer
+end

src/common/ReservoirTraces.jl

@@ -0,0 +1,20 @@
+using MacroTools: @forward
+
+mutable struct ReservoirTraces{T,R}
+    traces::T
+    n::Int
+    capacity::Int
+    rng::R
+end
+
+@forward ReservoirTrajectory.buffer sample, Base.keys, Base.haskey, Base.getindex, Base.view, Base.length, Base.setindex!, Base.lastindex, Base.firstindex
+
+function Base.push!(t::ReservoirTraces, x)
+    if t.n < t.capacity
+        push!(t.traces, x)
+        t.n += 1
+    else
+        i = rand(t.rng, 1:length(t))
+        t.traces[i] = x
+    end
+end