add more traces

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

src/traces.jl

@@ -1,4 +1,4 @@
-export Trace, Traces, sample
+export Trace, Traces
 
 """
     Trace(data)

src/trajectory.jl

@@ -1,7 +1,10 @@
-export Trajectory
+export Trajectory, TrajectoryStyle, SyncTrajectoryStyle, AsyncTrajectoryStyle
 
 using Base.Threads
 
+struct AsyncTrajectoryStyle end
+struct SyncTrajectoryStyle end
+
 
 """
     Trajectory(container, sampler, controler)
@@ -53,6 +56,16 @@ Base.@kwdef struct Trajectory{C,S,T}
     end
 end
 
+TrajectoryStyle(::Trajectory) = SyncTrajectoryStyle()
+TrajectoryStyle(::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}) = AsyncTrajectoryStyle()
+
+Base.bind(::Trajectory, ::Task) = nothing
+
+function Base.bind(t::Trajectory{<:Any,<:Any,<:AsyncInsertSampleRatioControler}, task)
+    bind(t.controler.ch_in, task)
+    bind(t.controler.ch_out, task)
+end
+
 
 Base.push!(t::Trajectory; kw...) = push!(t, values(kw))