Merge pull request #45 from JuliaReinforcementLearning/jpsl/type-stability-metaprogramming

Use metaprogramming to improve type stability...

Author: jeremiahpslewis

src/episodes.jl

@@ -85,21 +85,48 @@ ispartial_insert(traces::Traces, xs) = length(xs) < length(traces.traces) #this
 ispartial_insert(es::EpisodesBuffer, xs) = ispartial_insert(es.traces, xs)
 ispartial_insert(traces::CircularPrioritizedTraces, xs) = ispartial_insert(traces.traces, xs)
 
-function fill_multiplex(es::EpisodesBuffer)
-    for trace in es.traces.traces
-        if !(trace isa MultiplexTraces)
-            push!(trace, last(trace)) #push a duplicate of last element as a dummy element, should never be sampled.
-        end
-    end
+function pad!(trace::Trace)
+    pad!(trace.parent)
+    return nothing
 end
-function fill_multiplex(es::EpisodesBuffer{<:Any,<:Any,<:CircularPrioritizedTraces})
-    for trace in es.traces.traces.traces
-        if !(trace isa MultiplexTraces)
-            push!(trace, last(trace)) #push a duplicate of last element as a dummy element, should never be sampled.
+
+pad!(buf::CircularVectorBuffer{T}) where {T} = push!(buf, zero(T))
+pad!(vect::Vector{T}) where {T} = push!(vect, zero(T))
+
+#push a duplicate of last element as a dummy element for all 'trace' objects, ignores multiplex traces, should never be sampled.
+@generated function fill_multiplex(trace_tuple::Traces{names,Trs,N,E}) where {names,Trs,N,E}
+    traces_signature = Trs
+    ex = :()
+    i = 1
+
+    if traces_signature <: NamedTuple
+        # Handle 'simple' (non-multiplexed) Traces
+        for tr in traces_signature.parameters[1]
+            ex = :($ex; pad!(trace_tuple.traces[$i])) # pad everything
+            i += 1
         end
+    elseif traces_signature <: Tuple
+        traces_signature = traces_signature.parameters
+        
+    
+        for tr in traces_signature
+            if !(tr <: MultiplexTraces)
+                #push a duplicate of last element as a dummy element, should never be sampled.
+                ex = :($ex; pad!(trace_tuple.traces[$i]))
+            end
+            i += 1
+        end
+    else
+        error("Traces store is neither a tuple nor a named tuple!")
     end
+
+    return :($ex)
 end
 
+fill_multiplex(es::EpisodesBuffer) = fill_multiplex(es.traces)
+
+fill_multiplex(es::EpisodesBuffer{<:Any,<:Any,<:CircularPrioritizedTraces}) = fill_multiplex(es.traces.traces)
+
 function Base.push!(eb::EpisodesBuffer, xs::NamedTuple)
     push!(eb.traces, xs)
     partial = ispartial_insert(eb, xs)

src/samplers.jl

@@ -49,7 +49,7 @@ StatsBase.sample(s::BatchSampler{names}, t::AbstractTraces) where {names} = Stat
 
 function StatsBase.sample(s::BatchSampler, t::AbstractTraces, names, weights = StatsBase.UnitWeights{Int}(length(t)))
     inds = StatsBase.sample(s.rng, 1:length(t), weights, s.batch_size)
-    NamedTuple{names}(map(x -> collect(t[x][inds]), names))
+    NamedTuple{names}(map(x -> collect(t[Val(x)][inds]), names))
 end
 
 function StatsBase.sample(s::BatchSampler, t::EpisodesBuffer, names)
