Log more parts of the dataframe

main.py

@@ -6,16 +6,28 @@ import rerun as rr
 from datasets import load_dataset
 
 # 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)
 
-# load all frames in RAM in PIL format
-frames = ds_subset["observation.image"]
-
+print("Starting Rerun…")
 rr.init("rerun_example_lerobot", spawn=True)
 
-for i, frame in enumerate(frames):
-    rr.set_time_sequence("frame", i)
-    rr.log("observation/image", rr.Image(frame))
+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))