adding an EpisodesBuffer

Project.toml

@@ -1,15 +1,17 @@
 name = "ReinforcementLearningTrajectories"
 uuid = "6486599b-a3cd-4e92-a99a-2cea90cc8c3c"
-version = "0.1.10"
+version = "0.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StackViews = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 Adapt = "3"

README.md

@@ -12,7 +12,8 @@ The relationship of several concepts provided in this package:
 ┌───────────────────────────────────┐
 │ Trajectory                        │
 │ ┌───────────────────────────────┐ │
-│ │ AbstractTraces                │ │
+│ │ EpisodesBuffer wrapping a     | |
+| | AbstractTraces                │ │
 │ │             ┌───────────────┐ │ │
 │ │ :trace_A => │ AbstractTrace │ │ │
 │ │             └───────────────┘ │ │
@@ -61,8 +62,7 @@ julia> for batch in t
 - `Traces`
 - `MultiplexTraces`
 - `CircularSARTTraces`
-- `Episode`
-- `Episodes`
+- `NormalizedTraces`
 
 **Samplers**
 

src/ReinforcementLearningTrajectories.jl

@@ -2,13 +2,14 @@ module ReinforcementLearningTrajectories
 
 const RLTrajectories = ReinforcementLearningTrajectories
 export RLTrajectories
+import StatsBase
 
 include("patch.jl")
 include("traces.jl")
 include("common/common.jl")
+include("episodes.jl")
 include("samplers.jl")
 include("controllers.jl")
 include("trajectory.jl")
 include("normalization.jl")
-
 end

src/common/CircularArraySARTTraces.jl → src/common/CircularArraySARTSATraces.jl

@@ -1,8 +1,8 @@
-export CircularArraySARTTraces
+export CircularArraySARTSATraces
 
-import CircularArrayBuffers
+import CircularArrayBuffers.CircularArrayBuffer
 
