add RandomStartPolicy & restrict QBasedPolicy to return index of acti…

…on (#189)
JuliaReinforcementLearning · Dec 19, 2020 · 0cc7c7e · 0cc7c7e · findmyway · Dec 19, 2020
1 parent 135e691
commit 0cc7c7e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/policies/policies.jl b/src/policies/policies.jl
@@ -2,3 +2,4 @@ include("base.jl")
 include("agents/agents.jl")
 include("q_based_policies/q_based_policies.jl")
 include("random_policy.jl")
+include("random_start_policy.jl")
diff --git a/src/policies/q_based_policies/q_based_policy.jl b/src/policies/q_based_policies/q_based_policy.jl
@@ -19,9 +19,8 @@ end
 Flux.functor(x::QBasedPolicy) = (learner = x.learner,), y -> @set x.learner = y.learner
 
 (π::QBasedPolicy)(env) = π(env, ActionStyle(env))
-(π::QBasedPolicy)(env, ::MinimalActionSet) = action_space(env)[π.explorer(π.learner(env))]
-(π::QBasedPolicy)(env, ::FullActionSet) =
-    action_space(env)[π.explorer(π.learner(env), legal_action_space_mask(env))]
+(π::QBasedPolicy)(env, ::MinimalActionSet) = π.explorer(π.learner(env))
+(π::QBasedPolicy)(env, ::FullActionSet) = π.explorer(π.learner(env), legal_action_space_mask(env))
 
 RLBase.prob(p::QBasedPolicy, env) = prob(p, env, ActionStyle(env))
 RLBase.prob(p::QBasedPolicy, env, ::MinimalActionSet) =