Chapter 8.2 Dyna: Integrated Planning, Acting, and Learning

To demonstrate the flexibility of ReinforcementLearning.jl, the DynaAgent is also included and we'll explore its performance in this notebook.

12.8 μs

xxxxxxxxxx
 
using ReinforcementLearning

49.3 s

The Maze Environment

In this chapter, the authors introduced a specific maze environment. So let's define it by implementing the interfaces in ReinforcementLearning.jl.

8.7 μs

CartesianIndex{2}1

CartesianIndex(0, -1)

CartesianIndex(0, 1)

CartesianIndex(-1, 0)

CartesianIndex(1, 0)

xxxxxxxxxx
 
const LRUD = [
    CartesianIndex(0, -1),  # left
    CartesianIndex(0, 1),   # right
    CartesianIndex(-1, 0),  # up
    CartesianIndex(1, 0),   # down
]

40.4 ms

xxxxxxxxxx
 
begin
    mutable struct MazeEnv <: AbstractEnv
        walls::Set{CartesianIndex{2}}
        position::CartesianIndex{2}
        start::CartesianIndex{2}
        goal::CartesianIndex{2}
        NX::Int
        NY::Int
    end
​
    function MazeEnv()
        walls = Set([
            [CartesianIndex(i, 3) for i = 2:4]
            CartesianIndex(5, 6)
            [CartesianIndex(j, 8) for j = 1:3]
        ])
        start = CartesianIndex(3, 1)
        goal = CartesianIndex(1, 9)
        MazeEnv(walls, start, start, goal, 6, 9)
    end
    
    function (env::MazeEnv)(a::Int)
        p = env.position + LRUD[a]
        if p == env.goal
            env.position = env.goal
        elseif !(p ∈ env.walls)
            env.position = CartesianIndex(min(max(p[1], 1), env.NX), min(max(p[2], 1), env.NY))
        end
        nothing
    end
end

4.2 ms

xxxxxxxxxx
 
RLBase.state_space(env::MazeEnv) = Base.OneTo(env.NX * env.NY)

29.7 μs

xxxxxxxxxx
 
RLBase.action_space(env::MazeEnv) = Base.OneTo(length(LRUD))

25.6 μs

xxxxxxxxxx
 
RLBase.reward(env::MazeEnv) = Float64(env.position == env.goal)

21.6 μs

xxxxxxxxxx
 
RLBase.state(env::MazeEnv) = (env.position[2] - 1) * env.NX + env.position[1]

26.3 μs

xxxxxxxxxx
 
RLBase.is_terminated(env::MazeEnv) = env.position == env.goal

21.4 μs

xxxxxxxxxx
 
RLBase.reset!(env::MazeEnv) = env.position = env.start

19.9 μs

* (generic function with 861 methods)

xxxxxxxxxx
 
begin
    import Base: *
    
    function extend(p::CartesianIndex{2}, n::Int)
        x, y = Tuple(p)
        [CartesianIndex(n * (x - 1) + i, n * (y - 1) + j) for i = 1:n for j = 1:n]
    end
​
    function remap(p::CartesianIndex{2}, n::Int)
        x, y = Tuple(p)
        CartesianIndex((x - 1) * n + 1, (y - 1) * n + 1)
    end
​
    function *(env::MazeEnv, n::Int)
        walls = Set{CartesianIndex{2}}(ww for w in env.walls for ww in extend(w, n))
        start, position, goal = remap(env.start, n), remap(env.position, n), remap(env.goal, n)
        NX, NY = env.NX * n, env.NY * n
        MazeEnv(walls, position, start, goal, NX, NY)
    end
end

449 μs

# MazeEnv

## Traits

| Trait Type        |                                            Value |
|:----------------- | ------------------------------------------------:|
| NumAgentStyle     |          ReinforcementLearningBase.SingleAgent() |
| DynamicStyle      |           ReinforcementLearningBase.Sequential() |
| InformationStyle  | ReinforcementLearningBase.ImperfectInformation() |
| ChanceStyle       |           ReinforcementLearningBase.Stochastic() |
| RewardStyle       |           ReinforcementLearningBase.StepReward() |
| UtilityStyle      |           ReinforcementLearningBase.GeneralSum() |
| ActionStyle       |     ReinforcementLearningBase.MinimalActionSet() |
| StateStyle        |     ReinforcementLearningBase.Observation{Any}() |
| DefaultStateStyle |     ReinforcementLearningBase.Observation{Any}() |

## Is Environment Terminated?

No

## State Space

`Base.OneTo(54)`

## Action Space

`Base.OneTo(4)`

## Current State

```
3
```

xxxxxxxxxx
 
x = MazeEnv()

336 ms

Figure 8.2

5.5 μs

xxxxxxxxxx
 
using Flux

26.4 ms

plan_step (generic function with 1 method)

xxxxxxxxxx
 
function plan_step(n)
    env = MazeEnv()
    ns = length(state_space(env))
    na = length(action_space(env))
​
    agent = DynaAgent(
        policy=QBasedPolicy(
            learner=TDLearner(
                approximator=TabularQApproximator(
                    n_state=ns,
                    n_action=na, 
                    opt=Descent(0.1)
                    ),
                γ=0.1,
                method=:SARS
                ),
            explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
            ),
        model=ExperienceBasedSamplingModel(),
        trajectory=VectorSARTTrajectory(),
        plan_step=n
    )
