Updated for compat with RLBase 0.11 (#11)

* updated for compatabilty with RLBase 0.11

* bug fix for envpool

Author: dylan-asmar

Project.toml

@@ -1,13 +1,14 @@
 name = "MPOPIS"
 uuid = "e8a75bc8-90e1-4072-945a-20230e5738f6"
 authors = ["Dylan Asmar <[email protected]>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CovarianceEstimation = "587fd27a-f159-11e8-2dae-1979310e6154"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -18,3 +19,16 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ReinforcementLearning = "158674fc-8238-5cab-b5ba-03dfc80d1318"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+
+[compat]
+CSV = "0.10"
+CovarianceEstimation = "0.2"
+Distributions = "0.25"
+DomainSets = "0.7"
+IntervalSets = "0.7"
+Plots = "1"
+ProgressMeter = "1"
+PyCall = "1.96"
+Reexport = "1"
+ReinforcementLearning = "0.11"
+StatsBase = "0.34"

src/MPOPIS.jl

@@ -16,6 +16,7 @@ import CovarianceEstimation.SimpleCovariance
 @reexport using ReinforcementLearning
 import ReinforcementLearning.AbstractEnv
 import ReinforcementLearning.RLBase
+using DomainSets
 using Plots
 @reexport import Plots.plot
 using ProgressMeter

src/envs/car_racing.jl

@@ -25,10 +25,8 @@ Base.show(io::IO, params::CarRacingEnvParams) = print(
     join(["$p=$(getfield(params, p))" for p in fieldnames(CarRacingEnvParams)], ","),
 )
 
-mutable struct CarRacingEnv{A,T,R<:AbstractRNG} <: AbstractEnv
+mutable struct CarRacingEnv{T,R<:AbstractRNG} <: AbstractEnv
     params::CarRacingEnvParams{T}
-    action_space::A
-    observation_space::Space{Vector{ClosedInterval{T}}}
     state::Vector{T}
     done::Bool
     t::Int
@@ -136,25 +134,8 @@ function CarRacingEnv(
     rng=Random.GLOBAL_RNG
 )
 
-    action_space = ClosedInterval{Vector{T}}(
-        [-1.0, -1.0],
-        [1.0, 1.0],
-    )
-    observation_space = Space([
-        -Inf .. Inf,                          # X position in XY plane (x = north, y = west)
-        -Inf .. Inf,                          # Y position in XY plane (x = north, y = west)
-        -π .. π,                              # yaw (rotation from x axis toward y axis [north to west])
-        -Inf .. Inf,                          # Longitudinal velocity
-        -Inf .. Inf,                          # Lateral velocity
-        -Inf .. Inf,                          # yaw rate
-        -params.δ_max .. params.δ_max,        # steering angle
-        -1.0 .. 1.0,                          # acceleration/brake amount [-1, 1]
-    ])
-
     env = CarRacingEnv(
         params,
-        action_space,
-        observation_space,
         zeros(T, 8),
         false,
         0,
@@ -172,8 +153,25 @@ CarRacingEnv{T}(; kwargs...) where {T} = CarRacingEnv(; T=T, kwargs...)
 
 Random.seed!(env::CarRacingEnv, seed) = Random.seed!(env.rng, seed)
 
-RLBase.action_space(env::CarRacingEnv) = env.action_space
-RLBase.state_space(env::CarRacingEnv{T}) where {T} = env.observation_space
+function RLBase.action_space(::CarRacingEnv{T}) where {T}
+    action_space = ClosedInterval{Vector{T}}([-1.0, -1.0], [1.0, 1.0])
+    return action_space
+end
+
+function RLBase.state_space(env::CarRacingEnv)
+    state_space = ArrayProductDomain([
+        -Inf .. Inf,                          # X position in XY plane (x = north, y = west)
+        -Inf .. Inf,                          # Y position in XY plane (x = north, y = west)
+        -π .. π,                              # yaw (rotation from x axis toward y axis [north to west])
+        -Inf .. Inf,                          # Longitudinal velocity
+        -Inf .. Inf,                          # Lateral velocity
+        -Inf .. Inf,                          # yaw rate
+        -env.params.δ_max .. env.params.δ_max,        # steering angle
+        -1.0 .. 1.0,                          # acceleration/brake amount [-1, 1]
+    ])
+    return state_space
+end
+
 RLBase.is_terminated(env::CarRacingEnv) = env.done
 RLBase.state(env::CarRacingEnv) = env.state
 
@@ -214,7 +212,7 @@ function RLBase.reward(env::CarRacingEnv{T}) where {T}
     return rew
 end
 
-function RLBase.reset!(env::CarRacingEnv{A,T}) where {A,T}
+function RLBase.reset!(env::CarRacingEnv{T}) where {T}
     ss_size = length(env.state)
     env.state = zeros(T, ss_size)
     env.state[3] = deg2rad(90)
@@ -224,7 +222,7 @@ function RLBase.reset!(env::CarRacingEnv{A,T}) where {A,T}
     nothing
 end
 
-function RLBase.reset!(env::CarRacingEnv{A,T}, state::Vector{T}) where {A,T}
+function RLBase.reset!(env::CarRacingEnv{T}, state::Vector{T}) where {T}
     env.state = state
     env.t = 0
     env.done = false
@@ -237,16 +235,16 @@ end
     a[1] = Turn angle [-max turn angle, max turn angle] (-1 right turn, +1 left turn)
     a[2] = Pedal amount (-1 = full brake, 1 = full throttle)
 """
-function (env::CarRacingEnv{<:ClosedInterval})(a::Vector{Float64})
-    a in env.action_space || error("Action is not in action space")
+function (env::CarRacingEnv)(a::Vector{Float64})
+    a in action_space(env) || error("Action is not in action space")
     _step!(env, a)
 end
 
-function (env::CarRacingEnv{<:ClosedInterval})(a::Vector{Int})
+function (env::CarRacingEnv)(a::Vector{Int})
     env(Float64.(a))
 end
 
-function (env::CarRacingEnv{<:ClosedInterval})(a::Matrix{Float64})
+function (env::CarRacingEnv)(a::Matrix{Float64})
     size(a)[2] == 1 || error("Only implented for one step")
     env(vec(a))
 end

src/envs/envpool_env.jl

@@ -1,10 +1,8 @@
 using PyCall
 
-mutable struct EnvpoolEnv{A,T,R<:AbstractRNG} <: AbstractEnv
+mutable struct EnvpoolEnv{T,R<:AbstractRNG} <: AbstractEnv
     task::String
     py_env::PyObject
-    action_space::A
-    observation_space::Space{Vector{ClosedInterval{T}}}
     num_states::Int
     num_envs::Int
     info::Dict
@@ -62,26 +60,11 @@ function EnvpoolEnv(
 
     py_env = py"get_envs_ep"(task, "gym", num_envs, frame_skip)
     env_data = py_env.reset()
-    py_action_space = py_env.action_space
-
-    action_space = ClosedInterval{Vector{T}}(
-        py_action_space.low,
-        py_action_space.high,
-    )
-
-    py_observation_space = py_env.observation_space
-    py_obs_len = py_observation_space.shape[1]
-    py_obs_low = py_observation_space.low
-    py_obs_high = py_observation_space.high
-
-    observation_vec = [py_obs_low[ii] .. py_obs_high[ii] for ii in 1:py_obs_len]
-    observation_space = Space(observation_vec)
-
+    py_obs_len = py_env.observation_space.shape[1]
+    
     env = EnvpoolEnv(
         task,
         py_env,
-        action_space,
-        observation_space,
         py_obs_len,
         num_envs,
         env_data[end],
@@ -98,16 +81,35 @@ function EnvpoolEnv(
 end
 
 Random.seed!(env::EnvpoolEnv, seed) = Random.seed!(env.rng, seed)
-RLBase.action_space(env::EnvpoolEnv) = env.action_space
-RLBase.state_space(env::EnvpoolEnv) = env.observation_space
+
+function RLBase.action_space(env::EnvpoolEnv{T}) where {T}
+    py_action_space = env.py_env.action_space
+    action_space = ClosedInterval{Vector{T}}(
+        py_action_space.low,
+        py_action_space.high,
+    )   
+    return action_space 
+end
+
+function RLBase.state_space(env::EnvpoolEnv{T}) where {T}
+    py_obs_len = env.py_env.observation_space.shape[1]
+    py_obs_low = env.py_env.observation_space.low
+    py_obs_high = env.py_env.observation_space.high
+
+    observation_vec = [py_obs_low[ii] .. py_obs_high[ii] for ii in 1:py_obs_len]
+    observation_space = ArrayProductDomain(observation_vec)
+    
+    return observation_space
+end
+
 RLBase.is_terminated(env::EnvpoolEnv) = env.done
 RLBase.state(env::EnvpoolEnv) = env.state
 RLBase.reward(env::EnvpoolEnv) = env.rews
 
 """
     The keywork argument `restore` is used to restore the environments based on `acts`
 """
-function RLBase.reset!(env::EnvpoolEnv{A,T}; restore=false) where {A,T}
+function RLBase.reset!(env::EnvpoolEnv{T}; restore=false) where {T}
     env_data = env.py_env.reset()
     env.info = env_data[end]
     env.rews = zeros(T, env.num_envs)

src/envs/multi-car_racing.jl

@@ -1,9 +1,7 @@
 
-mutable struct MultiCarRacingEnv{A,T,R<:AbstractRNG} <: AbstractEnv
+mutable struct MultiCarRacingEnv{T,R<:AbstractRNG} <: AbstractEnv
     N::Int
     envs::Vector{CarRacingEnv}
-    action_space::A
-    observation_space::Space{Vector{A}}
     state::Vector{T}
     done::Bool
     t::Int
@@ -39,41 +37,26 @@ function MultiCarRacingEnv(N=2;
     envs = Vector{CarRacingEnv}(undef, N)
     for ii in 1:N
         if length(car_params) >= ii
-            cre = CarRacingEnv(car_params, T=T, dt=dt, δt=δt, track=track, rng=rng)
+            cre = CarRacingEnv(car_params; T=T, dt=dt, δt=δt, track=track, rng=rng)
         else
-            cre = CarRacingEnv(T=T, dt=dt, δt=δt, track=track, rng=rng)
+            cre = CarRacingEnv(; T=T, dt=dt, δt=δt, track=track, rng=rng)
         end
         envs[ii] = cre
     end
-
-
-    endpts_l = []
-    endpts_r = []
-    single_state_size = length(RLBase.state_space(envs[1]))
-    obs_space_vec = Vector{ClosedInterval}(undef, N * single_state_size)
+    
+    single_state_size = length(RLBase.state_space(envs[1]).domains)
     state = zeros(T, N * single_state_size)
-    for (idx, en) in enumerate(envs)
-        endpts_l = [endpts_l; leftendpoint(RLBase.action_space(en))]
-        endpts_r = [endpts_r; rightendpoint(RLBase.action_space(en))]
-        start_idx = single_state_size * (idx - 1) + 1
-        end_idx = single_state_size * idx
-        obs_space_vec[start_idx:end_idx] = RLBase.state_space(en)[:]
-    end
-    action_space = ClosedInterval{Vector{T}}(endpts_l, endpts_r)
-    observation_space = Space(obs_space_vec)
-
+    
     env = MultiCarRacingEnv(
         N,
         envs,
-        action_space,
-        observation_space,
         state,
         false,
         0,
         dt,
         δt,
         Track(track),
-        rng,
+        rng
     )
 
     reset!(env)
@@ -89,14 +72,38 @@ function Random.seed!(env::MultiCarRacingEnv, seed)
     end
 end
 
-RLBase.action_space(env::MultiCarRacingEnv) = env.action_space
-RLBase.state_space(env::MultiCarRacingEnv{T}) where {T} = env.observation_space
+function RLBase.action_space(env::MultiCarRacingEnv{T}) where {T}
+    endpts_l = []
+    endpts_r = []
+    for en in env.envs
+        endpts_l = [endpts_l; leftendpoint(RLBase.action_space(en))]
+        endpts_r = [endpts_r; rightendpoint(RLBase.action_space(en))]
+    end
+    action_space = ClosedInterval{Vector{T}}(endpts_l, endpts_r)
+    return action_space
+end
+
+function RLBase.state_space(env::MultiCarRacingEnv{T}) where {T}
+    envs = env.envs
+    single_state_size = length(RLBase.state_space(envs[1]).domains)
+    obs_space_vec = Vector{ClosedInterval}(undef, env.N * single_state_size)
+    for (idx, en) in enumerate(envs)
+        start_idx = single_state_size * (idx - 1) + 1
+        end_idx = single_state_size * idx
+        obs_space_vec[start_idx:end_idx] = RLBase.state_space(en)[:]
+    end
+    observation_space = ArrayProductDomain(obs_space_vec)
+    return observation_space
+end
+
+
+
 RLBase.is_terminated(env::MultiCarRacingEnv) = env.done
 RLBase.state(env::MultiCarRacingEnv) = env.state
 
 function _update_states_env2envs(env::MultiCarRacingEnv)
     for (idx, en) in enumerate(env.envs)
-        ss_size = length(RLBase.state_space(en))
+        ss_size = length(RLBase.state_space(en).domains)
         start_idx = ss_size * (idx - 1) + 1
         end_idx = ss_size * idx
         en.state = env.state[start_idx:end_idx]
@@ -105,7 +112,7 @@ end
 
 function _update_states_envs2env(env::MultiCarRacingEnv)
     for (idx, en) in enumerate(env.envs)
-        ss_size = length(RLBase.state_space(en))
+        ss_size = length(RLBase.state_space(en).domains)
         start_idx = ss_size * (idx - 1) + 1
         end_idx = ss_size * idx
         env.state[start_idx:end_idx] = en.state
@@ -150,7 +157,7 @@ function RLBase.reward(env::MultiCarRacingEnv{T}) where {T}
     return rew
 end
 
-function RLBase.reset!(env::MultiCarRacingEnv{A,T}) where {A,T}
+function RLBase.reset!(env::MultiCarRacingEnv{T}) where {T}
     ss_size = length(env.state)
     ind_ss_size = round(Int, length(env.state) / env.N)
     env.envs[1].state = zeros(T, ind_ss_size)
@@ -172,7 +179,7 @@ function RLBase.reset!(env::MultiCarRacingEnv{A,T}) where {A,T}
     nothing
 end
 
-function RLBase.reset!(env::MultiCarRacingEnv{A,T}, state::Vector{T}) where {A,T}
+function RLBase.reset!(env::MultiCarRacingEnv{T}, state::Vector{T}) where {T}
     env.state = state
     _update_states_env2envs(env)
     env.t = 0

src/examples/cartpole_example.jl

@@ -1,12 +1,8 @@
 
 # Modifications to the RLBase functions to work with different GMPPI algorithms
-function (env::CartPoleEnv{<:Base.OneTo{Int}})(a::Vector)
+function (env::CartPoleEnv)(a)
     length(a) == 1 || error("Only implented for 1 step")
-    env(a[1])
-end
-function (env::CartPoleEnv{<:ClosedInterval})(a::Vector)
-    length(a) == 1 || error("Only implented for 1 step")
-    env(a[1])
+    RLBase.act!(env, a[1])
 end
 
 """ 
@@ -189,4 +185,3 @@ function simulate_cartpole(;
         gif(anim, gif_name, fps=10)
     end
 end
-

src/examples/mountaincar_example.jl

@@ -1,13 +1,9 @@
 
 
 # Modifications to the RLBase functions to work with different GMPPI algorithms
-function (env::MountainCarEnv{<:ClosedInterval})(a::Vector)
+function (env::MountainCarEnv)(a)
     length(a) == 1 || error("Only implented for 1 step")
-    env(a[1])
-end
-function (env::MountainCarEnv{<:Base.OneTo{Int}})(a::Vector)
-    length(a) == 1 || error("Only implented for 1 step")
-    env(a[1])
+    RLBase.act!(env, a[1])
 end
 
 # Modified MountainCar reward function