Drop Functors and use Flux.@layer (#1048)

Author: jeremiahpslewis

src/ReinforcementLearningCore/Project.toml

@@ -11,7 +11,6 @@ Crayons = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
@@ -33,7 +32,6 @@ Crayons = "4"
 Distributions = "0.25"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
 Flux = "0.14"
-Functors = "0.1, 0.2, 0.3, 0.4"
 GPUArrays = "8, 9, 10"
 Metal = "1.0"
 ProgressMeter = "1"

src/ReinforcementLearningCore/src/policies/agent/agent_base.jl

@@ -1,8 +1,7 @@
 export Agent
 
 using Base.Threads: @spawn
-
-using Functors: @functor
+using Flux
 import Base.push!
 
 abstract type AbstractAgent <: AbstractPolicy end
@@ -41,7 +40,7 @@ RLBase.optimise!(::AsyncTrajectoryStyle, agent::AbstractAgent, stage::S) where {
 #by default, optimise does nothing at all stage
 function RLBase.optimise!(policy::AbstractPolicy, stage::AbstractStage, trajectory::Trajectory) end
 
-@functor Agent (policy,)
+Flux.@layer Agent trainable=(policy,)
 
 function Base.push!(agent::Agent, ::PreEpisodeStage, env::AbstractEnv)
     push!(agent.trajectory, (state = state(env),))

src/ReinforcementLearningCore/src/policies/agent/offline_agent.jl

@@ -1,5 +1,7 @@
 export OfflineAgent, OfflineBehavior
 
+using Flux
+
 """
     OfflineBehavior(; agent:: Union{<:Agent, Nothing}, steps::Int, reset_condition)
 
@@ -49,7 +51,7 @@ struct OfflineAgent{P<:AbstractPolicy,T<:Trajectory,B<:OfflineBehavior} <: Abstr
 end
 
 OfflineAgent(; policy, trajectory, offline_behavior=OfflineBehavior()) = OfflineAgent(policy, trajectory, offline_behavior)
-@functor OfflineAgent (policy,)
+Flux.@layer OfflineAgent trainable=(policy,)
 
 Base.push!(::OfflineAgent{P,T,<:OfflineBehavior{Nothing}}, ::PreExperimentStage, env::AbstractEnv) where {P,T} = nothing
 #fills the trajectory with interactions generated with the behavior_agent at the PreExperimentStage.

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

@@ -1,7 +1,6 @@
 export AbstractLearner, Approximator
 
 using Flux
-using Functors: @functor
 
 abstract type AbstractLearner end
 

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

@@ -36,9 +36,7 @@ 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,)
+Flux.@layer Approximator trainable=(model,)
 
 forward(A::Approximator, args...; kwargs...) = A.model(args...; kwargs...)
 forward(A::Approximator, env::E) where {E <: AbstractEnv} = env |> state |> (x -> forward(A, x))

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

@@ -1,7 +1,6 @@
 export Approximator, TargetNetwork, target, model
 
-using Flux: gpu
-
+using Flux
 
 target(ap::Approximator) = ap.model #see TargetNetwork
 model(ap::Approximator) = ap.model #see TargetNetwork
@@ -61,9 +60,7 @@ function TargetNetwork(network::Approximator; sync_freq = 1, ρ = 0f0, use_gpu =
     return TargetNetwork(network, target, sync_freq, ρ, 0)
 end
 
-@functor TargetNetwork (network, target)
-
-Flux.trainable(model::TargetNetwork) = (model.network,)
+Flux.@layer TargetNetwork trainable=(network,)
 
 forward(tn::TargetNetwork, args...) = forward(tn.network, args...)
 

src/ReinforcementLearningCore/src/policies/q_based_policy.jl

@@ -1,6 +1,6 @@
 export QBasedPolicy
 
-using Functors: @functor
+using Flux
 
 """
     QBasedPolicy(;learner, explorer)
@@ -17,7 +17,7 @@ Base.@kwdef mutable struct QBasedPolicy{L,E} <: AbstractPolicy
     explorer::E
 end
 
-@functor QBasedPolicy (learner,)
+Flux.@layer QBasedPolicy trainable=(learner,)
 
 function RLBase.plan!(p::QBasedPolicy{L,Ex}, env::E) where {Ex<:AbstractExplorer,L<:AbstractLearner,E<:AbstractEnv}
     RLBase.plan!(p.explorer, p.learner, env)

src/ReinforcementLearningCore/src/utils/networks.jl

@@ -1,4 +1,3 @@
-using Functors: @functor
 import Flux
 import Flux.onehotbatch
 using ChainRulesCore: ignore_derivatives
@@ -18,7 +17,7 @@ Base.@kwdef struct ActorCritic{A,C,O}
     critic::C
 end
 
-@functor ActorCritic
+Flux.@layer ActorCritic
 
 #####
 # GaussianNetwork
@@ -53,7 +52,7 @@ end
 
 GaussianNetwork(pre, μ, σ; squash = identity) = GaussianNetwork(pre, μ, σ, 0.0f0, Inf32, squash)
 
-@functor GaussianNetwork
+Flux.@layer GaussianNetwork
 
 """
 This function is compatible with a multidimensional action space.
@@ -142,7 +141,7 @@ end
 
 SoftGaussianNetwork(pre, μ, σ) = SoftGaussianNetwork(pre, μ, σ, 0.0f0, Inf32)
 
-@functor SoftGaussianNetwork
+Flux.@layer SoftGaussianNetwork
 
 """
 This function is compatible with a multidimensional action space.
@@ -225,7 +224,7 @@ Base.@kwdef mutable struct CovGaussianNetwork{P,U,S}
     Σ::S
 end
 
-@functor CovGaussianNetwork
+Flux.@layer CovGaussianNetwork
 
 """
     (model::CovGaussianNetwork)(rng::AbstractRNG, state::AbstractArray{<:Any, 3}; is_sampling::Bool=false, is_return_log_prob::Bool=false)
@@ -407,7 +406,7 @@ mutable struct CategoricalNetwork{P}
     model::P
 end
 
-@functor CategoricalNetwork
+Flux.@layer CategoricalNetwork
 
 function (model::CategoricalNetwork)(rng::AbstractRNG, state::AbstractArray; is_sampling::Bool=false, is_return_log_prob::Bool = false)
     logits = model.model(state) #may be 1-3 dimensional
@@ -514,7 +513,7 @@ Base.@kwdef struct DuelingNetwork{B,V,A}
     adv::A
 end
 
-Flux.@functor DuelingNetwork
+Flux.@layer DuelingNetwork
 
 function (m::DuelingNetwork)(state)
     x = m.base(state)
@@ -544,7 +543,7 @@ Base.@kwdef struct PerturbationNetwork{N}
     ϕ::Float32 = 0.05f0
 end
 
-Flux.@functor PerturbationNetwork
+Flux.@layer PerturbationNetwork
 
 """
 This function accepts `state` and `action`, and then outputs actions after disturbance.
@@ -570,7 +569,7 @@ Base.@kwdef struct VAE{E,D}
     latent_dims::Int
 end
 
-Flux.@functor VAE
+Flux.@layer VAE
 
 function (model::VAE)(rng::AbstractRNG, state, action)
     μ, σ = model.encoder(vcat(state, action))

src/ReinforcementLearningCore/test/core/hooks.jl

@@ -32,7 +32,7 @@ end
 
 function test_run!(hook::AbstractHook)
     hook_ = deepcopy(hook)
-    run(RandomPolicy(), RandomWalk1D(), StopAfterNEpisodes(10), hook_)
+    run(RandomPolicy(), RandomWalk1D(), StopAfterNEpisodes(100), hook_)
     return hook_
 end
 
@@ -49,7 +49,7 @@ end
 
     for h in (h_1, h_2, h_3, h_4, h_5)
         h_ = test_run!(h)
-        @test length(h_.rewards) == 10
+        @test length(h_.rewards) == 100
         @test sum(h_.rewards .== 1) > 0
         @test sum(h_.rewards .== -1) > 0        
 
@@ -77,11 +77,11 @@ end
     h_1 = TimePerStep()
     h_2 = TimePerStep{Float32}()
 
-    sleep_vect = [0.01, 0.02, 0.03]
+    sleep_vect = [0.05, 0.05, 0.05]
     for h in (h_1, h_2)
         push!(h, PostActStage(), 1, 1)
         [(sleep(i); push!(h, PostActStage(), 1, 1)) for i in sleep_vect]
-        @test all(0.1 .> h.times[2:end] .> 0)
+        @test all(0.2 .> h.times[2:end] .> 0)
         test_noop!(h, stages=[PreActStage(), PreEpisodeStage(), PostEpisodeStage(), PreExperimentStage(), PostExperimentStage()])
     end
 end
@@ -115,8 +115,8 @@ end
 
     for h in (h_1, h_2, h_3)
         h_ = test_run!(h)
-        @test length(h_.rewards) == 10
-        @test sum(abs.(sum.(h_.rewards))) == 10
+        @test length(h_.rewards) == 100
+        @test sum(abs.(sum.(h_.rewards))) == 100
         @test length(unique(length.(h_.rewards))) > 1
         test_noop!(h, stages=[PreActStage(), PostEpisodeStage(), PreExperimentStage(), PostExperimentStage()])
     end

src/ReinforcementLearningCore/test/utils/networks.jl

@@ -19,11 +19,11 @@ import ReinforcementLearningBase: RLBase
     @testset "NeuralNetworkApproximator" begin
         NN = NeuralNetworkApproximator(; model = Dense(2, 3), optimizer = Descent())
 
-        q_values = NN(rand(2))
+        q_values = NN(rand(Float32, 2))
         @test size(q_values) == (3,)
 
         gs = gradient(params(NN)) do
-            sum(NN(rand(2, 5)))
+            sum(NN(rand(Float32, 2, 5)))
         end
 
         old_params = deepcopy(collect(params(NN).params))
@@ -47,15 +47,15 @@ import ReinforcementLearningBase: RLBase
         D = ac.actor.model |> gpu |> device
         @test D === device(ac) === device(ac.actor) == device(ac.critic)
 
-        A = send_to_device(D, rand(3))
+        A = send_to_device(D, rand(Float32, 3))
         ac.actor(A)
         ac.critic(A)
     end=#
 
     @testset "GaussianNetwork" begin
         @testset "On CPU" begin
             gn = GaussianNetwork(Dense(20,15), Dense(15,10), Dense(15,10, softplus))
-            state = rand(Float32,20,3) #batch of 3 states
+            state = rand(Float32, 20, 3) #batch of 3 states
             @testset "Correctness of outputs" begin
                 m, L = gn(state)
                 @test size(m) == size(L) == (10,3)
@@ -115,7 +115,7 @@ import ReinforcementLearningBase: RLBase
             if (@isdefined CUDA) && CUDA.functional()
                 CUDA.allowscalar(false)
                 gn = GaussianNetwork(Dense(20,15), Dense(15,10), Dense(15,10, softplus)) |> gpu
-                state = rand(20,3)  |> gpu #batch of 3 states
+                state = rand(Float32, 20,3)  |> gpu #batch of 3 states
                 @testset "Forward pass compatibility" begin
                     @test Flux.params(gn) == Flux.Params([gn.pre.weight, gn.pre.bias, gn.μ.weight, gn.μ.bias, gn.σ.weight, gn.σ.bias])
                     m, L = gn(state)