Fix #251, ppo multidim action eval (#177)

JuliaReinforcementLearning · Apr 27, 2021 · d78f327 · d78f327
1 parent 52a9c85
commit d78f327
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/algorithms/policy_gradient/ppo.jl b/src/algorithms/policy_gradient/ppo.jl
@@ -173,11 +173,11 @@ RLBase.prob(p::PPOPolicy, env::MultiThreadEnv) = prob(p, state(env))
 function RLBase.prob(p::PPOPolicy, env::AbstractEnv)
     s = state(env)
     s = Flux.unsqueeze(s, ndims(s) + 1)
-    prob(p, s)[1]
+    prob(p, s)
 end
 
 (p::PPOPolicy)(env::MultiThreadEnv) = rand.(p.rng, prob(p, env))
-(p::PPOPolicy)(env::AbstractEnv) = rand(p.rng, prob(p, env))
+(p::PPOPolicy)(env::AbstractEnv) = rand.(p.rng, prob(p, env))
 
 function (agent::Agent{<:PPOPolicy})(env::MultiThreadEnv)
     dist = prob(agent.policy, env)