@@ -74,12 +74,12 @@ function StatsBase.sample(s::BatchSampler, e::EpisodesBuffer{<:Any, <:Any, <:Cir
     st = deepcopy(t.priorities)
     st .*= e.sampleable_inds[1:end-1] #temporary sumtree that puts 0 priority to non sampleable indices.
     inds, priorities = rand(s.rng, st, s.batch_size)
-    NamedTuple{(:key, :priority, names...)}((t.keys[inds], priorities, map(x -> collect(t.traces[x][inds]), names)...))
+    NamedTuple{(:key, :priority, names...)}((t.keys[inds], priorities, map(x -> collect(t.traces[Val(x)][inds]), names)...))
 end
 
 function StatsBase.sample(s::BatchSampler, t::CircularPrioritizedTraces, names)
     inds, priorities = rand(s.rng, t.priorities, s.batch_size)
-    NamedTuple{(:key, :priority, names...)}((t.keys[inds], priorities, map(x -> collect(t.traces[x][inds]), names)...))
+    NamedTuple{(:key, :priority, names...)}((t.keys[inds], priorities, map(x -> collect(t.traces[Val(x)][inds]), names)...))
 end
 
 #####

src/traces.jl

@@ -128,15 +128,18 @@ end
 
 Adapt.adapt_structure(to, t::MultiplexTraces{names}) where {names} = MultiplexTraces{names}(Adapt.adapt_structure(to, t.trace))
 
-function Base.getindex(t::MultiplexTraces{names}, k::Symbol) where {names}
-    a, b = names
-    if k == a
-        RelativeTrace{0,-1}(convert(AbstractTrace, t.trace))
-    elseif k == b
-        RelativeTrace{1,0}(convert(AbstractTrace, t.trace))
+Base.getindex(t::MultiplexTraces{names}, k::Symbol) where {names} = _getindex(t, Val(k))
+
+@generated function _getindex(t::MultiplexTraces{names}, ::Val{k}) where {names,k}
+    ex = :()
+    if QuoteNode(names[1]) == QuoteNode(k)
+        ex = :(RelativeTrace{0,-1}(t.trace))
+    elseif QuoteNode(names[2]) == QuoteNode(k)
+        ex = :(RelativeTrace{1,0}(t.trace))
     else
-        throw(ArgumentError("unknown trace name: $k"))
+        ex = :(throw(ArgumentError("unknown trace name: $k")))
     end
+    return :($ex)
 end
 
 Base.getindex(t::MultiplexTraces{names}, I::Int) where {names} = NamedTuple{names}((t.trace[I], t.trace[I+1]))
@@ -163,83 +166,144 @@ end
 
 struct Traces{names,T,N,E} <: AbstractTraces{names,E}
     traces::T
-    inds::NamedTuple{names,NTuple{N,Int}}
 end
 
 function Adapt.adapt_structure(to, t::Traces{names,T,N,E}) where {names,T,N,E}
     data = Adapt.adapt_structure(to, t.traces)
     # FIXME: `E` is not adapted here
-    Traces{names,typeof(data),length(names),E}(data, t.inds)
+    Traces{names,typeof(data),length(names),E}(data)
 end
 
 function Traces(; kw...)
     data = map(x -> convert(AbstractTrace, x), values(kw))
     names = keys(data)
-    inds = NamedTuple(k => i for (i, k) in enumerate(names))
-    Traces{names,typeof(data),length(names),typeof(values(data))}(data, inds)
+    Traces{names,typeof(data),length(names),typeof(values(data))}(data)
 end
 
 
-function Base.getindex(ts::Traces, s::Symbol)
-    t = ts.traces[ts.inds[s]]
+Base.getindex(ts::Traces, s::Symbol) = Base.getindex(ts::Traces, Val(s))
+
+function Base.getindex(ts::Traces, ::Val{s}) where {s}
+    t = _gettrace(ts, Val(s))
     if t isa AbstractTrace
         t
+    elseif t isa MultiplexTraces
+        _getindex(t, Val(s))
     else
-        t[s]
+        throw(ArgumentError("unknown trace name: $s"))
     end
 end
 
-Base.getindex(t::Traces{names}, i) where {names} = NamedTuple{names}(map(k -> t[k][i], names))
+@generated function _gettrace(ts::Traces{names,Trs,N,E}, ::Val{k}) where {names,Trs,N,E,k}
+    index_ = build_trace_index(names, Trs)
+    # Generate code, i.e. find the correct index for a given key
+    ex = :()
+    
+    for name in names
+        if QuoteNode(name) == QuoteNode(k)
+            index_element = index_[k]
+            ex = :(ts.traces[$index_element])
+            break
+        end
+    end
+
+    return :($ex)
+end
+
+@generated function Base.getindex(t::Traces{names}, i) where {names}
+    ex = :(NamedTuple{$(names)}($(Expr(:tuple))))
+    for k in names
+        push!(ex.args[2].args, :(t[Val($(QuoteNode(k)))][i]))
+    end
+    return ex
+end
 
 function Base.:(+)(t1::AbstractTraces{k1,T1}, t2::AbstractTraces{k2,T2}) where {k1,k2,T1,T2}
     ks = (k1..., k2...)
     ts = (t1, t2)
-    inds = (; (k => 1 for k in k1)..., (k => 2 for k in k2)...)
-    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts, inds)
+    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts)
 end
 
 function Base.:(+)(t1::AbstractTraces{k1,T1}, t2::Traces{k2,T,N,T2}) where {k1,T1,k2,T,N,T2}
     ks = (k1..., k2...)
     ts = (t1, t2.traces...)
