From edacc1d2a0309ad232db71e18ec807c4c5e98786 Mon Sep 17 00:00:00 2001
From: Cadene <re.cadene@gmail.com>
Date: Tue, 26 Mar 2024 10:40:06 +0000
Subject: [PATCH] Add root in example

---
 examples/1_visualize_dataset.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/1_visualize_dataset.py b/examples/1_visualize_dataset.py
index ac62757..f52ab76 100644
--- a/examples/1_visualize_dataset.py
+++ b/examples/1_visualize_dataset.py
@@ -1,3 +1,5 @@
+import os
+
 from torchrl.data.replay_buffers import SamplerWithoutReplacement
 
 import lerobot
@@ -10,7 +12,7 @@ print(lerobot.available_datasets)
 # we use this sampler to sample 1 frame after the other
 sampler = SamplerWithoutReplacement(shuffle=False)
 
-dataset = AlohaDataset("aloha_sim_transfer_cube_human", sampler=sampler)
+dataset = AlohaDataset("aloha_sim_transfer_cube_human", sampler=sampler, root=os.environ.get("DATA_DIR"))
 
 video_paths = render_dataset(
     dataset,