[Feature] init_random_frames to prefill replay buffers

facebookresearch · Dec 6, 2023 · c73e94a · c73e94a
1 parent 0d54682
commit c73e94a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/benchmarl/conf/experiment/base_experiment.yaml b/benchmarl/conf/experiment/base_experiment.yaml
@@ -67,6 +67,8 @@ off_policy_n_optimizer_steps: 1000
 off_policy_train_batch_size: 128
 # Maximum number of frames to keep in replay buffer memory for off-policy algorithms
 off_policy_memory_size: 1_000_000
+# Number of random action frames to prefill the replay buffer with
+off_policy_init_random_frames: 0
 
 
 evaluation: True

diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
@@ -80,6 +80,7 @@ class ExperimentConfig:
     off_policy_n_optimizer_steps: int = MISSING
     off_policy_train_batch_size: int = MISSING
     off_policy_memory_size: int = MISSING
+    off_policy_init_random_frames: int = MISSING
 
     evaluation: bool = MISSING
     render: bool = MISSING
@@ -447,6 +448,9 @@ def _setup_collector(self):
             storing_device=self.config.train_device,
             frames_per_batch=self.config.collected_frames_per_batch(self.on_policy),
             total_frames=self.config.get_max_n_frames(self.on_policy),
+            init_random_frames=self.config.off_policy_init_random_frames
+            if not self.on_policy
+            else 0,
         )
 
     def _setup_name(self):