-const CircularArraySARTTraces = Traces{
+const CircularArraySARTSATraces = Traces{
     SS′AA′RT,
     <:Tuple{
         <:MultiplexTraces{SS′,<:Trace{<:CircularArrayBuffer}},
@@ -12,7 +12,7 @@ const CircularArraySARTTraces = Traces{
     }
 }
 
-function CircularArraySARTTraces(;
+function CircularArraySARTSATraces(;
     capacity::Int,
     state=Int => (),
     action=Int => (),
@@ -24,12 +24,12 @@ function CircularArraySARTTraces(;
     reward_eltype, reward_size = reward
     terminal_eltype, terminal_size = terminal
 
-    MultiplexTraces{SS′}(CircularArrayBuffer{state_eltype}(state_size..., capacity + 1)) +
-    MultiplexTraces{AA′}(CircularArrayBuffer{action_eltype}(action_size..., capacity + 1)) +
+    MultiplexTraces{SS′}(CircularArrayBuffer{state_eltype}(state_size..., capacity+1)) +
+    MultiplexTraces{AA′}(CircularArrayBuffer{action_eltype}(action_size..., capacity+1)) +
     Traces(
         reward=CircularArrayBuffer{reward_eltype}(reward_size..., capacity),
         terminal=CircularArrayBuffer{terminal_eltype}(terminal_size..., capacity),
     )
 end
 
-CircularArrayBuffers.capacity(t::CircularArraySARTTraces) = CircularArrayBuffers.capacity(t.traces[end])
+CircularArrayBuffers.capacity(t::CircularArraySARTSATraces) = CircularArrayBuffers.capacity(minimum(map(capacity,t.traces)))

src/common/CircularArraySARTSTraces.jl

@@ -0,0 +1,35 @@
+export CircularArraySARTSTraces
+
+import CircularArrayBuffers.CircularArrayBuffer
+
+const CircularArraySARTSTraces = Traces{
+    SS′ART,
+    <:Tuple{
+        <:MultiplexTraces{SS′,<:Trace{<:CircularArrayBuffer}},
+        <:Trace{<:CircularArrayBuffer},
+        <:Trace{<:CircularArrayBuffer},
+        <:Trace{<:CircularArrayBuffer},
+    }
+}
+
+function CircularArraySARTSTraces(;
+    capacity::Int,
+    state=Int => (),
+    action=Int => (),
+    reward=Float32 => (),
+    terminal=Bool => ())
+
+    state_eltype, state_size = state
+    action_eltype, action_size = action
+    reward_eltype, reward_size = reward
+    terminal_eltype, terminal_size = terminal
+
+    MultiplexTraces{SS′}(CircularArrayBuffer{state_eltype}(state_size..., capacity+1)) +
+    Traces(
+        action = CircularArrayBuffer{action_eltype}(action_size..., capacity),
+        reward=CircularArrayBuffer{reward_eltype}(reward_size..., capacity),
+        terminal=CircularArrayBuffer{terminal_eltype}(terminal_size..., capacity),
+    )
+end
+
+CircularArrayBuffers.capacity(t::CircularArraySARTSTraces) = CircularArrayBuffers.capacity(minimum(map(capacity,t.traces)))

src/common/CircularArraySLARTTraces.jl

@@ -33,3 +33,5 @@ function CircularArraySLARTTraces(;
         terminal=CircularArrayBuffer{terminal_eltype}(terminal_size..., capacity),
     )
 end
+
+CircularArrayBuffers.capacity(t::CircularArraySLARTTraces) = CircularArrayBuffers.capacity(minimum(map(capacity,t.traces)))

src/common/CircularPrioritizedTraces.jl

@@ -1,6 +1,6 @@
 export CircularPrioritizedTraces
 
-using CircularArrayBuffers: capacity, CircularVectorBuffer
+using CircularArrayBuffers: CircularVectorBuffer
 
 struct CircularPrioritizedTraces{T,names,Ts} <: AbstractTraces{names,Ts}
     keys::CircularVectorBuffer{Int,Vector{Int}}

src/common/common.jl

@@ -1,7 +1,5 @@
 export SS′, LL′, AA′, RT, SS′ART, SS′AA′RT, SS′L′ART, SS′LL′AA′RT
 
-using CircularArrayBuffers
-
 const SS′ = (:state, :next_state)
 const LL′ = (:legal_actions_mask, :next_legal_actions_mask)
 const AA′ = (:action, :next_action)
@@ -12,7 +10,8 @@ const SS′L′ART = (SS′..., :next_legal_actions_mask, :action, RT...)
 const SS′LL′AA′RT = (SS′..., LL′..., AA′..., RT...)
 
 include("sum_tree.jl")
-include("CircularArraySARTTraces.jl")
+include("CircularArraySARTSTraces.jl")
+include("CircularArraySARTSATraces.jl")
 include("CircularArraySLARTTraces.jl")
 include("CircularPrioritizedTraces.jl")
 include("ElasticArraySARTTraces.jl")

src/episodes.jl

@@ -0,0 +1,138 @@
+export EpisodesBuffer, PartialNamedTuple
+import DataStructures.CircularBuffer
+
+"""
+    EpisodesBuffer(traces::AbstractTraces)
+
+Wraps an `AbstractTraces` object, usually the container of a `Trajectory`. 
+`EpisodesBuffer` tracks the indexes of the `traces` object that belong to the same episodes.
+To that end, it stores 
+1. an vector `sampleable_inds` of Booleans that determine whether an index in Traces is legally sampleable
+(i.e., it is not the index of a last state of an episode);
+2. a vector `episodes_lengths` that contains the total duration of the episode that each step belong to;
+3. an vector `step_numbers` that contains the index within the episode of the corresponding step.
+
+This information is used to correctly sample the traces. For example, if we have an episode that lasted 10 steps, the buffer stores 11 states. sampleable_inds[i] will
+be true for the index of the first ten steps, and 0 at the index of state number 11. episodes_lengths will be 11
+consecutive 10s (the episode saw 11 states but 10 steps occured). step_numbers will be 1 to 11.
+
+If `traces` is a capacitated buffer, such as a CircularArraySARTSTraces, then these three vectors will also be circular.
+
+EpisodesBuffer assumes that individual transitions are `push!`ed. Appending is not yet supported.
+"""
+
+mutable struct EpisodesBuffer{names, E, T<:AbstractTraces{names, E},B,S} <: AbstractTraces{names,E}
+    traces::T
+    sampleable_inds::S
+    step_numbers::B
+    episodes_lengths::B
+end
+
+"""
+    PartialNamedTuple(::NamedTuple)
+
+Wraps a NamedTuple to signal an EpisodesBuffer that it is pushed into that it should 
+ignore the fact that this is a partial insertion. Used at the end of an episode to
+complete multiplex traces before moving to the next episode.
+"""
+struct PartialNamedTuple{T}
+    namedtuple::T
+end
+
+function EpisodesBuffer(traces::AbstractTraces)
+    cap = any(t->t isa MultiplexTraces, traces.traces) ? capacity(traces) + 1 : capacity(traces)
+    @assert isempty(traces) "EpisodesBuffer must be initialized with empty traces."
+    if !isinf(cap)
+        legalinds =  CircularBuffer{Bool}(cap)
+        step_numbers = CircularBuffer{Int}(cap)
+        eplengths = deepcopy(step_numbers)
+        EpisodesBuffer(traces, legalinds, step_numbers, eplengths)
+    else
+        legalinds =  BitVector()
+        step_numbers = Vector{Int}()
+        eplengths = deepcopy(step_numbers)
+        EpisodesBuffer(traces, legalinds, step_numbers, eplengths)
+    end
+end
+
+Base.getindex(es::EpisodesBuffer, idx...) = getindex(es.traces, idx...)
+Base.setindex!(es::EpisodesBuffer, idx...) = setindex!(es.traces, idx...)
+Base.size(es::EpisodesBuffer) = size(es.traces)
+Base.length(es::EpisodesBuffer) = length(es.traces)
+Base.keys(es::EpisodesBuffer) = keys(es.traces)
+Base.keys(es::EpisodesBuffer{<:Any,<:Any,<:CircularPrioritizedTraces}) = keys(es.traces.traces)
+function Base.show(io::IO, m::MIME"text/plain", eb::EpisodesBuffer{names}) where {names}
+    s = nameof(typeof(eb))
+    t = eb.traces
+    println(io, "$s containing")
+    show(io, m::MIME"text/plain", t)
+end
+
+ispartial_insert(traces::Traces, xs) = length(xs) < length(traces.traces) #this is the number of traces it contains not the number of steps.
+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
+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.
+        end
+    end
+end
+
+function Base.push!(eb::EpisodesBuffer, xs::NamedTuple)
+    push!(eb.traces, xs)
+    partial = ispartial_insert(eb, xs)
+    if length(eb.traces) == 0 && partial #first push should be partial
+        push!(eb.step_numbers, 1)
+        push!(eb.episodes_lengths, 0)
+        push!(eb.sampleable_inds, 0)
+    elseif !partial #typical inserting
+        if length(eb.traces) < length(eb) && length(eb) > 2 #case when PartialNamedTuple is used. Steps are indexable one step later
+            eb.sampleable_inds[end-1] = 1 
+        else #case when we don't, length of traces and eb will match.
+            eb.sampleable_inds[end] = 1 #previous step is now indexable
+        end
+        push!(eb.sampleable_inds, 0) #this one is no longer
+        ep_length = last(eb.step_numbers)
+        push!(eb.episodes_lengths, ep_length)
+        startidx = max(1,length(eb.step_numbers) - last(eb.step_numbers))
+        eb.episodes_lengths[startidx:end] .= ep_length
+        push!(eb.step_numbers, ep_length + 1)
+    elseif partial
+        fill_multiplex(eb)
+        eb.sampleable_inds[end] = 0 #previous step is not indexable because it contains the last state
+        push!(eb.sampleable_inds, 0) #this one isn't either
+        push!(eb.step_numbers, 1)
+        push!(eb.episodes_lengths, 0)
+    end
+    return nothing
+end
+
+function Base.push!(eb::EpisodesBuffer, xs::PartialNamedTuple) #wrap a NamedTuple to push without incrementing the step number. 
+    push!(eb.traces, xs.namedtuple)
+    eb.sampleable_inds[end-1] = 1 #completes the episode trajectory.
+end
+
+for f in (:pop!, :popfirst!)
+    @eval function Base.$f(es::EpisodesBuffer)
+        $f(es.episodes_lengths)
+        $f(es.sampleable_inds)
+        $f(es.step_numbers)
+        $f(es.traces)
+    end
+end
+
+function Base.empty!(es::EpisodesBuffer)
+    empty!(es.traces)
+    empty!(es.episodes_lengths)
+    empty!(es.sampleable_inds)
+    empty!(es.step_numbers)
+end

src/normalization.jl

@@ -133,7 +133,7 @@ arrays (see [`array_normalizer`](@ref)).
 
 # Examples
 ```
-t = CircularArraySARTTraces(capacity = 10, state = Float64 => (5,))
+t = CircularArraySARTSTraces(capacity = 10, state = Float64 => (5,))
 nt = NormalizedTraces(t, reward = scalar_normalizer(), state = array_normalizer((5,)))
 # :next_state will also be normalized. 
 traj = Trajectory(
@@ -180,7 +180,7 @@ function Base.show(io::IO, ::MIME"text/plain", t::NormalizedTraces{names,T}) whe
     end
 end
 
-@forward NormalizedTraces.traces Base.length, Base.size, Base.lastindex, Base.firstindex, Base.getindex, Base.view, Base.pop!, Base.popfirst!, Base.empty!, Base.parent
+@forward NormalizedTraces.traces Base.length, Base.size, Base.lastindex, Base.firstindex, Base.view, Base.pop!, Base.popfirst!, Base.empty!, Base.parent
 
 for f in (:push!, :pushfirst!, :append!, :prepend!)
     @eval function Base.$f(nt::NormalizedTraces, x::T) where T
@@ -191,8 +191,19 @@ for f in (:push!, :pushfirst!, :append!, :prepend!)
     end
 end
 
-function sample(s::BatchSampler, nt::NormalizedTraces, names)
-    inds = rand(s.rng, 1:length(nt), s.batch_size)
+function StatsBase.sample(s::BatchSampler, nt::NormalizedTraces, names, weights = StatsBase.UnitWeights{Int}(length(nt)))
+    inds = StatsBase.sample(s.rng, 1:length(nt), weights, s.batch_size)
     maybe_normalize(data, key) = key in keys(nt.normalizers) ? normalize(nt.normalizers[key], data) : data
     NamedTuple{names}(collect(maybe_normalize(nt[x][inds], x)) for x in names)
 end
+
+function Base.getindex(nt::NormalizedTraces, inds)
+    maybe_normalize(data, key) = key in keys(nt.normalizers) ? normalize(nt.normalizers[key], data) : data
+    NamedTuple{keys(nt.traces)}(collect(maybe_normalize(nt.traces[x][inds], x)) for x in keys(nt.traces))
+end
+
+function Base.getindex(nt::NormalizedTraces, s::Symbol)
+    getindex(nt.traces, s)
+end
+
+ispartial_insert(traces::NormalizedTraces, xs) = ispartial_insert(traces.traces, xs)