Merge pull request #60 from CasBex/59_sumtree

Author: jeremiahpslewis

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,28 @@ function Base.empty!(t::SumTree)
     t
 end
 
+"""
+    correct_sample(t::SumTree, leaf_ind)
+Check whether the sampled leaf is valid and if not return another valid leaf close to it. Used to correct samples with zero priority which may occur due to numerical errors with floats.
+"""
+function correct_sample(t::SumTree, leaf_ind)
+    p = t.tree[leaf_ind]
+    # walk backwards until p != 0 or until leftmost leaf reached
+    tmp_ind = leaf_ind
+    while iszero(p) && (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
+    iszero(p) && (tmp_ind = leaf_ind)
+    while iszero(p) && (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 +174,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_sample(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
@@ -172,4 +194,4 @@ function Random.rand(rng::AbstractRNG, t::SumTree{T}, n::Int) where {T}
     inds, priorities
 end
 
-Random.rand(t::SumTree, n::Int) = rand(Random.GLOBAL_RNG, t, n)
+Random.rand(t::SumTree, n::Int) = rand(Random.GLOBAL_RNG, t, n)

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,63 @@
+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
+
+@testset "SumTree" begin
+    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))
+end