-    inds = merge(NamedTuple(k => 1 for k in k1), map(v -> v + 1, t2.inds))
-    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts, inds)
+    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts)
 end
 
 
 function Base.:(+)(t1::Traces{k1,T,N,T1}, t2::AbstractTraces{k2,T2}) where {k1,T,N,T1,k2,T2}
     ks = (k1..., k2...)
     ts = (t1.traces..., t2)
-    inds = merge(t1.inds, (; (k => length(ts) for k in k2)...))
-    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts, inds)
+    Traces{ks,typeof(ts),length(ks),Tuple{T1.types...,T2.types...}}(ts)
 end
 
 function Base.:(+)(t1::Traces{k1,T1,N1,E1}, t2::Traces{k2,T2,N2,E2}) where {k1,T1,N1,E1,k2,T2,N2,E2}
     ks = (k1..., k2...)
     ts = (t1.traces..., t2.traces...)
-    inds = merge(t1.inds, map(x -> x + length(t1.traces), t2.inds))
-    Traces{ks,typeof(ts),length(ks),Tuple{E1.types...,E2.types...}}(ts, inds)
+    Traces{ks,typeof(ts),length(ks),Tuple{E1.types...,E2.types...}}(ts)
 end
 
 Base.size(t::Traces) = (mapreduce(length, min, t.traces),)
-capacity(t::Traces) = minimum(map(idx->capacity(t.traces[idx]),t.inds))
 
-for f in (:push!, :pushfirst!)
-    @eval function Base.$f(ts::Traces, xs::NamedTuple)
-        for (k, v) in pairs(xs)
-            $f(ts, Val(k), v)
+function capacity(t::Traces{names,Trs,N,E}) where {names,Trs,N,E}
+    minimum(map(idx->capacity(t[idx]), names))
+end
+
+@generated function Base.push!(ts::Traces, xs::NamedTuple{N,T}) where {N,T}
+    ex = :()
+    for n in N
+        ex = :($ex; push!(ts, Val($(QuoteNode(n))), xs.$n))
+    end
+    return :($ex)
+end
+
+@generated function Base.pushfirst!(ts::Traces, xs::NamedTuple{N,T}) where {N,T}
+    ex = :()
+    for n in N
+        ex = :($ex; pushfirst!(ts, Val($(QuoteNode(n))), xs.$n))
+    end
+    return :($ex)
+end
+
+@generated function Base.pushfirst!(ts::Traces{names,Trs,N,E}, ::Val{k}, v) where {names,Trs,N,E,k}
+    index_ = build_trace_index(names, Trs)
+    # Generate code, i.e. find the correct index for a given key
+    ex = :()
+    
+    for name in names
+        if QuoteNode(name) == QuoteNode(k)
+            index_element = index_[k]
+            ex = :(pushfirst!(ts.traces[$index_element], Val($(QuoteNode(k))), v))
+            break
         end
     end
 
-    @eval function Base.$f(ts::Traces, ::Val{k}, v) where {k}
-        $f(ts.traces[ts.inds[k]], Val(k), v)
+    return :($ex)
+end
+
+@generated function Base.push!(ts::Traces{names,Trs,N,E}, ::Val{k}, v) where {names,Trs,N,E,k}
+    index_ = build_trace_index(names, Trs)
+    # Generate code, i.e. find the correct index for a given key
+    ex = :()
+    
+    for name in names
+        if QuoteNode(name) == QuoteNode(k)
+            index_element = index_[k]
+            ex = :(push!(ts.traces[$index_element], Val($(QuoteNode(k))), v))
+            break
+        end
     end
 
+    return :($ex)
+end
+
+for f in (:push!, :pushfirst!)
     @eval function Base.$f(t::AbstractTrace, ::Val{k}, v) where {k}
         $f(t, v)
     end
 
     @eval function Base.$f(t::Trace, ::Val{k}, v) where {k}
