From c37b1d45b6bf5b26bf9e507aead77fb4839fa11a Mon Sep 17 00:00:00 2001
From: Ermano Arruda <ermano.arruda@gmail.com>
Date: Thu, 13 Mar 2025 09:28:29 +0000
Subject: [PATCH] parametrise tolerance_s in visualize_dataset scripts (#716)

---
 lerobot/scripts/visualize_dataset.py      | 14 +++++++++++++-
 lerobot/scripts/visualize_dataset_html.py | 18 +++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/lerobot/scripts/visualize_dataset.py b/lerobot/scripts/visualize_dataset.py
index 11feb1af..cdfea6b8 100644
--- a/lerobot/scripts/visualize_dataset.py
+++ b/lerobot/scripts/visualize_dataset.py
@@ -265,13 +265,25 @@ def main():
         ),
     )
 
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(repo_id, root=root)
+    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
 
     visualize_dataset(dataset, **vars(args))
 
diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py
index f944144a..0fc21a8f 100644
--- a/lerobot/scripts/visualize_dataset_html.py
+++ b/lerobot/scripts/visualize_dataset_html.py
@@ -446,15 +446,31 @@ def main():
         help="Delete the output directory if it exists already.",
     )
 
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     load_from_hf_hub = kwargs.pop("load_from_hf_hub")
     root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
 
     dataset = None
     if repo_id:
-        dataset = LeRobotDataset(repo_id, root=root) if not load_from_hf_hub else get_dataset_info(repo_id)
+        dataset = (
+            LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+            if not load_from_hf_hub
+            else get_dataset_info(repo_id)
+        )
 
     visualize_dataset_html(dataset, **vars(args))