Update RLBase to the latest version (#188)

* fix RLBase imported functions

* rename functions in RLBase

* update dependency

Author: findmyway

Project.toml

@@ -1,7 +1,7 @@
 name = "ReinforcementLearningCore"
 uuid = "de1b191a-4ae0-4afa-a27b-92d07f46b2d6"
 authors = ["Jun Tian <tianjun.cpp@gmail.com>"]
-version = "0.6.0"
+version = "0.6.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -16,7 +16,6 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 ImageTransformations = "02fcd773-0e25-5acc-982a-7f6622650795"
-JLD = "4138dd39-2aa7-5051-a626-17a0bb65d9c8"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
@@ -39,10 +38,9 @@ Flux = "0.11.1"
 Functors = "0.1"
 GPUArrays = "5, 6.0"
 ImageTransformations = "0.8"
-JLD = "0.10, 0.11"
 MacroTools = "0.5"
 ProgressMeter = "1.2"
-ReinforcementLearningBase = "0.8.2"
+ReinforcementLearningBase = "0.9"
 Setfield = "0.6, 0.7"
 StatsBase = "0.32, 0.33"
 Zygote = "0.5"

src/core/hooks.jl

@@ -110,11 +110,7 @@ function (hook::RewardsPerEpisode)(::PreEpisodeStage, agent, env)
 end
 
 function (hook::RewardsPerEpisode)(::PostActStage, agent, env)
-    push!(hook.rewards[end], get_reward(env))
-end
-
-function (hook::RewardsPerEpisode)(::PostActStage, agent, env::RewardOverriddenEnv)
-    push!(hook.rewards[end], get_reward(env.env))
+    push!(hook.rewards[end], reward(env))
 end
 
 #####
@@ -125,9 +121,6 @@ end
     TotalRewardPerEpisode(; rewards = Float64[], reward = 0.0)
 
 Store the total rewards of each episode in the field of `rewards`.
-
-!!! note
-    If the environment is a [`RewardOverriddenenv`](@ref), then the original reward is recorded.
 """
 Base.@kwdef mutable struct TotalRewardPerEpisode <: AbstractHook
     rewards::Vector{Float64} = Float64[]
@@ -136,58 +129,11 @@ end
 
 Base.getindex(h::TotalRewardPerEpisode) = h.rewards
 
-(hook::TotalRewardPerEpisode)(s::AbstractStage, agent, env) =
-    hook(s, agent, env, RewardStyle(env), NumAgentStyle(env))
-(hook::TotalRewardPerEpisode)(::AbstractStage, agent, env, ::Any, ::Any) = nothing
-
-(hook::TotalRewardPerEpisode)(
-    ::PostEpisodeStage,
-    agent,
-    env,
-    ::TerminalReward,
-    ::SingleAgent,
-) = push!(hook.rewards, get_reward(env))
-(hook::TotalRewardPerEpisode)(
-    ::PostEpisodeStage,
-    agent,
-    env,
-    ::TerminalReward,
-    ::MultiAgent,
-) = push!(hook.rewards, get_reward(env, get_role(agent)))
-(hook::TotalRewardPerEpisode)(::PostActStage, agent, env, ::StepReward, ::SingleAgent) =
-    hook.reward += get_reward(env)
-(hook::TotalRewardPerEpisode)(::PostActStage, agent, env, ::StepReward, ::MultiAgent) =
-    hook.reward += get_reward(env, get_role(agent))
-(hook::TotalRewardPerEpisode)(
-    ::PostEpisodeStage,
-    agent,
-    env::RewardOverriddenEnv,
-    ::TerminalReward,
-    ::SingleAgent,
-) = push!(hook.rewards, get_reward(env.env))
-(hook::TotalRewardPerEpisode)(
-    ::PostEpisodeStage,
-    agent,
-    env::RewardOverriddenEnv,
-    ::TerminalReward,
-    ::MultiAgent,
-) = push!(hook.rewards, get_reward(env.env, get_role(agent)))
-(hook::TotalRewardPerEpisode)(
-    ::PostActStage,
-    agent,
-    env::RewardOverriddenEnv,
-    ::StepReward,
-    ::SingleAgent,
-) = hook.reward += get_reward(env.env)
-(hook::TotalRewardPerEpisode)(
-    ::PostActStage,
-    agent,
-    env::RewardOverriddenEnv,
-    ::StepReward,
-    ::MultiAgent,
-) = hook.reward += get_reward(env.env, get_role(agent))
-
-function (hook::TotalRewardPerEpisode)(::PostEpisodeStage, agent, env, ::StepReward, ::Any)
+function (hook::TotalRewardPerEpisode)(::PostActStage, agent, env)
+    hook.reward += reward(env)
+end
+
+function (hook::TotalRewardPerEpisode)(::PostEpisodeStage, agent, env)
     push!(hook.rewards, hook.reward)
     hook.reward = 0
 end
@@ -205,33 +151,27 @@ Base.getindex(h::TotalBatchRewardPerEpisode) = h.rewards
 """
     TotalBatchRewardPerEpisode(batch_size::Int)
 
-Similar to [`TotalRewardPerEpisode`](@ref), but will record total rewards per episode in [`MultiThreadEnv`](@ref).
-
-!!! note
-    If the environment is a [`RewardOverriddenEnv`](@ref), then the original reward is recorded.
+Similar to [`TotalRewardPerEpisode`](@ref), but is specific to environments
+which return a `Vector` of rewards (a typical case with `MultiThreadEnv`).
 """
 function TotalBatchRewardPerEpisode(batch_size::Int)
     TotalBatchRewardPerEpisode([Float64[] for _ in 1:batch_size], zeros(batch_size))
 end
 
-function (hook::TotalBatchRewardPerEpisode)(
-    ::PostActStage,
-    agent,
-    env::MultiThreadEnv{T},
-) where {T}
-    for i in 1:length(env)
-        if T <: RewardOverriddenEnv
-            hook.reward[i] += get_reward(env[i].env)
-        else
-            hook.reward[i] += get_reward(env[i])
-        end
-        if get_terminal(env[i])
+function (hook::TotalBatchRewardPerEpisode)(::PostActStage, agent, env)
+    for (i, (t, r)) in enumerate(zip(is_terminated(env), reward(env)))
+        hook.reward[i] += r
+        if t
             push!(hook.rewards[i], hook.reward[i])
             hook.reward[i] = 0.0
         end
     end
 end
 
+#####
+# BatchStepsPerEpisode
+#####
+
 struct BatchStepsPerEpisode <: AbstractHook
     steps::Vector{Vector{Int}}
     step::Vector{Int}
@@ -242,16 +182,17 @@ Base.getindex(h::BatchStepsPerEpisode) = h.steps
 """
     BatchStepsPerEpisode(batch_size::Int; tag = "TRAINING")
 
-Similar to [`StepsPerEpisode`](@ref), but only work for [`MultiThreadEnv`](@ref)
+Similar to [`StepsPerEpisode`](@ref), but is specific to environments
+which return a `Vector` of rewards (a typical case with `MultiThreadEnv`).
 """
 function BatchStepsPerEpisode(batch_size::Int)
     BatchStepsPerEpisode([Int[] for _ in 1:batch_size], zeros(Int, batch_size))
 end
 
-function (hook::BatchStepsPerEpisode)(::PostActStage, agent, env::MultiThreadEnv)
-    for i in 1:length(env)
+function (hook::BatchStepsPerEpisode)(::PostActStage, agent, env)
+    for (i, t) in enumerate(is_terminated(env))
         hook.step[i] += 1
-        if get_terminal(env[i])
+        if t
             push!(hook.steps[i], hook.step[i])
             hook.step[i] = 0
         end
@@ -266,24 +207,20 @@ end
     CumulativeReward(rewards::Vector{Float64} = [0.0])
 
 Store cumulative rewards since the beginning to the field of `rewards`.
-
-!!! note
-    If the environment is a [`RewardOverriddenEnv`](@ref), then the original reward is recorded instead.
 """
 Base.@kwdef struct CumulativeReward <: AbstractHook
-    rewards::Vector{Float64} = [0.0]
+    rewards::Vector{Vector{Float64}} = [[0.0]]
 end
 
 Base.getindex(h::CumulativeReward) = h.rewards
 
-function (hook::CumulativeReward)(::PostActStage, agent, env::T) where {T}
-    if T <: RewardOverriddenEnv
-        r = get_reward(env.env)
-    else
-        r = get_reward(env)
-    end
-    push!(hook.rewards, r + hook.rewards[end])
-    @debug hook.tag CUMULATIVE_REWARD = hook.rewards[end]
+function (hook::CumulativeReward)(::PostEpisodeStage, agent, env)
+    push!(hook.rewards, [0.0])
+end
+
+function (hook::CumulativeReward)(::PostActStage, agent, env)
+    r = reward(env)
+    push!(hook.rewards[end], r + hook.rewards[end][end])
 end
 
 #####
@@ -363,7 +300,7 @@ Base.@kwdef mutable struct UploadTrajectoryEveryNStep{M,S} <: AbstractHook
     sealer::S = deepcopy
 end
 
-function (hook::UploadTrajectoryEveryNStep)(::PostActStage, agent, env)
+function (hook::UploadTrajectoryEveryNStep)(::PostActStage, agent::Agent, env)
     hook.t += 1
     if hook.t > 0 && hook.t % hook.n == 0
         put!(hook.mailbox, hook.sealer(agent.trajectory))

src/core/run.jl

@@ -33,7 +33,7 @@ function _run(
         policy(PRE_EPISODE_STAGE, env)
         hook(PRE_EPISODE_STAGE, policy, env)
 
-        while !get_terminal(env) # one episode
+        while !is_terminated(env) # one episode
             action = policy(PRE_ACT_STAGE, env)
             hook(PRE_ACT_STAGE, policy, env, action)
 
@@ -53,29 +53,3 @@ function _run(
     end
     hook
 end
-
-function _run(
-    ::Sequential,
-    ::SingleAgent,
-    policy::AbstractPolicy,
-    env::MultiThreadEnv,
-    stop_condition,
-    hook::AbstractHook = EmptyHook(),
-)
-
-    while true
-        reset!(env)
-        action = policy(PRE_ACT_STAGE, env)
-        hook(PRE_ACT_STAGE, policy, env, action)
-
-        env(action)
-        policy(POST_ACT_STAGE, env)
-        hook(POST_ACT_STAGE, policy, env)
-
-        if stop_condition(policy, env)
-            policy(PRE_ACT_STAGE, env)  # let the policy see the last observation
-            break
-        end
-    end
-    hook
-end

src/core/stop_conditions.jl

@@ -95,7 +95,7 @@ function StopAfterEpisode(episode; cur = 0, is_show_progress = true)
 end
 
 function (s::StopAfterEpisode)(agent, env)
-    if get_terminal(env)
+    if is_terminated(env)
         s.cur += 1
         if !isnothing(s.progress)
             next!(s.progress;)
@@ -105,10 +105,6 @@ function (s::StopAfterEpisode)(agent, env)
     s.cur >= s.episode
 end
 
-(s::StopAfterEpisode)(agent, env::MultiThreadEnv) =
-    @error "MultiThreadEnv is not supported!"
-
-
 #####
 # StopAfterNoImprovement
 #####
@@ -128,7 +124,7 @@ Parameters:
 
 fn: a closure, return a scalar value, which indicates the performance of the policy (the higher the better)
 e.g. 
-1. () -> get_reward(env)
+1. () -> reward(env)
 1. () -> total_reward_per_episode.reward
 
 patience: Number of epochs with no improvement after which training will be stopped.
@@ -142,7 +138,7 @@ function StopAfterNoImprovement(fn, patience::Int, δ::T = 0.0f0) where {T<:Numb
 end
 
 function (s::StopAfterNoImprovement)(agent, env)::Bool
-    get_terminal(env) || return false # post episode stage
+    is_terminated(env) || return false # post episode stage
     val = s.fn()
     improved = isfull(s.buffer) ? all(s.buffer .< (val - s.δ)) : true
     push!(s.buffer, val)
@@ -160,7 +156,7 @@ Return `true` if the environment is terminated.
 """
 struct StopWhenDone end
 
-(s::StopWhenDone)(agent, env) = get_terminal(env)
+(s::StopWhenDone)(agent, env) = is_terminated(env)
 
 #####
 # StopSignal

src/extensions/ReinforcementLearningBase.jl

@@ -1,13 +1,5 @@
-using CUDA
-using Distributions: pdf
-using Random
-using Flux
 using AbstractTrees
 
-RLBase.update!(p::RandomPolicy, x) = nothing
-
-Random.rand(s::MultiContinuousSpace{<:CuArray}) = rand(CUDA.CURAND.generator(), s)
-
 Base.show(io::IO, p::AbstractPolicy) =
     AbstractTrees.print_tree(io, StructTree(p), get(io, :max_depth, 10))
 

src/policies/agents/agent.jl

@@ -85,7 +85,7 @@ function RLBase.update!(
     ::Union{PreActStage,PostEpisodeStage},
 )
     action = policy(env)
-    push!(trajectory[:state], get_state(env))
+    push!(trajectory[:state], state(env))
     push!(trajectory[:action], action)
     action
 end
@@ -100,9 +100,9 @@ function RLBase.update!(
     ::Union{PreActStage,PostEpisodeStage},
 )
     action = policy(env)
-    push!(trajectory[:state], get_state(env))
+    push!(trajectory[:state], state(env))
     push!(trajectory[:action], action)
-    push!(trajectory[:legal_actions_mask], get_legal_actions_mask(env))
+    push!(trajectory[:legal_actions_mask], legal_action_space_mask(env))
     action
 end
 
@@ -112,6 +112,6 @@ function RLBase.update!(
     env::AbstractEnv,
     ::PostActStage,
 )
-    push!(trajectory[:reward], get_reward(env))
-    push!(trajectory[:terminal], get_terminal(env))
+    push!(trajectory[:reward], reward(env))
+    push!(trajectory[:terminal], is_terminated(env))
 end

src/policies/policies.jl

@@ -1,3 +1,4 @@
 include("base.jl")
 include("agents/agents.jl")
 include("q_based_policies/q_based_policies.jl")
+include("random_policy.jl")

src/policies/q_based_policies/explorers/abstract_explorer.jl

@@ -14,15 +14,15 @@ function (p::AbstractExplorer)(x) end
 function (p::AbstractExplorer)(x, mask) end
 
 """
-    get_prob(p::AbstractExplorer, x) -> AbstractDistribution
+    prob(p::AbstractExplorer, x) -> AbstractDistribution
 
 Get the action distribution given action values.
 """
-function RLBase.get_prob(p::AbstractExplorer, x) end
+function RLBase.prob(p::AbstractExplorer, x) end
 
 """
-    get_prob(p::AbstractExplorer, x, mask)
+    prob(p::AbstractExplorer, x, mask)
 
-Similart to `get_prob(p::AbstractExplorer, x)`, but here only the `mask`ed elements are considered.
+Similart to `prob(p::AbstractExplorer, x)`, but here only the `mask`ed elements are considered.
 """
-function RLBase.get_prob(p::AbstractExplorer, x, mask) end
+function RLBase.prob(p::AbstractExplorer, x, mask) end