add random loading

benchmarks/openx.py

@@ -13,9 +13,9 @@
 # Constants
 DEFAULT_EXP_DIR = "/mnt/data/fog_x/"
 DEFAULT_NUMBER_OF_TRAJECTORIES = -1 # Load all trajectories
-DEFAULT_DATASET_NAMES = ["nyu_door_opening_surprising_effectiveness", "berkeley_cable_routing", "berkeley_autolab_ur5", "bridge"]
-#["nyu_door_opening_surprising_effectiveness"]
-CACHE_DIR = "/mnt/data/fog_x/cache/"
+# DEFAULT_DATASET_NAMES = ["nyu_door_opening_surprising_effectiveness", "berkeley_cable_routing", "berkeley_autolab_ur5", "bridge"]
+DEFAULT_DATASET_NAMES = ["nyu_door_opening_surprising_effectiveness"]
+CACHE_DIR = "/tmp/fog_x/cache/"
 DEFAULT_LOG_FREQUENCY = 20
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -103,6 +103,25 @@ def measure_loading_time(self):
                 print(f"RLDS - Loaded {i} trajectories, Time: {elapsed_time:.2f} s")
         return time.time() - start_time
 
+    def measure_random_loading_time(self, num_loads):
+        start_time = time.time()
+        loader = RLDSLoader(self.dataset_dir, split="train")
+        dataset_size = len(loader)
+        num_loads = num_loads * dataset_size
+
+        loader.ds = loader.ds.shuffle(buffer_size=num_loads)
+        # shuffled_ds = shuffled_ds.take(num_loads)
+
+        for i, data in enumerate(loader):
+            self._recursively_load_data(data)
+
+            elapsed_time = time.time() - start_time
+            self.write_result(f"RLDS-RandomLoad", elapsed_time, i)
+            if i % self.log_frequency == 0:
+                print(f"RLDS-RandomLoad - Loaded {i} random trajectories, Time: {elapsed_time:.2f} s")
+
+        return time.time() - start_time
+
 class VLAHandler(DatasetHandler):
     def __init__(self, exp_dir, dataset_name, num_trajectories, log_frequency=DEFAULT_LOG_FREQUENCY):
         super().__init__(exp_dir, dataset_name, num_trajectories, dataset_type="vla", log_frequency=log_frequency)
@@ -124,6 +143,26 @@ def measure_loading_time(self, mode="no_cache"):
                 print(f"Failed to load data: {e}")
         return time.time() - start_time
 
+    def measure_random_loading_time(self, num_loads):
+        start_time = time.time()
+        loader = VLALoader(self.dataset_dir, cache_dir=CACHE_DIR)
+        dataset_size = len(loader)
+        num_loads = num_loads * dataset_size
+
+        for i in range(num_loads):
+            random_index = np.random.randint(0, dataset_size)
+            data = loader[random_index]
+            try:
+                self._recursively_load_data(data.load(mode="cache"))
+                elapsed_time = time.time() - start_time
+                self.write_result(f"VLA-RandomLoad", elapsed_time, i + 1)
+                if (i + 1) % self.log_frequency == 0:
+                    print(f"VLA-RandomLoad - Loaded {i + 1} random trajectories, Time: {elapsed_time:.2f} s")
+            except Exception as e:
+                print(f"Failed to load data: {e}")
+
+        return time.time() - start_time
+
 class FFV1Handler(DatasetHandler):
     def __init__(self, exp_dir, dataset_name, num_trajectories, log_frequency=DEFAULT_LOG_FREQUENCY):
         super().__init__(exp_dir, dataset_name, num_trajectories, dataset_type="ffv1", log_frequency=log_frequency)
@@ -145,6 +184,26 @@ def measure_loading_time(self, mode="no_cache"):
                 print(f"Failed to load data: {e}")
         return time.time() - start_time
 
