Add default checks (#191)

* add default checks

* simplify condition

Author: findmyway

src/policies/agents/agent.jl

@@ -23,6 +23,17 @@ functor(x::Agent) = (policy = x.policy,), y -> @set x.policy = y.policy
 
 (agent::Agent)(env) = agent.policy(env)
 
+function check(agent::Agent, env::AbstractEnv)
+    if ActionStyle(env) === FULL_ACTION_SET && !haskey(agent.trajectory, :legal_actions_mask)
+        @warn "The env[$(nameof(env))] is of FULL_ACTION_SET, but I can not find a trace named :legal_actions_mask in the trajectory"
+    end
+    check(agent.policy, env)
+end
+
+#####
+# update!
+#####
+
 function (agent::Agent)(stage::AbstractStage, env::AbstractEnv)
     update!(agent.trajectory, agent.policy, env, stage)
     update!(agent.policy, agent.trajectory, env, stage)

src/policies/q_based_policies/q_based_policy.jl

@@ -33,6 +33,19 @@ RLBase.prob(p::QBasedPolicy, env, ::FullActionSet) =
 RLBase.update!(p::QBasedPolicy, trajectory::AbstractTrajectory) =
     update!(p.learner, trajectory)
 
+function check(p::QBasedPolicy, env::AbstractEnv)
+    A = action_space(env)
+    if (A isa AbstractVector && A == 1:length(A)) ||
+        (A isa Tuple && A == Tuple(1:length(A)))
+        # this is expected
+    else
+        @warn "Applying a QBasedPolicy to an environment with a unknown action space. Maybe convert the environment with `discrete2standard_discrete` in ReinforcementLearningEnvironments.jl first or redesign the environment."
+    end
+
+    check(p.learner, env)
+    check(p.explorer, env)
+end
+
 #####
 # TabularRandomPolicy
 #####