⚡ Pluto.jl ⚡

xxxxxxxxxx
 
begin
    using ReinforcementLearning
    using Flux
    using Statistics
    using Plots
end

45.1 s

To describe the Grid World in Example 3.5, we'll create a distributional environment model. Here the distributional means, given a state-action paire, we can predict the possible next state, reward, termination info and the corresponding probability.

12.5 μs

xxxxxxxxxx
 
begin
    function nextstep(s::CartesianIndex{2}, a::CartesianIndex{2})
        if s == CartesianIndex(1, 2)
            r, s′ =  10., CartesianIndex(5, 2)
        elseif s == CartesianIndex(1, 4)
            r, s′ = 5., CartesianIndex(3, 4)
        else
            s′ = s + a
            if 1 ≤ s′[1] ≤ 5 && 1 ≤ s′[2] ≤ 5
                r = 0.
            else
                r = -1.
                s′ = s
            end
        end
        [(r, false, LinearIndices((5,5))[s′]) => 1.0]
    end
​
    ACTIONS = (
        CartesianIndex(-1, 0),
        CartesianIndex(1,0),
        CartesianIndex(0, 1),
        CartesianIndex(0, -1)
    )
​
    struct GridWorldModel <: AbstractEnvironmentModel
    end
​
    function (m::GridWorldModel)(s, a)
        nextstep(CartesianIndices((5,5))[s], ACTIONS[a])
    end
​
    RLBase.state_space(m::GridWorldModel) = Base.OneTo(5*5)
    RLBase.action_space(m::GridWorldModel) = Base.OneTo(length(ACTIONS))
end

24.4 ms

TabularApproximator{1,Array{Float64,1},Descent}tableFloat641

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

optimizerDescenteta

1.0

xxxxxxxxxx
 
V = TabularVApproximator(;n_state=25, opt=Descent(1.0))

21.9 ms

TabularApproximator{1,Array{Float64,1},Descent}tableFloat641

3.30943

1.52199

0.0512075

-0.973217

-1.85733

8.78964

2.99266

0.738502

-0.435172

-1.34491

4.42795

2.25045

0.673411

-0.354592

-1.22898

5.32267

1.90786

0.358465

-0.585334

-1.42265

1.49249

0.54769

-0.40287

-1.18281

-1.97492

optimizerDescenteta

1.0

xxxxxxxxxx
 
policy_evaluation!(
    V = V,
    π=RandomPolicy(Base.OneTo(4)),
    model=GridWorldModel(),
    γ=0.9
)

456 ms

table

5×5 Array{Float64,2}:
  3.30943     8.78964    4.42795    5.32267    1.49249
  1.52199     2.99266    2.25045    1.90786    0.54769
  0.0512075   0.738502   0.673411   0.358465  -0.40287
 -0.973217   -0.435172  -0.354592  -0.585334  -1.18281
 -1.85733    -1.34491   -1.22898   -1.42265   -1.97492

xxxxxxxxxx
 
table = reshape(V.table, 5, 5)

1.8 μs

xxxxxxxxxx
 
heatmap(table,yflip=true)

16.7 s