Sumtree sampling

Project.toml

@@ -16,17 +16,18 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 [compat]
 Adapt = "3"
 CircularArrayBuffers = "0.1"
+DataStructures = "0.18"
 ElasticArrays = "1"
 MacroTools = "0.5"
 OnlineStats = "1"
 StackViews = "0.1"
-julia = "1.9"
-DataStructures = "0.18"
 StatsBase = "0.34"
+julia = "1.9"
 
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "CUDA"]
+test = ["Test", "CUDA", "StableRNGs"]

src/common/sum_tree.jl

@@ -131,6 +131,24 @@ function Base.empty!(t::SumTree)
     t
 end
 
+function correct_priority(t::SumTree, leaf_ind)
+    p = t.tree[leaf_ind]
+    # walk backwards until p != 0 or until leftmost leaf reached
+    tmp_ind = leaf_ind
+    while p == 0f0 && (tmp_ind-1)*2 > length(t.tree)
+        tmp_ind -= 1
+        p = t.tree[tmp_ind]
+    end
+    # walk forwards until p != 0 or until rightmost leaf reached
+    p == 0f0 && (tmp_ind = leaf_ind)
+    while p == 0f0 && (tmp_ind - t.nparents) <= t.length
+        tmp_ind += 1
+        p = t.tree[tmp_ind]
+    end
+    return p, tmp_ind
+end
+
+
 function Base.get(t::SumTree, v)
     parent_ind = 1
     leaf_ind = parent_ind
@@ -152,7 +170,7 @@ function Base.get(t::SumTree, v)
     if leaf_ind <= t.nparents
         leaf_ind += t.capacity
     end
-    p = t.tree[leaf_ind]
+    p, leaf_ind = correct_priority(t, leaf_ind)
     ind = leaf_ind - t.nparents
     real_ind = ind >= t.first ? ind - t.first + 1 : ind + t.capacity - t.first + 1
     real_ind, p

test/runtests.jl

@@ -1,9 +1,12 @@
 using ReinforcementLearningTrajectories
 using CircularArrayBuffers, DataStructures
+using StableRNGs
 using Test
 using CUDA
 using Adapt
+using Random
 import ReinforcementLearningTrajectories.StatsBase.sample
+import StatsBase.countmap
 
 struct TestAdaptor end
 
@@ -13,6 +16,7 @@ Adapt.adapt_storage(to::TestAdaptor, x) = CUDA.functional() ? CUDA.cu(x) : x
 
 @testset "ReinforcementLearningTrajectories.jl" begin
     include("traces.jl")
+    include("sum_tree.jl")
     include("common.jl")
     include("samplers.jl")
     include("controllers.jl")

test/sum_tree.jl

@@ -0,0 +1,61 @@
+function gen_rand_sumtree(n, seed, type::DataType=Float32)
+    rng = StableRNG(seed)
+    a = SumTree(type, n)
+    append!(a, rand(rng, type, n))
+    return a
+end
+
+function gen_sumtree_with_zeros(n, seed, type::DataType=Float32)
+    a = gen_rand_sumtree(n, seed, type)
+    b = rand(StableRNG(seed), Bool, n)
+    return copy_multiply(a, b)
+end 
+
+function copy_multiply(stree, m)
+    new_tree = deepcopy(stree)
+    new_tree .*= m
+    return new_tree
+end
+
+function sumtree_nozero(t::SumTree, rng::AbstractRNG, iters=1)
+    for _ in iters
+        (_, p) = rand(rng, t)
+        p == 0 && return false
+    end
+    return true
+end
+sumtree_nozero(n::Integer, seed::Integer, iters=1) = sumtree_nozero(gen_sumtree_with_zeros(n, seed), StableRNG(seed), iters)
+sumtree_nozero(n, seeds::AbstractVector, iters=1) = all(sumtree_nozero(n, seed, iters) for seed in seeds)
+
+
+function sumtree_distribution!(indices, priorities, t::SumTree, rng::AbstractRNG, iters=1000*t.length)
+    for i = 1:iters
+        indices[i], priorities[i] = rand(rng, t)
+    end
+    imap = countmap(indices)
+    est_pdf = Dict(k=>v/length(indices) for (k, v) in imap)
+    ex_pdf = Dict(k=>v/t.tree[1] for (k, v) in Dict(1:length(t) .=> t))
+    abserrs = [est_pdf[k] - ex_pdf[k] for k in keys(est_pdf)]
+    return abserrs
+end
+sumtree_distribution!(indices, priorities, n, seed, iters=1000*n) = sumtree_distribution!(indices, priorities, gen_rand_sumtree(n, seed), StableRNG(seed), iters)
+function sumtree_distribution(n, seeds::AbstractVector, iters=1000*n)
+    p = [zeros(Float32, iters) for _ = 1:Threads.nthreads()]
+    i = [zeros(Float32, iters) for _ = 1:Threads.nthreads()]
+    results = Vector{Vector{Float64}}(undef, length(seeds))
+    Threads.@threads for ix = 1:length(seeds)
+        results[ix] = sumtree_distribution!(i[Threads.threadid()], p[Threads.threadid()], gen_rand_sumtree(n, seeds[ix]), StableRNG(seeds[ix]), iters)
+    end
+    return results
+end
+
+n = 1024
+seeds = 1:100
+nozero_iters=1024
+distr_iters=1024*10_000
+abstol = 0.05
+maxerr=0.01
+
+@test sumtree_nozero(n, seeds, nozero_iters)
+@test all(x->all(x .< maxerr) && sum(abs2, x) < abstol,
+          sumtree_distribution(n, seeds, distr_iters))