Add TotalRewardPerEpisodeLastN hook (#1053)

Author: jeremiahpslewis

src/ReinforcementLearningFarm/Project.toml

@@ -3,13 +3,15 @@ uuid = "14eff660-7080-4cec-bba2-cfb12cd77ac3"
 version = "0.0.1"
 
 [deps]
+CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ReinforcementLearningBase = "e575027e-6cd6-5018-9292-cdc6200d2b44"
 ReinforcementLearningCore = "de1b191a-4ae0-4afa-a27b-92d07f46b2d6"
 
 [compat]
+CircularArrayBuffers = "0.1.12"
 ReinforcementLearningBase = "0.12"
 ReinforcementLearningCore = "0.14"
 ReinforcementLearningEnvironments = "0.8"

src/ReinforcementLearningFarm/src/ReinforcementLearningFarm.jl

@@ -6,5 +6,6 @@ const RLFarm = ReinforcementLearningFarm
 export RLFarm
 
 include("algorithms/algorithms.jl")
+include("hooks/hooks.jl")
 
 end # module

src/ReinforcementLearningFarm/src/hooks/hooks.jl

@@ -0,0 +1 @@
+include("total_reward_per_last_n_episodes.jl")

src/ReinforcementLearningFarm/src/hooks/total_reward_per_last_n_episodes.jl

@@ -0,0 +1,46 @@
+using ReinforcementLearningCore
+using ReinforcementLearningBase
+import Base.push!
+import Base.getindex
+using CircularArrayBuffers: CircularVectorBuffer, CircularArrayBuffer
+
+"""
+TotalRewardPerLastNEpisodes{F}(; max_episodes = 100)
+
+A hook that keeps track of the total reward per episode for the last `max_episodes` episodes.
+"""
+struct TotalRewardPerLastNEpisodes{B} <: AbstractHook where {B<:CircularArrayBuffer}
+    rewards::B
+
+    function TotalRewardPerLastNEpisodes(; max_episodes = 100)
+        buffer = CircularVectorBuffer{Float64}(max_episodes)
+        new{typeof(buffer)}(buffer)
+    end
+end
+
+Base.getindex(h::TotalRewardPerLastNEpisodes{B}, inds...) where {B<:CircularArrayBuffer} =
+    getindex(h.rewards, inds...)
+
+Base.push!(
+    h::TotalRewardPerLastNEpisodes{B},
+    ::PostActStage,
+    agent::P,
+    env::E,
+    player::Symbol,
+) where {P<:AbstractPolicy,E<:AbstractEnv,B<:CircularArrayBuffer} =
+    h.rewards[end] += reward(env, player)
+
+Base.push!(
+    hook::TotalRewardPerLastNEpisodes{B},
+    ::PreEpisodeStage,
+    agent,
+    env,
+) where {B<:CircularArrayBuffer} = Base.push!(hook.rewards, 0.0)
+
+Base.push!(
+    hook::TotalRewardPerLastNEpisodes{B},
+    stage::Union{PreEpisodeStage,PostEpisodeStage,PostExperimentStage},
+    agent,
+    env,
+    player::Symbol,
+) where {B<:CircularArrayBuffer} = Base.push!(hook, stage, agent, env)

src/ReinforcementLearningFarm/test/hooks/hooks.jl

@@ -0,0 +1 @@
+include("total_reward_per_last_n_episodes.jl")

src/ReinforcementLearningFarm/test/hooks/total_reward_per_last_n_episodes.jl

@@ -0,0 +1,29 @@
+using ReinforcementLearningFarm: TotalRewardPerLastNEpisodes
+
+@testset "TotalRewardPerLastNEpisodes" begin
+    @testset "Single Agent" begin
+        hook = TotalRewardPerLastNEpisodes(max_episodes = 10)
+        env = TicTacToeEnv()
+        agent = RandomPolicy()
+
+        for i = 1:15
+            push!(hook, PreEpisodeStage(), agent, env)
+            push!(hook, PostActStage(), agent, env)
+            @test length(hook.rewards) == min(i, 10)
+            @test hook.rewards[min(i, 10)] == reward(env)
+        end
+    end
+
+    @testset "MultiAgent" begin
+        hook = TotalRewardPerLastNEpisodes(max_episodes = 10)
+        env = TicTacToeEnv()
+        agent = RandomPolicy()
+
+        for i = 1:15
+            push!(hook, PreEpisodeStage(), agent, env, :Cross)
+            push!(hook, PostActStage(), agent, env, :Cross)
+            @test length(hook.rewards) == min(i, 10)
+            @test hook.rewards[min(i, 10)] == reward(env, :Cross)
+        end
+    end
+end

src/ReinforcementLearningFarm/test/runtests.jl

@@ -19,4 +19,5 @@ using ReinforcementLearningFarm
 
 @testset "ReinforcementLearningFarm.jl" begin
     include("algorithms/algorithms.jl")
+    include("hooks/hooks.jl")
 end