fix atari related experiments (#145)

Author: findmyway

src/algorithms/policy_gradient/ppo.jl

@@ -85,6 +85,7 @@ mutable struct PPOPolicy{A<:ActorCritic,D,R} <: AbstractPolicy
     critic_loss_weight::Float32
     entropy_loss_weight::Float32
     rng::R
+    n_random_start::Int
     update_freq::Int
     update_step::Int
     # for logging
@@ -98,6 +99,7 @@ end
 function PPOPolicy(;
     approximator,
     update_freq,
+    n_random_start = 0,
     update_step = 0,
     γ = 0.99f0,
     λ = 0.95f0,
@@ -123,6 +125,7 @@ function PPOPolicy(;
         critic_loss_weight,
         entropy_loss_weight,
         rng,
+        n_random_start,
         update_freq,
         update_step,
         zeros(Float32, n_microbatches, n_epochs),
@@ -137,17 +140,25 @@ function RLBase.prob(
     p::PPOPolicy{<:ActorCritic{<:GaussianNetwork},Normal},
     state::AbstractArray,
 )
-    p.approximator.actor(send_to_device(device(p.approximator), state)) |>
-    send_to_host |>
-    StructArray{Normal}
+    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}
+    end
 end
 
 function RLBase.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)]
+    if p.update_step < p.n_random_start
+        [Categorical(fill(1/length(x), length(x)); check_args = false) for x in eachcol(logits)]
+    else
+        [Categorical(x; check_args = false) for x in eachcol(logits)]
+    end
 end
 
 RLBase.prob(p::PPOPolicy, env::MultiThreadEnv) = prob(p, state(env))
@@ -161,29 +172,14 @@ end
 (p::PPOPolicy)(env::MultiThreadEnv) = rand.(p.rng, prob(p, env))
 (p::PPOPolicy)(env::AbstractEnv) = rand(p.rng, prob(p, env))
 
-function (agent::Agent{<:PPOPolicy})(env::AbstractEnv)
-    dist = prob(agent.policy, env)
-    a = rand(agent.policy.rng, dist)
-    EnrichedAction(a; action_log_prob=logpdf(dist, a))
-end
-
 function (agent::Agent{<:PPOPolicy})(env::MultiThreadEnv)
     dist = prob(agent.policy, env)
     action = rand.(agent.policy.rng, dist)
     EnrichedAction(action; action_log_prob=logpdf.(dist, action))
 end
 
-function (agent::Agent{<:RandomStartPolicy{<:PPOPolicy}})(env::AbstractEnv)
-    a = agent.policy(env)
-    if a isa EnrichedAction
-        a
-    else
-        EnrichedAction(a; action_log_prob=logpdf(prob(agent.policy, a)))
-    end
-end
-
 function RLBase.update!(p::PPOPolicy, t::Union{PPOTrajectory, MaskedPPOTrajectory}, ::AbstractEnv, ::PreActStage)
-    length(t) == 0 && return  # in the first update, only state & action is inserted into trajectory
+    length(t) == 0 && return  # in the first update, only state & action are inserted into trajectory
     p.update_step += 1
     if p.update_step % p.update_freq == 0
         _update!(p, t)
