From fd83c13333964d4ea7d2321b877eec8a1ec650ff Mon Sep 17 00:00:00 2001
From: Mo8it <mo8it@protonmail.com>
Date: Sun, 30 Jan 2022 04:38:57 +0100
Subject: [PATCH] Added reward shaping plot

---
 analysis/mean_squared_displacement.jl         |  2 +-
 .../radial_distribution_function.jl           |  2 +-
 analysis/reward_discount_analysis.jl          |  2 +-
 graphics/pontential.jl                        |  4 ++-
 graphics/reward_shaping.jl                    | 29 +++++++++++++++++++
 5 files changed, 35 insertions(+), 4 deletions(-)
 create mode 100644 graphics/reward_shaping.jl

diff --git a/analysis/mean_squared_displacement.jl b/analysis/mean_squared_displacement.jl
index 78db318..9166c74 100644
--- a/analysis/mean_squared_displacement.jl
+++ b/analysis/mean_squared_displacement.jl
@@ -133,7 +133,7 @@ function plot_mean_sq_displacement_with_expectation(
 
     fig = gen_figure()
 
-    ax = Axis(fig[1, 1]; xlabel=L"t", ylabel=L"\mathbf{MSD}", xscale=log10, yscale=log10)
+    ax = Axis(fig[1, 1]; xlabel=L"t", ylabel=L"\mathbf{MSD}(t)", xscale=log10, yscale=log10)
 
     t_linrange = LinRange(ts[1], ts[end], 1000)
 
diff --git a/analysis/radial_distribution_function/radial_distribution_function.jl b/analysis/radial_distribution_function/radial_distribution_function.jl
index 954420d..d146b3b 100644
--- a/analysis/radial_distribution_function/radial_distribution_function.jl
+++ b/analysis/radial_distribution_function/radial_distribution_function.jl
@@ -164,7 +164,7 @@ function plot_radial_distributions(;
         xticks=0:(2 * particle_radius):floor(Int64, max_lower_radius),
         yticks=0:ceil(Int64, max_g),
         xlabel=L"r / d",
-        ylabel=L"g",
+        ylabel=L"g(r)",
         limits=(min_lower_radius - 0.03, max_lower_radius + 0.03, min_g, max_g * 1.03),
     )
 
diff --git a/analysis/reward_discount_analysis.jl b/analysis/reward_discount_analysis.jl
index cd8ed07..937ed06 100644
--- a/analysis/reward_discount_analysis.jl
+++ b/analysis/reward_discount_analysis.jl
@@ -14,7 +14,7 @@ function run_reward_discount_analysis()
         γ = γs[γ_ind]
         env_helper = ReCo.run_rl(;
             EnvType=ReCo.OriginEnv,
-            n_episodes=3,
+            n_episodes=400,
             episode_duration=15.0,
             n_particles=150,
             update_actions_at=0.08,
diff --git a/graphics/pontential.jl b/graphics/pontential.jl
index 687ac53..bc79863 100644
--- a/graphics/pontential.jl
+++ b/graphics/pontential.jl
@@ -28,7 +28,9 @@ function plot_potentials()
     max_y = 1.05
     min_y = -max_y
 
-    ax = Axis(fig[1, 1]; xlabel=L"r/σ", ylabel=L"U/ϵ", limits=(0.88, max_x, min_y, max_y))
+    ax = Axis(
+        fig[1, 1]; xlabel=L"r / σ", ylabel=L"U(r) / ϵ", limits=(0.88, max_x, min_y, max_y)
+    )
 
     r_σ_ratio = LinRange(0.95, max_x, 1000)
 
diff --git a/graphics/reward_shaping.jl b/graphics/reward_shaping.jl
new file mode 100644
index 0000000..e185cb5
--- /dev/null
+++ b/graphics/reward_shaping.jl
@@ -0,0 +1,29 @@
+using CairoMakie
+using LaTeXStrings: @L_str
+
+using ReCo: ReCo
+
+includet("../src/Visualization/common_CairoMakie.jl")
+
+function plot_reward_function()
+    init_cairomakie!()
+
+    fig = gen_figure()
+
+    min_x = 0.0
+    max_x = 1.15
+
+    ax = Axis(
+        fig[1, 1]; xlabel=L"x / x_{\max}", ylabel=L"R(x)", limits=(min_x, max_x, 0.0, 1.05)
+    )
+
+    x = LinRange(min_x, max_x, 1000)
+
+    lines!(ax, x, ReCo.RL.minimizing_reward.(x, 1.0))
+
+    set_gaps!(fig)
+
+    save_fig("reward_shaping.pdf", fig)
+
+    return nothing
+end
\ No newline at end of file