+    def measure_random_loading_time(self, num_loads):
+        start_time = time.time()
+        loader = VLALoader(self.dataset_dir, cache_dir=CACHE_DIR)
+        dataset_size = len(loader)
+        num_loads = num_loads * dataset_size
+
+        for i in range(num_loads):
+            random_index = np.random.randint(0, dataset_size)
+            data = loader[random_index]
+            try:
+                self._recursively_load_data(data.load(mode="cache"))
+                elapsed_time = time.time() - start_time
+                self.write_result(f"FFV1-RandomLoad", elapsed_time, i + 1)
+                if (i + 1) % self.log_frequency == 0:
+                    print(f"FFV1-RandomLoad - Loaded {i + 1} random trajectories, Time: {elapsed_time:.2f} s")
+            except Exception as e:
+                print(f"Failed to load data: {e}")
+
+        return time.time() - start_time
+
 
 class HDF5Handler(DatasetHandler):
     def __init__(self, exp_dir, dataset_name, num_trajectories, log_frequency=DEFAULT_LOG_FREQUENCY):
@@ -164,6 +223,24 @@ def measure_loading_time(self):
                 print(f"HDF5 - Loaded {i} trajectories, Time: {elapsed_time:.2f} s")
         return time.time() - start_time
 
+    def measure_random_loading_time(self, num_loads):
+        start_time = time.time()
+        loader = HDF5Loader(path=os.path.join(self.dataset_dir, "*.h5"))
+        dataset_size = len(loader)
+        num_loads = num_loads * dataset_size
+
+        for i in range(num_loads):
+            random_index = np.random.randint(0, dataset_size)
+            data = loader[random_index]
+            self._recursively_load_data(data)
+
+            elapsed_time = time.time() - start_time
+            self.write_result(f"HDF5-RandomLoad", elapsed_time, i + 1)
+            if (i + 1) % self.log_frequency == 0:
+                print(f"HDF5-RandomLoad - Loaded {i + 1} random trajectories, Time: {elapsed_time:.2f} s")
+
+        return time.time() - start_time
+
 def prepare(args):
     # Clear the cache directory
     if os.path.exists(CACHE_DIR):
@@ -194,39 +271,48 @@ def evaluation(args):
             handler.clear_os_cache()
 
             avg_traj_size = handler.measure_average_trajectory_size()
-            loading_time = handler.measure_loading_time()
+            # loading_time = handler.measure_loading_time()
+
+            # new_results.append({
+            #     'Dataset': dataset_name,
+            #     'Format': handler.dataset_type.upper(),
+            #     'AverageTrajectorySize(MB)': avg_traj_size,
+            #     'LoadingTime(s)': loading_time,
+            # })
+
+            # print(f"{handler.dataset_type.upper()} - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {loading_time:.2f} s")
 
+            random_load_time = handler.measure_random_loading_time(args.random_loads)
             new_results.append({
                 'Dataset': dataset_name,
-                'Format': handler.dataset_type.upper(),
+                'Format': f"{handler.dataset_type.upper()}-RandomLoad",
                 'AverageTrajectorySize(MB)': avg_traj_size,
-                'LoadingTime(s)': loading_time,
+                'LoadingTime(s)': random_load_time,
             })
+            print(f"{handler.dataset_type.upper()}-RandomLoad - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {random_load_time:.2f} s")
+
+        # # Additional VLA measurements
+        # vla_handler = handlers[1]
+        # vla_handler.clear_cache()
+        # vla_handler.clear_os_cache()
+        # cold_cache_time = vla_handler.measure_loading_time(mode="cache")
+        # hot_cache_time = vla_handler.measure_loading_time(mode="cache")
+
+        # new_results.append({
+        #     'Dataset': dataset_name,
+        #     'Format': 'VLA-ColdCache',
+        #     'AverageTrajectorySize(MB)': avg_traj_size,
+        #     'LoadingTime(s)': cold_cache_time,
+        # })
 
-            print(f"{handler.dataset_type.upper()} - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {loading_time:.2f} s")
-
-        # Additional VLA measurements
-        vla_handler = handlers[1]
-        vla_handler.clear_cache()
-        vla_handler.clear_os_cache()
-        cold_cache_time = vla_handler.measure_loading_time(mode="cache")
-        hot_cache_time = vla_handler.measure_loading_time(mode="cache")
-
-        new_results.append({
-            'Dataset': dataset_name,
-            'Format': 'VLA-ColdCache',
-            'AverageTrajectorySize(MB)': avg_traj_size,
-            'LoadingTime(s)': cold_cache_time,
-        })
-
-        new_results.append({
-            'Dataset': dataset_name,
-            'Format': 'VLA-HotCache',
-            'AverageTrajectorySize(MB)': avg_traj_size,
-            'LoadingTime(s)': hot_cache_time,
-        })
-        print(f"VLA-ColdCache - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {cold_cache_time:.2f} s")
-        print(f"VLA-HotCache - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {hot_cache_time:.2f} s")
+        # new_results.append({
+        #     'Dataset': dataset_name,
+        #     'Format': 'VLA-HotCache',
+        #     'AverageTrajectorySize(MB)': avg_traj_size,
+        #     'LoadingTime(s)': hot_cache_time,
+        # })
+        # print(f"VLA-ColdCache - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {cold_cache_time:.2f} s")
+        # print(f"VLA-HotCache - Average Trajectory Size: {avg_traj_size:.2f} MB, Loading Time: {hot_cache_time:.2f} s")
 
         # Combine existing and new results
         all_results = existing_results + new_results
