Merge branch 'main' into fix-prio

Author: HenriDeh

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,14 @@
 using ReinforcementLearningTrajectories
 using CircularArrayBuffers, DataStructures
+using StableRNGs
 using Test
 import ReinforcementLearningTrajectories.StatsBase.sample
 using CUDA
 using Adapt
+using Random
+import ReinforcementLearningTrajectories.StatsBase.sample
+import StatsBase.countmap
+
 
 struct TestAdaptor end
 
@@ -13,6 +18,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