[BugFix] Fix init_random_frames=0

ghstack-source-id: 38a544ea15631f9affb4c385c09e7c4df94af55d Pull Request resolved: #2645
pytorch · Dec 12, 2024 · 19dfefc · 19dfefc
1 parent b840a77
commit 19dfefc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/test/test_collector.py b/test/test_collector.py
@@ -1345,7 +1345,7 @@ def make_env():
         functools.partial(MultiSyncDataCollector, cat_results="stack"),
     ],
 )
-@pytest.mark.parametrize("init_random_frames", [50])  # 1226: faster execution
+@pytest.mark.parametrize("init_random_frames", [0, 50])  # 1226: faster execution
 @pytest.mark.parametrize(
     "explicit_spec,split_trajs", [[True, True], [False, False]]
 )  # 1226: faster execution

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -712,10 +712,10 @@ def __init__(
         )
         self.reset_at_each_iter = reset_at_each_iter
         self.init_random_frames = (
-            int(init_random_frames) if init_random_frames is not None else 0
+            int(init_random_frames) if init_random_frames not in (None, -1) else 0
         )
         if (
-            init_random_frames is not None
+            init_random_frames not in (-1, None, 0)
             and init_random_frames % frames_per_batch != 0
             and RL_WARNINGS
         ):