[CI] Fix dreamer run in SOTA tests

ghstack-source-id: dfe3ab6fe0d29fcdcaf57f31f84d04e07e36bad3 Pull Request resolved: #2627
pytorch · Dec 3, 2024 · aed03fd · aed03fd
1 parent 3da76f0
commit aed03fd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/.github/unittest/linux_sota/scripts/test_sota.py b/.github/unittest/linux_sota/scripts/test_sota.py
@@ -190,12 +190,12 @@
   logger.backend=
 """,
     "dreamer": """python sota-implementations/dreamer/dreamer.py \
-  collector.total_frames=200 \
+  collector.total_frames=600 \
   collector.init_random_frames=10 \
   collector.frames_per_batch=200 \
   env.n_parallel_envs=1 \
   optimization.optim_steps_per_batch=1 \
-  logger.video=True \
+  logger.video=False \
   logger.backend=csv \
   replay_buffer.buffer_size=120 \
   replay_buffer.batch_size=24 \

diff --git a/sota-implementations/dreamer/dreamer.py b/sota-implementations/dreamer/dreamer.py
@@ -321,8 +321,10 @@ def compile_rssms(module):
 
         t_collect_init = time.time()
 
-    test_env.close()
-    train_env.close()
+    if not test_env.is_closed:
+        test_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     collector.shutdown()
 
     del test_env