Fix interaction between CircularPrioritizedTraces and NStepBatchSampler

Project.toml

@@ -1,6 +1,6 @@
 name = "ReinforcementLearningTrajectories"
 uuid = "6486599b-a3cd-4e92-a99a-2cea90cc8c3c"
-version = "0.2"
+version = "0.2.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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)

src/samplers.jl

@@ -209,8 +209,12 @@ function StatsBase.sample(s::NStepBatchSampler, ts, ::Val{SS′L′ART}, inds)
     NamedTuple{SSLART}(map(collect, (s, s′, l, a, r, t)))
 end
 
-function StatsBase.sample(s::NStepBatchSampler{names}, t::CircularPrioritizedTraces) where {names}
-    inds, priorities = rand(s.rng, t.priorities, s.batch_size)
+function StatsBase.sample(s::NStepBatchSampler{names}, e::EpisodesBuffer{<:Any, <:Any, <:CircularPrioritizedTraces}) where {names}
+    t = e.traces
+    st = deepcopy(t.priorities)
+    st .*= e.sampleable_inds[1:end-1] #temporary sumtree that puts 0 priority to non sampleable indices.
+    inds, priorities = rand(s.rng, st, s.batch_size)
+
     merge(
         (key=t.keys[inds], priority=priorities),
         StatsBase.sample(s, t.traces, Val(names), inds)

test/samplers.jl

@@ -136,4 +136,36 @@ end
     @test xs3.reward[2] ≈ 5 + γ * (6 + γ * 7)
     @test xs3.reward[3] ≈ 7 + γ * (8 + γ * 9)
 end
-#! format: on
+#! format: on
+
+@testset "Trajectory with CircularArraySARTSTraces and NStepBatchSampler" begin
+    n=1
+    γ=0.99f0
+
+    t = Trajectory(
+        container=CircularPrioritizedTraces(
+            CircularArraySARTSTraces(
+                capacity=5,
+                state=Float32 => (4,),
+            );
+            default_priority=100.0f0
+        ),
+        sampler=NStepBatchSampler{SS′ART}(
+            n=n,
+            γ=γ,
+            batch_size=32,
+        ),
+        controller=InsertSampleRatioController(
+            threshold=100,
+            n_inserted=-1
+        )
+    )
+
+    push!(t, (state = 1, action = true))
+    for i = 1:9
+        push!(t, (state = i+1, action = true, reward = i, terminal = false))
+    end
+
+    b = RLTrajectories.StatsBase.sample(t)
+    @test haskey(b, :priority)
+end