fix(buffer): Fixed warm method bug when num_envs=1.

Vectorized environments don't like having a single environment. Added condition to ensure warm method requires a minimum of '2' envs.

Author: Ryan Partridge

tests/test_buffer.py

@@ -264,10 +264,19 @@ def test_buffer_warm(self, replay_buffer: ReplayBuffer):
         assert len(replay_buffer) == 0
 
         n_samples = 10
-        model.buffer.warm(model, n_samples, 2)
+        model.buffer.warm(model, n_samples)
 
         assert len(model.buffer) >= n_samples
 
+    def test_buffer_warm_single_env(self, replay_buffer: ReplayBuffer):
+        model = NeuroFlowCT("InvertedPendulum-v5", 8, 16, device=torch.device("cpu"))
+        assert len(replay_buffer) == 0
+
+        n_samples = 10
+
+        with pytest.raises(ValueError):
+            model.buffer.warm(model, n_samples, 1)
+
     def test_add_multi(self, replay_buffer: ReplayBuffer):
         """Test adding multiple experiences at once using add_multi."""
         # Create test data: batch of 5 experiences

velora/buffer/replay.py

@@ -95,8 +95,12 @@ def warm(self, agent: "RLModuleAgent", n_samples: int, num_envs: int = 8) -> Non
         Parameters:
             agent (Any): the agent to generate samples with
             n_samples (int): the maximum number of samples to generate
-            num_envs (int, optional): number of vectorized environments
+            num_envs (int, optional): number of vectorized environments. Cannot
+                be smaller than `2`
         """
+        if num_envs < 2:
+            raise ValueError(f"'{num_envs=}' cannot be smaller than 2.")
+
         envs = gym.make_vec(
             agent.env.spec.id,
             num_envs=num_envs,
@@ -115,6 +119,7 @@ def warm(self, agent: "RLModuleAgent", n_samples: int, num_envs: int = 8) -> Non
             dones = terminated | truncated
 
             self.add_multi(states, actions, rewards, next_states, dones, hidden)
+
             states = next_states
 
         envs.close()

velora/models/nf/agent.py

@@ -321,7 +321,7 @@ def train(
         )
 
         if warmup_steps > 0:
-            self.buffer.warm(self, warmup_steps, 1 if warmup_steps < 8 else 8)
+            self.buffer.warm(self, warmup_steps, 2 if warmup_steps < 8 else 8)
 
         with TrainHandler(
             self, n_episodes, max_steps, log_freq, window_size, callbacks
@@ -716,7 +716,7 @@ def train(
         )
 
         if warmup_steps > 0:
-            self.buffer.warm(self, warmup_steps, 1 if warmup_steps < 8 else 8)
+            self.buffer.warm(self, warmup_steps, 2 if warmup_steps < 8 else 8)
 
         with TrainHandler(
             self, n_episodes, max_steps, log_freq, window_size, callbacks