Unify the definition of AbstractTraces

.gitignore

@@ -1,4 +1,7 @@
+coverage/
 *.jl.*.cov
 *.jl.cov
 *.jl.mem
 /Manifest.toml
+
+.DS_Store

src/LastDimSlices.jl

@@ -0,0 +1,5 @@
+export LastDimSlices
+
+using MacroTools: @forward
+
+# See also https://github.com/JuliaLang/julia/pull/32310

src/Trajectories.jl

@@ -1,11 +1,10 @@
 module Trajectories
 
-include("samplers.jl")
-include("controlers.jl")
 include("traces.jl")
 include("episodes.jl")
+include("samplers.jl")
+include("controlers.jl")
 include("trajectory.jl")
-include("rendering.jl")
 include("common/common.jl")
 
 end

src/common/CircularArraySARTTraces.jl

@@ -1,16 +1,5 @@
 export CircularArraySARTTraces
 
-const CircularArraySARTTraces = Traces{
- SART,
- <:Tuple{
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer}
- }
-}
-
-
 function CircularArraySARTTraces(;
  capacity::Int,
  state=Int => (),
@@ -23,32 +12,10 @@ function CircularArraySARTTraces(;
  reward_eltype, reward_size = reward
  terminal_eltype, terminal_size = terminal
 
+ MultiplexTraces{(:state, :next_state)}(CircularArrayBuffer{state_eltype}(state_size..., capacity + 1)) +
+ MultiplexTraces{(:action, :next_action)}(CircularArrayBuffer{action_eltype}(action_size..., capacity + 1)) +
  Traces(
- state=CircularArrayBuffer{state_eltype}(state_size..., capacity + 1), # !!! state is one step longer
- action=CircularArrayBuffer{action_eltype}(action_size..., capacity + 1), # !!! action is one step longer
  reward=CircularArrayBuffer{reward_eltype}(reward_size..., capacity),
  terminal=CircularArrayBuffer{terminal_eltype}(terminal_size..., capacity),
  )
 end
-
-function Random.rand(s::BatchSampler, t::CircularArraySARTTraces)
- inds = rand(s.rng, 1:length(t), s.batch_size)
- inds′ = inds .+ 1
- (
- state=t[:state][inds],
- action=t[:action][inds],
- reward=t[:reward][inds],
- terminal=t[:terminal][inds],
- next_state=t[:state][inds′],
- next_action=t[:state][inds′]
- ) |> s.transformer
-end
-
-function Base.push!(t::CircularArraySARTTraces, x::NamedTuple{SA})
- if length(t[:state]) == length(t[:terminal]) + 1
- pop!(t[:state])
- pop!(t[:action])
- end
- push!(t[:state], x[:state])
- push!(t[:action], x[:action])
-end

src/common/CircularArraySLARTTraces.jl

@@ -1,17 +1,5 @@
 export CircularArraySLARTTraces
 
-const CircularArraySLARTTraces = Traces{
- SLART,
- <:Tuple{
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer},
- <:Trace{<:CircularArrayBuffer}
- }
-}
-
-
 function CircularArraySLARTTraces(;
  capacity::Int,
  state=Int => (),
@@ -26,37 +14,11 @@ function CircularArraySLARTTraces(;
  reward_eltype, reward_size = reward
  terminal_eltype, terminal_size = terminal
 
+ MultiplexTraces{(:state, :next_state)}(CircularArrayBuffer{state_eltype}(state_size..., capacity + 1)) +
+ MultiplexTraces{(:legal_actions_mask, :next_legal_actions_mask)}(CircularArrayBuffer{legal_actions_mask_eltype}(legal_actions_mask_size..., capacity + 1)) +
+ MultiplexTraces{(:action, :next_action)}(CircularArrayBuffer{action_eltype}(action_size..., capacity + 1)) +
  Traces(
- state=CircularArrayBuffer{state_eltype}(state_size..., capacity + 1), # !!! state is one step longer
- legal_actions_mask=CircularArrayBuffer{legal_actions_mask_eltype}(legal_actions_mask_size..., capacity + 1), # !!! legal_actions_mask is one step longer
- action=CircularArrayBuffer{action_eltype}(action_size..., capacity + 1), # !!! action is one step longer
  reward=CircularArrayBuffer{reward_eltype}(reward_size..., capacity),
  terminal=CircularArrayBuffer{terminal_eltype}(terminal_size..., capacity),
  )
-end
-
-function sample(s::BatchSampler, t::CircularArraySLARTTraces)
- inds = rand(s.rng, 1:length(t), s.batch_size)
- inds′ = inds .+ 1
- (
- state=t[:state][inds],
- legal_actions_mask=t[:legal_actions_mask][inds],
- action=t[:action][inds],
- reward=t[:reward][inds],
- terminal=t[:terminal][inds],
- next_state=t[:state][inds′],
- next_legal_actions_mask=t[:legal_actions_mask][inds′],
- next_action=t[:state][inds′]
- ) |> s.transformer
-end
-
-function Base.push!(t::CircularArraySLARTTraces, x::NamedTuple{SLA})
- if length(t[:state]) == length(t[:terminal]) + 1
- pop!(t[:state])
- pop!(t[:legal_actions_mask])
- pop!(t[:action])
- end
- push!(t[:state], x[:state])
- push!(t[:legal_actions_mask], x[:legal_actions_mask])
- push!(t[:action], x[:action])
-end
+end

src/common/common.jl

@@ -1,13 +1,5 @@
 using CircularArrayBuffers
 
-const SA = (:state, :action)
-const SLA = (:state, :legal_actions_mask, :action)
-const RT = (:reward, :terminal)
-const SART = (:state, :action, :reward, :terminal)
-const SARTSA = (:state, :action, :reward, :terminal, :next_state, :next_action)
-const SLART = (:state, :legal_actions_mask, :action, :reward, :terminal)
-const SLARTSLA = (:state, :legal_actions_mask, :action, :reward, :terminal, :next_state, :next_legal_actions_mask, :next_action)
-
 include("sum_tree.jl")
 include("CircularArraySARTTraces.jl")
 include("CircularArraySLARTTraces.jl")

src/episodes.jl

@@ -1,54 +1,45 @@
 export Episode, Episodes
 
-using MLUtils: batch
 
 """
  Episode(traces)
 
 An `Episode` is a wrapper around [`Traces`](@ref). You can use `(e::Episode)[]`
 to check/update whether the episode reaches a terminal or not.
 """
-struct Episode{T}
+struct Episode{T,E} <: AbstractVector{E}
  traces::T
- is_done::Ref{Bool}
+ is_terminated::Ref{Bool}
 end
 
-Base.getindex(e::Episode, s::Symbol) = getindex(e.traces, s)
-Base.keys(e::Episode) = keys(e.traces)
+Base.getindex(e::Episode, I) = getindex(e.traces, I)
+Base.getindex(e::Episode) = getindex(e.is_terminated)
+Base.setindex!(e::Episode, x::Bool) = setindex!(e.is_terminated, x)
 
-Base.getindex(e::Episode) = getindex(e.is_done)
-Base.setindex!(e::Episode, x::Bool) = setindex!(e.is_done, x)
+Base.size(e::Episode) = size(e.traces)
 
-Base.length(e::Episode) = length(e.traces)
+Episode(t::T) where {T<:AbstractTraces} = Episode{T,eltype(t)}(t, Ref(false))
 
-Episode(t::Traces) = Episode(t, Ref(false))
-
-function Base.push!(t::Episode, x)
- if t.is_done[]
- throw(ArgumentError("The episode is already flagged as done!"))
- else
- push!(t.traces, x)
- end
-end
-
-function Base.append!(t::Episode, x)
- if t.is_done[]
- throw(ArgumentError("The episode is already flagged as done!"))
- else
- append!(t.traces, x)
+for f in (:push!, :pushfirst!, :append!, :prepend!)
+ @eval function Base.$f(t::Episode, x)
+ if t.is_terminated[]
+ throw(ArgumentError("The episode is already flagged as done!"))
+ else
+ $f(t.traces, x)
+ end
  end
 end
 
 function Base.pop!(t::Episode)
  pop!(t.traces)
- t.is_done[] = false
+ t.is_terminated[] = false
 end
 
 Base.popfirst!(t::Episode) = popfirst!(t.traces)
 
 function Base.empty!(t::Episode)
  empty!(t.traces)
- t.is_done[] = false
+ t.is_terminated[] = false
 end
 
 #####
@@ -58,13 +49,14 @@ end
 
 A container for multiple [`Episode`](@ref)s. `init` is a parameterness function which return an [`Episode`](@ref).
 """
-struct Episodes
+struct Episodes <: AbstractVector{Episode}
  init::Any
  episodes::Vector{Episode}
  inds::Vector{Tuple{Int,Int}}
 end
 
-Base.length(e::Episodes) = length(e.inds)
+Base.size(e::Episodes) = size(e.inds)
+Base.getindex(e::Episodes, I) = getindex(e.episodes, I)
 
 function Base.push!(e::Episodes, x::Episode)
  push!(e.episodes, x)
@@ -98,10 +90,3 @@ function Base.append!(e::Episodes, x)
  push!(e.inds, (lengthe.episodes, i))
  end
 end
-
-##
-
-function sample(s::BatchSampler, e::Episodes)
- inds = rand(s.rng, 1:length(t), s.batch_size)
- batch([@view(s.episodes[e.inds[i][1]][e.inds[i][2]]) for i in inds]) |> s.transformer
-end