From 10706ed75344916c70af814ce39b5e58a73eab9d Mon Sep 17 00:00:00 2001
From: Harsimrat Sandhawalia <10599550+sandhawalia@users.noreply.github.com>
Date: Thu, 6 Mar 2025 10:27:29 +0100
Subject: [PATCH] Support for discrete actions (#810)

---
 lerobot/scripts/visualize_dataset_html.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py
index d5825aa6..f944144a 100644
--- a/lerobot/scripts/visualize_dataset_html.py
+++ b/lerobot/scripts/visualize_dataset_html.py
@@ -234,7 +234,7 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     This file will be loaded by Dygraph javascript to plot data in real time."""
     columns = []
 
-    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] == "float32"]
+    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
     selected_columns.remove("timestamp")
 
     ignored_columns = []