fix(vault): handle buffer full state when current_index resets to 0

flashbax/utils.py

@@ -92,3 +92,16 @@ def get_timestep_count(buffer_state: chex.ArrayTree) -> int:
     )
     timestep_count: int = b_size * t_size
     return timestep_count
+
+
+def get_max_index(buffer_state: chex.ArrayTree) -> int:
+    """Utility to compute the maximum index of the buffer state.
+
+    Args:
+        buffer_state: the buffer state to compute the maximum index for.
+
+    Returns:
+        int: the maximum index of the buffer state.
+    """
+    assert hasattr(buffer_state, "experience")
+    return get_tree_shape_prefix(buffer_state.experience, n_axes=2)[1]

flashbax/vault/vault.py

@@ -15,6 +15,7 @@
 import asyncio
 import json
 import os
+import warnings
 from ast import literal_eval as make_tuple
 from datetime import datetime
 from typing import Optional, Tuple, Union
@@ -359,7 +360,23 @@ def write(
 
         # By default, we read from `last received` to `current index`
         if source_interval == (0, 0):
-            source_interval = (self._last_received_fbx_index, fbx_current_index)
+            # Special case: if buffer is full and current_index is 0, and this is the first write
+            # (last_received_fbx_index is 0), we need to write the entire buffer
+            if (
+                fbx_state.is_full
+                and fbx_current_index == 0
+                and self._last_received_fbx_index == 0
+            ):
+                warnings.warn(
+                    "Writing nothing: Buffer is full and current_index has reset to 0, but this "
+                    "is the first write to vault (last_received_fbx_index=0). No data will be "
+                    "written. To write the entire buffer, specify source_interval explicitly: "
+                    "v.write(buffer_state, source_interval=(0, get_max_index(buffer_state)))",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            else:
+                source_interval = (self._last_received_fbx_index, fbx_current_index)
 
         # By default, we continue writing from the current vault index
         dest_start = self.vault_index if dest_start is None else dest_start