Complain if the policy cannot think

AlignmentResearch · Feb 27, 2024 · e216715 · e216715
1 parent 8008ca8
commit e216715
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/stable_baselines3/common/evaluation.py b/stable_baselines3/common/evaluation.py
@@ -102,6 +102,9 @@ def evaluate_policy(
         with th.no_grad():
             if hasattr(model, "think_for_n_steps"):
                 states = model.think_for_n_steps(n_steps_to_think, observations, states, episode_starts)
+            else:
+                if n_steps_to_think > 0:
+                    raise TypeError(f"Policy {model} cannot think for longer than 0 steps.")
             actions, states = model.predict(
                 observations,  # type: ignore[arg-type]
                 state=states,