move to NormalizedTraces

Author: HenriDeh

src/ReinforcementLearningTrajectories.jl

@@ -4,11 +4,11 @@ const RLTrajectories = ReinforcementLearningTrajectories
 export RLTrajectories
 
 include("patch.jl")
-
 include("traces.jl")
 include("samplers.jl")
 include("controllers.jl")
 include("trajectory.jl")
+include("normalization.jl")
 include("common/common.jl")
 
 end

src/normalization.jl

@@ -1,11 +1,11 @@
 import OnlineStats: OnlineStats, Group, Moments, fit!, OnlineStat, Weight, EqualWeight, mean, std
-export scalar_normalizer, array_normalizer, NormalizedTrace, Normalizer
+export scalar_normalizer, array_normalizer, NormalizedTraces, Normalizer
 import MacroTools.@forward
 
 """
     Normalizer(::OnlineStat)
 
-Wraps an OnlineStat to be used by a [`NormalizedTrajectory`](@ref).
+Wraps an OnlineStat to be used by a [`NormalizedTraces`](@ref).
 """
 struct Normalizer{OS<:OnlineStat}
     os::OS
@@ -87,25 +87,32 @@ t = Trajectory(
 )
 
 """
-struct NormalizedTrace{T <: Trace, N <: Normalizer}
-    trace::T
-    normalizer::N
+struct NormalizedTraces{T <: AbstractTraces, names, N} #<: AbstractTraces
+    traces::T
+    normalizers::NamedTuple{names, N}
+    function NormalizedTraces(traces::AbstractTraces; trace_normalizer_pairs...)
+        for key in keys(trace_normalizer_pairs)
+            @assert key in keys(traces) "Traces do not have key $key, valid keys are $(keys(traces))."
+        end
+        nt = (; trace_normalizer_pairs...)
+        new{typeof(traces), keys(nt), typeof(values(nt))}(traces, nt)
+    end
 end 
 
-NormalizedTrace(x, normalizer) = NormalizedTrace(convert(Trace, x), normalizer)
-
-@forward NormalizedTrace.trace Base.length, Base.lastindex, Base.firstindex, Base.getindex, Base.view, Base.pop!, Base.popfirst!, Base.empty!
-
-Base.convert(::Type{Trace}, x::NormalizedTrace) = x #ignore conversion to Trace
+@forward NormalizedTraces.traces Base.length, Base.lastindex, Base.firstindex, Base.getindex, Base.view, Base.pop!, Base.popfirst!, Base.empty!, Base.prepend!, Base.pushfirst!
 
-function Base.push!(nt::NormalizedTrace, x)
-    fit!(nt.normalizer, x)
-    push!(nt.trace, x)
+function Base.push!(nt::NormalizedTraces, x)
+    for key in intersect(keys(nt.normalizers), keys(x))
+        fit!(nt.normalizers[key], x[key])
+    end
+    push!(nt.traces, x)
 end
 
-function Base.append!(nt::NormalizedTrace, x)
-    fit!(nt.normalizer, x)
-    append!(nt.trace, x)
+function Base.append!(nt::NormalizedTraces, x)
+    for key in intersect(keys(nt.normalizers), keys(x))
+        fit!(nt.normalizers[key], x[key])
+    end
+    append!(nt.traces, x)
 end
 
 """
