Chapter 13 Short Corridor

begin

    using ReinforcementLearning

    using Flux

    using Statistics

    using Plots

    using LinearAlgebra:dot

end

begin

    Base.@kwdef mutable struct ShortCorridorEnv <: AbstractEnv

        position::Int = 1

    end

​

    RLBase.state_space(env::ShortCorridorEnv) = Base.OneTo(4)

    RLBase.action_space(env::ShortCorridorEnv) = Base.OneTo(2)

​

    function (env::ShortCorridorEnv)(a)

        if env.position == 1 && a == 2

            env.position += 1

        elseif env.position == 2

            env.position += a == 1 ? 1 : -1

        elseif env.position == 3

            env.position += a == 1 ? -1 : 1

        end

        nothing

    end

​

    function RLBase.reset!(env::ShortCorridorEnv)

        env.position = 1

        nothing

    end

​

    RLBase.state(env::ShortCorridorEnv) = env.position

    RLBase.is_terminated(env::ShortCorridorEnv) = env.position == 4

    RLBase.reward(env::ShortCorridorEnv) =  env.position == 4 ? 0.0 : -1.0

end

world = ShortCorridorEnv()

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

function run_once(A)

    avg_rewards = []

    for ϵ in A

        p = TabularRandomPolicy(;table=Dict(s => [1-ϵ, ϵ] for s in 1:ns))

        env = ShortCorridorEnv()

        hook=TotalRewardPerEpisode()

        run(p, env, StopAfterEpisode(1000),hook)

        push!(avg_rewards, mean(hook.rewards[end-100:end]))

    end

    avg_rewards

end

X = 0.05:0.05:0.95

plot(X, mean([run_once(X) for _ in 1:10]), legend=nothing)

begin

    Base.@kwdef struct LinearPreferenceApproximator{F,O} <: AbstractApproximator

        weight::Vector{Float64}

        feature_func::F

        actions::Int

        opt::O

    end

​

    function (A::LinearPreferenceApproximator)(s)

        h = [dot(A.weight, A.feature_func(s, a)) for a in 1:A.actions]

        softmax(h)

    end

​

    function RLBase.update!(A::LinearPreferenceApproximator, correction::Pair)

        (s, a), Δ = correction

        w, x = A.weight, A.feature_func

        w̄ = -Δ .* (x(s,a) .- sum(A(s) .* [x(s, b) for b in 1:A.actions]))

        Flux.Optimise.update!(A.opt, w, w̄)

    end

end

begin

    Base.@kwdef struct ReinforcePolicy{A<:AbstractApproximator} <: AbstractPolicy

        approximator::A

        γ::Float64

    end

​

    (p::ReinforcePolicy)(env::AbstractEnv) = prob(p, state(env)) |> WeightedExplorer(;is_normalized=true)

​

    RLBase.prob(p::ReinforcePolicy, s) = p.approximator(s)

​

    function RLBase.update!(

        p::ReinforcePolicy,

        t::AbstractTrajectory,

        ::AbstractEnv,

        ::PostEpisodeStage

    )

        S, A, R = t[:state], t[:action], t[:reward]

        Q, γ = p.approximator, p.γ

        G = 0.

​

        for i in 1:length(R)

            s,a,r = S[end-i], A[end-i], R[end-i+1]

            G = r + γ*G

​

            update!(Q, (s, a) => G)

        end

    end

​

    function RLBase.update!(

        t::AbstractTrajectory,

        ::ReinforcePolicy,

        ::AbstractEnv,

        ::PreEpisodeStage

    )

        empty!(t)

    end

end

function run_once_RL(α)

    agent = Agent(

        policy=ReinforcePolicy(

            approximator=LinearPreferenceApproximator(

                weight=[-1.47, 1.47],  # init_weight

                feature_func=(s,a) -> a == 1 ? [0, 1] : [1, 0],

                actions=na,

                opt=Descent(α)

            ),

            γ=1.0

        ),

        trajectory=VectorSARTTrajectory()

    )

    env = ShortCorridorEnv()

    hook = TotalRewardPerEpisode()

    run(agent,env,StopAfterEpisode(1000;is_show_progress=false),hook)

    hook.rewards

end

begin

    fig_13_1 = plot(legend=:bottomright)

    for x in [-13, -14]  # for -12, it seems not that easy to converge in short time

        plot!(fig_13_1, mean(run_once_RL(2. ^ x) for _ in 1:50), label="alpha = 2^$x")

    end

    fig_13_1

end