fix bugs with DefaultStateStyleEnv

Author: findmyway

src/environments/wrappers/DefaultStateStyle.jl

@@ -14,7 +14,7 @@ DefaultStateStyleEnv{S}(env::E) where {S,E} = DefaultStateStyleEnv{S,E}(env)
 RLBase.DefaultStateStyle(::DefaultStateStyleEnv{S}) where {S} = S
 
 for f in vcat(RLBase.ENV_API, RLBase.MULTI_AGENT_ENV_API)
-    if f ∉ (:DefaultStateStyle,)
+    if f ∉ (:DefaultStateStyle, :state, :state_space)
         @eval RLBase.$f(x::DefaultStateStyleEnv, args...; kwargs...) =
             $f(x.env, args...; kwargs...)
     end
@@ -23,5 +23,10 @@ end
 (env::DefaultStateStyleEnv)(args...; kwargs...) = env.env(args...; kwargs...)
 
 RLBase.state(env::DefaultStateStyleEnv, ss::RLBase.AbstractStateStyle) = state(env.env, ss)
+RLBase.state(env::DefaultStateStyleEnv, ss::RLBase.AbstractStateStyle, p) = state(env.env, ss, p)
+
 RLBase.state_space(env::DefaultStateStyleEnv, ss::RLBase.AbstractStateStyle) =
     state_space(env.env, ss)
+
+RLBase.state_space(env::DefaultStateStyleEnv, ss::RLBase.AbstractStateStyle, p) =
+    state_space(env.env, ss, p)