@@ -289,7 +285,7 @@ end
 
 function RLBase.update!(
     trajectory::Union{PPOTrajectory,MaskedPPOTrajectory},
-    policy::Union{PPOPolicy,RandomStartPolicy{<:PPOPolicy}},
+    ::PPOPolicy,
     env::MultiThreadEnv,
     ::PreActStage,
     action::EnrichedAction

src/algorithms/policy_gradient/run.jl

@@ -13,7 +13,7 @@ function RLCore._run(
 )
 
     while true
-        reset!(env)  # this is a soft reset!, only environments reached the end will get reset.
+        reset!(env)  # this is a soft reset!, only environments reached the end will be reset.
         action = policy(env)
         policy(PRE_ACT_STAGE, env, action)
         hook(PRE_ACT_STAGE, policy, env, action)
@@ -23,9 +23,11 @@ function RLCore._run(
         hook(POST_ACT_STAGE, policy, env)
 
         if stop_condition(policy, env)
-            policy(PRE_ACT_STAGE, env)  # let the policy see the last observation
             break
         end
     end
-    hook
+    action = policy(env)
+    policy(PRE_ACT_STAGE, env, action)  # let the policy see the last observation
+    hook(PRE_ACT_STAGE, policy, env, action)
+    nothing
 end

src/experiments/atari/Dopamine_DQN_Atari.jl

@@ -71,7 +71,7 @@ function RLCore.Experiment(
     MAX_EPISODE_STEPS_EVAL = 27_000
     N_CHECKPOINTS = 3
 
-    total_reward_per_episode = TotalRewardPerEpisode()
+    total_reward_per_episode = TotalOriginalRewardPerEpisode()
     time_per_step = TimePerStep()
     steps_per_episode = StepsPerEpisode()
     hook = ComposedHook(
@@ -93,7 +93,7 @@ function RLCore.Experiment(
             @info "evaluating agent at $t step..."
             p = agent.policy
             p = @set p.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
-            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
+            h = ComposedHook(TotalOriginalRewardPerEpisode(), StepsPerEpisode())
             s = @elapsed run(
                 p,
                 atari_env_factory(

src/experiments/atari/Dopamine_IQN_Atari.jl

@@ -86,7 +86,7 @@ function RLCore.Experiment(
     MAX_EPISODE_STEPS_EVAL = 27_000
     N_CHECKPOINTS = 3
 
-    total_reward_per_episode = TotalRewardPerEpisode()
+    total_reward_per_episode = TotalOriginalRewardPerEpisode()
     time_per_step = TimePerStep()
     steps_per_episode = StepsPerEpisode()
     hook = ComposedHook(
@@ -108,7 +108,7 @@ function RLCore.Experiment(
             @info "evaluating agent at $t step..."
             p = agent.policy
             p = @set p.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
-            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
+            h = ComposedHook(TotalOriginalRewardPerEpisode(), StepsPerEpisode())
             s = @elapsed run(
                 p,
                 atari_env_factory(

src/experiments/atari/Dopamine_Rainbow_Atari.jl

@@ -74,7 +74,7 @@ function RLCore.Experiment(
     MAX_EPISODE_STEPS_EVAL = 27_000
     N_CHECKPOINTS = 3
 
-    total_reward_per_episode = TotalRewardPerEpisode()
+    total_reward_per_episode = TotalOriginalRewardPerEpisode()
     time_per_step = TimePerStep()
     steps_per_episode = StepsPerEpisode()
     hook = ComposedHook(
@@ -96,7 +96,7 @@ function RLCore.Experiment(
             @info "evaluating agent at $t step..."
             p = agent.policy
             p = @set p.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
-            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
+            h = ComposedHook(TotalOriginalRewardPerEpisode(), StepsPerEpisode())
             s = @elapsed run(
                 p,
                 atari_env_factory(

src/experiments/atari/atari.jl

@@ -1,5 +1,7 @@
 using .ArcadeLearningEnvironment
 using .ReinforcementLearningEnvironments
+using BSON
+using Flux:Chain
 
 function atari_env_factory(
     name,
@@ -8,8 +10,9 @@ function atari_env_factory(
     max_episode_steps = 100_000;
     seed = nothing,
     repeat_action_probability = 0.25,
+    n_replica = 1
 )
-    AtariEnv(;
+    init(seed) = AtariEnv(;
         name = string(name),
         grayscale_obs = true,
         noop_max = 30,
@@ -21,14 +24,62 @@ function atari_env_factory(
         full_action_space = false,
         seed = seed,
     ) |>
-    StateOverriddenEnv(
-        ResizeImage(state_size...),  # this implementation is different from cv2.resize https://github.com/google/dopamine/blob/e7d780d7c80954b7c396d984325002d60557f7d1/dopamine/discrete_domains/atari_lib.py#L629
-        StackFrames(state_size..., n_frames),
+    env -> StateOverriddenEnv(
+        env,
+        Chain(ResizeImage(state_size...), StackFrames(state_size..., n_frames))
     ) |>
     StateCachedEnv |>
-    RewardOverriddenEnv(r -> clamp(r, -1, 1))
+    env -> RewardOverriddenEnv(env, r -> clamp(r, -1, 1))
+
+    if n_replica == 1
+        init(seed)
+    else
+        envs = [init(hash(seed+i)) for i in 1:n_replica]
+        states = Flux.batch(state.(envs))
+        rewards = reward.(envs)
+        terminals = is_terminated.(envs)
+        A = Space([action_space(x) for x in envs])
+        S = Space(fill(0..255, size(states)))
+        MultiThreadEnv(envs, states, rewards, terminals, A, S, nothing)
+    end
+end
+
+"Total reward per episode before reward reshaping"
+Base.@kwdef mutable struct TotalOriginalRewardPerEpisode <: AbstractHook
+    rewards::Vector{Float64} = Float64[]
+    reward::Float64 = 0.0
+end
+
+function (hook::TotalOriginalRewardPerEpisode)(::PostActStage, agent, env::RewardOverriddenEnv)
+    hook.reward += reward(env.env)
+end
+
+function (hook::TotalOriginalRewardPerEpisode)(::PostEpisodeStage, agent, env)
+    push!(hook.rewards, hook.reward)
+    hook.reward = 0
 end
 
+"Total reward of each inner env per episode before reward reshaping"
+struct TotalBatchOriginalRewardPerEpisode <: AbstractHook
+    rewards::Vector{Vector{Float64}}
+    reward::Vector{Float64}
+end
+
+function TotalBatchOriginalRewardPerEpisode(batch_size::Int)
+    TotalBatchOriginalRewardPerEpisode([Float64[] for _ in 1:batch_size], zeros(batch_size))
+end
+
+function (hook::TotalBatchOriginalRewardPerEpisode)(::PostActStage, agent, env::MultiThreadEnv{<:RewardOverriddenEnv})
+    for (i, e) in enumerate(env.envs)
+        hook.reward[i] += reward(e.env)
+        if is_terminated(e)
+            push!(hook.rewards[i], hook.reward[i])
+            hook.reward[i] = 0.0
+        end
+    end
+end
+
+
 for f in readdir(@__DIR__)
     if f != splitdir(@__FILE__)[2]
         include(f)

src/experiments/atari/rlpyt_A2C_Atari.jl

@@ -19,15 +19,14 @@ function RLCore.Experiment(
     UPDATE_FREQ = 5
     N_FRAMES = 4
     STATE_SIZE = (80, 104)
-    env = MultiThreadEnv([
-        atari_env_factory(
-            name,
-            STATE_SIZE,
-            N_FRAMES;
-            repeat_action_probability = 0,
-            seed = hash(seed + i),
-        ) for i in 1:N_ENV
-    ])
+    env = atari_env_factory(
+        name,
+        STATE_SIZE,
+        N_FRAMES;
+        repeat_action_probability = 0,
+        seed = seed,
+        n_replica = N_ENV
+    ) 
     N_ACTIONS = length(action_space(env[1]))
 
     init = orthogonal(rng)
@@ -77,7 +76,7 @@ function RLCore.Experiment(
     N_CHECKPOINTS = 3
     stop_condition = StopAfterStep(N_TRAINING_STEPS)
 
-    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
+    total_batch_reward_per_episode = TotalBatchOriginalRewardPerEpisode(N_ENV)
     batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
     evaluation_result = []
 
@@ -112,19 +111,18 @@ function RLCore.Experiment(
         end,
         DoEveryNStep(EVALUATION_FREQ) do t, agent, env
             @info "evaluating agent at $t step..."
-            h = TotalBatchRewardPerEpisode(N_ENV)
+            h = TotalBatchOriginalRewardPerEpisode(N_ENV)
             s = @elapsed run(
                 agent.policy,
-                MultiThreadEnv([
-                    atari_env_factory(
-                        name,
-                        STATE_SIZE,
-                        N_FRAMES,
-                        MAX_EPISODE_STEPS_EVAL;
-                        repeat_action_probability = 0,
-                        seed = hash(seed + t + i),
-                    ) for i in 1:N_ENV
-                ]),
+                atari_env_factory(
+                    name,
+                    STATE_SIZE,
+                    N_FRAMES,
+                    MAX_EPISODE_STEPS_EVAL;
+                    repeat_action_probability = 0,
+                    seed = seed + t,
+                    n_replica = 4
+                ),
                 StopAfterStep(27_000; is_show_progress = false),
                 h,
             )