fix: reset before collect on watch

Author: li-plus

examples/atari/atari_c51.py

@@ -174,7 +174,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_dqn.py

@@ -216,7 +216,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_fqf.py

@@ -187,7 +187,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_iqn.py

@@ -184,7 +184,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_ppo.py

@@ -244,7 +244,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_qrdqn.py

@@ -178,7 +178,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_rainbow.py

@@ -219,7 +219,7 @@ def watch() -> None:
                 beta=args.beta,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/atari/atari_sac.py

@@ -227,7 +227,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/vizdoom/vizdoom_c51.py

@@ -180,7 +180,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)

examples/vizdoom/vizdoom_ppo.py

@@ -246,7 +246,7 @@ def watch() -> None:
                 stack_num=args.frames_stack,
             )
             collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
-            result = collector.collect(n_step=args.buffer_size)
+            result = collector.collect(n_step=args.buffer_size, reset_before_collect=True)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size
             buffer.save_hdf5(args.save_buffer_name)