From c2d655599ee8754f1c2a4f5b73cb0a59121a8480 Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 12 Sep 2025 00:06:19 +0800 Subject: [PATCH 1/7] fix(vault): handle buffer full state when current_index resets to 0 - Add condition that ensure writing full buffer on first write - Ensure data is not lost when buffer becomes full and index resets --- flashbax/vault/vault.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index 290c84a..d4ed493 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -359,7 +359,13 @@ 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: + fbx_max_index = get_tree_shape_prefix(fbx_state.experience, n_axes=2)[1] + source_interval = (0, fbx_max_index) + 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 From 0a7b6a8f463a3808b68522f1a4db58c652d0aeb9 Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 12 Sep 2025 00:26:09 +0800 Subject: [PATCH 2/7] style: format code with black - Apply black formatting to vault.py - Fix long condition statement formatting --- flashbax/vault/vault.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index d4ed493..ac8871c 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -361,7 +361,11 @@ def write( if source_interval == (0, 0): # 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: + if ( + fbx_state.is_full + and fbx_current_index == 0 + and self._last_received_fbx_index == 0 + ): fbx_max_index = get_tree_shape_prefix(fbx_state.experience, n_axes=2)[1] source_interval = (0, fbx_max_index) else: From 59e5dbb93bd4d7d3f7c31b6805845685ba20e273 Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 12 Sep 2025 18:56:23 +0800 Subject: [PATCH 3/7] add warning for first full-buffer write. add get_max_index in utils. --- flashbax/utils.py | 12 ++++++++++++ flashbax/vault/vault.py | 10 ++++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/flashbax/utils.py b/flashbax/utils.py index bc34a15..c104b7f 100644 --- a/flashbax/utils.py +++ b/flashbax/utils.py @@ -92,3 +92,15 @@ 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 (BufferStateTypes): the buffer state to compute the maximum index for. + + Returns: + int: the maximum index of the buffer state. + """ + return get_tree_shape_prefix(buffer_state.experience, n_axes=2)[1] diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index ac8871c..3b62321 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -18,6 +18,7 @@ from ast import literal_eval as make_tuple from datetime import datetime from typing import Optional, Tuple, Union +import warnings import jax import jax.numpy as jnp @@ -366,8 +367,13 @@ def write( and fbx_current_index == 0 and self._last_received_fbx_index == 0 ): - fbx_max_index = get_tree_shape_prefix(fbx_state.experience, n_axes=2)[1] - source_interval = (0, fbx_max_index) + warnings.warn( + "Writing nothing: buffer is full and current_index is 0 and first write " + "(last_received_fbx_index=0). If you want to write the entire buffer, " + "you should specify source_interval explicitly.", + UserWarning, + stacklevel=2, + ) else: source_interval = (self._last_received_fbx_index, fbx_current_index) From 40f1aabd91cbe305105257c08646a67b0b2f2021 Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 12 Sep 2025 19:13:57 +0800 Subject: [PATCH 4/7] fix mypy type checking --- flashbax/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flashbax/utils.py b/flashbax/utils.py index c104b7f..3f3823d 100644 --- a/flashbax/utils.py +++ b/flashbax/utils.py @@ -98,9 +98,10 @@ def get_max_index(buffer_state: chex.ArrayTree) -> int: """Utility to compute the maximum index of the buffer state. Args: - buffer_state (BufferStateTypes): the buffer state to compute the maximum index for. + 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] From 9e9a1cb9a0f8e72052dfc49ef319ae6e30c8387f Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 12 Sep 2025 19:20:54 +0800 Subject: [PATCH 5/7] style: fix import order with isort --- flashbax/vault/vault.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index 3b62321..8a844d0 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -15,10 +15,10 @@ 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 -import warnings import jax import jax.numpy as jnp From 0cc080fe09327c0ee9756a8f243db20586358aab Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 19 Sep 2025 17:42:34 +0800 Subject: [PATCH 6/7] add example in warning to write entire buffer --- flashbax/vault/vault.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index 8a844d0..92c65ea 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -368,9 +368,10 @@ def write( and self._last_received_fbx_index == 0 ): warnings.warn( - "Writing nothing: buffer is full and current_index is 0 and first write " - "(last_received_fbx_index=0). If you want to write the entire buffer, " - "you should specify source_interval explicitly.", + "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, ) From 2649710e24ff608784aa9ac36eddead49bd55897 Mon Sep 17 00:00:00 2001 From: Dcyaprogrammer Date: Fri, 19 Sep 2025 23:48:51 +0800 Subject: [PATCH 7/7] fix flake8 format --- flashbax/vault/vault.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flashbax/vault/vault.py b/flashbax/vault/vault.py index 92c65ea..1ac1452 100644 --- a/flashbax/vault/vault.py +++ b/flashbax/vault/vault.py @@ -368,9 +368,9 @@ def write( 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: " + "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,