1
0
Fork 0
mirror of https://gitlab.rlp.net/mobitar/ReCo.jl.git synced 2024-09-17 18:51:17 +00:00

Added RewardsPlot

This commit is contained in:
Mo8it 2022-01-31 17:50:55 +01:00
parent d6900d9a92
commit 63a1506fad
2 changed files with 45 additions and 0 deletions

View file

@ -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

View file

@ -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