Chapter12 Random Walk

begin

    using ReinforcementLearning

    using Flux

    using Statistics

    using Plots

end

N = 21

true_values = -1:0.1:1

Base.@kwdef struct RecordRMS <: AbstractHook

    rms::Vector{Float64}=[]

end

(h::RecordRMS)(::PostEpisodeStage, agent, env) = push!(h.rms, sqrt(mean((agent.policy.learner.approximator.table[2:end-1] - true_values[2:end-1]).^2)))

function create_agent_env(α, λ)

    env = RandomWalk1D(N=21)

    ns, na =  length(state_space(env)),  length(action_space(env))

    agent = Agent(

        policy=VBasedPolicy(

            learner=TDλReturnLearner(

                approximator=TabularVApproximator(;n_state=ns, opt=Descent(α)),

                γ=1.0,

                λ=λ

            ),

            mapping = (env, V) -> rand(1:na)

        ),

        trajectory=VectorSARTTrajectory()

    )

    agent, env

end

function records(α, λ, nruns=10)

    rms = []

    for _ in 1:nruns

        hook = RecordRMS()

        run(create_agent_env(α, λ)..., StopAfterEpisode(10, is_show_progress=false),hook)

        push!(rms, mean(hook.rms))

    end

    mean(rms)

end

begin

    As = [0:0.1:1, 0:0.1:1, 0:0.1:1, 0:0.1:1, 0:0.1:1, 0:0.05:0.5, 0:0.02:0.2, 0:0.01:0.1]

    Λ = [0., 0.4, .8, 0.9, 0.95, 0.975, 0.99, 1.]

    p = plot(legend=:topright)

    for (A, λ) in zip(As, Λ)

        plot!(p, A, [records(α, λ) for α in A], label="lambda = $λ")

    end

    ylims!(p, (0.25, 0.55))

    p

end