fix render issue

lerobot/common/datasets/abstract.py

@@ -44,9 +44,9 @@ class AbstractExperienceReplay(TensorDictReplayBuffer):
                 f"The version of the dataset ({self.version}) is not enforced when root is provided ({self.root})."
             )
 
-        # HACK
-        if dataset_id == "xarm_lift_medium":
-            self.data_dir = self.root / self.dataset_id
+        # HACK: to remove before merge
+        self.data_dir = self.root / self.dataset_id
+        if not (self.data_dir / "replay_buffer").exists():
             storage = self._download_and_preproc_obsolete()
         else:
             storage = self._download_or_load_dataset()

lerobot/common/datasets/simxarm.py

@@ -113,7 +113,11 @@ class SimxarmExperienceReplay(AbstractExperienceReplay):
 
             if episode_id == 0:
                 # hack to initialize tensordict data structure to store episodes
-                td_data = episode[0].expand(total_frames).memmap_like(self.root / f"{self.dataset_id}")
+                td_data = (
+                    episode[0]
+                    .expand(total_frames)
+                    .memmap_like(self.root / f"{self.dataset_id}" / "replay_buffer")
+                )
 
             td_data[idx0:idx1] = episode