@@ -149,12 +156,22 @@ function normalize(os::Group{<:AbstractVector{<:Moments}}, x::AbstractVector{<:A
     return xn
 end
 
-function fetch(nt::NormalizedTrace, inds)
-    batch = fetch(nt.trace, inds)
-    normalize(nt.normalizer.os, batch)
+function fetch(nt::NormalizedTraces, inds)
+    batch = fetch(nt.traces, inds)
+    return (; (key => (key in keys(nt.normalizers) ? normalize(nt.normalizers[key].os, data) : data) for (key, data) in pairs(batch))...)
+end
+
+#=
+abstract type AbstractFoo{T} end
+
+struct Foo{T} <: AbstractFoo{T}
+    x::T
+end
+
+struct Baz{F <: Foo} <: supertype(Type(F))
+    foo::F
 end
 
-function sample(s, nt::NormalizedTrace)
-    batch = sample(s, nt.trace)
-    normalize(nt.normalizer.os, batch)
-end
+struct Bal{F <: AbstractFoo{T where T}} <: AbstractFoo{T} 
+    foo::F{T}
+end=#

src/samplers.jl

@@ -27,7 +27,7 @@ sample(s::BatchSampler{names}, t::AbstractTraces) where {names} = sample(s, t, n
 
 function sample(s::BatchSampler, t::AbstractTraces, names)
     inds = rand(s.rng, 1:length(t), s.batch_size)
-    NamedTuple{names}(s.transformer(t[x][inds]) for x in names)
+    NamedTuple{names}(s.transformer(fetch(t[x], inds) for x in names))
 end
 
 """

src/traces.jl

@@ -33,14 +33,12 @@ Base.setindex!(s::Trace, v, I) = setindex!(s.parent, v, ntuple(i -> i == ndims(s
 
 @forward Trace.parent Base.parent, Base.pushfirst!, Base.push!, Base.append!, Base.prepend!, Base.pop!, Base.popfirst!, Base.empty!
 
-fetch(t::Trace, inds) = t[inds]
+fetch(t::AbstractTrace, inds) = t[inds]
 #####
 
 """
 For each concrete `AbstractTraces`, we have the following assumption:
 
-
-##
 1. Every inner trace is an `AbstractVector`
 1. Support partial updating
 1. Return *View* by default when getting elements.
@@ -58,6 +56,8 @@ end
 Base.keys(t::AbstractTraces{names}) where {names} = names
 Base.haskey(t::AbstractTraces{names}, k::Symbol) where {names} = k in names
 
+#use fetch instead of getindex when sampling to retain compatibility
+fetch(t::AbstractTraces, inds) = t[inds]
 #####
 
 """
@@ -84,7 +84,7 @@ end
 
 function MultiplexTraces{names}(t) where {names}
     if length(names) != 2
-        throw(ArgumentError("MultiplexTraces has exactly two sub traces, got $length(names) trace names"))
+        throw(ArgumentError("MultiplexTraces has exactly two sub traces, got $(length(names)) trace names"))
     end
     trace = convert(AbstractTrace, t)
     MultiplexTraces{names,typeof(trace),eltype(trace)}(trace)
@@ -301,13 +301,6 @@ end
 
 Base.size(t::Traces) = (mapreduce(length, min, t.traces),)
 
-function sample(s::BatchSampler, t::Traces)
-    inds = rand(s.rng, 1:length(t), s.batch_size)
-    map(t.traces) do x
-        fetch(x, inds)
-    end |> s.transformer
-end
-
 for f in (:push!, :pushfirst!)
     @eval function Base.$f(ts::Traces, xs::NamedTuple)
         for (k, v) in pairs(xs)

test/normalization.jl

@@ -1,46 +1,43 @@
 using Test
-using Trajectories
-import Trajectories.normalize
-import OnlineStats: fit!, mean, std
+using ReinforcementLearningTrajectories
+import ReinforcementLearningTrajectories: fetch, sample
+import OnlineStats: mean, std
 
 @testset "normalization.jl" begin
-    #scalar normalization
-    rewards = [1.:10;]
-    rn = scalar_normalizer()
-    fit!(rn, rewards)
-    batch_reward = normalize(rn, [6.,5.,10.])
-    @test batch_reward ≈ ([6.,5.,10.] .- mean(1:10))./std(1:10)
-    #vector normalization
-    states = reshape([1:50;], 5, 10)
-    sn = array_normalizer((5,))
-    fit!(sn, states)
-    @test [mean(stat) for stat in sn] == [mean((1:5:46) .+i) for i in 0:4]
-    batch_states =normalize(sn, reshape(repeat(5.:-1:1, 5), 5,5))
-    @test all(length(unique(x)) == 1 for x in eachrow(batch_states))
-    #array normalization
-    states = reshape(1.:250, 5,5,10)
-    sn = array_normalizer((5,5))
-    fit!(sn, states)
-    batch_states = normalize(sn, collect(states))
-    
-    #NormalizedTrace
-    t = Trajectory(
-        container=Traces(
-            a= NormalizedTrace(Float32[], scalar_normalizer()),
-            b=Int[],
-            c=NormalizedTrace(Vector{Float32}[], array_normalizer((10,))) #TODO check with ElasticArrays and Episodes
-        ),
-        sampler=BatchSampler(300000),
-        controler=InsertSampleRatioControler(Inf, 0)
+    t = CircularArraySARTTraces(capacity = 10, state = Float64 => (5,))
+    nt = NormalizedTraces(t, reward = scalar_normalizer(), state = array_normalizer((5,)))
+    m = mean(0:4)
+    s = std(0:4)
+
+    for i in 0:4
+        r = ((1.0:5.0) .+ i) .% 5
+        push!(nt, (state = [r;], action = 1, reward = Float32(i), terminal = false))
+    end
+    push!(nt, (next_state = fill(m, 5), next_action = 1)) #does not update because next_state is not in keys of normlizers. Is this desirable or not ?
+
+    @test mean(nt.normalizers[:reward].os) == m && std(nt.normalizers[:reward].os) == s
+    @test all(nt.normalizers[:state].os) do moments
+        mean(moments) == m && std(moments) == s
+    end
+
+    unnormalized_batch = fetch(t, [1:5;])
+    @test unnormalized_batch[:reward] == [0:4;]
+    @test extrema(unnormalized_batch[:state]) == (0, 4)
+    normalized_batch = fetch(nt, [1:5;])
+    @test normalized_batch[:reward] ≈ ([0:4;] .- m)./s
+    @test all(extrema(normalized_batch[:state]) .≈ ((0, 4) .- m)./s)
+    @test normalized_batch[:state][:,5] ≈ ([0:4;] .- m)./s
+    #check for no mutation
+    unnormalized_batch = fetch(t, [1:5;])
+    @test unnormalized_batch[:reward] == [0:4;]
+    @test extrema(unnormalized_batch[:state]) == (0, 4)
+    #=
+    traj = Trajectory(
+        container = nt,
+        sampler = BatchSampler(10),
+        controller = InsertSampleRatioController(ratio = Inf, threshold = 0)
     )
-    append!(t, a = [1,2,3], b = [1,2,3], c = eachcol(reshape(1f0:30, 10,3)))
-    push!(t, a = 2, b = 2, c = fill(mean(1:30), 10))
-    @test mean(t.container[:a].trace.x) ≈ 2.
-    @test std(t.container[:a].trace.x) ≈ std([1,2,2,3])
-    a,b,c = take!(t)
-    @test eltype(a) == Float32
-    @test mean(a) ≈ 0 atol = 0.01
-    @test mean(b) ≈ 2 atol = 0.01 #b is not normalized
-    @test eltype(first(c)) == Float32
-    @test all(isapprox(0f0, atol = 0.01), vec(mean(reduce(hcat,c), dims = 2)))
+
+    batch = sample(traj)=#
+
 end