Fix CircularPrioritizedTraces with SARTSA

Author: Dharanish

src/common/CircularPrioritizedTraces.jl

@@ -12,7 +12,11 @@ end
 function CircularPrioritizedTraces(traces::AbstractTraces{names,Ts}; default_priority) where {names,Ts}
     new_names = (:key, :priority, names...)
     new_Ts = Tuple{Int,Float32,Ts.parameters...}
-    c = capacity(traces)
+    if traces isa CircularArraySARTSATraces
+        c = capacity(traces) - 1
+    else
+        c = capacity(traces)
+    end
     CircularPrioritizedTraces{typeof(traces),new_names,new_Ts}(
         CircularVectorBuffer{Int}(c),
         SumTree(c),
@@ -34,6 +38,22 @@ function Base.push!(t::CircularPrioritizedTraces, x)
     end
 end
 
+function Base.push!(t::CircularPrioritizedTraces{<:CircularArraySARTSATraces}, x)
+    initial_length = length(t.traces)
+    push!(t.traces, x)
+    if length(t.traces) == 1
+        push!(t.keys, 1)
+        push!(t.priorities, t.default_priority)
+    elseif length(t.traces) > 1 && (initial_length < length(t.traces) || initial_length == capacity(t.traces)-1 ) 
+        # only add a key if the length changes after insertion of the tuple
+        # or if the trace is already at capacity
+        push!(t.keys, t.keys[end] + 1)
+        push!(t.priorities, t.default_priority)
+    else
+        # may be partial inserting at the first step, ignore it
+    end
+end
+
 function Base.setindex!(t::CircularPrioritizedTraces, vs, k::Symbol, keys)
     if k === :priority
         @assert length(vs) == length(keys)
@@ -48,6 +68,7 @@ function Base.setindex!(t::CircularPrioritizedTraces, vs, k::Symbol, keys)
 end
 
 Base.size(t::CircularPrioritizedTraces) = size(t.traces)
+max_length(t::CircularPrioritizedTraces) = max_length(t.traces)
 
 function Base.getindex(ts::CircularPrioritizedTraces, s::Symbol)
     if s === :priority

src/episodes.jl

@@ -184,6 +184,20 @@ function Base.push!(eb::EpisodesBuffer{<:Any,<:Any,<:CircularArraySARTSATraces},
     eb.sampleable_inds[end-1] = 1 #completes the episode trajectory.
 end
 
+function Base.push!(eb::EpisodesBuffer{<:Any,<:Any,<:CircularPrioritizedTraces{<:CircularArraySARTSATraces}}, xs::PartialNamedTuple{@NamedTuple{action::Int64}})
+    if max_length(eb) == capacity(eb.traces)
+        addition = (name => zero(eltype(eb.traces[name])) for name in [:state, :reward, :terminal])
+        xs = (xs.namedtuple, addition)
+        push!(eb.traces, xs)
+        pop!(eb.traces[:state].trace)
+        pop!(eb.traces[:reward])
+        pop!(eb.traces[:terminal])
+    else
+        push!(eb.traces, xs.namedtuple)
+        eb.sampleable_inds[end-1] = 1
+    end
+end
+
 for f in (:pop!, :popfirst!)
     @eval function Base.$f(eb::EpisodesBuffer)
         $f(eb.episodes_lengths)