Normalizer Wrapper

Project.toml

@@ -7,13 +7,15 @@ version = "0.1.0"
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Term = "22787eb5-b846-44ae-b979-8e399b8463ab"
 
 [compat]
 CircularArrayBuffers = "0.1"
 Term = "0.3"
 julia = "1.6"
+OnlineStats = "1.0"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/Trajectories.jl

@@ -5,6 +5,7 @@ include("controlers.jl")
 include("traces.jl")
 include("episodes.jl")
 include("trajectory.jl")
+include("normalization.jl")
 include("rendering.jl")
 include("common/common.jl")
 

src/normalization.jl

@@ -0,0 +1,160 @@
+import OnlineStats: OnlineStats, Group, Moments, fit!, OnlineStat, Weight, EqualWeight, mean, std
+export scalar_normalizer, array_normalizer, NormalizedTrace, Normalizer
+import MacroTools.@forward
+
+"""
+    Normalizer(::OnlineStat)
+
+Wraps an OnlineStat to be used by a [`NormalizedTrajectory`](@ref).
+"""
+struct Normalizer{OS<:OnlineStat}
+    os::OS
+end
+
+@forward Normalizer.os OnlineStats.mean, OnlineStats.std, Base.iterate, normalize, Base.length
+
+
+
+#Treats last dim as batch dim
+function OnlineStats.fit!(n::Normalizer, data::AbstractArray)
+    for d in eachslice(data, dims = ndims(data))
+        fit!(n.os, vec(d))
+    end
+    n
+end
+
+function OnlineStats.fit!(n::Normalizer{<:Group}, y::AbstractVector)
+    fit!(n.os, y)
+    n
+end
+
+function OnlineStats.fit!(n::Normalizer, y)
+    for yi in y
+        fit!(n.os, vec(yi))
+    end
+    n
+end
+
+function OnlineStats.fit!(n::Normalizer{<:Moments}, y::AbstractVector{<:Number})
+    for yi in y
+        fit!(n.os, yi)
+    end
+    n
+end
+
+function OnlineStats.fit!(n::Normalizer, data::Number)
+    fit!(n.os, data)
+    n
+end
+
+"""
+    scalar_normalizer(;weights = OnlineStats.EqualWeight())
+
+Returns preconfigured normalizer for scalar traces such as rewards. By default, all samples have equal weights in the computation of the moments.
+See the [OnlineStats documentation](https://joshday.github.io/OnlineStats.jl/stable/weights/) to use variants such as exponential weights to favor the most recent observations.
+"""
+scalar_normalizer(; weight::Weight = EqualWeight()) = Normalizer(Moments(weight = weight))
+
+"""
+    array_normalizer(size::Tuple{Int}; weights = OnlineStats.EqualWeight())
+
+Returns preconfigured normalizer for array traces such as vector or matrix states. 
+`size` is a tuple containing the dimension sizes of a state. E.g. `(10,)` for a 10-elements vector, or `(252,252)` for a square image.
+By default, all samples have equal weights in the computation of the moments.
+See the [OnlineStats documentation](https://joshday.github.io/OnlineStats.jl/stable/weights/) to use variants such as exponential weights to favor the most recent observations.
+"""
+array_normalizer(size::NTuple{N,Int}; weight::Weight = EqualWeight()) where N = Normalizer(Group([Moments(weight = weight) for _ in 1:prod(size)]))
+
+
+"""
+    NormalizedTrace(trace::Trace, normalizer::Normalizer)
+
+Wraps a [`Trace`](@ref) and a [`Normalizer`](@ref). When pushing new elements to the trace, a `NormalizedTrace` will first update a running estimate of the moments of that trace.
+When sampling a normalized trace, it will first normalize the samples using to zero mean and unit variance.
+
+preconfigured normalizers are provided for scalar (see [`scalar_normalizer`](@ref)) and arrays (see [`array_normalizer`](@ref))
+
+#Example
+t = Trajectory(
+    container=Traces(
+        a_scalar_trace = NormalizedTrace(Float32[], scalar_normalizer()),
+        a_non_normalized_trace=Bool[],
+        a_vector_trace = NormalizedTrace(Vector{Float32}[], array_normalizer((10,))),
+        a_matrix_trace = NormalizedTrace(Matrix{Float32}[], array_normalizer((252,252), weight = OnlineStats.ExponientialWeight(0.9f0)))
+    ),
+    sampler=BatchSampler(3),
+    controler=InsertSampleRatioControler(0.25, 4)
+)
+
+"""
+struct NormalizedTrace{T <: Trace, N <: Normalizer}
+    trace::T
+    normalizer::N
+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
+
+function Base.push!(nt::NormalizedTrace, x)
+    fit!(nt.normalizer, x)
+    push!(nt.trace, x)
+end
+
+function Base.append!(nt::NormalizedTrace, x)
+    fit!(nt.normalizer, x)
+    append!(nt.trace, x)
+end
+
+"""
+    normalize!(os::Moments, x)
+
+Given an Moments estimate of the elements of x, a vector of scalar traces,
+normalizes x elementwise to zero mean, and unit variance. 
+"""
+function normalize(os::Moments, x::AbstractVector)
+    T = eltype(x)
+    m, s = T(mean(os)), T(std(os))
+    return (x .- m) ./ s
+end
+
+"""
+    normalize!(os::Group{<:AbstractVector{<:Moments}}, x)
+
+Given an os::Group{<:Tuple{Moments}}, that is, a multivariate estimator of the moments of each element of x,
+normalizes each element of x to zero mean, and unit variance. Treats the last dimension as a batch dimension if `ndims(x) >= 2`.
+"""
+function normalize(os::Group{<:AbstractVector{<:Moments}}, x::AbstractVector)
+    T = eltype(x)
+    m = [T(mean(stat)) for stat in os]
+    s = [T(std(stat)) for stat in os]
+    return (x .- m) ./ s
+end
+
+function normalize(os::Group{<:AbstractVector{<:Moments}}, x::AbstractArray) 
+    xn = similar(x)
+    for (i, slice) in enumerate(eachslice(x, dims = ndims(x)))
+        xn[repeat([:], ndims(x)-1)..., i] .= reshape(normalize(os, vec(slice)), size(x)[1:end-1]...) 
+    end
+    return xn
+end
+
+function normalize(os::Group{<:AbstractVector{<:Moments}}, x::AbstractVector{<:AbstractArray})
+    xn = similar(x)
+    for (i,el) in enumerate(x)
+        xn[i] = normalize(os, vec(el))
+    end
+    return xn
+end
+
+function fetch(nt::NormalizedTrace, inds)
+    batch = deepcopy(fetch(nt.trace, inds))
+    normalize(nt.normalizer.os, batch)
+end
+
+function sample(s, nt::NormalizedTrace)
+    batch = deepcopy(sample(s, nt.trace))
+    normalize(nt.normalizer.os, batch)
+end

