Merge pull request #33 from JuliaPOMDP/abstractarrays

Support AbstractArray in DQExperience

Author: MaximeBouton

Project.toml

@@ -1,10 +1,11 @@
 name = "DeepQLearning"
 uuid = "de0a67f4-c691-11e8-0034-5fc6e16e22d3"
 repo = "https://github.com/JuliaPOMDP/DeepQLearning.jl"
-version = "0.4.5"
+version = "0.4.6"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 POMDPModelTools = "08074719-1b2a-587c-a292-00f91cc44415"
@@ -21,13 +22,14 @@ TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
 Flux = "0.10"
 POMDPPolicies = "0.2.1"
 POMDPs = "0.7.3, 0.8"
-RLInterface = "0.3.2"
+RLInterface = "0.3.6"
 julia = "1"
 
 [extras]
 POMDPModels = "355abbd5-f08e-5560-ac9e-8b5f2592a0ca"
 POMDPSimulators = "e0d0a172-29c6-5d4e-96d0-f262df5d01fd"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["POMDPModels", "POMDPSimulators", "Test"]
+test = ["POMDPModels", "POMDPSimulators", "StaticArrays", "Test"]

src/DeepQLearning.jl

@@ -12,6 +12,7 @@ using POMDPPolicies
 using RLInterface
 using LinearAlgebra
 using TensorBoardLogger: TBLogger, log_value
+using EllipsisNotation
 
 export DeepQLearningSolver,
        AbstractNNPolicy,

src/episode_replay.jl

@@ -1,6 +1,6 @@
 # Replay buffer that store full episodes
 
-mutable struct EpisodeReplayBuffer{N<:Integer, T<:AbstractFloat, CI, Q}
+mutable struct EpisodeReplayBuffer{N<:Integer, T<:Real, CI, Q<:AbstractArray{T},A<:AbstractArray{T}}
     max_size::Int64
     batch_size::Int64
     trace_length::Int64
@@ -9,43 +9,43 @@ mutable struct EpisodeReplayBuffer{N<:Integer, T<:AbstractFloat, CI, Q}
     _idx::Int64
     _experience::Vector{Vector{DQExperience{N,T,Q}}}
 
-    _s_batch::Vector{Array{T}}
+    _s_batch::Vector{A}
     _a_batch::Vector{Vector{CI}}
-    _r_batch::Vector{Array{T}}
-    _sp_batch::Vector{Array{T}}
-    _done_batch::Vector{Array{T}}
-    _trace_mask::Vector{Array{N}}
+    _r_batch::Vector{Vector{T}}
+    _sp_batch::Vector{A}
+    _done_batch::Vector{Vector{T}}
+    _trace_mask::Vector{Vector{N}}
     _episode::Vector{DQExperience{N,T,Q}}
 end
 
-function EpisodeReplayBuffer(env::AbstractEnvironment,
+function EpisodeReplayBuffer(env::AbstractEnvironment{OV},
                         max_size::Int64,
                         batch_size::Int64,
                         trace_length::Int64,
-                        rng::AbstractRNG = MersenneTwister(0))
+                        rng::AbstractRNG = MersenneTwister(0)) where {OV}
     s_dim = obs_dimensions(env)
     Q = length(s_dim)
-    experience = Vector{Vector{DQExperience{Int32, Float32, Q}}}(undef, max_size)
+    experience = Vector{Vector{DQExperience{Int32, Float32, OV}}}(undef, max_size)
     _s_batch = [zeros(Float32, s_dim..., batch_size) for i=1:trace_length]
     _a_batch = [[CartesianIndex(1,1) for i=1:batch_size] for i=1:trace_length]
     _r_batch = [zeros(Float32, batch_size) for i=1:trace_length]
     _sp_batch = [zeros(Float32, s_dim..., batch_size) for i=1:trace_length]
     _done_batch = [zeros(Float32, batch_size) for i=1:trace_length]
     _trace_mask = [zeros(Int32, batch_size) for i=1:trace_length]
