add REMDQN (#708)

Author: findmyway

src/ReinforcementLearningCore/src/policies/agent.jl

@@ -2,7 +2,7 @@ export Agent
 
 using Base.Threads: @spawn
 
-import Functors
+using Functors: @functor
 
 """
     Agent(;policy, trajectory)
@@ -20,7 +20,7 @@ mutable struct Agent{P,T} <: AbstractPolicy
     trajectory::T
     cache::NamedTuple # trajectory do not support partial inserting
 
-    function Agent(policy::P, trajectory::T, cache = NamedTuple()) where {P,T}
+    function Agent(policy::P, trajectory::T, cache=NamedTuple()) where {P,T}
         agent = new{P,T}(policy, trajectory, cache)
         if TrajectoryStyle(trajectory) === AsyncTrajectoryStyle()
             bind(trajectory, @spawn(optimise!(p, t)))
@@ -29,7 +29,7 @@ mutable struct Agent{P,T} <: AbstractPolicy
     end
 end
 
-Agent(; policy, trajectory, cache = NamedTuple()) = Agent(policy, trajectory, cache)
+Agent(; policy, trajectory, cache=NamedTuple()) = Agent(policy, trajectory, cache)
 
 RLBase.optimise!(agent::Agent) = optimise!(TrajectoryStyle(agent.trajectory), agent)
 RLBase.optimise!(::SyncTrajectoryStyle, agent::Agent) =
@@ -44,21 +44,20 @@ function RLBase.optimise!(policy::AbstractPolicy, trajectory::Trajectory)
     end
 end
 
-Functors.functor(x::Agent) =
-    (policy = x.policy,), y -> Agent(y.policy, x.trajectory, x.cache)
+@functor Agent (policy,)
 
 # !!! TODO: In async scenarios, parameters of the policy may still be updating
 # (partially), which will result to incorrect action. This should be addressed
 # in Oolong.jl with a wrapper
 function (agent::Agent)(env::AbstractEnv)
     action = agent.policy(env)
-    push!(agent.trajectory, (agent.cache..., action = action))
+    push!(agent.trajectory, (agent.cache..., action=action))
     agent.cache = (;)
     action
 end
 
 (agent::Agent)(::PreActStage, env::AbstractEnv) =
-    agent.cache = (agent.cache..., state = state(env))
+    agent.cache = (agent.cache..., state=state(env))
 
 (agent::Agent)(::PostActStage, env::AbstractEnv) =
-    agent.cache = (agent.cache..., reward = reward(env), terminal = is_terminated(env))
+    agent.cache = (agent.cache..., reward=reward(env), terminal=is_terminated(env))

src/ReinforcementLearningCore/src/policies/learners.jl

@@ -1,7 +1,7 @@
 export AbstractLearner, Approximator
 
 import Flux
-import Functors
+using Functors: @functor
 
 abstract type AbstractLearner end
 
@@ -12,7 +12,7 @@ Base.@kwdef mutable struct Approximator{M,O}
     optimiser::O
 end
 
-Functors.functor(x::Approximator) = (model=x.model,), y -> Approximator(y.model, x.state)
+@functor Approximator (model,)
 
 (A::Approximator)(x) = A.model(x)
 

src/ReinforcementLearningCore/src/policies/q_based_policy.jl

@@ -3,15 +3,14 @@ export QBasedPolicy
 include("learners.jl")
 include("explorers/explorers.jl")
 
-import Functors
+using Functors: @functor
 
 Base.@kwdef mutable struct QBasedPolicy{L,E} <: AbstractPolicy
     learner::L
     explorer::E
 end
 
-Functors.functor(x::QBasedPolicy) =
-    (learner = x.learner,), y -> QBasedPolicy(y.learner, x.explorer)
+@functor QBasedPolicy (learner,)
 
 (p::QBasedPolicy)(env) = p.explorer(p.learner(env), legal_action_space_mask(env))
 

src/ReinforcementLearningCore/src/utils/networks.jl

@@ -1,8 +1,6 @@
-import Functors
+using Functors: @functor
 import Flux
 
-using Setfield: @set
-
 #####
 # ActorCritic
 #####
@@ -18,7 +16,7 @@ Base.@kwdef struct ActorCritic{A,C,O}
     critic::C
 end
 
-Functors.@functor ActorCritic
+@functor ActorCritic
 
 #####
 # GaussianNetwork
@@ -44,7 +42,7 @@ end
 
 GaussianNetwork(pre, μ, logσ, normalizer=tanh) = GaussianNetwork(pre, μ, logσ, 0.0f0, Inf32, normalizer)
 
-Functors.@functor GaussianNetwork
+@functor GaussianNetwork
 
 """
 This function is compatible with a multidimensional action space. When outputting an action, it uses the `normalizer` function to normalize it elementwise.
@@ -138,7 +136,7 @@ end
 
 CovGaussianNetwork(pre, m, s) = CovGaussianNetwork(pre, m, s, tanh)
 
-Functors.@functor CovGaussianNetwork
+@functor CovGaussianNetwork
 
 """
     (model::CovGaussianNetwork)(rng::AbstractRNG, state; is_sampling::Bool=false, is_return_log_prob::Bool=false)
@@ -397,7 +395,7 @@ end
 
 TwinNetwork(x; kw...) = TwinNetwork(; source=x, target=deepcopy(x), kw...)
 
-Functors.functor(x::TwinNetwork) = (; source=x.source), y -> @set x.source = y.source
+@functor TwinNetwork (source,)
 
 (model::TwinNetwork)(x) = model.source(x)
 

src/ReinforcementLearningExperiments/deps/experiments/experiments/DQN/JuliaRL_REMDQN_CartPole.jl

@@ -2,7 +2,7 @@
 # title: JuliaRL\_REMDQN\_CartPole
 # cover: assets/JuliaRL_REMDQN_CartPole.png
 # description: REMDQN applied to CartPole
-# date: 2021-05-22
+# date: 2021-06-25
 # author: "[Jun Tian](https://github.com/findmyway)"
 # ---
 
@@ -16,61 +16,70 @@ function RL.Experiment(
     ::Val{:JuliaRL},
     ::Val{:REMDQN},
     ::Val{:CartPole},
-    ::Nothing;
-    seed = 123,
+    ; seed=123,
+    ensemble_num=16
 )
     rng = StableRNG(seed)
 
-    env = CartPoleEnv(; T = Float32, rng = rng)
+    env = CartPoleEnv(; T=Float32, rng=rng)
     ns, na = length(state(env)), length(action_space(env))
-    ensemble_num = 16
+
+    n = 1
+    γ = 0.99f0
 
     agent = Agent(
-        policy = QBasedPolicy(
-            learner = REMDQNLearner(
-                approximator = NeuralNetworkApproximator(
-                    model = Chain(
-                        ## Multi-head method, please refer to "https://github.com/google-research/batch_rl/tree/b55ba35ebd2381199125dd77bfac9e9c59a64d74/batch_rl/multi_head".
-                        Dense(ns, 128, relu; init = glorot_uniform(rng)),
-                        Dense(128, 128, relu; init = glorot_uniform(rng)),
-                        Dense(128, na * ensemble_num; init = glorot_uniform(rng)),
-                    ) |> gpu,
-                    optimizer = ADAM(),
-                ),
-                target_approximator = NeuralNetworkApproximator(
-                    model = Chain(
-                        Dense(ns, 128, relu; init = glorot_uniform(rng)),
-                        Dense(128, 128, relu; init = glorot_uniform(rng)),
-                        Dense(128, na * ensemble_num; init = glorot_uniform(rng)),
-                    ) |> gpu,
+        policy=QBasedPolicy(
+            learner=REMDQNLearner(
+                approximator=Approximator(
+                    model=TwinNetwork(
+                        Chain(
+                            ## Multi-head method, please refer to "https://github.com/google-research/batch_rl/tree/b55ba35ebd2381199125dd77bfac9e9c59a64d74/batch_rl/multi_head".
+                            Dense(ns, 128, relu; init=glorot_uniform(rng)),
+                            Dense(128, 128, relu; init=glorot_uniform(rng)),
+                            Dense(128, na * ensemble_num; init=glorot_uniform(rng)),
+                        ),
+                        sync_freq=100
+                    ),
+                    optimiser=ADAM(),
                 ),
-                loss_func = huber_loss,
-                stack_size = nothing,
-                batch_size = 32,
-                update_horizon = 1,
-                min_replay_history = 100,
-                update_freq = 1,
-                target_update_freq = 100,
-                ensemble_num = ensemble_num,
-                ensemble_method = :rand, 
-                rng = rng,
+                n=n,
+                γ=γ,
+                loss_func=huber_loss,
+                ensemble_num=ensemble_num,
+                ensemble_method=:rand,
+                rng=rng,
             ),
-            explorer = EpsilonGreedyExplorer(
-                kind = :exp,
-                ϵ_stable = 0.01,
-                decay_steps = 500,
-                rng = rng,
+            explorer=EpsilonGreedyExplorer(
+                kind=:exp,
+                ϵ_stable=0.01,
+                decay_steps=500,
+                rng=rng,
             ),
         ),
-        trajectory = CircularArraySARTTrajectory(
-            capacity = 1000,
-            state = Vector{Float32} => (ns,),
-        ),
+        trajectory=Trajectory(
+            container=CircularArraySARTTraces(
+                capacity=1000,
+                state=Float32 => (ns,),
+            ),
+            sampler=NStepBatchSampler{SS′ART}(
+                n=n,
+                γ=γ,
+                batch_size=32,
+                rng=rng
+            ),
+            controller=InsertSampleRatioController(
+                threshold=100,
+                n_inserted=-1
+            )
+        )
     )
 
     stop_condition = StopAfterStep(10_000, is_show_progress=!haskey(ENV, "CI"))
     hook = TotalRewardPerEpisode()
-    Experiment(agent, env, stop_condition, hook, "")
+
+    ## !!! note that REMDQN is used in offline RL
+    ## TODO: use DQN to collect experiences and then optimise the REMDQN
+    Experiment(agent, env, stop_condition, hook)
 end
 
 #+ tangle=false

src/ReinforcementLearningExperiments/src/ReinforcementLearningExperiments.jl

@@ -12,6 +12,7 @@ include(joinpath(EXPERIMENTS_DIR, "JuliaRL_BasicDQN_CartPole.jl"))
 include(joinpath(EXPERIMENTS_DIR, "JuliaRL_DQN_CartPole.jl"))
 include(joinpath(EXPERIMENTS_DIR, "JuliaRL_PrioritizedDQN_CartPole.jl"))
 include(joinpath(EXPERIMENTS_DIR, "JuliaRL_QRDQN_CartPole.jl"))
+include(joinpath(EXPERIMENTS_DIR, "JuliaRL_REMDQN_CartPole.jl"))
 
 # dynamic loading environments
 function __init__() end

src/ReinforcementLearningExperiments/test/runtests.jl

@@ -7,9 +7,9 @@ run(E`JuliaRL_BasicDQN_CartPole`)
 run(E`JuliaRL_DQN_CartPole`)
 run(E`JuliaRL_PrioritizedDQN_CartPole`)
 run(E`JuliaRL_QRDQN_CartPole`)
+run(E`JuliaRL_REMDQN_CartPole`)
 # run(E`JuliaRL_BC_CartPole`)
 # run(E`JuliaRL_Rainbow_CartPole`)
-# run(E`JuliaRL_REMDQN_CartPole`)
 # run(E`JuliaRL_IQN_CartPole`)
 # run(E`JuliaRL_VMPO_CartPole`)
 # run(E`JuliaRL_VPG_CartPole`)

src/ReinforcementLearningZoo/src/algorithms/dqns/basic_dqn.jl

@@ -3,8 +3,7 @@ export BasicDQNLearner
 using Flux: gradient, params
 using Zygote: ignore
 using Setfield: @set
-
-import Functors
+using Functors: @functor
 
 """
     BasicDQNLearner(;kwargs...)
@@ -32,7 +31,7 @@ Base.@kwdef mutable struct BasicDQNLearner{Q} <: AbstractLearner
     loss::Float32 = 0.0f0
 end
 
-Functors.functor(x::BasicDQNLearner) = (Q=x.approximator,), y -> @set x.approximator = y.Q
+@functor BasicDQNLearner (approximator,)
 
 (L::BasicDQNLearner)(s::AbstractArray) = L.approximator(s)
 

src/ReinforcementLearningZoo/src/algorithms/dqns/dqn.jl

@@ -1,8 +1,7 @@
 export DQNLearner
 
-using Setfield: @set
 using Random: AbstractRNG, GLOBAL_RNG
-import Functors
+using Functors: @functor
 
 Base.@kwdef mutable struct DQNLearner{A<:Approximator{<:TwinNetwork}} <: AbstractLearner
     approximator::A
@@ -17,7 +16,7 @@ end
 
 (L::DQNLearner)(s::AbstractArray) = L.approximator(s)
 
-Functors.functor(x::DQNLearner) = (; approximator=x.approximator), y -> @set x.approximator = y.approximator
+@functor DQNLearner (approximator,)
 
 function RLBase.optimise!(learner::DQNLearner, batch::Union{NamedTuple{SS′ART},NamedTuple{SS′L′ART}})
     A = learner.approximator

src/ReinforcementLearningZoo/src/algorithms/dqns/dqns.jl

@@ -2,7 +2,7 @@ include("basic_dqn.jl")
 include("dqn.jl")
 include("prioritized_dqn.jl")
 include("qr_dqn.jl")
-# include("rem_dqn.jl")
+include("rem_dqn.jl")
 # include("rainbow.jl")
 # include("iqn.jl")
 # include("common.jl")