​
    hook = StepsPerEpisode()
    run(agent, env, StopAfterEpisode(50),hook)
    hook.steps
end

48.5 μs

xxxxxxxxxx
 
using Plots

8.9 s

xxxxxxxxxx
 
using Statistics

26.2 ms

xxxxxxxxxx
 
begin
    fig_8_2 = plot(legend=:topright)
    for n in [0, 5, 50]
        plot!(fig_8_2, mean(plan_step(n) for _ in 1:30), label="plan_step = $n")
    end
    fig_8_2
end

75.1 s

Figure 8.4

4.9 μs

cumulative_dyna_reward (generic function with 1 method)

xxxxxxxxxx
 
function cumulative_dyna_reward(model, walls, nstep1, change, nstep2)
    env = MazeEnv(
        walls,
        CartesianIndex(6, 4),
        CartesianIndex(6, 4),
        CartesianIndex(1, 9),
        6,
        9
    )
    ns = length(state_space(env))
    na = length(action_space(env))
    agent = DynaAgent(
        policy=QBasedPolicy(
            learner=TDLearner(
                approximator=TabularQApproximator(
                    n_state=ns,
                    n_action=na, 
                    opt=Descent(1.)
                    ),
                γ=0.95,
                method=:SARS
                ),
            explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
            ),
        model=model,
        trajectory=VectorSARTTrajectory(),
        plan_step=10
    )
    
    hook = StepsPerEpisode()
    run(agent, env, StopAfterStep(nstep1;is_show_progress=false),hook)
    change(env.walls)
    run(agent, env, StopAfterStep(nstep2;is_show_progress=false),hook)
    
    cumulative_reward = []
    for (i, n) in enumerate(hook.steps)
        for _ in 1:n
            push!(cumulative_reward, i)
        end
    end
    for _ in (nstep1+nstep2):-1:length(cumulative_reward)
        push!(cumulative_reward, length(hook.steps))
    end
    cumulative_reward
end

79.6 μs

walls (generic function with 1 method)

xxxxxxxxxx
 
walls() = Set([CartesianIndex(4, j) for j in 1:8])

47.1 μs

change_walls (generic function with 1 method)

xxxxxxxxxx
 
function change_walls(walls)
    pop!(walls, CartesianIndex(4,1))
    push!(walls, CartesianIndex(4,9))
end

20.3 μs

xxxxxxxxxx
 
begin
    fig_8_4 = plot(legend=:topleft)
    plot!(fig_8_4, mean(cumulative_dyna_reward(ExperienceBasedSamplingModel(), walls(), 1000, change_walls, 2000) for _ in 1:30), label="Dyna-Q")
    plot!(fig_8_4, mean(cumulative_dyna_reward(TimeBasedSamplingModel(;n_actions=4), walls(), 1000, change_walls, 2000) for _ in 1:30), label="Dyna-Q+")
    fig_8_4
end

36.8 s

Figure 8.5

5.6 μs

new_walls (generic function with 1 method)

xxxxxxxxxx
 
new_walls() = Set([CartesianIndex(4, j) for j in 2:9])

47.7 μs

new_change_walls (generic function with 1 method)

xxxxxxxxxx
 
function new_change_walls(walls)
    pop!(walls, CartesianIndex(4,9))
end

20.2 μs

xxxxxxxxxx
 
begin
    fig_8_5 = plot(legend=:topleft)
    plot!(fig_8_5, mean(cumulative_dyna_reward(ExperienceBasedSamplingModel(), new_walls(), 3000, new_change_walls, 3000) for _ in 1:50), label="dyna-Q")
    plot!(fig_8_5, mean(cumulative_dyna_reward(TimeBasedSamplingModel(n_actions=4, κ = 1e-3), new_walls(), 3000, new_change_walls, 3000) for _ in 1:50), label="dyna-Q+")
    fig_8_5
end

120 s

Example 8.4

5.3 μs

run_once (generic function with 2 methods)

xxxxxxxxxx
 
function run_once(model, ratio=1)
    env = MazeEnv() * ratio
    ns = length(state_space(env))
    na = length(action_space(env))
    agent = DynaAgent(
        policy=QBasedPolicy(
            learner=TDLearner(
                approximator=TabularQApproximator(
                    n_state=ns,
                    n_action=na,
                    opt=Descent(0.5)),
                γ=0.95,
                method=:SARS
                ),
            explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
            ),
        model=model,
        trajectory=VectorSARTTrajectory(),
        plan_step=5
    )
    hook = StepsPerEpisode()
    run(agent, env, (args...) -> length(hook.steps) > 0 && hook.steps[end] <= 14 * ratio * 1.2,hook)
    model.sample_count
end

83.2 μs

xxxxxxxxxx
 
begin
    p = plot(legend=:topleft)
    plot!(mean([run_once(ExperienceBasedSamplingModel(), ratio) for ratio in 1:6] for _ in 1:5), label="Dyna", yscale=:log10)
    plot!(mean([run_once(PrioritizedSweepingSamplingModel(), ratio) for ratio in 1:6] for _ in 1:5), label="Prioritized", yscale=:log10)
    p
end

136 s