@@ -243,6 +329,7 @@ def evaluation(args):
     parser.add_argument("--dataset_names", nargs="+", default=DEFAULT_DATASET_NAMES, help="List of dataset names to evaluate.")
     parser.add_argument("--prepare", action="store_true", help="Prepare the datasets before evaluation.")
     parser.add_argument("--log_frequency", type=int, default=DEFAULT_LOG_FREQUENCY, help="Frequency of logging results.")
+    parser.add_argument("--random_loads", type=int, default=2, help="Number of random loads to perform for each loader.")
     args = parser.parse_args()
 
     if args.prepare:

fog_x/loader/hdf5.py

@@ -30,6 +30,9 @@ def __init__(self, path, split = None):
         self.index = 0
         self.files = glob.glob(self.path, recursive=True)
 
+    def __getitem__(self, idx):
+        return self._read_hdf5(self.files[idx])
+
     def _read_hdf5(self, data_path):
 
         with h5py.File(data_path, "r") as f:
@@ -52,4 +55,6 @@ def __next__(self):
             self.index += 1
             return self._read_hdf5(file_path)
         raise StopIteration
-
+
+    def __len__(self):
+        return len(self.files)

fog_x/loader/rlds.py

@@ -23,6 +23,12 @@ def __init__(self, path, split):
         self.index = 0
 
     def __len__(self):
+        try:
+            import tensorflow as tf
+            import tensorflow_datasets as tfds
+        except ImportError:
+            raise ImportError("Please install tensorflow and tensorflow_datasets to use rlds loader")
+
         return tf.data.experimental.cardinality(self.ds).numpy()
 
     def __iter__(self):
@@ -48,4 +54,7 @@ def __next__(self):
         except StopIteration:
             self.index = 0
             self.iterator = iter(self.ds)
-            raise StopIteration
+            raise StopIteration
+
+    def __getitem__(self, idx):
+        return next(iter(self.ds.skip(idx).take(1)))

fog_x/loader/vla.py

@@ -50,5 +50,8 @@ def __next__(self):
     def __len__(self):
         return len(self.files)
 
-    def peak(self, index):
-        return self._read_vla(self.files[index])
+    def __getitem__(self, index):
+        return self._read_vla(self.files[index])
+
+    def peak(self):
+        return self._read_vla(self.files[self.index])

openx_to_vla.sh

@@ -31,7 +31,7 @@
 
 # nyu_door_opening_surprising_effectiveness dataset
 # python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name nyu_door_opening_surprising_effectiveness --destination_dir /mnt/data/fog_x/vla --version 0.1.0 --split train[0:] --max_workers 4
-# python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name nyu_door_opening_surprising_effectiveness --destination_dir /mnt/data/fog_x/ffv1 --version 0.1.0 --split train[0:] --max_workers 4  --lossless
+python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name nyu_door_opening_surprising_effectiveness --destination_dir /mnt/data/fog_x/ffv1 --version 0.1.0 --split train[0:] --max_workers 4  --lossless
 
 # python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name bridge --destination_dir /mnt/data/fog_x/vla --version 0.1.0 --split train[0:] --max_workers 4
-python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name bridge --destination_dir /mnt/data/fog_x/ffv1 --version 0.1.0 --split train[0:] --max_workers 4 --lossless
+# python examples/openx_loader.py --data_dir /home/kych/datasets/rtx --dataset_name bridge --destination_dir /mnt/data/fog_x/ffv1 --version 0.1.0 --split train[0:] --max_workers 4 --lossless