From fa998c631bbca75af2b148fc6f70866455a6f85c Mon Sep 17 00:00:00 2001
From: Callum Tilbury <c.tilbury@instadeep.com>
Date: Wed, 28 Feb 2024 10:56:47 +0200
Subject: [PATCH] fix: correct current index.

---
 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 09cb4a7..f366dd5 100644
--- a/flashbax/vault/vault.py
+++ b/flashbax/vault/vault.py
@@ -478,6 +478,6 @@ def read(
         # Return the read result as a fbx buffer state
         return TrajectoryBufferState(
             experience=read_result,
-            current_index=jnp.array(self.vault_index, dtype=int),
+            current_index=jnp.array(0, dtype=int),
             is_full=jnp.array(True, dtype=bool),
         )