Soft Actor Critic (#71)

* inital SAC implementation

* PR review fixes

Author: rbange

Project.toml

@@ -8,6 +8,7 @@ AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"

README.md

@@ -1,4 +1,4 @@
-<div align="center"> 
+<div align="center">
 <a href="https://en.wikipedia.org/wiki/Tangram"> <img src="https://upload.wikimedia.org/wikipedia/commons/7/7a/Tangram-man.svg" width="200"> </a>
 <p> <a href="https://wiki.c2.com/?MakeItWorkMakeItRightMakeItFast">"Make It Work Make It Right Make It Fast"</a></p>
 <p>― <a href="https://wiki.c2.com/?KentBeck">KentBeck</a></p>
@@ -23,6 +23,7 @@ This project aims to provide some implementations of the most typical reinforcem
 - A2C
 - PPO
 - DDPG
+- SAC
 
 If you are looking for tabular reinforcement learning algorithms, you may refer [ReinforcementLearningAnIntroduction.jl](https://github.com/JuliaReinforcementLearning/ReinforcementLearningAnIntroduction.jl).
 
@@ -41,6 +42,7 @@ Some built-in experiments are exported to help new users to easily run benchmark
 - ``E`JuliaRL_A2CGAE_CartPole` `` (Thanks to [@sriram13m](https://github.com/sriram13m))
 - ``E`JuliaRL_PPO_CartPole` ``
 - ``E`JuliaRL_DDPG_Pendulum` ``
+- ``E`JuliaRL_SAC_Pendulum` `` (Thanks to [@rbange](https://github.com/rbange))
 - ``E`JuliaRL_BasicDQN_MountainCar` `` (Thanks to [@felixchalumeau](https://github.com/felixchalumeau))
 - ``E`JuliaRL_DQN_MountainCar` `` (Thanks to [@felixchalumeau](https://github.com/felixchalumeau))
 - ``E`Dopamine_DQN_Atari(pong)` ``

src/algorithms/policy_gradient/policy_gradient.jl

@@ -2,3 +2,4 @@ include("A2C.jl")
 include("ppo.jl")
 include("A2CGAE.jl")
 include("ddpg.jl")
+include("sac.jl")

src/algorithms/policy_gradient/sac.jl

@@ -0,0 +1,186 @@
+export SACPolicy, SACPolicyNetwork
+
+using Random
+using Flux
+using Flux.Losses: mse
+using Distributions: Normal, logpdf
+
+# Define SAC Actor
+struct SACPolicyNetwork
+    pre::Chain
+    mean::Chain
+    log_std::Chain
+end
+Flux.@functor SACPolicyNetwork
+(m::SACPolicyNetwork)(state) = (x = m.pre(state); (m.mean(x), m.log_std(x)))
+
+mutable struct SACPolicy{
+    BA<:NeuralNetworkApproximator,
+    BC1<:NeuralNetworkApproximator,
+    BC2<:NeuralNetworkApproximator,
+    P,
+    R<:AbstractRNG,
+} <: AbstractPolicy
+
+    policy::BA
+    qnetwork1::BC1
+    qnetwork2::BC2
+    target_qnetwork1::BC1
+    target_qnetwork2::BC2
+    γ::Float32
+    ρ::Float32
+    α::Float32
+    batch_size::Int
+    start_steps::Int
+    start_policy::P
+    update_after::Int
+    update_every::Int
+    step::Int
+    rng::R
+end
+
+"""
+    SACPolicy(;kwargs...)
+
+# Keyword arguments
+
+- `policy`,
+- `qnetwork1`,
+- `qnetwork2`,
+- `target_qnetwork1`,
+- `target_qnetwork2`,
+- `start_policy`,
+- `γ = 0.99f0`,
+- `ρ = 0.995f0`,
+- `α = 0.2f0`,
+- `batch_size = 32`,
+- `start_steps = 10000`,
+- `update_after = 1000`,
+- `update_every = 50`,
+- `step = 0`,
+- `rng = Random.GLOBAL_RNG`,
+"""
+function SACPolicy(;
+    policy,
+    qnetwork1,
+    qnetwork2,
+    target_qnetwork1,
+    target_qnetwork2,
+    start_policy,
+    γ = 0.99f0,
+    ρ = 0.995f0,
+    α = 0.2f0,
+    batch_size = 32,
+    start_steps = 10000,
+    update_after = 1000,
+    update_every = 50,
+    step = 0,
+    rng = Random.GLOBAL_RNG,
+)
+    copyto!(qnetwork1, target_qnetwork1)  # force sync
+    copyto!(qnetwork2, target_qnetwork2)  # force sync
+    SACPolicy(
+        policy,
+        qnetwork1,
+        qnetwork2,
+        target_qnetwork1,
+        target_qnetwork2,
+        γ,
+        ρ,
+        α,
+        batch_size,
+        start_steps,
+        start_policy,
+        update_after,
+        update_every,
+        step,
+        rng,
+    )
+end
+
+# TODO: handle Training/Testing mode
+function (p::SACPolicy)(env)
+    p.step += 1
+
+    if p.step <= p.start_steps
+        p.start_policy(env)
+    else
+        D = device(p.policy)
+        s = get_state(env)
+        s = Flux.unsqueeze(s, ndims(s) + 1)
+        # trainmode:
+        action = evaluate(p, s)[1][] # returns action as scalar
+
+        # testmode:
+        # if testing dont sample an action, but act deterministically by
+        # taking the "mean" action
+        # action = p.policy(s)[1][] # returns action as scalar
+    end
+end
+
+"""
+This function is compatible with a multidimensional action space.
+"""
+function evaluate(p::SACPolicy, state)
+    μ, log_σ = p.policy(state)
+    π_dist = Normal.(μ, exp.(log_σ))
+    z = rand.(p.rng, π_dist)
+    logp_π = sum(logpdf.(π_dist, z), dims = 1)
+    logp_π -= sum((2f0 .* (log(2f0) .- z - softplus.(-2f0 * z))), dims = 1)
+    return tanh.(z), logp_π
+end
+
+function RLBase.update!(p::SACPolicy, traj::CircularCompactSARTSATrajectory)
+    length(traj[:terminal]) > p.update_after || return
+    p.step % p.update_every == 0 || return
+
+    inds = rand(p.rng, 1:(length(traj[:terminal])-1), p.batch_size)
+    s = select_last_dim(traj[:state], inds)
+    a = select_last_dim(traj[:action], inds)
+    r = select_last_dim(traj[:reward], inds)
+    t = select_last_dim(traj[:terminal], inds)
+    s′ = select_last_dim(traj[:next_state], inds)
+
+    γ, ρ, α = p.γ, p.ρ, p.α
+
+    # !!! we have several assumptions here, need revisit when we have more complex environments
+    # state is vector
+    # action is scalar
+    a′, log_π = evaluate(p, s′)
+    q′_input = vcat(s′, a′)
+    q′ = min.(p.target_qnetwork1(q′_input), p.target_qnetwork2(q′_input))
+
+    y = r .+ γ .* (1 .- t) .* vec((q′ .- α .* log_π))
+
+    # Train Q Networks
+    a = Flux.unsqueeze(a, 1)
+    q_input = vcat(s, a)
+
+    q_grad_1 = gradient(Flux.params(p.qnetwork1)) do
+        q1 = p.qnetwork1(q_input) |> vec
+        mse(q1, y)
+    end
+    update!(p.qnetwork1, q_grad_1)
+    q_grad_2 = gradient(Flux.params(p.qnetwork2)) do
+        q2 = p.qnetwork1(q_input) |> vec
+        mse(q2, y)
+    end
+    update!(p.qnetwork2, q_grad_2)
+
+    # Train Policy
+    p_grad = gradient(Flux.params(p.policy)) do
+        a, log_π = evaluate(p, s)
+        q_input = vcat(s, a)
+        q = min.(p.qnetwork1(q_input), p.qnetwork2(q_input))
+        mean(α .* log_π .- q)
+    end
+    update!(p.policy, p_grad)
+
+    # polyak averaging
+    for (dest, src) in zip(
+        Flux.params([p.target_qnetwork1, p.target_qnetwork2]),
+        Flux.params([p.qnetwork1, p.qnetwork2]),
+    )
+        dest .= ρ .* dest .+ (1 - ρ) .* src
+    end
+end

src/experiments/rl_envs.jl

@@ -954,3 +954,75 @@ function RLCore.Experiment(
 
     Experiment(agent, env, stop_condition, hook, description)
 end
+
+function RLCore.Experiment(
+    ::Val{:JuliaRL},
+    ::Val{:SAC},
+    ::Val{:Pendulum},
+    ::Nothing;
+    seed = 123,
+)
+    rng = MersenneTwister(seed)
+    inner_env = PendulumEnv(T = Float32, rng = rng)
+    action_space = get_actions(inner_env)
+    low = action_space.low
+    high = action_space.high
+    ns = length(get_state(inner_env))
+
+    env = inner_env |> ActionTransformedEnv(x -> low + (x + 1) * 0.5 * (high - low))
+    init = glorot_uniform(rng)
+
+    create_policy_net() = NeuralNetworkApproximator(
+        model = SACPolicyNetwork(
+            Chain(Dense(ns, 30, relu), Dense(30, 30, relu)),
+            Chain(Dense(30, 1, initW = init)),
+            Chain(Dense(
+                30,
+                1,
+                x -> min(max(x, typeof(x)(-20)), typeof(x)(2)),
+                initW = init,
+            )),
+        ),
+        optimizer = ADAM(0.003),
+    )
+
+    create_q_net() = NeuralNetworkApproximator(
+        model = Chain(
+            Dense(ns + 1, 30, relu; initW = init),
+            Dense(30, 30, relu; initW = init),
+            Dense(30, 1; initW = init),
+        ),
+        optimizer = ADAM(0.003),
+    )
+
+    agent = Agent(
+        policy = SACPolicy(
+            policy = create_policy_net(),
+            qnetwork1 = create_q_net(),
+            qnetwork2 = create_q_net(),
+            target_qnetwork1 = create_q_net(),
+            target_qnetwork2 = create_q_net(),
+            γ = 0.99f0,
+            ρ = 0.995f0,
+            α = 0.2f0,
+            batch_size = 64,
+            start_steps = 1000,
+            start_policy = RandomPolicy(ContinuousSpace(-1.0, 1.0); rng = rng),
+            update_after = 1000,
+            update_every = 1,
+            rng = rng,
+        ),
+        trajectory = CircularCompactSARTSATrajectory(
+            capacity = 10000,
+            state_type = Float32,
+            state_size = (ns,),
+            action_type = Float32,
+        ),
+    )
+
+    description = """
+    # Play Pendulum with SAC
+    """
+
+    Experiment(agent, env, StopAfterStep(10000), TotalRewardPerEpisode(), description)
+end

test/runtests.jl

@@ -41,8 +41,11 @@ using Random
                     mean(Iterators.flatten(res.hook.rewards))
             end
 
-            res = run(E`JuliaRL_DDPG_Pendulum`)
-            @info "stats for DDPG Pendulum" avg_reward = mean(res.hook.rewards)
+            for method in (:DDPG, :SAC)
+                res = run(Experiment(Val(:JuliaRL), Val(method), Val(:Pendulum), nothing))
+                @info "stats for $method" avg_reward =
+                    mean(Iterators.flatten(res.hook.rewards))
+            end
         end
     end