move AcrobotEnv optional to reduce first loading time (#139)

Author: findmyway

Project.toml

@@ -8,7 +8,6 @@ GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReinforcementLearningBase = "e575027e-6cd6-5018-9292-cdc6200d2b44"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -18,7 +17,6 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 GR = "0.46, 0.47, 0.48, 0.49, 0.50, 0.51, 0.52, 0.53, 0.54, 0.55"
 IntervalSets = "0.5"
 MacroTools = "0.5"
-OrdinaryDiffEq = "5"
 ReinforcementLearningBase = "0.9.2"
 Requires = "1.0"
 StatsBase = "0.32, 0.33"

src/ReinforcementLearningEnvironments.jl

@@ -5,7 +5,7 @@ using Random
 using GR
 using Requires
 using IntervalSets
-using Base.Threads: @spawn
+using Base.Threads:@spawn
 using Markdown
 
 const RLEnvs = ReinforcementLearningEnvironments
@@ -29,6 +29,11 @@ function __init__()
     @require SnakeGames = "34dccd9f-48d6-4445-aa0f-8c2e373b5429" include(
         "environments/3rd_party/snake.jl",
     )
+    @require OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" include(
+        "environments/3rd_party/AcrobotEnv.jl",
+    )
+
+
 end
 
 end # module

src/environments/examples/AcrobotEnv.jl renamed to src/environments/3rd_party/AcrobotEnv.jl