-        $f(t, v)
+        $f(t.parent, v)
     end
 
     @eval function Base.$f(ts::MultiplexTraces, ::Val{k}, v) where {k}
@@ -251,7 +315,7 @@ end
 for f in (:append!, :prepend!)
     @eval function Base.$f(ts::Traces, xs::Traces)
         for k in keys(xs)
-            t = ts.traces[ts.inds[k]]
+            t = _gettrace(ts, Val(k))
             $f(t, xs[k])
         end
     end
@@ -264,3 +328,38 @@ for f in (:pop!, :popfirst!, :empty!)
         end
     end
 end
+
+
+"""
+    build_trace_index(names::NTuple, traces_signature::DataType)
+
+Take type signature from `Traces` and build a mapping from trace name to trace index
+"""
+function build_trace_index(names::NTuple, traces_signature::DataType)
+    # Build index
+    index_ = Dict()
+
+    if traces_signature <: NamedTuple
+        # Handle simple Traces
+        index_ = Dict(name => i for (name, i) ∈ zip(names, 1:length(names)))
+    elseif traces_signature <: Tuple
+        # Handle MultiplexTracesup
+        i = 1
+        j = 1
+        trace_list = traces_signature.parameters
+        for tr in trace_list
+            if tr <: MultiplexTraces
+                index_[names[i]] = j
+                i += 1
+                index_[names[i]] = j
+            else
+                index_[names[i]] = j
+            end
+            i += 1
+            j += 1
+        end
+    else
+        error("Traces store is neither a tuple nor a named tuple!")
+    end
+    return index_
+end

test/episodes.jl

@@ -28,7 +28,7 @@ using Test
         @test eb.episodes_lengths[end] == 0
         @test eb.step_numbers[end] == 1
         @test eb.sampleable_inds == [1,1,1,1,1,0,0]
-        @test eb[6][:reward] == 5 #6 is not a valid index, the reward there is dummy duplicate of previous (5)
+        @test eb[6][:reward] == 0 #6 is not a valid index, the reward there is filled as zero
         ep2_len = 0
         for (j,i) = enumerate(8:11)
             ep2_len += 1
@@ -113,7 +113,7 @@ using Test
         @test eb.episodes_lengths[end] == 0
         @test eb.step_numbers[end] == 1
         @test eb.sampleable_inds == [1,1,1,1,1,0,0]
-        @test eb[6][:reward] == 5 #6 is not a valid index, the reward there is dummy duplicate of previous (5)
+        @test eb[6][:reward] == 0 #6 is not a valid index, the reward there is dummy, filled as zero
         ep2_len = 0
         for (j,i) = enumerate(8:11)
             ep2_len += 1

test/traces.jl

@@ -115,3 +115,47 @@ end
     t8_view.a[1] = 0
     @test t8[:a][2] == 0
 end
+
+using ReinforcementLearningTrajectories: build_trace_index
+
+@testset "build_trace_index" begin
+    t1 = CircularArraySARTSATraces(;
+        capacity=3,
+        state=Float32 => (2, 3),
+        action=Float32 => (2,),
+        reward=Float32 => (),
+        terminal=Bool => ()
+    )
+    @test build_trace_index(typeof(t1).parameters[1], typeof(t1).parameters[2]) == Dict(:reward => 3,
+        :next_state => 1,
+        :state => 1,
+        :action => 2,
+        :next_action => 2,
+        :terminal => 4)
+
+    t2 = Traces(; a=[2, 3], b=[false, true])
+    build_trace_index(typeof(t2).parameters[1], typeof(t2).parameters[2])
+end
+
+@testset "push!(ts::Traces{names,Trs,N,E}, ::Val{k}, v)" begin
+    t1 = CircularArraySARTSATraces(;
+        capacity=3,
+        state=Float32 => (2, 3),
+        action=Float32 => (2,),
+        reward=Float32 => (),
+        terminal=Bool => ()
+    )
+    push!(t1, Val(:reward), 5)
+    @test t1[:reward][1] == 5
+
+    @test size(Base.getindex(t1, :reward)) == (1,)
+    @test size(Base.getindex(t1, 1).state) == (2,3)
+
+
+    t2 = Traces(; a=[2, 3], b=[false, true])
+    push!(t2, Val(:a), 5)
+    @test t2[:a][3] == 5
+
+    @test size(Base.getindex(t2, :a)) == (3,)
+    @test Base.getindex(t2, 1) == (; a = 2, b= false)
+end