-    _episode = Vector{DQExperience{Int32, Float32, Q}}()
-    return EpisodeReplayBuffer{Int32, Float32, CartesianIndex{2}, Q}(max_size, batch_size, trace_length, rng, 0, 1, experience,
+    _episode = Vector{DQExperience{Int32, Float32, OV}}()
+    return EpisodeReplayBuffer(max_size, batch_size, trace_length, rng, 0, 1, experience,
                 _s_batch, _a_batch, _r_batch, _sp_batch, _done_batch, _trace_mask, _episode)
 end
 
 is_full(r::EpisodeReplayBuffer) = r._curr_size == r.max_size
 
 max_size(r::EpisodeReplayBuffer) = r.max_size
 
-function add_exp!(r::EpisodeReplayBuffer{N, T}, exp::DQExperience) where {N, T}
+function add_exp!(r::EpisodeReplayBuffer{N,T,CI,Q}, exp::DQExperience) where {N,T,CI,Q}
     push!(r._episode, exp)
     if exp.done
         add_episode!(r, r._episode)
-        r._episode = Vector{DQExperience{N, T}}()
+        r._episode = Vector{DQExperience{N,T,Q}}()
     end
 end
 
@@ -73,30 +73,22 @@ function StatsBase.sample(r::EpisodeReplayBuffer)
     sample_indices = sample(r.rng, 1:r._curr_size, r.batch_size, replace=false)
     @assert length(sample_indices) == size(r._s_batch[1])[end]
     s_batch_size = size(first(r._s_batch))
-    for t=1:r.trace_length
-        r._s_batch[t] = reshape(r._s_batch[t], (:, r.batch_size))
-        r._sp_batch[t] = reshape(r._sp_batch[t], (:, r.batch_size))
-    end
     for (i, idx) in enumerate(sample_indices)
         ep = r._experience[idx]
         # randomized start TODO add as an option of the buffer
         ep_start = rand(r.rng, 1:length(ep))
         t = 1
         for j=ep_start:min(length(ep), r.trace_length)
             expe = ep[t]
-            r._s_batch[t][:, i] = vec(expe.s)
+            r._s_batch[t][.., i] = vec(expe.s)
             r._a_batch[t][i] = CartesianIndex(expe.a, i)
             r._r_batch[t][i] = expe.r
-            r._sp_batch[t][:, i] = vec(expe.sp)
+            r._sp_batch[t][.., i] = vec(expe.sp)
             r._done_batch[t][i] = expe.done
             r._trace_mask[t][i] = 1
             t += 1
         end
     end
-    for t=1:r.trace_length
-        r._s_batch[t] = reshape(r._s_batch[t], s_batch_size)
-        r._sp_batch[t] = reshape(r._sp_batch[t], s_batch_size)
-    end
     return r._s_batch, r._a_batch, r._r_batch, r._sp_batch, r._done_batch, r._trace_mask
 end
 
@@ -112,9 +104,9 @@ function populate_replay_buffer!(r::EpisodeReplayBuffer,
     @assert r._curr_size >= r.batch_size
 end
 
-function generate_episode(env::AbstractEnvironment, action_indices; max_steps::Int64 = 100)
+function generate_episode(env::AbstractEnvironment{OV}, action_indices; max_steps::Int64 = 100) where OV
     s_dim = obs_dimensions(env)
-    episode = DQExperience{Int32, Float32, length(s_dim)}[]
+    episode = DQExperience{Int32, Float32, OV}[]
     sizehint!(episode, max_steps)
     # start simulation
     o = reset!(env)

src/prioritized_experience_replay.jl

@@ -1,22 +1,22 @@
 # Naive implementation
 
-struct DQExperience{N <: Real,T <: Real, Q}
-    s::Array{T, Q}
+struct DQExperience{N <: Real,T <: Real, A<:AbstractArray{T}}
+    s::A
     a::N
     r::T
-    sp::Array{T, Q}
+    sp::A
     done::Bool
 end
 
 function Base.convert(::Type{DQExperience{Int32, Float32, C}}, x::DQExperience{A, B, C}) where {A, B, C}
-    return DQExperience{Int32, Float32, C}(convert(Array{Float32, C}, x.s),
-                 convert(Int32, x.a),
-                 convert(Float32, x.r),
-                 convert(Array{Float32, C}, x.sp),
-                 x.done)
+    return DQExperience{Int32, Float32, C}(convert(C, x.s),
+                                            convert(Int32, x.a),
+                                            convert(Float32, x.r),
+                                            convert(C, x.sp),
+                                            x.done)
 end
 
-mutable struct PrioritizedReplayBuffer{N<:Integer, T<:AbstractFloat,CI, Q}
+mutable struct PrioritizedReplayBuffer{N<:Integer, T<:AbstractFloat,CI,Q,A<:AbstractArray{T}}
     max_size::Int64
     batch_size::Int64
     rng::AbstractRNG
@@ -28,23 +28,23 @@ mutable struct PrioritizedReplayBuffer{N<:Integer, T<:AbstractFloat,CI, Q}
     _priorities::Vector{T}
     _experience::Vector{DQExperience{N,T,Q}}
 
-    _s_batch::Array{T}
+    _s_batch::A
     _a_batch::Vector{CI}
     _r_batch::Vector{T}
-    _sp_batch::Array{T}
+    _sp_batch::A
     _done_batch::Vector{T}
     _weights_batch::Vector{T}
 end
 
-function PrioritizedReplayBuffer(env::AbstractEnvironment,
+function PrioritizedReplayBuffer(env::AbstractEnvironment{OV},
                                 max_size::Int64,
                                 batch_size::Int64;
                                 rng::AbstractRNG = MersenneTwister(0),
                                 α::Float32 = 6f-1,
                                 β::Float32 = 4f-1,
-                                ϵ::Float32 = 1f-3)
+                                ϵ::Float32 = 1f-3) where {OV}
     s_dim = obs_dimensions(env)
-    experience = Vector{DQExperience{Int32, Float32, length(s_dim)}}(undef, max_size)
+    experience = Vector{DQExperience{Int32, Float32, OV}}(undef, max_size)
     priorities = Vector{Float32}(undef, max_size)
     _s_batch = zeros(Float32, s_dim..., batch_size)
     _a_batch = [CartesianIndex(0,0) for i=1:batch_size]
@@ -87,21 +87,16 @@ end
 
 function get_batch(r::PrioritizedReplayBuffer, sample_indices::Vector{Int64})
     @assert length(sample_indices) == size(r._s_batch)[end]
-    s_batch_size = size(r._s_batch)
-    r._s_batch = reshape(r._s_batch, (:, r.batch_size))
-    r._sp_batch = reshape(r._sp_batch, (:, r.batch_size))
     for (i, idx) in enumerate(sample_indices)
         @inbounds begin
-            r._s_batch[:, i] = vec(r._experience[idx].s)
+            r._s_batch[.., i] = vec(r._experience[idx].s)
             r._a_batch[i] = CartesianIndex(r._experience[idx].a, i)
             r._r_batch[i] = r._experience[idx].r
-            r._sp_batch[:, i] = vec(r._experience[idx].sp)
+            r._sp_batch[.., i] = vec(r._experience[idx].sp)
             r._done_batch[i] = r._experience[idx].done
             r._weights_batch[i] = r._priorities[idx]
         end
     end
-    r._s_batch = reshape(r._s_batch, s_batch_size)
-    r._sp_batch = reshape(r._sp_batch, s_batch_size)
     p = r._weights_batch ./ sum(r._priorities[1:r._curr_size])
     weights = (r._curr_size * p).^(-r.β)
     return r._s_batch, r._a_batch, r._r_batch, r._sp_batch, r._done_batch, sample_indices, weights

test/prototype.jl

@@ -8,7 +8,7 @@ using Flux
 using DeepQLearning
 include("test/test_env.jl")
 
-# mdp = TestMDP((5,5), 4, 6)
+mdp = TestMDP((5,5), 4, 6)
 # mdp = SimpleGridWorld()
 rng = MersenneTwister(1)
 mdp = TestMDP((5,5), 1, 6)
@@ -20,8 +20,12 @@ solver = DeepQLearningSolver(batch_size = 128, eval_freq = 10_000, save_freq=10_
 
 @btime policy = solve($solver, $mdp)
 
+
 policy = solve(solver, mdp)
 
+env = MDPEnvironment(mdp)
+o = reset!(env)
+
 using RLInterface
 using LinearAlgebra
 

test/runtests.jl

@@ -5,6 +5,7 @@ using POMDPPolicies
 using Flux
 using Random
 using RLInterface
+using StaticArrays
 using Test
 Random.seed!(7)
 GLOBAL_RNG = MersenneTwister(1) # for test consistency
@@ -110,3 +111,30 @@ end
     policy = solve(solver, pomdp)
     @test size(actionvalues(policy, true)) == (length(actions(pomdp)),)
 end
+
+mutable struct StaticArrayMDP <: MDP{typeof(SVector(1)), Int64}
+    state::typeof(SVector(1))
+end
+POMDPs.discount(::StaticArrayMDP) = 0.95f0
+POMDPs.initialstate(m::StaticArrayMDP, rng::AbstractRNG) = m.state 
+
+function POMDPs.gen(m::StaticArrayMDP, s, a, rng::AbstractRNG)
+    return (sp=s + SVector(a), r=m.state[1]^2)
+end
+
+POMDPs.isterminal(::StaticArrayMDP, s) = s[1] >= 3
+POMDPs.actions(::StaticArrayMDP) = [0,1]
+
+
+@testset "Static Array Env" begin
+    mdp = StaticArrayMDP(SVector(1))
+
+    model = Chain(Dense(1, 32), Dense(32, length(actions(mdp))))
+
+    solver = DeepQLearningSolver(qnetwork = model, max_steps=10, 
+                                learning_rate=0.005,log_freq=500,
+                                recurrence=false,double_q=true, dueling=true, prioritized_replay=true)
+    policy = solve(solver, mdp)
+
+    @test evaluate(mdp, policy, GLOBAL_RNG) > 1.0
+end