Add test and fix for sampler dispatch issue

Author: jeremiahpslewis

src/samplers.jl

@@ -209,8 +209,21 @@ 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})
+    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)
+
+    NamedTuple{(:key, :priority, names...)}((t.keys[inds], priorities, map(x -> collect(t.traces[x][inds]), names)...))
+end
+
+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

@@ -161,3 +161,35 @@ end
     end    
 end
 #! format: on
+
+@testset "Trajectory with CircularArraySARTSTraces and NStepBatchSampler" begin
+    n=1
+    γ=0.99f0
+
+    t = Trajectory(
+        container=CircularPrioritizedTraces(
+            CircularArraySARTSTraces(
+                capacity=1000,
+                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