Add dueling network (#171)

* Add dueling network

* Add docs

* Some adjustment

Author: pilgrimygy

src/algorithms/dqns/common.jl

@@ -35,3 +35,22 @@ function RLBase.update!(
     push!(trajectory[:terminal], is_terminated(env))
     push!(trajectory[:priority], p.learner.default_priority)
 end
+
+"""
+    DuelingNetwork(;base, val, adv)
+    
+Dueling network automatically produces separate estimates of the state value function network and advantage function network. The expected output size of val is 1, and adv is the size of the action space.
+"""
+struct DuelingNetwork{B,V,A}
+    base::B
+    val::V
+    adv::A
+end
+
+Flux.@functor DuelingNetwork
+
+function (m::DuelingNetwork)(state)
+    x = m.base(state)
+    val = m.val(x)
+    return val .+ m.adv(x) .- mean(m.adv(x), dims=1)
+end

src/experiments/rl_envs/JuliaRL_DQN_CartPole.jl

@@ -16,24 +16,21 @@ function RLCore.Experiment(
 
     env = CartPoleEnv(; T = Float32, rng = rng)
     ns, na = length(state(env)), length(action_space(env))
+    base_model = Chain(
+        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
+        Dense(128, 128, relu; initW = glorot_uniform(rng)),
+        Dense(128, na; initW = glorot_uniform(rng))
+        )
 
     agent = Agent(
         policy = QBasedPolicy(
             learner = DQNLearner(
                 approximator = NeuralNetworkApproximator(
-                    model = Chain(
-                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
-                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
-                        Dense(128, na; initW = glorot_uniform(rng)),
-                    ) |> cpu,
+                    model = build_dueling_network(base_model) |> cpu,
                     optimizer = ADAM(),
                 ),
                 target_approximator = NeuralNetworkApproximator(
-                    model = Chain(
-                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
-                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
-                        Dense(128, na; initW = glorot_uniform(rng)),
-                    ) |> cpu,
+                    model = build_dueling_network(base_model) |> cpu,
                 ),
                 loss_func = huber_loss,
                 stack_size = nothing,

src/experiments/rl_envs/JuliaRL_REMDQN_CartPole.jl

@@ -16,7 +16,7 @@ function RLCore.Experiment(
 
     env = CartPoleEnv(; T = Float32, rng = rng)
     ns, na = length(state(env)), length(action_space(env))
-    ensemble_num = 6
+    ensemble_num = 16
 
     agent = Agent(
         policy = QBasedPolicy(

src/experiments/rl_envs/rl_envs.jl

@@ -5,3 +5,16 @@ for f in readdir(@__DIR__)
         include(f)
     end
 end
+
+# Build Dueling Network
+function build_dueling_network(network::Chain)
+    lm = length(network)
+    if !(network[lm] isa Dense) || !(network[lm-1] isa Dense) 
+        error("The Qnetwork provided is incompatible with dueling.")
+    end
+    base = Chain([deepcopy(network[i]) for i=1:lm-2]...)
+    last_layer_dims = size(network[lm].W, 2)
+    val = Chain(deepcopy(network[lm-1]), Dense(last_layer_dims, 1))
+    adv = Chain([deepcopy(network[i]) for i=lm-1:lm]...)
+    return DuelingNetwork(base, val, adv)
+end