Fix GaussianNetwork stddev and replace SACPolicyNetwork (#172)

* Switch sigma to log_sigma

* Replace SAC network with gaussian network

* Missed a logsigma spot...

* Remove unwanted prints

* Remove na from example

* Remove StructArray

* Update src/algorithms/policy_gradient/sac.jl

Co-authored-by: Jun Tian <[email protected]>

* Add more missed logsigma spots

Co-authored-by: Jun Tian <[email protected]>

Author: albheim

src/algorithms/policy_gradient/ppo.jl

@@ -148,9 +148,8 @@ function RLBase.prob(
     if p.update_step < p.n_random_start
         @error "todo"
     else
-        p.approximator.actor(send_to_device(device(p.approximator), state)) |>
-        send_to_host |>
-        StructArray{Normal}
+        μ, logσ = p.approximator.actor(send_to_device(device(p.approximator), state)) |> send_to_host 
+        StructArray{Normal}((μ, exp.(logσ)))
     end
 end
 
@@ -265,13 +264,13 @@ function _update!(p::PPOPolicy, t::AbstractTrajectory)
             gs = gradient(ps) do
                 v′ = AC.critic(s) |> vec
                 if AC.actor isa GaussianNetwork
-                    μ, σ = AC.actor(s)
+                    μ, logσ = AC.actor(s)
                     if ndims(a) == 2
-                        log_p′ₐ = vec(sum(normlogpdf(μ, σ, a), dims = 1))
+                        log_p′ₐ = vec(sum(normlogpdf(μ, exp.(logσ), a), dims = 1))
                     else
-                        log_p′ₐ = normlogpdf(μ, σ, a)
+                        log_p′ₐ = normlogpdf(μ, exp.(logσ), a)
                     end
-                    entropy_loss = mean(size(σ, 1) * (log(2.0f0π) + 1) .+ sum(log, σ; dims = 1)) / 2
+                    entropy_loss = mean(size(logσ, 1) * (log(2.0f0π) + 1) .+ sum(logσ; dims = 1)) / 2
                 else
                     # actor is assumed to return discrete logits
                     logit′ = AC.actor(s)

src/algorithms/policy_gradient/sac.jl

@@ -1,13 +1,4 @@
-export SACPolicy, SACPolicyNetwork
-
-# 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)))
+export SACPolicy
 
 mutable struct SACPolicy{
     BA<:NeuralNetworkApproximator,
@@ -54,6 +45,10 @@ end
 - `update_every = 50`,
 - `step = 0`,
 - `rng = Random.GLOBAL_RNG`,
+
+`policy` is expected to output a tuple `(μ, logσ)` of mean and
+log standard deviations for the desired action distributions, this
+can be implemented using a `GaussianNetwork` in a `NeuralNetworkApproximator`.
 """
 function SACPolicy(;
     policy,
@@ -117,8 +112,8 @@ end
 This function is compatible with a multidimensional action space.
 """
 function evaluate(p::SACPolicy, state)
-    μ, log_σ = p.policy(state)
-    π_dist = Normal.(μ, exp.(log_σ))
+    μ, logσ = p.policy(state)
+    π_dist = Normal.(μ, exp.(logσ))
     z = rand.(p.rng, π_dist)
     logp_π = sum(logpdf.(π_dist, z), dims = 1)
     logp_π -= sum((2.0f0 .* (log(2.0f0) .- z - softplus.(-2.0f0 * z))), dims = 1)

src/algorithms/policy_gradient/vpg.jl

@@ -1,21 +1,23 @@
 export VPGPolicy, GaussianNetwork
 
 """
-    GaussianNetwork(;pre=identity, μ, σ)
+    GaussianNetwork(;pre=identity, μ, logσ)
 
-`σ` should return the log of std, `exp` will be applied to it automatically.
+Returns `μ` and `logσ` when called. 
+Create a distribution to sample from 
+using `Normal.(μ, exp.(logσ))`.
 """
 Base.@kwdef struct GaussianNetwork{P,U,S}
     pre::P = identity
     μ::U
-    σ::S
+    logσ::S
 end
 
 Flux.@functor GaussianNetwork
 
 function (m::GaussianNetwork)(S)
     x = m.pre(S)
-    m.μ(x), m.σ(x) .|> exp
+    m.μ(x), m.logσ(x) 
 end
 
 """

src/experiments/rl_envs/JuliaRL_PPO_Pendulum.jl

@@ -37,7 +37,7 @@ function RLCore.Experiment(
                         Dense(64, 64, relu; initW = glorot_uniform(rng)),
                     ),
                     μ = Chain(Dense(64, 1, tanh; initW = glorot_uniform(rng)), vec),
-                    σ = Chain(Dense(64, 1; initW = glorot_uniform(rng)), vec),
+                    logσ = Chain(Dense(64, 1; initW = glorot_uniform(rng)), vec),
                 ),
                 critic = Chain(
                     Dense(ns, 64, relu; initW = glorot_uniform(rng)),

src/experiments/rl_envs/JuliaRL_SAC_Pendulum.jl

@@ -26,12 +26,13 @@ function RLCore.Experiment(
     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 -> clamp(x, typeof(x)(-2), typeof(x)(2)), initW = init),
+        model = GaussianNetwork(
+            pre = Chain(
+                Dense(ns, 30, relu), 
+                Dense(30, 30, relu),
             ),
+            μ = Chain(Dense(30, 1, initW = init)),
+            logσ = Chain(Dense(30, 1, x -> clamp.(x, typeof(x)(-10), typeof(x)(2)), initW = init)),
         ),
         optimizer = ADAM(0.003),
     )