From 63a1506fad2f9065677ec6b7047fe09bc839fafe Mon Sep 17 00:00:00 2001
From: Mo8it <mo8it@protonmail.com>
Date: Mon, 31 Jan 2022 17:50:55 +0100
Subject: [PATCH] Added RewardsPlot

---
 src/ReCo.jl                      |  4 ++++
 src/Visualization/RewardsPlot.jl | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 src/Visualization/RewardsPlot.jl

diff --git a/src/ReCo.jl b/src/ReCo.jl
index e51d7cd..ff02a9e 100644
--- a/src/ReCo.jl
+++ b/src/ReCo.jl
@@ -5,6 +5,7 @@ export init_sim,
     run_rl,
     animate,
     plot_snapshot,
+    plot_rewards,
     LocalCOMWithAdditionalShapeRewardEnv,
     LocalCOMWithAdditionalShapeRewardEnv2,
     OriginEnv,
@@ -48,4 +49,7 @@ using .Animation
 include("Visualization/SnapshotPlot.jl")
 using .SnapshotPlot
 
+include("Visualization/RewardsPlot.jl")
+using .RewardsPlot
+
 end # module
\ No newline at end of file
diff --git a/src/Visualization/RewardsPlot.jl b/src/Visualization/RewardsPlot.jl
new file mode 100644
index 0000000..924bf35
--- /dev/null
+++ b/src/Visualization/RewardsPlot.jl
@@ -0,0 +1,41 @@
+module RewardsPlot
+
+export plot_rewards
+
+using CairoMakie
+using JLD2: JLD2
+
+using ReCo: ReCo
+
+include("common_CairoMakie.jl")
+
+function plot_rewards_from_env_helper(; env_helper::ReCo.RL.EnvHelper, rl_dir::String)
+    rewards = env_helper.shared.hook.rewards
+    n_episodes = length(rewards)
+
+    init_cairomakie!()
+
+    fig = gen_figure()
+
+    ax = Axis(
+        fig[1, 1]; xlabel="Episode", ylabel="Reward", limits=((0, n_episodes), nothing)
+    )
+
+    lines!(ax, 1:n_episodes, rewards)
+
+    set_gaps!(fig)
+
+    save_fig("rewards.pdf", fig; parent_dir=rl_dir)
+
+    return nothing
+end
+
+function plot_rewards(rl_dir::String, env_helper_file_name::String="env_helper.jld2")
+    env_helper::ReCo.RL.EnvHelper = JLD2.load_object("$rl_dir/$env_helper_file_name")
+
+    plot_rewards_from_env_helper(; env_helper, rl_dir)
+
+    return nothing
+end
+
+end # module
\ No newline at end of file