From 2d71d0295de154f55bc4d5753226cd9d55ce2548 Mon Sep 17 00:00:00 2001 From: Albin Cassirer Date: Thu, 18 Feb 2021 07:54:49 -0800 Subject: [PATCH] Fix bug where `TrajectoryWriter.history` didn't work for unstructured data (i.e a single value per step). PiperOrigin-RevId: 358177758 Change-Id: I75d5a0d9ccbdb90219f0509ef00f292ca50007c0 --- reverb/trajectory_writer.py | 2 +- reverb/trajectory_writer_test.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) 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})