Fix buffer size issue

Author: jeremiahpslewis

src/common/CircularPrioritizedTraces.jl

@@ -59,4 +59,6 @@ function Base.getindex(ts::CircularPrioritizedTraces, s::Symbol)
     end
 end
 
-Base.getindex(t::CircularPrioritizedTraces{<:Any,names}, i) where {names} = NamedTuple{names}(map(k -> t[k][i], names))
+Base.getindex(t::CircularPrioritizedTraces{<:Any,names}, i) where {names} = NamedTuple{names}(map(k -> t[k][i], names))
+
+capacity(t::CircularPrioritizedTraces) = ReinforcementLearningTrajectories.capacity(t.traces)

src/episodes.jl

@@ -39,8 +39,21 @@ struct PartialNamedTuple{T}
     namedtuple::T
 end
 
+# Capacity of an EpisodesBuffer is the capacity of the underlying traces + 1 for certain cases
+function is_capacity_plus_one(traces::AbstractTraces)
+    if any(t->t isa MultiplexTraces, traces.traces)
+        # MultiplexTraces buffer next_state and next_action, so we need to add one to the capacity
+        return true
+    elseif traces isa CircularPrioritizedTraces
+        # CircularPrioritizedTraces buffer next_state and next_action, so we need to add one to the capacity
+        return true
+    else
+        false
+    end
+end
+
 function EpisodesBuffer(traces::AbstractTraces)
-    cap = any(t->t isa MultiplexTraces, traces.traces) ? capacity(traces) + 1 : capacity(traces)
+    cap = is_capacity_plus_one(traces) ? capacity(traces) + 1 : capacity(traces)
     @assert isempty(traces) "EpisodesBuffer must be initialized with empty traces."
     if !isinf(cap)
         legalinds =  CircularBuffer{Bool}(cap)

test/samplers.jl

@@ -145,7 +145,7 @@ end
     t = Trajectory(
         container=CircularPrioritizedTraces(
             CircularArraySARTSTraces(
-                capacity=1000,
+                capacity=5,
                 state=Float32 => (4,),
             );
             default_priority=100.0f0