Prepare for the next release of RLCore (#129)

* sync

* fix experiments in rl_env

* fix experiments of Atari

* bugfix with ppo

* fix tests

* add compat

* create Project.toml in test

* only run tests on ubuntu

Author: findmyway

.github/workflows/ci.yml

@@ -16,11 +16,8 @@ jobs:
       matrix:
         version:
           - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
-          - 'nightly'
         os:
           - ubuntu-latest
-          - macOS-latest
-          - windows-latest
         arch:
           - x64
     steps:

Project.toml

@@ -7,6 +7,7 @@ version = "0.2.2"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
@@ -34,7 +35,7 @@ Distributions = "0.24"
 Flux = "0.11"
 MacroTools = "0.5"
 ReinforcementLearningBase = "0.8.4"
-ReinforcementLearningCore = "0.5"
+ReinforcementLearningCore = "0.6"
 Requires = "1"
 Setfield = "0.6, 0.7"
 StableRNGs = "1.0"
@@ -43,11 +44,3 @@ StructArrays = "0.4"
 TensorBoardLogger = "0.1"
 Zygote = "0.5"
 julia = "1.4"
-
-[extras]
-OpenSpiel = "ceb70bd2-fe3f-44f0-b81f-41608acaf2f2"
-ReinforcementLearningEnvironments = "25e41dd2-4622-11e9-1641-f1adca772921"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test", "ReinforcementLearningEnvironments", "OpenSpiel"]

src/ReinforcementLearningZoo.jl

@@ -7,20 +7,21 @@ using ReinforcementLearningBase
 using ReinforcementLearningCore
 using Setfield: @set
 using StableRNGs
+using Logging
+using Flux.Losses
+using Dates
 
 include("patch.jl")
 include("algorithms/algorithms.jl")
-include("utils.jl")
 
 using Requires
 
 # dynamic loading environments
 function __init__()
     @require ReinforcementLearningEnvironments = "25e41dd2-4622-11e9-1641-f1adca772921" begin
-        include("experiments/rl_envs.jl")
-        @require ArcadeLearningEnvironment = "b7f77d8d-088d-5e02-8ac0-89aab2acc977" include("experiments/atari.jl")
-        @require SnakeGames = "34dccd9f-48d6-4445-aa0f-8c2e373b5429" include("experiments/snake.jl")
-        @require OpenSpiel = "ceb70bd2-fe3f-44f0-b81f-41608acaf2f2" include("experiments/open_spiel.jl")
+        include("experiments/rl_envs/rl_envs.jl")
+        @require ArcadeLearningEnvironment = "b7f77d8d-088d-5e02-8ac0-89aab2acc977" include("experiments/atari/atari.jl")
+        # @require OpenSpiel = "ceb70bd2-fe3f-44f0-b81f-41608acaf2f2" include("experiments/open_spiel/open_spiel.jl")
     end
 end
 

src/algorithms/dqns/basic_dqn.jl

@@ -12,9 +12,11 @@ This is the very basic implementation of DQN. Compared to the traditional Q lear
 in the updating step it uses a batch of transitions sampled from an experience buffer instead of current transition.
 And the `approximator` is usually a [`NeuralNetworkApproximator`](@ref).
 You can start from this implementation to understand how everything is organized and how to write your own customized algorithm.
+
 # Keywords
+
 - `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
-- `loss_func`: the loss function to use. TODO: provide a default [`huber_loss`](@ref)?
+- `loss_func`: the loss function to use.
 - `γ::Float32=0.99f0`: discount rate.
 - `batch_size::Int=32`
 - `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
@@ -24,9 +26,10 @@ mutable struct BasicDQNLearner{Q,F,R} <: AbstractLearner
     approximator::Q
     loss_func::F
     γ::Float32
-    batch_size::Int
+    sampler::BatchSampler
     min_replay_history::Int
     rng::R
+    # for debugging
     loss::Float32
 end
 
@@ -38,14 +41,13 @@ end
 (learner::BasicDQNLearner)(env) =
     env |>
     get_state |>
-    x ->
-        send_to_device(device(learner.approximator), x) |>
-        learner.approximator |>
-        send_to_host
+    x -> send_to_device(device(learner), x) |>
+    learner.approximator |>
+    send_to_host
 
 function BasicDQNLearner(;
     approximator::Q,
-    loss_func::F,
+    loss_func::F = huber_loss,
     γ = 0.99f0,
     batch_size = 32,
     min_replay_history = 32,
@@ -55,45 +57,28 @@ function BasicDQNLearner(;
         approximator,
         loss_func,
         γ,
-        batch_size,
+        BatchSampler{SARTS}(batch_size),
         min_replay_history,
         rng,
         0.0,
     )
 end
 
-function RLBase.update!(learner::BasicDQNLearner, T::AbstractTrajectory)
-    length(T[:terminal]) < learner.min_replay_history && return
-
-    inds = rand(learner.rng, 1:length(T[:terminal]), learner.batch_size)
-
-    batch = (
-        state = consecutive_view(T[:state], inds),
-        action = consecutive_view(T[:action], inds),
-        reward = consecutive_view(T[:reward], inds),
-        terminal = consecutive_view(T[:terminal], inds),
-        next_state = consecutive_view(T[:next_state], inds),
-    )
-
-    update!(learner, batch)
+function RLBase.update!(learner::BasicDQNLearner, traj::AbstractTrajectory)
+    if length(traj) >= learner.min_replay_history
+        inds, batch = sample(learner.rng, traj, learner.sampler)
+        update!(learner, batch)
+    end
 end
 
-function RLBase.update!(learner::BasicDQNLearner, batch::NamedTuple)
+function RLBase.update!(learner::BasicDQNLearner, batch::NamedTuple{SARTS})
 
     Q = learner.approximator
-    D = device(Q)
     γ = learner.γ
     loss_func = learner.loss_func
 
-    batch_size = nframes(batch.terminal)
-
-    s = send_to_device(D, batch.state)
-    a = batch.action
-    r = send_to_device(D, batch.reward)
-    t = send_to_device(D, batch.terminal)
-    s′ = send_to_device(D, batch.next_state)
-
-    a = CartesianIndex.(a, 1:batch_size)
+    s, a, r, t, s′ = send_to_device(device(Q), batch)
+    a = CartesianIndex.(a, 1:length(a))
 
     gs = gradient(params(Q)) do
         q = Q(s)[a]

src/algorithms/dqns/common.jl

@@ -4,73 +4,7 @@
 
 const PERLearners = Union{PrioritizedDQNLearner,RainbowLearner,IQNLearner}
 
-function extract_experience(t::AbstractTrajectory, learner::PERLearners)
-    s = learner.stack_size
-    h = learner.update_horizon
-    n = learner.batch_size
-    γ = learner.γ
-
-    # 1. sample indices based on priority
-    valid_ind_range =
-        isnothing(s) ? (1:(length(t[:terminal])-h)) : (s:(length(t[:terminal])-h))
-    if haskey(t, :priority)
-        inds = Vector{Int}(undef, n)
-        priorities = Vector{Float32}(undef, n)
-        for i in 1:n
-            ind, p = sample(learner.rng, t[:priority])
-            while ind ∉ valid_ind_range
-                ind, p = sample(learner.rng, t[:priority])
-            end
-            inds[i] = ind
-            priorities[i] = p
-        end
-    else
-        inds = rand(learner.rng, valid_ind_range, n)
-        priorities = nothing
-    end
-
-    next_inds = inds .+ h
-
-    # 2. extract SARTS
-    states = consecutive_view(t[:state], inds; n_stack = s)
-    actions = consecutive_view(t[:action], inds)
-    next_states = consecutive_view(t[:state], next_inds; n_stack = s)
-
-    if haskey(t, :legal_actions_mask)
-        legal_actions_mask = consecutive_view(t[:legal_actions_mask], inds)
-        next_legal_actions_mask = consecutive_view(t[:next_legal_actions_mask], inds)
-    else
-        legal_actions_mask = nothing
-        next_legal_actions_mask = nothing
-    end
-
-    consecutive_rewards = consecutive_view(t[:reward], inds; n_horizon = h)
-    consecutive_terminals = consecutive_view(t[:terminal], inds; n_horizon = h)
-    rewards, terminals = zeros(Float32, n), fill(false, n)
-
-    rewards = discount_rewards_reduced(
-        consecutive_rewards,
-        γ;
-        terminal = consecutive_terminals,
-        dims = 1,
-    )
-    terminals = mapslices(any, consecutive_terminals; dims = 1) |> vec
-
-    inds,
-    (
-        states = states,
-        legal_actions_mask = legal_actions_mask,
-        actions = actions,
-        rewards = rewards,
-        terminals = terminals,
-        next_states = next_states,
-        next_legal_actions_mask = next_legal_actions_mask,
-        priorities = priorities,
-    )
-end
-
-function RLBase.update!(p::QBasedPolicy{<:PERLearners}, t::AbstractTrajectory)
-    learner = p.learner
+function RLBase.update!(learner::Union{DQNLearner, PERLearners}, t::AbstractTrajectory)
     length(t[:terminal]) < learner.min_replay_history && return
 
     learner.update_step += 1
@@ -81,20 +15,18 @@ function RLBase.update!(p::QBasedPolicy{<:PERLearners}, t::AbstractTrajectory)
 
     learner.update_step % learner.update_freq == 0 || return
 
-    inds, experience = extract_experience(t, p.learner)
+    inds, batch = sample(learner.rng, t, learner.sampler)
 
-    if haskey(t, :priority)
-        priorities = update!(p.learner, experience)
+    if t isa PrioritizedTrajectory
+        priorities = update!(learner, batch)
         t[:priority][inds] .= priorities
     else
-        update!(p.learner, experience)
+        update!(learner,batch)
     end
 end
 
-function (agent::Agent{<:QBasedPolicy{<:PERLearners}})(::RLCore.Training{PostActStage}, env)
-    push!(agent.trajectory; reward = get_reward(env), terminal = get_terminal(env))
-    if haskey(agent.trajectory, :priority)
-        push!(agent.trajectory; priority = agent.policy.learner.default_priority)
-    end
-    nothing
+function RLBase.update!(trajectory::PrioritizedTrajectory, p::QBasedPolicy{<:PERLearners}, env::AbstractEnv, ::PostActStage)
+    push!(trajectory[:reward], get_reward(env))
+    push!(trajectory[:terminal], get_terminal(env))
+    push!(trajectory[:priority], p.learner.default_priority)
 end