sync

Author: findmyway

notebooks/Chapter08_Maze.ipynb

@@ -40,8 +40,13 @@ function RLBase.extract_experience(
     end
 end
 
-function RLBase.update!(m::PrioritizedSweepingSampleModel, transition, P)
-    s, a, r, d, s′ = transition
+function RLBase.update!(m::PrioritizedSweepingSampleModel, t::AbstractTrajectory, p::AbstractPolicy)
+    experience = extract_experience(t, m)
+    isnothing(experience) || update!(m, (experience..., get_priority(p, experience)))
+end
+
+function RLBase.update!(m::PrioritizedSweepingSampleModel, transition::Tuple)
+    s, a, r, d, s′, P = transition
     m.experiences[(s, a)] = (r, d, s′)
     if P >= m.θ
         m.PQueue[(s, a)] = P

notebooks/Chapter08_Trajectory_Sampling.ipynb

@@ -30,7 +30,7 @@ mutable struct TimeBasedSampleModel <: AbstractEnvironmentModel
         )
 end
 
-function extract_transitions(t::AbstractTrajectory, m::TimeBasedSampleModel)
+function RLBase.extract_experience(t::AbstractTrajectory, m::TimeBasedSampleModel)
     if length(t) > 0
         get_trace(t, :state)[end],
         get_trace(t, :action)[end],

src/extensions/environment_models/prioritized_sweeping_sample_model.jl

@@ -1,5 +1,5 @@
 include("utils/utils.jl")
+include("environment_models/environment_models.jl")
 include("learners/learners.jl")
 include("policies/policies.jl")
-include("environment_models/environment_models.jl")
 include("iteration_methods.jl")

src/extensions/environment_models/time_based_sample_model.jl

@@ -208,6 +208,88 @@ function RLBase.extract_experience(
     end
 end
 
+function RLBase.update!(
+    learner::TDLearner{<:AbstractApproximator,:SARS},
+    model::Union{TimeBasedSampleModel,ExperienceBasedSampleModel},
+    t::AbstractTrajectory,
+    plan_step::Int,
+)
+    @assert learner.n == 0 "n must be 0 here"
+    for _ = 1:plan_step
+        transitions = extract_experience(model, learner)
+        if !isnothing(transitions)
+            update!(learner, transitions)
+        end
+    end
+end
+
+function RLBase.extract_experience(
+    model::Union{ExperienceBasedSampleModel,TimeBasedSampleModel},
+    learner::TDLearner{<:AbstractApproximator,:SARS},
+)
+    if length(model.experiences) > 0
+        s = sample(model)
+        (
+            states = [s[1]],
+            actions = [s[2]],
+            rewards = [s[3]],
+            terminals = [s[4]],
+            next_states = [s[5]],
+        )
+    else
+        nothing
+    end
+end
+
+function RLBase.get_priority(learner::TDLearner{<:AbstractApproximator,:SARS}, transition::Tuple)
+    s, a, r, d, s′ = transition
+    γ, Q, opt = learner.γ, learner.approximator, learner.optimizer
+    error = d ? apply!(opt, (s, a), r - Q(s, a)) :
+            apply!(opt, (s, a), r + γ^(learner.n + 1) * maximum(Q(s′)) - Q(s, a))
+    abs(error)
+end
+
+function RLBase.update!(
+    learner::TDLearner{<:AbstractApproximator,:SARS},
+    model::PrioritizedSweepingSampleModel,
+    t::AbstractTrajectory,
+    plan_step::Int,
+)
+    for _ = 1:plan_step
+        # @assert learner.n == 0 "n must be 0 here"
+        transitions = extract_experience(model, learner)
+        if !isnothing(transitions)
+            update!(learner, transitions)
+            s, _, _, _, _ = transitions
+            s = s[]  # length(s) is assumed to be 1
+            for (s̄, ā, r̄, d̄) in model.predecessors[s]
+                P = get_priority(learner, (s̄, ā, r̄, d̄, s))
+                if P ≥ model.θ
+                    model.PQueue[(s̄, ā)] = P
+                end
+            end
+        end
+    end
+end
+
+function RLBase.extract_experience(
+    model::PrioritizedSweepingSampleModel,
+    learner::TDLearner{<:AbstractApproximator,:SARS},
+)
+    if length(model.PQueue) > 0
+        s = sample(model)
+        (
+            states = [s[1]],
+            actions = [s[2]],
+            rewards = [s[3]],
+            terminals = [s[4]],
+            next_states = [s[5]],
+        )
+    else
+        nothing
+    end
+end
+
 #####
 # SARS DoubleLearner
 #####