begin

    using ReinforcementLearning

    using Flux

    using Statistics

    using Plots

end

md"""

## The Baird Count Environment

"""

begin

    const DASH_SOLID = (:dashed, :solid)

​

    Base.@kwdef mutable struct BairdCounterEnv <: AbstractEnv

        current::Int = rand(1:7)

    end

​

    RLBase.state_space(env::BairdCounterEnv) = Base.OneTo(7)

    RLBase.action_space(env::BairdCounterEnv) = Base.OneTo(length(DASH_SOLID))

​

    function (env::BairdCounterEnv)(a)

        if DASH_SOLID[a] == :dashed

            env.current = rand(1:6)

        else

            env.current = 7

        end

        nothing

    end

​

    RLBase.reward(env::BairdCounterEnv) = 0.

    RLBase.is_terminated(env::BairdCounterEnv) = false

    RLBase.state(env::BairdCounterEnv) = env.current

    RLBase.reset!(env::BairdCounterEnv) = env.current = rand(1:6)

end

# Base.@kwdef struct OffPolicy{P,B} <: AbstractPolicy

#     π_target::P

#     π_behavior::B

# end

# (π::OffPolicy)(env) = π.π_behavior(env)

begin

    # const VectorWSARTTrajectory = Trajectory{<:NamedTuple{(:weight, SART...)}}

​

    # function VectorWSARTTrajectory(;weight=Float64, state=Int, action=Int, reward=Float32, terminal=Bool)

    #   VectorTrajectory(;weight=Float64, state=state, action=action, reward=reward, terminal=terminal)

    # end

    # function RLBase.update!(

    #   p::OffPolicy,

    #   t::VectorWSARTTrajectory,

    #   e::AbstractEnv,

    #   s::AbstractStage

    # )

    #   update!(p.π_target, t, e, s)

    # end

​

    # function RLBase.update!(

    #   t::VectorWSARTTrajectory,

    #   p::OffPolicy,

    #   env::AbstractEnv,

    #   s::PreActStage,

    #   a

    # )

    #   push!(t[:state], state(env))

    #   push!(t[:action], a)

​

    #   w = prob(p.π_target, s, a) / prob(p.π_behavior, s, a)

    #   push!(t[:weight], w)

    # end

​

    # function RLBase.update!(

    #   t::VectorWSARTTrajectory,

    #   p::OffPolicy{<:QBasedPolicy{<:TDLearner}},

    #   env::AbstractEnv,

    #   s::PreEpisodeStage,

    # )

    #   empty!(t)

    # end

​

    # function RLBase.update!(

    #   t::VectorWSARTTrajectory,

    #   p::OffPolicy{<:QBasedPolicy{<:TDLearner}},

    #   env::AbstractEnv,

    #   s::PostEpisodeStage,

    # )

    #   action = rand(action_space(env))

​

    #   push!(trajectory[:state], state(env))

    #   push!(trajectory[:action], action)

    #   push!(t[:weight], 1.0)

    # end

end

world = BairdCounterEnv()

begin

    Base.@kwdef struct RecordWeights <: AbstractHook

        weights::Vector{Vector{Float64}}=[]

    end

​

    (h::RecordWeights)(::PostActStage, agent, env) = push!(

        h.weights,

        agent.policy.π_target.learner.approximator.weights |> deepcopy

    )

end

NW = 8

INIT_WEIGHT = ones(8)

INIT_WEIGHT[7] = 10

STATE_MAPPING = zeros(NW, length(state_space(world)))

begin

    for i in 1:6

        STATE_MAPPING[i, i] = 2

        STATE_MAPPING[8, i] = 1

    end

    STATE_MAPPING[7, 7] = 1

    STATE_MAPPING[8, 7] = 2

end

STATE_MAPPING

π_b = x -> rand() < 6/7 ? 1 : 2

π_t = VBasedPolicy(

    learner=TDLearner(

        approximator=RLZoo.LinearApproximator(INIT_WEIGHT, Descent(0.01)),

        γ=0.99,

        n=0,

        method=:SRS

    ),

    mapping = (env, V) -> 2

)

prob_b = [6/7, 1/7]

prob_t = [0., 1.]

RLBase.prob(::typeof(π_b), s, a::Integer) = prob_b[a]

RLBase.prob(::typeof(π_t), s, a::Integer) = prob_t[a]

agent = Agent(

    policy=OffPolicy(

        π_target=π_t,

        π_behavior=π_b

    ),

    trajectory=VectorWSARTTrajectory(state=Any)

)

new_env = StateOverriddenEnv(

    BairdCounterEnv(),

    s -> @view STATE_MAPPING[:, s]

)

hook = RecordWeights()

run(agent, new_env, StopAfterStep(1000),hook)

begin

    p = plot(legend=:topleft)

    for i in 1:length(INIT_WEIGHT)

        plot!(p, [w[i] for w in hook.weights])

    end

    p

end