Update FluxModelApproximator references to FluxApproximator (#1051)

Co-authored-by: Jeremiah Lewis <--get>

Author: jeremiahpslewis

src/ReinforcementLearningCore/src/policies/learners/flux_approximator.jl

@@ -0,0 +1,47 @@
+export FluxApproximator
+
+using Flux
+
+"""
+    FluxApproximator(model, optimiser)
+
+Wraps a Flux trainable model and implements the `RLBase.optimise!(::FluxApproximator, ::Gradient)` 
+interface. See the RLCore documentation for more information on proper usage.
+"""
+struct FluxApproximator{M,O} <: AbstractLearner
+    model::M
+    optimiser_state::O
+end
+
+
+"""
+    FluxApproximator(; model, optimiser, usegpu=false)
+
+Constructs an `FluxApproximator` 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 `FluxApproximator` object.
+"""
+function FluxApproximator(; model, optimiser, use_gpu=false)
+    optimiser_state = Flux.setup(optimiser, model)
+    if use_gpu  # Pass model to GPU (if available) upon creation
+        return FluxApproximator(gpu(model), gpu(optimiser_state))
+    else
+        return FluxApproximator(model, optimiser_state)
+    end
+end
+
+FluxApproximator(model, optimiser::Flux.Optimise.AbstractOptimiser; use_gpu=false) = FluxApproximator(model=model, optimiser=optimiser, use_gpu=use_gpu)
+
+Flux.@layer FluxApproximator trainable=(model,)
+
+forward(A::FluxApproximator, args...; kwargs...) = A.model(args...; kwargs...)
+forward(A::FluxApproximator, env::E) where {E <: AbstractEnv} = env |> state |> (x -> forward(A, x))
+
+RLBase.optimise!(A::FluxApproximator, grad::NamedTuple) =
+    Flux.Optimise.update!(A.optimiser_state, A.model, grad.model)

src/ReinforcementLearningCore/src/policies/learners/flux_model_approximator.jl

@@ -1,5 +1,5 @@
 include("abstract_learner.jl")
-include("flux_model_approximator.jl")
+include("flux_approximator.jl")
 include("tabular_approximator.jl")
 include("td_learner.jl")
 include("target_network.jl")

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

@@ -2,13 +2,13 @@ export TargetNetwork, target, model
 
 using Flux
 
-target(ap::FluxModelApproximator) = ap.model #see TargetNetwork
-model(ap::FluxModelApproximator) = ap.model #see TargetNetwork
+target(ap::FluxApproximator) = ap.model #see TargetNetwork
+model(ap::FluxApproximator) = ap.model #see TargetNetwork
 
 """
-    TargetNetwork(network::FluxModelApproximator; sync_freq::Int = 1, ρ::Float32 = 0f0)
+    TargetNetwork(network::FluxApproximator; sync_freq::Int = 1, ρ::Float32 = 0f0)
 
-Wraps an FluxModelApproximator to hold a target network that is updated towards the model of the 
+Wraps an FluxApproximator to hold a target network that is updated towards the model of the 
 approximator. 
 - `sync_freq` is the number of updates of `network` between each update of the `target`. 
 - ρ (\rho) is "how much of the target is kept when updating it". 
@@ -21,11 +21,11 @@ Implements the `RLBase.optimise!(::TargetNetwork, ::Gradient)` interface to upda
 and the target with weights replacement or Polyak averaging.
 
 Note to developers: `model(::TargetNetwork)` will return the trainable Flux model 
-and `target(::TargetNetwork)` returns the target model and `target(::FluxModelApproximator)`
+and `target(::TargetNetwork)` returns the target model and `target(::FluxApproximator)`
 returns the non-trainable Flux model. See the RLCore documentation.
 """
 mutable struct TargetNetwork{M}
-    network::FluxModelApproximator{M}
+    network::FluxApproximator{M}
     target::M
     sync_freq::Int
     ρ::Float32
@@ -46,13 +46,13 @@ Constructs a target network for reinforcement learning.
 # Returns
 A `TargetNetwork` object.
 """
-function TargetNetwork(network::FluxModelApproximator; sync_freq = 1, ρ = 0f0, use_gpu = false)
+function TargetNetwork(network::FluxApproximator; sync_freq = 1, ρ = 0f0, use_gpu = false)
     @assert 0 <= ρ <= 1 "ρ must in [0,1]"
     ρ = Float32(ρ)
 
     if use_gpu
-        @assert typeof(gpu(network.model)) == typeof(network.model) "`FluxModelApproximator` model is not on GPU. Please set `use_gpu=false`` or ensure model is on GPU, by setting `use_gpu=true` when constructing `FluxModelApproximator`."
-        # NOTE: model is pushed to gpu in FluxModelApproximator, need to transfer to cpu before deepcopy, then push target model to gpu
+        @assert typeof(gpu(network.model)) == typeof(network.model) "`FluxApproximator` model is not on GPU. Please set `use_gpu=false`` or ensure model is on GPU, by setting `use_gpu=true` when constructing `FluxApproximator`."
+        # NOTE: model is pushed to gpu in FluxApproximator, need to transfer to cpu before deepcopy, then push target model to gpu
         target = gpu(deepcopy(cpu(network.model)))
     else
         target = deepcopy(network.model)

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

@@ -1,21 +1,21 @@
 using Test
 using Flux
 
-@testset "FluxModelApproximator Tests" begin
+@testset "FluxApproximator Tests" begin
     @testset "Creation, with use_gpu = true toggle" begin
         model = Chain(Dense(10, 5, relu), Dense(5, 2))
         optimiser = Adam()
-        approximator = FluxModelApproximator(model=model, optimiser=optimiser, use_gpu=true)
+        approximator = FluxApproximator(model=model, optimiser=optimiser, use_gpu=true)
 
-        @test approximator isa FluxModelApproximator
+        @test approximator isa FluxApproximator
         @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 = FluxModelApproximator(model=model, optimiser=optimiser, use_gpu=false)
+        approximator = FluxApproximator(model=model, optimiser=optimiser, use_gpu=false)
 
         input = rand(Float32, 10)
         output = RLCore.forward(approximator, input)
@@ -27,7 +27,7 @@ using Flux
     @testset "Forward to environment" begin
         model = Chain(Dense(4, 5, relu), Dense(5, 2))
         optimiser = Adam()
-        approximator = FluxModelApproximator(model=model, optimiser=optimiser, use_gpu=false)
+        approximator = FluxApproximator(model=model, optimiser=optimiser, use_gpu=false)
 
         env = CartPoleEnv(T=Float32)
         output = RLCore.forward(approximator, env)
@@ -38,7 +38,7 @@ using Flux
     @testset "Optimise" begin
         model = Chain(Dense(10, 5, relu), Dense(5, 2))
         optimiser = Adam()
-        approximator = FluxModelApproximator(model=model, optimiser=optimiser)
+        approximator = FluxApproximator(model=model, optimiser=optimiser)
 
         input = rand(Float32, 10)
 

src/ReinforcementLearningCore/test/policies/learners/flux_model_approximator.jl renamed to src/ReinforcementLearningCore/test/policies/learners/flux_approximator.jl

@@ -1,6 +1,6 @@
 @testset "approximators.jl" begin
     include("abstract_learner.jl")
-    include("flux_model_approximator.jl")
+    include("flux_approximator.jl")
     include("tabular_approximator.jl")
     include("target_network.jl")
     include("td_learner.jl")

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

@@ -7,12 +7,12 @@ using ReinforcementLearningCore
         model = Chain(Dense(10, 5, relu), Dense(5, 2))
         optimiser = Adam()
         if ((@isdefined CUDA) && CUDA.functional()) || ((@isdefined Metal) && Metal.functional())
-            @test_throws "AssertionError: `FluxModelApproximator` model is not on GPU." TargetNetwork(FluxModelApproximator(model, optimiser), use_gpu=true)
+            @test_throws "AssertionError: `FluxApproximator` model is not on GPU." TargetNetwork(FluxApproximator(model, optimiser), use_gpu=true)
         end
-        @test TargetNetwork(FluxModelApproximator(model=model, optimiser=optimiser, use_gpu=true), use_gpu=true) isa TargetNetwork
-        @test TargetNetwork(FluxModelApproximator(model, optimiser, use_gpu=true), use_gpu=true) isa TargetNetwork
+        @test TargetNetwork(FluxApproximator(model=model, optimiser=optimiser, use_gpu=true), use_gpu=true) isa TargetNetwork
+        @test TargetNetwork(FluxApproximator(model, optimiser, use_gpu=true), use_gpu=true) isa TargetNetwork
 
-        approx = FluxModelApproximator(model, optimiser, use_gpu=false)
+        approx = FluxApproximator(model, optimiser, use_gpu=false)
         target_network = TargetNetwork(approx, use_gpu=false)
 
 
@@ -26,7 +26,7 @@ using ReinforcementLearningCore
 
     @testset "Forward" begin
         model = Chain(Dense(10, 5, relu), Dense(5, 2))
-        target_network = TargetNetwork(FluxModelApproximator(model, Adam()))
+        target_network = TargetNetwork(FluxApproximator(model, Adam()))
 
         input = rand(Float32, 10)
         output = RLCore.forward(target_network, input)
@@ -38,7 +38,7 @@ using ReinforcementLearningCore
     @testset "Optimise" begin
         optimiser = Adam()
         model = Chain(Dense(10, 5, relu), Dense(5, 2))
-        approximator = FluxModelApproximator(model, optimiser)
+        approximator = FluxApproximator(model, optimiser)
         target_network = TargetNetwork(approximator)
         input = rand(Float32, 10)    
         grad = Flux.Zygote.gradient(target_network) do model
@@ -54,7 +54,7 @@ using ReinforcementLearningCore
 
     @testset "Sync" begin
         optimiser = Adam()
-        model = FluxModelApproximator(Chain(Dense(10, 5, relu), Dense(5, 2)), optimiser)
+        model = FluxApproximator(Chain(Dense(10, 5, relu), Dense(5, 2)), optimiser)
         target_network = TargetNetwork(model, sync_freq=2, ρ=0.5)
 
         input = rand(Float32, 10)
@@ -73,7 +73,7 @@ end
 
 @testset "TargetNetwork" begin 
     m = Chain(Dense(4,1))
-    app = FluxModelApproximator(model = m, optimiser = Flux.Adam(), use_gpu=true)
+    app = FluxApproximator(model = m, optimiser = Flux.Adam(), use_gpu=true)
     tn = TargetNetwork(app, sync_freq = 3, use_gpu=true)
     @test typeof(model(tn)) == typeof(target(tn))
     p1 = Flux.destructure(model(tn))[1]