@@ -1,46 +1,3 @@
-import OrdinaryDiffEq
-
-export AcrobotEnv
-
-struct AcrobotEnvParams{T}
-    link_length_a::T # [m]
-    link_length_b::T # [m]
-    link_mass_a::T #: [kg] mass of link 1
-    link_mass_b::T #: [kg] mass of link 2
-    #: [m] position of the center of mass of link 1
-    link_com_pos_a::T
-    #: [m] position of the center of mass of link 2
-    link_com_pos_b::T
-    #: Rotation related parameters
-    link_moi::T
-    max_torque_noise::T
-    #: [m/s] maximum velocity of link 1
-    max_vel_a::T
-    #: [m/s] maximum velocity of link 2
-    max_vel_b::T
-    #: [m/s2] acceleration due to gravity
-    g::T
-    #: [s] timestep
-    dt::T
-    #: maximum steps in episode
-    max_steps::Int
-end
-
-mutable struct AcrobotEnv{T,R<:AbstractRNG} <: AbstractEnv
-    params::AcrobotEnvParams{T}
-    state::Vector{T}
-    action::Int
-    done::Bool
-    t::Int
-    rng::R
-    reward::T
-    # difference in second link angular acceleration equation
-    # as per python gym
-    book_or_nips::String
-    # array of available torques based on actions
-    avail_torque::Vector{T}
-end
-
 """
 AcrobotEnv(;kwargs...)
 # Keyword arguments
@@ -61,23 +18,23 @@ AcrobotEnv(;kwargs...)
 - `avail_torque = [T(-1.), T(0.), T(1.)]`
 """
 function AcrobotEnv(;
-    T = Float64,
-    link_length_a = T(1.0),
-    link_length_b = T(1.0),
-    link_mass_a = T(1.0),
-    link_mass_b = T(1.0),
-    link_com_pos_a = T(0.5),
-    link_com_pos_b = T(0.5),
-    link_moi = T(1.0),
-    max_torque_noise = T(0.0),
-    max_vel_a = T(4 * π),
-    max_vel_b = T(9 * π),
-    g = T(9.8),
-    dt = T(0.2),
-    max_steps = 200,
-    rng = Random.GLOBAL_RNG,
-    book_or_nips = "book",
-    avail_torque = [T(-1.0), T(0.0), T(1.0)],
+    T=Float64,
+    link_length_a=T(1.0),
+    link_length_b=T(1.0),
+    link_mass_a=T(1.0),
+    link_mass_b=T(1.0),
+    link_com_pos_a=T(0.5),
+    link_com_pos_b=T(0.5),
+    link_moi=T(1.0),
+    max_torque_noise=T(0.0),
+    max_vel_a=T(4 * π),
+    max_vel_b=T(9 * π),
+    g=T(9.8),
+    dt=T(0.2),
+    max_steps=200,
+    rng=Random.GLOBAL_RNG,
+    book_or_nips="book",
+    avail_torque=[T(-1.0), T(0.0), T(1.0)],
 )
 
     params = AcrobotEnvParams{T}(
@@ -124,7 +81,7 @@ RLBase.is_terminated(env::AcrobotEnv) = env.done
 RLBase.state(env::AcrobotEnv) = acrobot_observation(env.state)
 RLBase.reward(env::AcrobotEnv) = env.reward
 
-function RLBase.reset!(env::AcrobotEnv{T}) where {T<:Number}
+function RLBase.reset!(env::AcrobotEnv{T}) where {T <: Number}
     env.state[:] = T(0.1) * rand(env.rng, T, 4) .- T(0.05)
     env.t = 0
     env.action = 2
@@ -133,7 +90,7 @@ function RLBase.reset!(env::AcrobotEnv{T}) where {T<:Number}
 end
 
 # governing equations as per python gym
-function (env::AcrobotEnv{T})(a) where {T<:Number}
+function (env::AcrobotEnv{T})(a) where {T <: Number}
     env.action = a
     env.t += 1
     torque = env.avail_torque[a]
@@ -178,7 +135,7 @@ function dsdt(du, s_augmented, env::AcrobotEnv, t)
 
     # extract action and state
     a = s_augmented[end]
-    s = s_augmented[1:end-1]
+    s = s_augmented[1:end - 1]
 
     # writing in standard form
     theta1 = s[1]
@@ -242,7 +199,7 @@ function wrap(x, m, M)
     while x < m
         x = x + diff
     end
-    return x
+return x
 end
 
 function bound(x, m, M)

src/environments/3rd_party/structs.jl

@@ -6,7 +6,7 @@ struct GymEnv{T,Ta,To,P} <: AbstractEnv
 end
 export GymEnv
 
-mutable struct AtariEnv{IsGrayScale,TerminalOnLifeLoss,N,S<:AbstractRNG} <: AbstractEnv
+mutable struct AtariEnv{IsGrayScale,TerminalOnLifeLoss,N,S <: AbstractRNG} <: AbstractEnv
     ale::Ptr{Nothing}
     name::String
     screens::Tuple{Array{UInt8,N},Array{UInt8,N}}  # for max-pooling
@@ -38,3 +38,46 @@ mutable struct SnakeGameEnv{A,N,G} <: AbstractEnv
     is_terminated::Bool
 end
 export SnakeGameEnv
+
+struct AcrobotEnvParams{T}
+    link_length_a::T # [m]
+    link_length_b::T # [m]
+    link_mass_a::T # : [kg] mass of link 1
+    link_mass_b::T # : [kg] mass of link 2
+    # : [m] position of the center of mass of link 1
+    link_com_pos_a::T
+    # : [m] position of the center of mass of link 2
+    link_com_pos_b::T
+    # : Rotation related parameters
+    link_moi::T
+    max_torque_noise::T
+    # : [m/s] maximum velocity of link 1
+    max_vel_a::T
+    # : [m/s] maximum velocity of link 2
+    max_vel_b::T
+    # : [m/s2] acceleration due to gravity
+    g::T
+    # : [s] timestep
+    dt::T
+    # : maximum steps in episode
+    max_steps::Int
+end
+
+export AcrobotEnvParams
+
+mutable struct AcrobotEnv{T,R <: AbstractRNG} <: AbstractEnv
+    params::AcrobotEnvParams{T}
+    state::Vector{T}
+    action::Int
+    done::Bool
+    t::Int
+    rng::R
+    reward::T
+    # difference in second link angular acceleration equation
+    # as per python gym
+    book_or_nips::String
+    # array of available torques based on actions
+    avail_torque::Vector{T}
+end
+
+export AcrobotEnv

src/environments/examples/examples.jl

@@ -7,7 +7,6 @@ include("TicTacToeEnv.jl")
 include("TinyHanabiEnv.jl")
 include("PigEnv.jl")
 include("KuhnPokerEnv.jl")
-include("AcrobotEnv.jl")
 include("CartPoleEnv.jl")
 include("MountainCarEnv.jl")
 include("PendulumEnv.jl")

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 ArcadeLearningEnvironment = "b7f77d8d-088d-5e02-8ac0-89aab2acc977"
 OpenSpiel = "ceb70bd2-fe3f-44f0-b81f-41608acaf2f2"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReinforcementLearningBase = "e575027e-6cd6-5018-9292-cdc6200d2b44"

test/runtests.jl

@@ -8,6 +8,7 @@ using OpenSpiel
 using Random
 using StableRNGs
 using Statistics
+using OrdinaryDiffEq
 
 @testset "ReinforcementLearningEnvironments" begin
     include("environments/environments.jl")