Fix coverage

HumanCompatibleAI · Jul 7, 2023 · 551fa7e · 551fa7e
1 parent 29cdbfa
commit 551fa7e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/algorithms/test_sqil.py b/tests/algorithms/test_sqil.py
@@ -6,6 +6,7 @@
 import stable_baselines3.common.vec_env as vec_env
 import stable_baselines3.dqn as dqn
 from stable_baselines3 import ppo
+from stable_baselines3.common import policies
 from stable_baselines3.common.evaluation import evaluate_policy
 
 from imitation.algorithms import base as algo_base
@@ -38,6 +39,8 @@ def test_sqil_demonstration_buffer(rng):
         policy=policy,
     )
 
+    assert isinstance(model.policy, policies.BasePolicy)
+
     assert isinstance(model.dqn.replay_buffer, sqil.SQILReplayBuffer)
     expert_buffer = model.dqn.replay_buffer.expert_buffer