Reduce training times in SQIL tests to make the test suite faster.

HumanCompatibleAI · Oct 9, 2023 · 9343344 · 9343344
1 parent 7b8b4bf
commit 9343344
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/algorithms/test_sqil.py b/tests/algorithms/test_sqil.py
@@ -90,7 +90,7 @@ def _test_sqil_no_crash(
         rl_algo_class=rl_algo_class,
         rl_kwargs=rl_kwargs,
     )
-    model.train(total_timesteps=5000)
+    model.train(total_timesteps=500)
 
 
 def test_sqil_no_crash_discrete(
@@ -104,7 +104,7 @@ def test_sqil_no_crash_discrete(
         cartpole_venv,
         "seals/CartPole-v0",
         rl_algo_class=dqn.DQN,
-        rl_kwargs=dict(learning_starts=1000),
+        rl_kwargs=dict(learning_starts=100),
     )
 
 
@@ -143,7 +143,7 @@ def _test_sqil_few_demonstrations(
         rl_algo_class=rl_algo_class,
         rl_kwargs=rl_kwargs,
     )
-    model.train(total_timesteps=1_000)
+    model.train(total_timesteps=1_00)
 
 
 def test_sqil_few_demonstrations_discrete(
@@ -203,7 +203,7 @@ def _test_sqil_performance(
         return_episode_rewards=True,
     )
 
-    model.train(total_timesteps=10_000)
+    model.train(total_timesteps=1_000)
 
     venv.seed(SEED)
     rewards_after, _ = evaluate_policy(