Chapter 8.6 Trajectory Sampling

The general function run(policy, env, stop_condition, hook) is very flexible and powerful. However, we are not restricted to use it only. In this notebook, we'll see how to use part of the components provided in ReinforcementLearning.jl to finish some specific experiments.

First, let's define the environment mentioned in Chapter 8.6:

13.5 μs

xxxxxxxxxx
 
begin
    using ReinforcementLearning
    using Flux
    using Statistics
    using Plots
end

59.1 s

xxxxxxxxxx
 
begin
    mutable struct TestEnv <: AbstractEnv
        transitions::Array{Int, 3}
        rewards::Array{Float64, 3}
        reward_table::Array{Float64, 2}
        terminate_prob::Float64
        # cache
        s_init::Int
        s::Int
        reward::Float64
        is_terminated::Bool
    end
​
    function TestEnv(;ns=1000, na=2, b=1, terminate_prob=0.1,init_state=1)
        transitions = rand(1:ns, b, na, ns)
        rewards = randn(b, na, ns)
        reward_table = randn(na, ns)
        TestEnv(
            transitions,
            rewards,
            reward_table, 
            terminate_prob,
            init_state,
            init_state,
            0.,
            false
        )
    end
​
    function (env::TestEnv)(a::Int)
        t = rand() < 0.1
        bᵢ = rand(axes(env.transitions, 1))
​
        env.is_terminated = t
        if t
            env.reward = env.reward_table[a, env.s]
        else
            env.reward = env.rewards[bᵢ, a, env.s]
        end
​
        env.s = env.transitions[bᵢ, a, env.s]
​
    end
​
    RLBase.state_space(env::TestEnv) = Base.OneTo(1:size(env.rewards, 3))
    RLBase.action_space(env::TestEnv) = Base.OneTo(1:size(env.rewards, 2))
​
    function RLBase.reset!(env::TestEnv)
        env.s = env.s_init
        env.reward = 0.0
        env.is_terminated = false
    end
​
    RLBase.is_terminated(env::TestEnv) = env.is_terminated
    RLBase.state(env::TestEnv) = env.s
    RLBase.reward(env::TestEnv) = env.reward
end

4.9 ms

Note that this environment is not described very clearly on the book. Part of the information are inferred from the lisp source code.

Info

Actually the lisp code is also not perfect, I spent a whole afternoon to figure out the code logic. So good luck if you also want to understand it.

The definitions above are just like any other environment we've defined before in previous chapters. Now we'll add an extra function to make it work for our planning purpose.

9.9 μs

Main.workspace46.successors

xxxxxxxxxx
 
"""
Return all the possible next states and corresponding reward.
"""
function successors(env::TestEnv, s, a)
    S = @view env.transitions[:, a, s]
    R = @view env.rewards[:, a, s]
    zip(R, S)
end

334 μs

0.9

xxxxxxxxxx
 
γ = 0.9

76.0 ns

n_sweep

xxxxxxxxxx
 
n_sweep=10

77.0 ns

Main.workspace46.eval_Q

xxxxxxxxxx
 
"""
Here we are only interested in the performance of Q
with env starting at state `1`. Note here we're calculating
the discounted reward.
"""
function eval_Q(Q, env;n_eval=100)
    R = 0.
    for _ in 1:n_eval
        reset!(env)
        i = 0
        while !is_terminated(env)
            a = Q(state(env)) |> argmax  # greedy
            env(a)
            R += reward(env) * γ^i
            i += 1
        end
    end
    R/n_eval
end

204 μs

Main.workspace46.gain

xxxxxxxxxx
 
"""
Calculate the expected gain.
"""
function gain(Q,env,s,a)
    p = env.terminate_prob
    r = env.reward_table[a, s]
    p * r + (1-p) * mean(r̄ + γ * maximum(Q(s′)) for (r̄, s′) in  successors(env, s, a))
end

119 μs

sweep (generic function with 1 method)

xxxxxxxxxx
 
function sweep(;b = 1, ns=1000)
​
    na = 2
    
    α=1.0
    p = 0.1
​
    env= TestEnv(;ns=ns, na=na, b=b, terminate_prob=p)
    Q = TabularQApproximator(;n_state=ns, n_action=na, opt=Descent(α))
​
    i = 1
    vals = [eval_Q(Q, env)]
    for _ in 1:n_sweep
        for s in 1:ns
            for a in 1:na
                G = gain(Q,env,s,a)
                update!(Q, (s,a) =>  Q(s, a) - G)
                if i % 100 == 0
                    push!(vals, eval_Q(Q, env))
                end
                i += 1
            end
        end
    end
    vals
end

81.2 μs

on_policy (generic function with 1 method)

xxxxxxxxxx
 
function on_policy(;b = 1, ns=1000)
​
    na = 2
    
    α=1.0
    p = 0.1
​
    env= TestEnv(;ns=ns, na=na, b=b, terminate_prob=p)
    Q = TabularQApproximator(;n_state=ns, n_action=na, opt=Descent(α))
​
    i = 1
    vals = [eval_Q(Q, env)]
    
    explorer = EpsilonGreedyExplorer(0.1)
    for i in 1:(n_sweep * ns * na)
        is_terminated(env) && reset!(env)
        s = state(env)
        a = Q(s) |> explorer
        env(a)
        G = gain(Q, env, s, a)
        update!(Q, (s,a) => Q(s,a) - G)
        if i % 100 == 0
            push!(vals, eval_Q(Q, env))
        end
    end
    
    vals
end

79.2 μs

xxxxxxxxxx
 
begin
    fig_8_8 = plot(legend=:bottomright)
    for b in [1, 3, 10]
        plot!(fig_8_8, mean(sweep(;b=b) for _ in 1:200), label="uniform b=$b")
        plot!(fig_8_8, mean(on_policy(;b=b) for _ in 1:200), label="on policy b=$b")
    end
    fig_8_8
end

88.5 s

xxxxxxxxxx
 
begin
    fig_8_8_2 = plot(legend=:bottomright)
    plot!(fig_8_8_2, mean(sweep(;ns=10_000) for _ in 1:200), label="uniform")
    plot!(fig_8_8_2, mean(on_policy(;ns=10_000) for _ in 1:200), label="on_policy")
    fig_8_8_2
end

173 s