Merge pull request #35 from JuliaReinforcementLearning/HenriDeh-patch-normalization

collect on normalized traces for performance.

Author: HenriDeh

Project.toml

@@ -1,6 +1,6 @@
 name = "ReinforcementLearningTrajectories"
 uuid = "6486599b-a3cd-4e92-a99a-2cea90cc8c3c"
-version = "0.1.8"
+version = "0.1.9"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/normalization.jl

@@ -194,5 +194,5 @@ end
 function sample(s::BatchSampler, nt::NormalizedTraces, names)
     inds = rand(s.rng, 1:length(nt), s.batch_size)
     maybe_normalize(data, key) = key in keys(nt.normalizers) ? normalize(nt.normalizers[key], data) : data
-    NamedTuple{names}(maybe_normalize(nt[x][inds], x) for x in names)
+    NamedTuple{names}(collect(maybe_normalize(nt[x][inds], x)) for x in names)
 end