fix size issue

Author: KeplerC

.gitignore

@@ -136,4 +136,5 @@ temp.gif
 
 *.vla
 *.mkv
-*.csv
+*.csv
+*.pdf

benchmarks/Visualization.ipynb

@@ -60,16 +60,23 @@ def __init__(
     def measure_average_trajectory_size(self):
         """Calculates the average size of trajectory files in the dataset directory."""
         total_size = 0
-        file_count = 0
         for dirpath, dirnames, filenames in os.walk(self.dataset_dir):
             for f in filenames:
-                if f.endswith(self.file_extension):
-                    file_path = os.path.join(dirpath, f)
-                    total_size += os.path.getsize(file_path)
-                    file_count += 1
-        if file_count == 0:
-            return 0
-        return (total_size / file_count) / (1024 * 1024)  # Convert to MB
+                file_path = os.path.join(dirpath, f)
+                total_size += os.path.getsize(file_path)
+        
+        print(f"total_size: {total_size} of directory {self.dataset_dir}")
+        # trajectory number 
+        traj_num = 0
+        if self.dataset_name == "nyu_door_opening_surprising_effectiveness":
+            traj_num = 435
+        if self.dataset_name == "berkeley_cable_routing":
+            traj_num = 1482
+        if self.dataset_name == "bridge":
+            traj_num = 25460
+        if self.dataset_name == "berkeley_autolab_ur5":
+            traj_num = 896
+        return (total_size / traj_num) / (1024 * 1024)  # Convert to MB
 
     def clear_cache(self):
         """Clears the cache directory."""
@@ -274,7 +281,7 @@ def __init__(
             exp_dir,
             dataset_name,
             num_batches,
-            dataset_type="lerobot",
+            dataset_type="hf",
             batch_size=batch_size,
             log_frequency=log_frequency,
         )

benchmarks/openx.py

@@ -1,11 +1,9 @@
 # ask for sudo access
 sudo echo "Use sudo access for clearning cache"
 
-rm *.csv
-
 # Define a list of batch sizes to iterate through
-batch_sizes=(1 2 4 6 8)
-num_batches=200
+batch_sizes=(1)
+num_batches=20
 # batch_sizes=(1 2)
 
 # batch_sizes=(2)

evaluation.sh

@@ -29,8 +29,12 @@ def _frame_to_numpy(frame):
                     # repeat
                     if self.episode_index >= len(self.dataset):
                         self.episode_index = 0
-                    from_idx = self.dataset.episode_data_index["from"][self.episode_index].item()
-                    to_idx = self.dataset.episode_data_index["to"][self.episode_index].item()
+                    try:
+                        from_idx = self.dataset.episode_data_index["from"][self.episode_index].item()
+                        to_idx = self.dataset.episode_data_index["to"][self.episode_index].item()
+                    except Exception as e:
+                        self.episode_index = 0
+                        continue
                     frames = [_frame_to_numpy(self.dataset[idx]) for idx in range(from_idx, to_idx)]
                     episode.extend(frames)
                     self.episode_index += 1