sample namespace

Author: HenriDeh

test/common.jl

@@ -145,7 +145,7 @@ end
 
     s = BatchSampler(5)
 
-    b = ReinforcementLearningTrajectories.StatsBase.sample(s, t)
+    b = sample(s, t)
 
     t[:priority, [1, 2]] = [0, 0]
 
@@ -154,7 +154,7 @@ end
 
     t[:priority, [3, 4, 5]] = [0, 1, 0]
 
-    b = ReinforcementLearningTrajectories.StatsBase.sample(s, t)
+    b = sample(s, t)
 
     @test b.key == [4, 4, 4, 4, 4] # the priority of the rest transitions are set to 0
 
@@ -176,7 +176,7 @@ end
         push!(eb, (state = i, action =i, reward = i-1, terminal = false))
     end
     s = BatchSampler(1000)
-    b = ReinforcementLearningTrajectories.StatsBase.sample(s, eb)
+    b = sample(s, eb)
     cm = counter(b[:state])
     @test !haskey(cm, 6)
     @test !haskey(cm, 11)

test/normalization.jl

@@ -29,7 +29,7 @@ import OnlineStats: mean, std
     @test extrema(unnormalized_batch[:state]) == (0, 4)
     normalized_batch = nt[[1:5;]]
 
-    normalized_batch = ReinforcementLearningTrajectories.StatsBase.sample(traj)
+    normalized_batch = sample(traj)
     @test all(extrema(normalized_batch[:state]) .≈ ((0, 4) .- m)./ss)
     @test all(extrema(normalized_batch[:next_state]) .≈ ((0, 4) .- m)./ss)
     @test all(extrema(normalized_batch[:reward]) .≈ ((0, 4) .- m)./s)

test/runtests.jl

@@ -3,6 +3,7 @@ using CircularArrayBuffers, DataStructures
 using Test
 using CUDA
 using Adapt
+import ReinforcementLearningTrajectories.StatsBase.sample
 
 struct TestAdaptor end
 

test/samplers.jl

@@ -6,7 +6,7 @@
         action=rand(1:4, 5),
     )
 
-    b = ReinforcementLearningTrajectories.StatsBase.sample(s, t)
+    b = sample(s, t)
 
     @test keys(b) == (:state, :action)
     @test size(b.state) == (3, 4, sz)
@@ -23,7 +23,7 @@
         push!(eb, (state = i, action =i, reward = i-1, terminal = false))
     end
     s = BatchSampler(1000)
-    b = ReinforcementLearningTrajectories.StatsBase.sample(s, eb)
+    b = sample(s, eb)
     cm = counter(b[:state])
     @test !haskey(cm, 6)
     @test !haskey(cm, 11)