refactor(tweaks): Minor changes to DDPG and buffer methods.

- Removed 'flatten()' from action return. Can cause issues in future with vector environments
- Changed 'buffer.push()' -> 'buffer.add()' for better UX
- Changes 'state.unsqueeze(0)' in predict method to conditional
- Added 'pragma: no cover' comment to 'except ImportError', required for Python 3.11 compatibility. Not needed to test when have 'tox'

Author: Achronus

docs/learn/customize/buffers.md

@@ -26,9 +26,9 @@ buffer = ReplayBuffer(capacity=100_000, device=device)
 
 ???+ api "API Docs"
 
-    [`velora.buffer.ReplayBuffer.push(exp)`](../reference/buffer.md#velora.buffer.BufferBase.push)
+    [`velora.buffer.ReplayBuffer.add(exp)`](../reference/buffer.md#velora.buffer.BufferBase.add)
 
-To add an item, we `push()` a set of `Experience` to it:
+To add an item, we `add()` a set of `Experience` to it:
 
 ```python
 from velora.buffer import Experience
@@ -42,7 +42,7 @@ exp = Experience(
     done=False,
 )
 
-buffer.push(exp)
+buffer.add(exp)
 ```
 
 `Experience` is a simple dataclass that holds the information of a single environment `timestep`. We'll talk about them in more detail later.
@@ -126,7 +126,7 @@ exp = Experience(
     next_state=torch.zeros(state_dim, device=device),
     done=False,
 )
-buffer.push(exp)
+buffer.add(exp)
 
 # Get a batch
 batch = buffer.sample(batch_size=5)
@@ -158,11 +158,11 @@ buffer = RolloutBuffer(capacity=10, device=device)
 
 ???+ api "API Docs"
 
-    [`velora.buffer.RolloutBuffer.push(exp)`](../reference/buffer.md#velora.buffer.RolloutBuffer.push)
+    [`velora.buffer.RolloutBuffer.add(exp)`](../reference/buffer.md#velora.buffer.RolloutBuffer.add)
 
     [`velora.buffer.RolloutBuffer.empty()`](../reference/buffer.md#velora.buffer.RolloutBuffer.empty)
 
-To add an item, we `push()` a set of `Experience` to it:
+To add an item, we `add()` a set of `Experience` to it:
 
 ```python
 from velora.buffer import Experience
@@ -176,7 +176,7 @@ exp = Experience(
     done=False,
 )
 
-buffer.push(exp)
+buffer.add(exp)
 ```
 
 Once the buffer is full, we need to `empty` it before we can add new samples:
@@ -228,8 +228,8 @@ exp = Experience(
     done=False,
 )
 
-buffer.push(exp)
-# buffer.push(exp)  # BufferError
+buffer.add(exp)
+# buffer.add(exp)  # BufferError
 
 batch = buffer.sample()
 
@@ -374,7 +374,7 @@ for i_ep in range(n_episodes):
         done = terminated or truncated
 
         # Add it to the buffer
-        buffer.push(
+        buffer.add(
             Experience(state, action.item(), reward, next_state, done),
         )
 

tests/models/test_ddpg.py

@@ -172,7 +172,7 @@ def test_train_step(self, ddpg: LiquidDDPG):
 
             # Create Experience object explicitly
             exp = Experience(state, action, reward, next_state, done)
-            ddpg.buffer.push(exp)
+            ddpg.buffer.add(exp)
 
         # Perform training step
         result = ddpg._train_step(batch_size, gamma)
@@ -190,7 +190,7 @@ def test_train_step_insufficient_buffer(self, ddpg: LiquidDDPG):
         for _ in range(batch_size - 1):
             state = torch.zeros(ddpg.state_dim)
             exp = Experience(state, 1.0, 2.0, state, False)
-            ddpg.buffer.push(exp)
+            ddpg.buffer.add(exp)
 
         # Should return None when buffer is insufficient
         result = ddpg._train_step(batch_size, gamma)
@@ -228,7 +228,7 @@ def test_save_load_with_buffer(self, ddpg: LiquidDDPG):
             next_state = torch.ones(ddpg.state_dim)
             done = i == 9
             exp = Experience(state, action, reward, next_state, done)
-            ddpg.buffer.push(exp)
+            ddpg.buffer.add(exp)
 
         with tempfile.TemporaryDirectory() as temp_dir:
             filepath = os.path.join(temp_dir, "model.pt")
@@ -415,7 +415,7 @@ def patched_init(self, dirname, **kwargs):
 
                             # Mock buffer.push to prevent storing experiences
                             with (
-                                patch.object(ddpg.buffer, "push"),
+                                patch.object(ddpg.buffer, "add"),
                                 patch.object(ddpg.buffer, "warm"),
                             ):
                                 # Mock _train_step to avoid network operations
@@ -468,7 +468,7 @@ def test_early_stopping(self, ddpg: LiquidDDPG, env: gym.Env):
         # Mock necessary methods to avoid actual training
         with (
             patch.object(ddpg.buffer, "warm"),
-            patch.object(ddpg.buffer, "push"),
+            patch.object(ddpg.buffer, "add"),
             patch.object(ddpg, "_train_step", return_value=(0.1, 0.2)),
             patch.object(
                 ddpg,

tests/test_buffer.py

@@ -86,7 +86,7 @@ def sample_experience(self) -> Experience:
     def filled_buffer(self, replay_buffer: ReplayBuffer) -> ReplayBuffer:
         """Fixture that returns a replay buffer with 10 experiences."""
         for i in range(10):
-            replay_buffer.push(
+            replay_buffer.add(
                 Experience(
                     state=torch.tensor([float(i), float(i + 1)]),
                     action=torch.tensor([i]),
@@ -109,7 +109,7 @@ def test_config(self, replay_buffer: ReplayBuffer):
     def test_push_experience(
         self, replay_buffer: ReplayBuffer, sample_experience: Experience
     ) -> None:
-        replay_buffer.push(sample_experience)
+        replay_buffer.add(sample_experience)
         assert len(replay_buffer) == 1
         assert isinstance(replay_buffer.buffer[0], Experience)
 
@@ -118,13 +118,13 @@ def test_buffer_capacity(
     ) -> None:
         # Fill buffer beyond capacity
         for _ in range(150):
-            replay_buffer.push(sample_experience)
+            replay_buffer.add(sample_experience)
         assert len(replay_buffer) == 100  # Should not exceed capacity
 
     def test_sample_insufficient_experiences(
         self, replay_buffer: ReplayBuffer, sample_experience: Experience
     ) -> None:
-        replay_buffer.push(sample_experience)
+        replay_buffer.add(sample_experience)
         with pytest.raises(ValueError):
             replay_buffer.sample(batch_size=2)
 
@@ -133,7 +133,7 @@ def test_sample_batch(
     ) -> None:
         # Fill buffer with multiple experiences
         for _ in range(10):
-            replay_buffer.push(sample_experience)
+            replay_buffer.add(sample_experience)
 
         batch_size = 5
         batch = replay_buffer.sample(batch_size)
@@ -149,7 +149,7 @@ def test_len_method(
         self, replay_buffer: ReplayBuffer, sample_experience: Experience
     ) -> None:
         assert len(replay_buffer) == 0
-        replay_buffer.push(sample_experience)
+        replay_buffer.add(sample_experience)
         assert len(replay_buffer) == 1
 
     def test_state_dict_empty_buffer(self, replay_buffer: ReplayBuffer) -> None:
@@ -276,7 +276,7 @@ def test_buffer_warm(self):
             next_state=torch.zeros(state_dim, device=device),
             done=False,
         )
-        buffer.push(exp)
+        buffer.add(exp)
 
         # Verify buffer length increases
         assert len(buffer) == n_samples + 1
@@ -318,7 +318,7 @@ def sample_experience(self) -> Experience:
     def filled_buffer(self, rollout_buffer: RolloutBuffer) -> RolloutBuffer:
         """Fixture that returns a filled rollout buffer with 3 experiences."""
         for i in range(3):
-            rollout_buffer.push(
+            rollout_buffer.add(
                 Experience(
                     state=torch.tensor([float(i), float(i + 1)]),
                     action=torch.tensor([i]),
@@ -341,7 +341,7 @@ def test_config(self, rollout_buffer: RolloutBuffer):
     def test_push_experience(
         self, rollout_buffer: RolloutBuffer, sample_experience: Experience
     ) -> None:
-        rollout_buffer.push(sample_experience)
+        rollout_buffer.add(sample_experience)
         assert len(rollout_buffer) == 1
         assert isinstance(rollout_buffer.buffer[0], Experience)
 
@@ -350,11 +350,11 @@ def test_buffer_capacity_error(
     ) -> None:
         # Fill buffer to capacity
         for _ in range(5):
-            rollout_buffer.push(sample_experience)
+            rollout_buffer.add(sample_experience)
 
         # Attempt to push when buffer is full
         with pytest.raises(BufferError):
-            rollout_buffer.push(sample_experience)
+            rollout_buffer.add(sample_experience)
 
     def test_sample_empty_buffer(self, rollout_buffer: RolloutBuffer) -> None:
         with pytest.raises(BufferError) as exc_info:
@@ -367,7 +367,7 @@ def test_sample_buffer(
         # Fill buffer with experiences
         num_experiences = 3
         for _ in range(num_experiences):
-            rollout_buffer.push(sample_experience)
+            rollout_buffer.add(sample_experience)
 
         batch = rollout_buffer.sample()
 
@@ -384,7 +384,7 @@ def test_clear_buffer(
     ) -> None:
         # Add some experiences
         for _ in range(3):
-            rollout_buffer.push(sample_experience)
+            rollout_buffer.add(sample_experience)
         assert len(rollout_buffer) == 3
 
         # Clear buffer
@@ -395,9 +395,9 @@ def test_len_method(
         self, rollout_buffer: RolloutBuffer, sample_experience: Experience
     ) -> None:
         assert len(rollout_buffer) == 0
-        rollout_buffer.push(sample_experience)
+        rollout_buffer.add(sample_experience)
         assert len(rollout_buffer) == 1
-        rollout_buffer.push(sample_experience)
+        rollout_buffer.add(sample_experience)
         assert len(rollout_buffer) == 2
         rollout_buffer.empty()
         assert len(rollout_buffer) == 0
@@ -514,7 +514,7 @@ def test_empty_after_save(self, filled_buffer: RolloutBuffer) -> None:
             assert len(loaded_buffer) == 3  # Original size before emptying
 
             # Add more experiences to emptied buffer
-            filled_buffer.push(
+            filled_buffer.add(
                 Experience(
                     state=torch.tensor([10.0, 11.0]),
                     action=10.0,
@@ -534,7 +534,7 @@ def test_load_and_continue_filling(self) -> None:
         # Create and fill a buffer
         buffer = RolloutBuffer(capacity=5)
         for i in range(3):
-            buffer.push(
+            buffer.add(
                 Experience(
                     state=torch.tensor([float(i), float(i + 1)]),
                     action=torch.tensor([i]),
@@ -556,7 +556,7 @@ def test_load_and_continue_filling(self) -> None:
             assert len(loaded_buffer) == 3
 
             # Add more experiences
-            loaded_buffer.push(
+            loaded_buffer.add(
                 Experience(
                     state=torch.tensor([10.0, 11.0]),
                     action=torch.tensor([10.0]),
@@ -569,7 +569,7 @@ def test_load_and_continue_filling(self) -> None:
             assert len(loaded_buffer) == 4
 
             # Try to add experiences up to capacity
-            loaded_buffer.push(
+            loaded_buffer.add(
                 Experience(
                     state=torch.tensor([11.0, 12.0]),
                     action=torch.tensor([11.0]),
@@ -583,7 +583,7 @@ def test_load_and_continue_filling(self) -> None:
 
             # Should raise error on next push
             with pytest.raises(BufferError, match="Buffer full!"):
-                loaded_buffer.push(
+                loaded_buffer.add(
                     Experience(
                         state=torch.tensor([12.0, 13.0]),
                         action=torch.tensor([12.0]),

velora/buffer/base.py

@@ -27,15 +27,24 @@ def __init__(self, capacity: int, *, device: torch.device | None = None) -> None
         self.buffer: Deque[Experience] = deque(maxlen=capacity)
         self.device = device
 
-    def push(self, exp: Experience) -> None:
+    def add(self, exp: Experience) -> None:
         """
-        Stores an experience in the buffer.
+        Adds a single experience to the buffer.
 
         Parameters:
-            exp (Experience): a single set of experience as an object
+            exp (Experience): a single set of experience
         """
         self.buffer.append(exp)
 
+    def add_multi(self, exp: List[Experience]) -> None:
+        """
+        Adds multiple experiences to the buffer.
+
+        Parameters:
+            exp (List[Experience]): a list of experience
+        """
+        self.buffer.extend(exp)
+
     def _batch(self, batch: List[Experience]) -> BatchExperience:
         """
         Helper method. Converts a `List[Experience]` into a `BatchExperience`.
@@ -163,7 +172,7 @@ def load(cls, filepath: str | Path) -> Self:
             data["next_states"],
             data["dones"],
         ):
-            buffer.push(
+            buffer.add(
                 Experience(
                     state=to_tensor(state, device=device),
                     action=to_tensor(action, device=device),

velora/buffer/replay.py

@@ -3,7 +3,7 @@
 
 try:
     from typing import override
-except ImportError:
+except ImportError:  # pragma: no cover
     from typing_extensions import override  # pragma: no cover
 
 import gymnasium as gym
@@ -84,7 +84,7 @@ def warm(self, agent: RLAgent, env_name: str, n_samples: int) -> None:
             next_state, reward, terminated, truncated, _ = env.step(action)
             done = terminated or truncated
 
-            self.push(Experience(state, action, reward, next_state, done))
+            self.add(Experience(state, action, reward, next_state, done))
 
             state = next_state
 

velora/buffer/rollout.py

@@ -1,6 +1,6 @@
 try:
     from typing import override
-except ImportError:
+except ImportError:  # pragma: no cover
     from typing_extensions import override  # pragma: no cover
 
 import torch
@@ -36,7 +36,7 @@ def config(self) -> BufferConfig:
         return BufferConfig(type="RolloutBuffer", capacity=self.capacity)
 
     @override
-    def push(self, exp: Experience) -> None:
+    def add(self, exp: Experience) -> None:
         """
         Stores an experience in the buffer.
 
@@ -46,7 +46,7 @@ def push(self, exp: Experience) -> None:
         if len(self.buffer) == self.capacity:
             raise BufferError("Buffer full! Use the 'empty()' method first.")
 
-        super().push(exp)
+        super().add(exp)
 
     @override
     def sample(self) -> BatchExperience:

velora/callbacks.py

@@ -5,7 +5,7 @@
 
 try:
     from typing import override
-except ImportError:
+except ImportError:  # pragma: no cover
     from typing_extensions import override  # pragma: no cover
 
 import gymnasium as gym