Corner case of first non-sampleable idx

Author: HenriDeh

src/samplers.jl

@@ -262,9 +262,13 @@ function StatsBase.sample(::EpisodesSampler, t::EpisodesBuffer, names)
     ranges = UnitRange{Int}[]
     idx = 1
     while idx < length(t)
-        last_state_idx = idx + t.episodes_lengths[idx] - t.step_numbers[idx] + 1
-        push!(ranges,idx:last_state_idx)
-        idx = last_state_idx + 1
+        if t.sampleable_inds[idx] == 1
+            last_state_idx = idx + t.episodes_lengths[idx] - t.step_numbers[idx] + 1
+            push!(ranges,idx:last_state_idx)
+            idx = last_state_idx + 1
+        else
+            idx += 1
+        end
     end
 
     return [Episode(NamedTuple{names}(map(x -> collect(t[Val(x)][r]), names))) for r in ranges]

test/samplers.jl

@@ -218,6 +218,14 @@
         @test length(b) == 2
         @test length(b[1][:state]) == 5
         @test length(b[2][:state]) == 6
+        
+        for (j,i) = enumerate(2:5)
+            push!(eb, (state = i, action =i, reward = i-1, terminal = false))
+        end
+        #only the last state of the first episode is still buffered. Should not be sampled.
+        b = sample(s, eb)
+        @test length(b) == 1
+        
 
         #with specified traces
         s = EpisodesSampler{(:state,)}()