diff --git a/examples/vault_demonstration.ipynb b/examples/vault_demonstration.ipynb
index b25aeca..e2ee7cf 100644
--- a/examples/vault_demonstration.ipynb
+++ b/examples/vault_demonstration.ipynb
@@ -1,354 +1,644 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Vault demonstration"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%capture\n",
-    "try:\n",
-    "    import flashbax as fbx\n",
-    "except ModuleNotFoundError:\n",
-    "    print('installing flashbax')\n",
-    "    %pip install -q flashbax\n",
-    "    import flashbax as fbx"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import jax\n",
-    "from typing import NamedTuple\n",
-    "import jax.numpy as jnp\n",
-    "from flashbax.vault import Vault\n",
-    "import flashbax as fbx\n",
-    "from chex import Array"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We create a simple timestep structure, with a corresponding flat buffer."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
+  "cells": [
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
-     ]
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KDbulg4KK1Jx"
+      },
+      "source": [
+        "# Vault demonstration"
+      ]
     },
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/claude/flashbax/flashbax/buffers/trajectory_buffer.py:473: UserWarning: Setting max_size dynamically sets the `max_length_time_axis` to be `max_size`//`add_batch_size = 5`.This allows one to control exactly how many timesteps are stored in the buffer.Note that this overrides the `max_length_time_axis` argument.\n",
-      "  warnings.warn(\n"
-     ]
-    }
-   ],
-   "source": [
-    "class FbxTransition(NamedTuple):\n",
-    "    obs: Array\n",
-    "\n",
-    "tx = FbxTransition(obs=jnp.zeros(shape=(2,)))\n",
-    "\n",
-    "buffer = fbx.make_flat_buffer(\n",
-    "    max_length=5,\n",
-    "    min_length=1,\n",
-    "    sample_batch_size=1,\n",
-    ")\n",
-    "buffer_state = buffer.init(tx)\n",
-    "buffer_add = jax.jit(buffer.add, donate_argnums=0)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The shape of this buffer is $(B = 1, T = 5, E = 2)$, meaning the buffer can hold 5 timesteps, where each observation is of shape $(2,)$."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Q4XRiiktK1Jy"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "try:\n",
+        "    import flashbax as fbx\n",
+        "except ModuleNotFoundError:\n",
+        "    print('installing flashbax')\n",
+        "    %pip install -q flashbax\n",
+        "    import flashbax as fbx"
+      ]
+    },
     {
-     "data": {
-      "text/plain": [
-       "(1, 5, 2)"
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8CGhMJSlK1Jz"
+      },
+      "outputs": [],
+      "source": [
+        "import jax\n",
+        "from typing import NamedTuple\n",
+        "import jax.numpy as jnp\n",
+        "from flashbax.vault import Vault\n",
+        "import flashbax as fbx\n",
+        "from chex import Array"
       ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "buffer_state.experience.obs.shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We create the vault, based on the buffer's experience structure."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
+    },
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "New vault created at /tmp/demo/20240205140817\n"
-     ]
-    }
-   ],
-   "source": [
-    "v = Vault(\n",
-    "    vault_name=\"demo\",\n",
-    "    experience_structure=buffer_state.experience,\n",
-    "    rel_dir=\"/tmp\"\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We now add 10 timesteps to the buffer, and write that buffer to the vault. We inspect the buffer and vault state after each timestep."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "teAD4hI1K1Jz"
+      },
+      "source": [
+        "We create a simple timestep structure, with a corresponding flat buffer."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "OCwy3AOOK1J0"
+      },
+      "outputs": [],
+      "source": [
+        "class FbxTransition(NamedTuple):\n",
+        "    obs: Array\n",
+        "\n",
+        "tx = FbxTransition(obs=jnp.zeros(shape=(2,)))\n",
+        "\n",
+        "buffer = fbx.make_flat_buffer(\n",
+        "    max_length=5,\n",
+        "    min_length=1,\n",
+        "    sample_batch_size=1,\n",
+        ")\n",
+        "buffer_state = buffer.init(tx)\n",
+        "buffer_add = jax.jit(buffer.add, donate_argnums=0)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ewG8Wpq0K1J0"
+      },
+      "source": [
+        "The shape of this buffer is $(B = 1, T = 5, E = 2)$, meaning the buffer can hold 5 timesteps, where each observation is of shape $(2,)$."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "slWCha6mK1J0",
+        "outputId": "fb3fe704-c8f6-46d5-bce4-42f0caedf37b"
+      },
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "(1, 5, 2)"
+            ]
+          },
+          "execution_count": 4,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "buffer_state.experience.obs.shape"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "J-LDhiqLK1J0"
+      },
+      "source": [
+        "We create the vault, based on the buffer's experience structure."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "6c9odXqbK1J1",
+        "outputId": "76bc3648-21a0-40f0-a8cd-ee2d92af6f24"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "New vault created at /tmp/demo/20250923012939\n",
+            "Since the provided buffer state has a temporal dimension of 5, you must write to the vault at least every 4 timesteps to avoid data loss.\n"
+          ]
+        }
+      ],
+      "source": [
+        "v = Vault(\n",
+        "    vault_name=\"demo\",\n",
+        "    experience_structure=buffer_state.experience,\n",
+        "    rel_dir=\"/tmp\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "65f0m-1XK1J1"
+      },
+      "source": [
+        "We now add 10 timesteps to the buffer, and write that buffer to the vault. We inspect the buffer and vault state after each timestep."
+      ]
+    },
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[1. 1.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [0. 0.]\n",
-      "  [0. 0.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [0. 0.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[6. 6.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]\n",
-      "  [6. 6.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[6. 6.]\n",
-      "  [7. 7.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]\n",
-      "  [6. 6.]\n",
-      "  [7. 7.]]]\n",
-      "------------------\n",
-      "------------------\n",
-      "Buffer state:\n",
-      "[[[6. 6.]\n",
-      "  [7. 7.]\n",
-      "  [8. 8.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]]]\n",
-      "\n",
-      "Vault state:\n",
-      "[[[1. 1.]\n",
-      "  [2. 2.]\n",
-      "  [3. 3.]\n",
-      "  [4. 4.]\n",
-      "  [5. 5.]\n",
-      "  [6. 6.]\n",
-      "  [7. 7.]\n",
-      "  [8. 8.]]]\n",
-      "------------------\n"
-     ]
+      "cell_type": "code",
+      "execution_count": 24,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "4VsxnEM_K1J1",
+        "outputId": "8e1159df-e2af-47c3-dff2-167c18609c27"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[1. 1.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [9. 9.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[6. 6.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]]]\n",
+            "------------------\n",
+            "------------------\n",
+            "Buffer state:\n",
+            "[[[6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]]]\n",
+            "\n",
+            "Vault state:\n",
+            "[[[1. 1.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]\n",
+            "  [9. 9.]\n",
+            "  [1. 1.]\n",
+            "  [2. 2.]\n",
+            "  [3. 3.]\n",
+            "  [4. 4.]\n",
+            "  [5. 5.]\n",
+            "  [6. 6.]\n",
+            "  [7. 7.]\n",
+            "  [8. 8.]]]\n",
+            "------------------\n"
+          ]
+        }
+      ],
+      "source": [
+        "for i in range(1, 10):\n",
+        "    print('------------------')\n",
+        "    print(\"Buffer state:\")\n",
+        "    print(buffer_state.experience.obs)\n",
+        "    print()\n",
+        "\n",
+        "    v.write(buffer_state)\n",
+        "\n",
+        "    print(\"Vault state:\")\n",
+        "    print(v.read().experience.obs)\n",
+        "    print('------------------')\n",
+        "\n",
+        "    buffer_state = buffer_add(\n",
+        "        buffer_state,\n",
+        "        FbxTransition(obs=i * jnp.ones(1))\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "50qFidTxK1J1"
+      },
+      "source": [
+        "Notice that when the buffer (implemented as a ring buffer) wraps around, the vault continues storing the data:\n",
+        "```\n",
+        "Buffer state:\n",
+        "[[[6. 6.]\n",
+        "  [2. 2.]\n",
+        "  [3. 3.]\n",
+        "  [4. 4.]\n",
+        "  [5. 5.]]]\n",
+        "\n",
+        "Vault state:\n",
+        "[[[1. 1.]\n",
+        "  [2. 2.]\n",
+        "  [3. 3.]\n",
+        "  [4. 4.]\n",
+        "  [5. 5.]\n",
+        "  [6. 6.]]]\n",
+        "```\n",
+        "\n",
+        "Note: the vault must be given the buffer state at least every `max_steps` number of timesteps (i.e. before stale data is overwritten in the ring buffer)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "oAo6U40wQG4v"
+      },
+      "source": [
+        "However, for offline learning scenario where you want to directly write in a full buffer, remember to specify the interval"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "53u-lEqzQ-jx"
+      },
+      "source": [
+        "Here we create a same buffer and add data until the buffer is full"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "9fTm-M1DN8YH",
+        "outputId": "e5989ac5-8f32-4e1e-fb38-d2b2555fb945"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "True\n"
+          ]
+        }
+      ],
+      "source": [
+        "MAX_LENGTH = 5\n",
+        "\n",
+        "fbuffer = fbx.make_flat_buffer(\n",
+        "    max_length=MAX_LENGTH,\n",
+        "    min_length=1,\n",
+        "    sample_batch_size=1,\n",
+        ")\n",
+        "fbuffer_state = fbuffer.init(tx)\n",
+        "fbuffer_add = jax.jit(fbuffer.add, donate_argnums=0)\n",
+        "\n",
+        "for i in range(MAX_LENGTH):\n",
+        "    fbuffer_state = fbuffer_add(\n",
+        "        fbuffer_state,\n",
+        "        FbxTransition(obs=i * jnp.ones(1))\n",
+        "    )\n",
+        "\n",
+        "print(fbuffer_state.is_full)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XYRoL8h1SDaK"
+      },
+      "source": [
+        "Now we create a new vault to write in the full buffer"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oDEbHhpQRT6x",
+        "outputId": "7abfcc50-e199-4dda-c728-faac7cf2151a"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "New vault created at /tmp/TestVault/20250923015646\n",
+            "Since the provided buffer state has a temporal dimension of 5, you must write to the vault at least every 4 timesteps to avoid data loss.\n",
+            "Vault data shape: (1, 0, 2)\n",
+            "Vault data is empty: True\n"
+          ]
+        }
+      ],
+      "source": [
+        "fv = Vault(\n",
+        "    vault_name=\"TestVault\",\n",
+        "    experience_structure=buffer_state.experience,\n",
+        "    rel_dir=\"/tmp\"\n",
+        ")\n",
+        "\n",
+        "fv.write(buffer_state)\n",
+        "vault_data = fv.read()\n",
+        "\n",
+        "print(\"Vault data shape:\", vault_data.experience.obs.shape)\n",
+        "print(\"Vault data is empty:\", jnp.all(vault_data.experience.obs == 0))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Sw_wMZ9tSgIK"
+      },
+      "source": [
+        "You can notice that in this way the write is unsuccessful, the correct way is to specify the interval:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "_Z0bTSQySqLO",
+        "outputId": "f570754b-e37c-41f4-bbc7-9e70bdfa3c8c"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "New vault created at /tmp/TestVault/20250923015840\n",
+            "Since the provided buffer state has a temporal dimension of 5, you must write to the vault at least every 4 timesteps to avoid data loss.\n",
+            "Vault data shape: (1, 5, 2)\n",
+            "Vault data is empty: False\n"
+          ]
+        }
+      ],
+      "source": [
+        "del fv\n",
+        "\n",
+        "fv = Vault(\n",
+        "    vault_name=\"TestVault\",\n",
+        "    experience_structure=buffer_state.experience,\n",
+        "    rel_dir=\"/tmp\"\n",
+        ")\n",
+        "\n",
+        "fv.write(buffer_state, source_interval=(0, MAX_LENGTH))\n",
+        "vault_data = fv.read()\n",
+        "\n",
+        "print(\"Vault data shape:\", vault_data.experience.obs.shape)\n",
+        "print(\"Vault data is empty:\", jnp.all(vault_data.experience.obs == 0))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bdKBIyU0TebU"
+      },
+      "source": [
+        "if you don't know the length of the buffer, you can use:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "qgwfqscFS10T",
+        "outputId": "b6e2cc02-1ae3-4376-9cd1-b199660f7327"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "5\n"
+          ]
+        }
+      ],
+      "source": [
+        "max_index = fbx.utils.get_tree_shape_prefix(fbuffer_state.experience, n_axes=2)[1]\n",
+        "print(max_index)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "flashbax",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.9.18"
     }