src/rendering.jl

@@ -19,7 +19,7 @@ function inner_convert(::Type{Term.AbstractRenderable}, x; style="gray1", width=
 end
 
 Base.convert(T::Type{Term.AbstractRenderable}, t::Trace{<:AbstractArray}; kw...) = convert(T, Trace(collect(eachslice(t.x, dims=ndims(t.x)))); kw..., type=typeof(t), subtitle="size: $(size(t.x))")
-
+Base.convert(T::Type{Term.AbstractRenderable}, t::NormalizedTrace; kw...) = convert(T, t.trace; kw..., type = typeof(t))
 function Base.convert(
     ::Type{Term.AbstractRenderable},
     t::Trace{<:AbstractVector};

src/traces.jl

@@ -32,6 +32,8 @@ Base.pop!(t::Trace) = pop!(t.x)
 Base.popfirst!(t::Trace) = popfirst!(t.x)
 Base.empty!(t::Trace) = empty!(t.x)
 
+fetch(t::Trace, inds) = t[inds]
+
 ##
 
 function sample(s::BatchSampler, t::Trace)
@@ -83,6 +85,6 @@ Base.empty!(t::Traces) = map(empty!, t.traces)
 function sample(s::BatchSampler, t::Traces)
     inds = rand(s.rng, 1:length(t), s.batch_size)
     map(t.traces) do x
-        x[inds]
+        fetch(x, inds)
     end |> s.transformer
 end

test/normalization.jl

@@ -0,0 +1,46 @@
+using Test
+using Trajectories
+import Trajectories.normalize
+import OnlineStats: fit!, 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, eachslice(states, dims = 3))
+    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(30000),
+        controler=InsertSampleRatioControler(Inf, 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)))
+end

test/runtests.jl

@@ -4,4 +4,5 @@ using Test
 @testset "Trajectories.jl" begin
     include("traces.jl")
     include("trajectories.jl")
+    include("normalization.jl")
 end