Rename PPOLearner to PPOPolicy and make it support continuous space (#93)

findmyway · web-flow · commit 2674ba187797 · 2020-09-29T10:02:00.000+08:00
* initial changes

* added experiment of PPO with Pendulum

* fix test errors

* update README

* bump version and update dependency

* fix conflict

* avoid calculating distribution twice!

* update dependency
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ReinforcementLearningZoo"
 uuid = "d607f57d-ee1e-4ba7-bcf2-7734c1e31854"
 authors = ["Jun Tian <tianjun.cpp@gmail.com>"]
-version = "0.1.7"
+version = "0.2.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -21,6 +21,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -31,8 +32,8 @@ CUDA = "1"
 Distributions = "0.23"
 Flux = "0.11"
 MacroTools = "0.5"
-ReinforcementLearningBase = "0.8"
-ReinforcementLearningCore = "0.4.2"
+ReinforcementLearningBase = "0.8.4"
+ReinforcementLearningCore = "0.4.5"
 Requires = "1"
 Setfield = "0.6, 0.7"
 StatsBase = "0.32, 0.33"
diff --git a/README.md b/README.md
@@ -50,6 +50,7 @@ Some built-in experiments are exported to help new users to easily run benchmark
 - ``E`JuliaRL_DDPG_Pendulum` ``
 - ``E`JuliaRL_TD3_Pendulum` `` (Thanks to [@rbange](https://github.com/rbange))
 - ``E`JuliaRL_SAC_Pendulum` `` (Thanks to [@rbange](https://github.com/rbange))
+- ``E`JuliaRL_PPO_Pendulum` ``
 - ``E`JuliaRL_BasicDQN_MountainCar` `` (Thanks to [@felixchalumeau](https://github.com/felixchalumeau))
 - ``E`JuliaRL_DQN_MountainCar` `` (Thanks to [@felixchalumeau](https://github.com/felixchalumeau))
 - ``E`JuliaRL_Minimax_OpenSpiel(tic_tac_toe)` ``
diff --git a/src/algorithms/policy_gradient/policy_gradient.jl b/src/algorithms/policy_gradient/policy_gradient.jl
@@ -1,7 +1,7 @@
+include("vpg.jl") 
 include("A2C.jl")
 include("ppo.jl")
 include("A2CGAE.jl")
 include("ddpg.jl")
 include("td3.jl")
 include("sac.jl")
-include("vpg.jl") 
diff --git a/src/algorithms/policy_gradient/ppo.jl b/src/algorithms/policy_gradient/ppo.jl
@@ -1,11 +1,13 @@
 include("ppo_trajectory.jl")
 
 using Random
+using Distributions: Categorical, Normal, logpdf
+using StructArrays
 
-export PPOLearner
+export PPOPolicy
 
 """
-    PPOLearner(;kwargs)
+    PPOPolicy(;kwargs)
 
 # Keyword arguments
 
@@ -19,9 +21,13 @@ export PPOLearner
 - `actor_loss_weight = 1.0f0`,
 - `critic_loss_weight = 0.5f0`,
 - `entropy_loss_weight = 0.01f0`,
+- `dist = Categorical`,
 - `rng = Random.GLOBAL_RNG`,
+
+By default, `dist` is set to `Categorical`, which means it will only works
+on environments of discrete actions. To work with environments of
 """
-mutable struct PPOLearner{A<:ActorCritic,R} <: AbstractLearner
+mutable struct PPOPolicy{A<:ActorCritic,D,R} <: AbstractPolicy
     approximator::A
     γ::Float32
     λ::Float32
@@ -41,7 +47,7 @@ mutable struct PPOLearner{A<:ActorCritic,R} <: AbstractLearner
     loss::Matrix{Float32}
 end
 
-function PPOLearner(;
+function PPOPolicy(;
     approximator,
     γ = 0.99f0,
     λ = 0.95f0,
@@ -52,9 +58,10 @@ function PPOLearner(;
     actor_loss_weight = 1.0f0,
     critic_loss_weight = 0.5f0,
     entropy_loss_weight = 0.01f0,
+    dist = Categorical,
     rng = Random.GLOBAL_RNG,
 )
-    PPOLearner(
+    PPOPolicy{typeof(approximator),dist,typeof(rng)}(
         approximator,
         γ,
         λ,
@@ -74,21 +81,33 @@ function PPOLearner(;
     )
 end
 
-function (learner::PPOLearner)(env::MultiThreadEnv)
-    learner.approximator.actor(send_to_device(
-        device(learner.approximator),
-        get_state(env),
-    )) |> send_to_host
+function RLBase.get_prob(p::PPOPolicy{<:ActorCritic{<:NeuralNetworkApproximator{<:GaussianNetwork}}, Normal}, state::AbstractArray)
+    p.approximator.actor(send_to_device(
+        device(p.approximator),
+        state,
+    )) |> send_to_host |> StructArray{Normal}
+end
+
+function RLBase.get_prob(p::PPOPolicy{<:ActorCritic, Categorical}, state::AbstractArray)
+    logits = p.approximator.actor(send_to_device(
+        device(p.approximator),
+        state,
+    )) |> softmax |> send_to_host
+    [Categorical(x;check_args=false) for x in eachcol(logits)]
 end
 
-function (learner::PPOLearner)(env)
+RLBase.get_prob(p::PPOPolicy, env::MultiThreadEnv) = get_prob(p, get_state(env))
+
+function RLBase.get_prob(p::PPOPolicy, env::AbstractEnv)
     s = get_state(env)
     s = Flux.unsqueeze(s, ndims(s) + 1)
-    s = send_to_device(device(learner.approximator), s)
-    learner.approximator.actor(s) |> vec |> send_to_host
+    get_prob(p, s)[1]
 end
 
-function RLBase.update!(learner::PPOLearner, t::PPOTrajectory)
+(p::PPOPolicy)(env::MultiThreadEnv) = rand.(p.rng, get_prob(p, env))
+(p::PPOPolicy)(env::AbstractEnv) = rand(p.rng, get_prob(p, env))
+
+function RLBase.update!(p::PPOPolicy, t::PPOTrajectory)
     isfull(t) || return
 
     states = t[:state]
@@ -98,16 +117,16 @@ function RLBase.update!(learner::PPOLearner, t::PPOTrajectory)
     terminals = t[:terminal]
     states_plus = t[:full_state]
 
-    rng = learner.rng
-    AC = learner.approximator
-    γ = learner.γ
-    λ = learner.λ
-    n_epochs = learner.n_epochs
-    n_microbatches = learner.n_microbatches
-    clip_range = learner.clip_range
-    w₁ = learner.actor_loss_weight
-    w₂ = learner.critic_loss_weight
-    w₃ = learner.entropy_loss_weight
+    rng = p.rng
+    AC = p.approximator
+    γ = p.γ
+    λ = p.λ
+    n_epochs = p.n_epochs
+    n_microbatches = p.n_microbatches
+    clip_range = p.clip_range
+    w₁ = p.actor_loss_weight
+    w₂ = p.critic_loss_weight
+    w₃ = p.entropy_loss_weight
     D = device(AC)
 
     n_envs, n_rollout = size(terminals)
@@ -142,60 +161,63 @@ function RLBase.update!(learner::PPOLearner, t::PPOTrajectory)
             ps = Flux.params(AC)
             gs = gradient(ps) do
                 v′ = AC.critic(s) |> vec
-                logit′ = AC.actor(s)
-                p′ = softmax(logit′)
-                log_p′ = logsoftmax(logit′)
-                log_p′ₐ = log_p′[CartesianIndex.(a, 1:length(a))]
+                if AC.actor isa NeuralNetworkApproximator{<:GaussianNetwork}
+                    μ, σ = AC.actor(s)
+                    log_p′ₐ = normlogpdf(μ, σ, a)
+                    entropy_loss = mean((log(2.0f0π)+1)/2 .+ log.(σ))
+                else
+                    # actor is assumed to return discrete logits
+                    logit′ = AC.actor(s)
+                    p′ = softmax(logit′)
+                    log_p′ = logsoftmax(logit′)
+                    log_p′ₐ = log_p′[CartesianIndex.(a, 1:length(a))]
+                    entropy_loss = -sum(p′ .* log_p′) * 1//size(p′, 2)
+                end
 
                 ratio = exp.(log_p′ₐ .- log_p)
                 surr1 = ratio .* adv
                 surr2 = clamp.(ratio, 1.0f0 - clip_range, 1.0f0 + clip_range) .* adv
 
                 actor_loss = -mean(min.(surr1, surr2))
                 critic_loss = mean((r .- v′) .^ 2)
-                entropy_loss = -sum(p′ .* log_p′) * 1//size(p′, 2)
                 loss = w₁ * actor_loss + w₂ * critic_loss - w₃ * entropy_loss
 
                 ignore() do
-                    learner.actor_loss[i, epoch] = actor_loss
-                    learner.critic_loss[i, epoch] = critic_loss
-                    learner.entropy_loss[i, epoch] = entropy_loss
-                    learner.loss[i, epoch] = loss
+                    p.actor_loss[i, epoch] = actor_loss
+                    p.critic_loss[i, epoch] = critic_loss
+                    p.entropy_loss[i, epoch] = entropy_loss
+                    p.loss[i, epoch] = loss
                 end
 
                 loss
             end
 
-            learner.norm[i, epoch] = clip_by_global_norm!(gs, ps, learner.max_grad_norm)
+            p.norm[i, epoch] = clip_by_global_norm!(gs, ps, p.max_grad_norm)
             update!(AC, gs)
         end
     end
 end
 
-function (π::QBasedPolicy{<:PPOLearner})(env::MultiThreadEnv)
-    action_values = π.learner(env)
-    logits = logsoftmax(action_values)
-    actions = π.explorer(action_values)
-    actions_log_prob = logits[CartesianIndex.(actions, 1:size(action_values, 2))]
-    actions, actions_log_prob
-end
+function (agent::Agent{<:Union{PPOPolicy, RandomStartPolicy{<:PPOPolicy}}})(::Training{PreActStage}, env::MultiThreadEnv)
+    state = get_state(env)
+    dist = get_prob(agent.policy, env)
 
-(π::QBasedPolicy{<:PPOLearner})(env) = env |> π.learner |> π.explorer
+    # currently RandomPolicy returns a Matrix instead of a (vector of) distribution.
+    if dist isa Matrix{<:Number}
+        dist = [Categorical(x;check_args=false) for x in eachcol(dist)]
+    elseif dist isa Vector{<:Vector{<:Number}}
+        dist = [Categorical(x;check_args=false) for x in dist]
+    end
 
-function (p::RandomStartPolicy{<:QBasedPolicy{<:PPOLearner}})(env::MultiThreadEnv)
-    p.num_rand_start -= 1
-    if p.num_rand_start < 0
-        p.policy(env)
-    else
-        a = p.random_policy(env)
-        log_p = log.(get_prob(p.random_policy, env, a))
-        a, log_p
+    # !!! a little ugly
+    rng = if agent.policy isa PPOPolicy
+        agent.policy.rng
+    elseif agent.policy isa RandomStartPolicy
+        agent.policy.policy.rng
     end
-end
 
-function (agent::Agent{<:AbstractPolicy,<:PPOTrajectory})(::Training{PreActStage}, env)
-    action, action_log_prob = agent.policy(env)
-    state = get_state(env)
+    action = [rand(rng, d) for d in dist]
+    action_log_prob = [logpdf(d, a) for (d, a) in zip(dist, action)]
     push!(
         agent.trajectory;
         state = state,
@@ -217,12 +239,3 @@ function (agent::Agent{<:AbstractPolicy,<:PPOTrajectory})(::Training{PreActStage
 
     action
 end
-
-function (agent::Agent{<:AbstractPolicy,<:PPOTrajectory})(::Training{PostActStage}, env)
-    push!(agent.trajectory; reward = get_reward(env), terminal = get_terminal(env))
-    nothing
-end
-
-function (agent::Agent{<:AbstractPolicy,<:PPOTrajectory})(::Testing{PreActStage}, env)
-    agent.policy(env)[1]  # ignore the log_prob of action
-end
diff --git a/src/algorithms/policy_gradient/vpg.jl b/src/algorithms/policy_gradient/vpg.jl
@@ -6,12 +6,19 @@ using ReinforcementLearningCore
 
 export VPGPolicy, GaussianNetwork
 
-struct GaussianNetwork
-    pre::Chain
-    μ::Chain
-    σ::Chain
+"""
+    GaussianNetwork(;pre=identity, μ, σ)
+
+`σ` should return the log of std, `exp` will be applied to it automatically.
+"""
+Base.@kwdef struct GaussianNetwork{P,U,S}
+    pre::P = identity
+    μ::U
+    σ::S
 end
+
 Flux.@functor GaussianNetwork
+
 function (m::GaussianNetwork)(S)
     x = m.pre(S)
     m.μ(x), m.σ(x) .|> exp
diff --git a/src/experiments/atari.jl b/src/experiments/atari.jl
@@ -753,25 +753,22 @@ function RLCore.Experiment(
         policy = RandomStartPolicy(
             num_rand_start = 1000,
             random_policy = RandomPolicy(get_actions(env); rng = rng),
-            policy = QBasedPolicy(
-                learner = PPOLearner(
-                    approximator = ActorCritic(
-                        actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
-                        critic = Chain(model, Dense(512, 1; initW = init)),
-                        optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
-                    ) |> gpu,
-                    γ = 0.99f0,
-                    λ = 0.98f0,
-                    clip_range = INIT_CLIP_RANGE,  # decrease with a hook
-                    max_grad_norm = 1.0f0,
-                    n_microbatches = 4,
-                    n_epochs = 4,
-                    actor_loss_weight = 1.0f0,
-                    critic_loss_weight = 0.5f0,
-                    entropy_loss_weight = 0.01f0,
-                    rng = rng,
-                ),
-                explorer = BatchExplorer(GumbelSoftmaxExplorer(; rng = rng)),
+            policy = PPOPolicy(
+                approximator = ActorCritic(
+                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
+                    critic = Chain(model, Dense(512, 1; initW = init)),
+                    optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
+                ) |> gpu,
+                γ = 0.99f0,
+                λ = 0.98f0,
+                clip_range = INIT_CLIP_RANGE,  # decrease with a hook
+                max_grad_norm = 1.0f0,
+                n_microbatches = 4,
+                n_epochs = 4,
+                actor_loss_weight = 1.0f0,
+                critic_loss_weight = 0.5f0,
+                entropy_loss_weight = 0.01f0,
+                rng = rng,
             ),
         ),
         trajectory = PPOTrajectory(;
@@ -803,19 +800,19 @@ function RLCore.Experiment(
         total_batch_reward_per_episode,
         batch_steps_per_episode,
         DoEveryNStep(UPDATE_FREQ) do t, agent, env
-            learner = agent.policy.policy.learner
+            p = agent.policy.policy
             with_logger(lg) do
-                @info "training" loss = mean(learner.loss) actor_loss =
-                    mean(learner.actor_loss) critic_loss = mean(learner.critic_loss) entropy_loss =
-                    mean(learner.entropy_loss) norm = mean(learner.norm) log_step_increment =
+                @info "training" loss = mean(p.loss) actor_loss =
+                    mean(p.actor_loss) critic_loss = mean(p.critic_loss) entropy_loss =
+                    mean(p.entropy_loss) norm = mean(p.norm) log_step_increment =
                     UPDATE_FREQ
             end
         end,
         DoEveryNStep(UPDATE_FREQ) do t, agent, env
             decay = (N_TRAINING_STEPS - t) / N_TRAINING_STEPS
-            agent.policy.policy.learner.approximator.optimizer.eta =
+            agent.policy.policy.approximator.optimizer.eta =
                 INIT_LEARNING_RATE * decay
-            agent.policy.policy.learner.clip_range = INIT_CLIP_RANGE * Float32(decay)
+            agent.policy.policy.clip_range = INIT_CLIP_RANGE * Float32(decay)
         end,
         DoEveryNStep() do t, agent, env
             with_logger(lg) do
diff --git a/src/experiments/rl_envs.jl b/src/experiments/rl_envs.jl