-   ],
-   "source": [
-    "for i in range(1, 10):\n",
-    "    print('------------------')\n",
-    "    print(\"Buffer state:\")\n",
-    "    print(buffer_state.experience.obs)\n",
-    "    print()\n",
-    "\n",
-    "    v.write(buffer_state)\n",
-    "\n",
-    "    print(\"Vault state:\")\n",
-    "    print(v.read().experience.obs)\n",
-    "    print('------------------')\n",
-    "\n",
-    "    buffer_state = buffer_add(\n",
-    "        buffer_state,\n",
-    "        FbxTransition(obs=i * jnp.ones(1))\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Notice that when the buffer (implemented as a ring buffer) wraps around, the vault continues storing the data:\n",
-    "```\n",
-    "Buffer state:\n",
-    "[[[6. 6.]\n",
-    "  [2. 2.]\n",
-    "  [3. 3.]\n",
-    "  [4. 4.]\n",
-    "  [5. 5.]]]\n",
-    "\n",
-    "Vault state:\n",
-    "[[[1. 1.]\n",
-    "  [2. 2.]\n",
-    "  [3. 3.]\n",
-    "  [4. 4.]\n",
-    "  [5. 5.]\n",
-    "  [6. 6.]]]\n",
-    "```\n",
-    "\n",
-    "Note: the vault must be given the buffer state at least every `max_steps` number of timesteps (i.e. before stale data is overwritten in the ring buffer)."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "flashbax",
-   "language": "python",
-   "name": "python3"
   },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.9.18"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
+  "nbformat": 4,
+  "nbformat_minor": 0
 }