Log any and all columns of the dataset

main.py

@@ -4,30 +4,62 @@ from __future__ import annotations
 
 import rerun as rr
 from datasets import load_dataset
+from PIL import Image
+from tqdm import tqdm
 
-# download/load dataset in pyarrow format
-print("Loading dataset…")
-dataset = load_dataset("lerobot/pusht", split="train")
-
-# select the frames belonging to episode number 5
-print("Select specific episode…")
-ds_subset = dataset.filter(lambda frame: frame["episode_id"] == 5)
-
-print("Starting Rerun…")
-rr.init("rerun_example_lerobot", spawn=True)
-
-print("Logging to Rerun…")
-for frame_id, timestamp, image, state, action, next_reward in zip(
-    ds_subset["frame_id"],
-    ds_subset["timestamp"],
-    ds_subset["observation.image"],
-    ds_subset["observation.state"],
-    ds_subset["action"],
-    ds_subset["next.reward"],
-):
-    rr.set_time_sequence("frame_id", frame_id)
-    rr.set_time_seconds("timestamp", timestamp)
-    rr.log("observation/image", rr.Image(image))
-    rr.log("observation/state", rr.BarChart(state))
-    rr.log("observation/action", rr.BarChart(action))
-    rr.log("next/reward", rr.Scalar(next_reward))
+
+def log_dataset_to_rerun(dataset) -> None:
+    # Special time-like columns
+    TIME_LIKE = {"index", "frame_id", "timestamp"}
+
+    # Ignore these columns
+    IGNORE = {"episode_data_index_from", "episode_data_index_to", "episode_id"}
+
+    num_rows = len(dataset)
+    for row_nr in tqdm(range(num_rows)):
+        row = dataset[row_nr]
+
+        # Handle time-like columns first, since they set a state (time is an index in Rerun):
+        for column_name in TIME_LIKE:
+            if column_name in row:
+                cell = row[column_name]
+                if isinstance(cell, int):
+                    rr.set_time_sequence(column_name, cell)
+                elif isinstance(cell, float):
+                    rr.set_time_seconds(column_name, cell)  # assume seconds
+                else:
+                    print(f"Unknown time-like column {column_name} with value {cell}")
+
+        # Now log actual data columns
+        for column_name in dataset.column_names:
+            if column_name in TIME_LIKE or column_name in IGNORE:
+                continue
+
+            cell = row[column_name]
+            if isinstance(cell, Image.Image):
+                rr.log(column_name, rr.Image(cell))
+            elif isinstance(cell, list):
+                rr.log(column_name, rr.BarChart(cell))
+            elif isinstance(cell, float) or isinstance(cell, int):
+                rr.log(column_name, rr.Scalar(cell))
+            else:
+                rr.log(column_name, rr.TextDocument(str(cell)))
+
+
+def main():
+    print("Loading dataset…")
+    # dataset = load_dataset("lerobot/pusht", split="train")
+    dataset = load_dataset("lerobot/aloha_sim_transfer_cube_human", split="train")
+
+    print("Selecting specific episode…")
+    ds_subset = dataset.filter(lambda frame: frame["episode_id"] == 3)
+
+    print("Starting Rerun…")
+    rr.init("rerun_example_lerobot", spawn=True)
+
+    print("Logging to Rerun…")
+    log_dataset_to_rerun(ds_subset)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -1,2 +1,4 @@
 datasets
+Pillow
 rerun-sdk>=0.15.0,<0.16.0
+tqdm