diff --git a/reverb/trajectory_writer.py b/reverb/trajectory_writer.py
index 301142d5..7e2b2dc6 100644
--- a/reverb/trajectory_writer.py
+++ b/reverb/trajectory_writer.py
@@ -120,7 +120,7 @@ def history(self):
     Raises:
       RuntimeError: If `append` hasn't been called at least once before.
     """
-    if self._structure is None:
+    if not self._column_history:
       raise RuntimeError(
           'history cannot be accessed before `append` is called at least once.')
 
diff --git a/reverb/trajectory_writer_test.py b/reverb/trajectory_writer_test.py
index ada59ac2..1171e949 100644
--- a/reverb/trajectory_writer_test.py
+++ b/reverb/trajectory_writer_test.py
@@ -55,6 +55,14 @@ def test_history_require_append_to_be_called_before(self):
     with self.assertRaises(RuntimeError):
       _ = self.writer.history
 
+  def test_history_contains_references_when_data_flat(self):
+    self.writer.append(1)
+    self.writer.append(2)
+    self.writer.append(3)
+
+    history = tree.map_structure(extract_data, self.writer.history)
+    self.assertListEqual(history, [1, 2, 3])
+
   def test_history_contains_structured_references(self):
     self.writer.append({'x': 1, 'y': 100})
     self.writer.append({'x': 2, 'y': 101})