Tabular Approximator fixes (pre v0.11 changes) (#1040)

* Fix tabular approx

* Expand tests for approx / learners

* Add target network docstring

* add more target_network tests

* add docstring, fix gpu toggle

* Fix naming

* Fix approximator tests

* Fix tests

* target network expanded tests pass

* final target network test passes

* delete excess file

* Drop gpu code from generic forward function

* Tab approx forward to env fixes

* Add missing test and methods for approximator

* Add missing env import

* Add RLEnv back to RLCore test dependencies

* only run gpu check when gpu is functional

* Fix dqn

* Try dqn fix

* Update src/ReinforcementLearningCore/src/policies/learners/target_network.jl

Co-authored-by: Henri Dehaybe <47037088+HenriDeh@users.noreply.github.com>

* Fix var naming style

---------

Co-authored-by: Henri Dehaybe <47037088+HenriDeh@users.noreply.github.com>

Author: jeremiahpslewis

src/ReinforcementLearningCore/Project.toml

@@ -55,9 +55,10 @@ DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReinforcementLearningEnvironments = "25e41dd2-4622-11e9-1641-f1adca772921"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [targets]
-test = ["CommonRLInterface", "CUDA", "cuDNN", "DomainSets", "Metal", "Preferences", "Test", "UUIDs"]
+test = ["CommonRLInterface", "CUDA", "cuDNN", "DomainSets", "Metal", "Preferences", "ReinforcementLearningEnvironments", "Test", "UUIDs"]

src/ReinforcementLearningCore/src/policies/learners/abstract_learner.jl

@@ -5,11 +5,11 @@ using Functors: @functor
 
 abstract type AbstractLearner end
 
-Base.show(io::IO, m::MIME"text/plain", L::AbstractLearner) = show(io, m, convert(AnnotatedStructTree, L))
+Base.show(io::IO, m::MIME"text/plain", learner::AbstractLearner) = show(io, m, convert(AnnotatedStructTree, learner))
 
 # Take Learner and Environment, get state, send to RLCore.forward(Learner, State)
-function forward(L::Le, env::E) where {Le <: AbstractLearner, E <: AbstractEnv}
-    env |> state |> Flux.gpu |> (x -> forward(L, x)) |> Flux.cpu
+function forward(learner::L, env::E) where {L <: AbstractLearner, E <: AbstractEnv}
+    env |> state |> (x -> forward(learner, x))
 end
 
 function RLBase.optimise!(::AbstractLearner, ::AbstractStage, ::Trajectory) end

src/ReinforcementLearningCore/src/policies/learners/approximator.jl

@@ -1,3 +1,5 @@
+using Flux
+
 """
     Approximator(model, optimiser)
 
@@ -9,20 +11,37 @@ struct Approximator{M,O} <: AbstractLearner
     optimiser_state::O
 end
 
-function Approximator(; model, optimiser, gpu=false)
+
+"""
+    Approximator(; model, optimiser, usegpu=false)
+
+Constructs an `Approximator` object for reinforcement learning.
+
+# Arguments
+- `model`: The model used for approximation.
+- `optimiser`: The optimizer used for updating the model.
+- `usegpu`: A boolean indicating whether to use GPU for computation. Default is `false`.
+
+# Returns
+An `Approximator` object.
+"""
+function Approximator(; model, optimiser::Flux.Optimise.AbstractOptimiser, use_gpu=false)
     optimiser_state = Flux.setup(optimiser, model)
-    if gpu  # Pass model to GPU (if available) upon creation
+    if use_gpu  # Pass model to GPU (if available) upon creation
         return Approximator(gpu(model), gpu(optimiser_state))
     else
         return Approximator(model, optimiser_state)
     end
 end
 
+Approximator(model, optimiser::Flux.Optimise.AbstractOptimiser; use_gpu=false) = Approximator(model=model, optimiser=optimiser, use_gpu=use_gpu)
+
 Base.show(io::IO, m::MIME"text/plain", A::Approximator) = show(io, m, convert(AnnotatedStructTree, A))
 
 @functor Approximator (model,)
 
 forward(A::Approximator, args...; kwargs...) = A.model(args...; kwargs...)
+forward(A::Approximator, env::E) where {E <: AbstractEnv} = env |> state |> (x -> forward(A, x))
 
-RLBase.optimise!(A::Approximator, grad) =
-    Flux.Optimise.update!(A.optimiser_state, A.model, grad)
+RLBase.optimise!(A::Approximator, grad::NamedTuple) =
+    Flux.Optimise.update!(A.optimiser_state, A.model, grad.model)

src/ReinforcementLearningCore/src/policies/learners/tabular_approximator.jl

@@ -26,22 +26,21 @@ TabularQApproximator(; n_state, n_action, init = 0.0, opt = InvDecay(1.0)) =
     TabularApproximator(fill(init, n_action, n_state), opt)
 
 # Take Learner and Environment, get state, send to RLCore.forward(Learner, State)
-function forward(L::Approximator{A, Any}, env::E) where {A <:AbstractArray, E <: AbstractEnv}
-    env |> state |> (x -> forward(L, x))
-end
+forward(L::TabularVApproximator, env::E) where {E <: AbstractEnv} = env |> state |> (x -> forward(L, x))
+forward(L::TabularQApproximator, env::E) where {E <: AbstractEnv} = env |> state |> (x -> forward(L, x))
 
 RLCore.forward(
-    app::Approximator{R,O},
+    app::TabularVApproximator{R,O},
     s::I,
-) where {R<:AbstractVector,O} = @views app.model[s]
+) where {R<:AbstractVector,O,I} = @views app.model[s]
 
 RLCore.forward(
-    app::Approximator{R,O},
+    app::TabularQApproximator{R,O},
     s::I,
-) where {R<:AbstractArray,O} = @views app.model[:, s]
+) where {R<:AbstractArray,O,I} = @views app.model[:, s]
 
 RLCore.forward(
-    app::Approximator{R,O},
+    app::TabularQApproximator{R,O},
     s::I1,
     a::I2,
-) where {R<:AbstractArray,O} = @views app.model[a, s]
+) where {R<:AbstractArray,O,I1,I2} = @views app.model[a, s]

src/ReinforcementLearningCore/src/policies/learners/target_network.jl

@@ -1,6 +1,6 @@
 export Approximator, TargetNetwork, target, model
 
-using Flux
+using Flux: gpu
 
 
 target(ap::Approximator) = ap.model #see TargetNetwork
@@ -33,11 +33,32 @@ mutable struct TargetNetwork{M}
     n_optimise::Int
 end
 
-function TargetNetwork(network; sync_freq = 1, ρ = 0f0)
+"""
+    TargetNetwork(network; sync_freq = 1, ρ = 0f0, use_gpu = false)
+
+Constructs a target network for reinforcement learning.
+
+# Arguments
+- `network`: The main network used for training.
+- `sync_freq`: The frequency (in number of calls to `optimise!`) at which the target network is synchronized with the main network. Default is 1.
+- `ρ`: The interpolation factor used for updating the target network. Must be in the range [0, 1]. Default is 0 (the old weights are completely replaced by the new ones).
+- `use_gpu`: Specifies whether to use GPU for the target network. Default is `false`.
+
+# Returns
+A `TargetNetwork` object.
+"""
+function TargetNetwork(network::Approximator; sync_freq = 1, ρ = 0f0, use_gpu = false)
     @assert 0 <= ρ <= 1 "ρ must in [0,1]"
-    # NOTE: model is pushed to gpu in Approximator, need to transfer to cpu before deepcopy, then push target model to gpu
-    target = gpu(deepcopy(cpu(network.model)))
-    TargetNetwork(network, target, sync_freq, ρ, 0)
+    ρ = Float32(ρ)
+    
+    if use_gpu
+        @assert typeof(gpu(network.model)) == typeof(network.model) "`Approximator` model is not on GPU. Please set `use_gpu=false`` or ensure model is on GPU, by setting `use_gpu=true` when constructing `Approximator`."
+        # NOTE: model is pushed to gpu in Approximator, need to transfer to cpu before deepcopy, then push target model to gpu
+        target = gpu(deepcopy(cpu(network.model)))
+    else
+        target = deepcopy(network.model)
+    end
+    return TargetNetwork(network, target, sync_freq, ρ, 0)
 end
 
 @functor TargetNetwork (network, target)
@@ -49,9 +70,9 @@ forward(tn::TargetNetwork, args...) = forward(tn.network, args...)
 model(tn::TargetNetwork) = model(tn.network)
 target(tn::TargetNetwork) = tn.target
 
-function RLBase.optimise!(tn::TargetNetwork, grad)
+function RLBase.optimise!(tn::TargetNetwork, grad::NamedTuple)
     A = tn.network
-    optimise!(A, grad)
+    optimise!(A, grad.network)
 
     tn.n_optimise += 1
 
@@ -62,4 +83,6 @@ function RLBase.optimise!(tn::TargetNetwork, grad)
         end
         tn.n_optimise = 0
     end
+
+    return
 end

src/ReinforcementLearningCore/test/policies/learners/abstract_learner.jl

@@ -0,0 +1,73 @@
+using Test
+using Flux
+
+@testset "AbstractLearner Tests" begin
+    @testset "Forward" begin
+        # Mock environment and learner
+        struct MockEnv <: AbstractEnv end
+        struct MockLearner <: AbstractLearner end
+
+        function RLCore.forward(::MockLearner, ::AbstractState)
+            return rand(2)
+        end
+
+        env = MockEnv()
+        learner = MockLearner()
+
+        output = forward(learner, env)
+
+        @test typeof(output) == Array{Float64,1}
+        @test length(output) == 2
+    end
+
+    @testset "Plan" begin
+        # Mock explorer, environment, and learner
+        struct MockExplorer <: AbstractExplorer end
+        struct MockEnv <: AbstractEnv end
+        struct MockLearner <: AbstractLearner end
+
+        function RLBase.plan!(::MockExplorer, ::AbstractState, ::AbstractActionSpace)
+            return rand(2)
+        end
+
+        env = MockEnv()
+        learner = MockLearner()
+        explorer = MockExplorer()
+
+        output = RLBase.plan!(explorer, learner, env)
+
+        @test typeof(output) == Array{Float64,1}
+        @test length(output) == 2
+    end
+
+    @testset "Plan with Player" begin
+        # Mock explorer, environment, and learner
+        struct MockExplorer <: AbstractExplorer end
+        struct MockEnv <: AbstractEnv end
+        struct MockLearner <: AbstractLearner end
+
+        function RLBase.plan!(::MockExplorer, ::AbstractState, ::AbstractActionSpace)
+            return rand(2)
+        end
+
+        env = MockEnv()
+        learner = MockLearner()
+        explorer = MockExplorer()
+        player = :player1
+
+        output = RLBase.plan!(explorer, learner, env, player)
+
+        @test typeof(output) == Array{Float64,1}
+        @test length(output) == 2
+    end
+
+    @testset "optimise!" begin
+        struct MockLearner <: AbstractLearner end
+        tr = Trajectory(
+                    CircularArraySARTSTraces(; capacity = 1_000),
+                    BatchSampler(1),
+                    InsertSampleRatioController(n_inserted = -1),
+                )
+        @test optimise!(MockLearner(), PreActStage(), tr) is nothing
+    end
+end

src/ReinforcementLearningCore/test/policies/learners/approximator.jl

@@ -0,0 +1,56 @@
+using Test
+using Flux
+using ReinforcementLearningEnvironments
+
+@testset "Approximator Tests" begin
+    @testset "Creation, with use_gpu = true toggle" begin
+        model = Chain(Dense(10, 5, relu), Dense(5, 2))
+        optimiser = Adam()
+        approximator = Approximator(model=model, optimiser=optimiser, use_gpu=true)
+
+        @test approximator isa Approximator
+        @test typeof(approximator.model) == typeof(gpu(model))
+        @test approximator.optimiser_state isa NamedTuple
+    end
+
+    @testset "Forward" begin
+        model = Chain(Dense(10, 5, relu), Dense(5, 2))
+        optimiser = Adam()
+        approximator = Approximator(model=model, optimiser=optimiser, use_gpu=false)
+
+        input = rand(Float32, 10)
+        output = RLCore.forward(approximator, input)
+
+        @test typeof(output) == Array{Float32,1}
+        @test length(output) == 2
+    end
+
+    @testset "Forward to environment" begin
+        model = Chain(Dense(4, 5, relu), Dense(5, 2))
+        optimiser = Adam()
+        approximator = Approximator(model=model, optimiser=optimiser, use_gpu=false)
+
+        env = CartPoleEnv()
+        output = RLCore.forward(approximator, env)
+        @test typeof(output) == Array{Float32,1}
+        @test length(output) == 2
+    end
+
+    @testset "Optimise" begin
+        model = Chain(Dense(10, 5, relu), Dense(5, 2))
+        optimiser = Adam()
+        approximator = Approximator(model=model, optimiser=optimiser)
+
+        input = rand(Float32, 10)
+        
+
+        grad = Flux.Zygote.gradient(approximator) do model
+            sum(RLCore.forward(model, input))
+        end
+    
+        @test approximator.model.layers[2].bias == [0, 0]
+        RLCore.optimise!(approximator, grad[1])
+
+        @test approximator.model.layers[2].bias != [0, 0]
+    end
+end

src/ReinforcementLearningCore/test/policies/learners/approximators.jl

@@ -1,2 +1,5 @@
-include("approximator.jl")
-include("tabular_approximator.jl")
+@testset "approximators.jl" begin
+    include("approximator.jl")
+    include("tabular_approximator.jl")
+    include("target_network.jl")
+end

src/ReinforcementLearningCore/test/policies/learners/learners.jl

@@ -1,6 +1,7 @@
 
 using Test
 using ReinforcementLearningCore
+using ReinforcementLearningEnvironments
 using Flux
 
 @testset "Constructors" begin
@@ -15,7 +16,11 @@ end
     v_approx = TabularVApproximator(n_state = 10)
     @test RLCore.forward(v_approx, 1) == 0.0
 
+    env = RockPaperScissorsEnv()
+    @test RLCore.forward(v_approx, env) == 0.0
+
     q_approx = TabularQApproximator(n_state = 5, n_action = 10)
     @test RLCore.forward(q_approx, 1) == zeros(Float64, 10)
     @test RLCore.forward(q_approx, 1, 5) == 0.0
-end
+    @test RLCore.forward(q_approx, env) == zeros(10)
+end