Extend MultiThreadEnv to accept a vector of environments (#108)

* extend MultiThreadEnv
JuliaReinforcementLearning · Dec 19, 2020 · 7d06c21 · 7d06c21
1 parent 1eb14f3
commit 7d06c21
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/environments/wrappers/MultiThreadEnv.jl b/src/environments/wrappers/MultiThreadEnv.jl
@@ -24,9 +24,15 @@ function Base.show(io::IO, t::MIME"text/markdown", env::MultiThreadEnv)
     show(io, t, Markdown.parse(s))
 end
 
-function MultiThreadEnv(f, n)
-    envs = [f() for _ in 1:n]
+"""
+    MultiThreadEnv(f, n::Int)
+
+`f` is a lambda function which creates an `AbstractEnv` by calling `f()`.
+"""
+MultiThreadEnv(f, n::Int) = MultiThreadEnv([f() for _ in 1:n])
 
+function MultiThreadEnv(envs::Vector{<:AbstractEnv})
+    n = length(envs)
     S = state_space(envs[1])
     s = state(envs[1])
     if S isa Space