diff --git a/vla/benchmarks/block_type_classification/config.py b/vla/benchmarks/block_type_classification/config.py
index 63cbaee..62c64d1 100644
--- a/vla/benchmarks/block_type_classification/config.py
+++ b/vla/benchmarks/block_type_classification/config.py
@@ -8,6 +8,9 @@ class LabelType(str, Enum):
 class ConfigPaths:
     path_to_raw_data = None
     feature_store_path = "./precomputed"
+    generalization_set_folder = ""
+    train_json = "train_los_dataset.json"
+    test_json = "test_los_dataset.json"
 
 class TrainConfig:
     learning_rate = 1e-3
@@ -18,6 +21,7 @@ class TrainConfig:
     layers_sizes = [512]
 
 def construct_configs(path_to_raw_data=None, feature_store_path=None,
+        generalization_set_folder="", train_json="train_los_dataset.json", test_json="test_los_dataset.json",
         learning_rate=1e-3, momentum=0.9, epochs=10, output_activation="none", num_layers=1, layers_sizes=[512]):
     TrainConfig.learning_rate = learning_rate
     TrainConfig.momentum = momentum
@@ -33,4 +37,7 @@ def construct_configs(path_to_raw_data=None, feature_store_path=None,
         warnings.warn("Feature store path is not set by config. If benchmark will be ran with use_precomputed_features "
                       "default folder ./precomputed will be used to store features")
     else:
-        ConfigPaths.feature_store_path = feature_store_path
\ No newline at end of file
+        ConfigPaths.feature_store_path = feature_store_path
+    ConfigPaths.generalization_set_folder = generalization_set_folder
+    ConfigPaths.train_json = train_json
+    ConfigPaths.test_json = test_json
\ No newline at end of file
diff --git a/vla/benchmarks/block_type_classification/dataset_index_creation_example.py b/vla/benchmarks/block_type_classification/dataset_index_creation_example.py
index fc2af68..a511a20 100644
--- a/vla/benchmarks/block_type_classification/dataset_index_creation_example.py
+++ b/vla/benchmarks/block_type_classification/dataset_index_creation_example.py
@@ -1,10 +1,15 @@
+import json
 from manager import DatasetManager
-from config import LabelType
+from config import LabelType, ConfigPaths, construct_configs
 
 
 if __name__ == "__main__":
+    with open("example_config.json", 'r') as f:
+        config_dict = json.load(f)
+    construct_configs(**config_dict)
+
     manager = DatasetManager(
-        directory="./2025_LOS/Day_clear/",
+        directory=ConfigPaths.path_to_raw_data,
         label_type=LabelType.DISTANCE,
         dataset_size=200,
         train_split=0.25,
diff --git a/vla/benchmarks/block_type_classification/example_config.json b/vla/benchmarks/block_type_classification/example_config.json
index e3d8639..b57509d 100644
--- a/vla/benchmarks/block_type_classification/example_config.json
+++ b/vla/benchmarks/block_type_classification/example_config.json
@@ -1,6 +1,9 @@
 {
   "path_to_raw_data" : "./2025_LOS",
   "feature_store_path" : "./precomputed",
+  "generalization_set_folder" : "./2025_LOS/Night_clear/Mountain_Range",
+  "train_json" : "train_los_dataset.json",
+  "test_json" : "test_los_dataset.json",
   "learning_rate" : 1e-3,
   "momentum" : 0.9,
   "epochs" : 10,
diff --git a/vla/benchmarks/block_type_classification/test_dino_example.py b/vla/benchmarks/block_type_classification/test_dino_example.py
index 740609e..3a66581 100644
--- a/vla/benchmarks/block_type_classification/test_dino_example.py
+++ b/vla/benchmarks/block_type_classification/test_dino_example.py
@@ -9,6 +9,5 @@
         preprocessor,
         label_type=LabelType.DISTANCE,
         random_seed=10,
-        generalization_set_folder="./2025_LOS/Night_clear/Mountain_Range",
         config_path="example_config.json"
     )
diff --git a/vla/benchmarks/block_type_classification/test_eagle_example.py b/vla/benchmarks/block_type_classification/test_eagle_example.py
index 98ddd11..9988350 100755
--- a/vla/benchmarks/block_type_classification/test_eagle_example.py
+++ b/vla/benchmarks/block_type_classification/test_eagle_example.py
@@ -43,8 +43,7 @@ def encode(self, x):
 benchmark(
     ModelWrapper(),
     preprocess,
-    label_type=LabelType.DISTANCE, #LabelType.TYPE_CLASSIFICATION, #
+    label_type=LabelType.DISTANCE,
     random_seed=10,
-    generalization_set_folder="./2025_LOS/Night_clear/Mountain_Range",
     config_path="example_config.json"
 )
diff --git a/vla/benchmarks/block_type_classification/test_llava_example.py b/vla/benchmarks/block_type_classification/test_llava_example.py
index d1bd8a8..6c555d3 100755
--- a/vla/benchmarks/block_type_classification/test_llava_example.py
+++ b/vla/benchmarks/block_type_classification/test_llava_example.py
@@ -35,9 +35,8 @@ def encode(self, x):
 benchmark(
     ModelWrapper(),
     preprocess,
-    label_type=LabelType.TYPE_CLASSIFICATION, #LabelType.DISTANCE, 
+    label_type=LabelType.TYPE_CLASSIFICATION,
     random_seed=10,
-    generalization_set_folder="./2025_LOS/Night_clear/Mountain_Range",
     config_path="example_config.json"
 )
  
\ No newline at end of file
diff --git a/vla/benchmarks/block_type_classification/train.py b/vla/benchmarks/block_type_classification/train.py
index d4bace5..ea32435 100644
--- a/vla/benchmarks/block_type_classification/train.py
+++ b/vla/benchmarks/block_type_classification/train.py
@@ -153,9 +153,8 @@ def score(model, classifier, loader, label_type: LabelType, device):
         raise ValueError(f"Unsupported label type {label_type}")
 
 
-def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_json="test_los_dataset.json",
-              label_type=LabelType.DISTANCE, use_precomputed_features=True, random_seed=None,
-              generalization_set_folder="", config_path="example_config.json"):
+def benchmark(model, preprocessor, label_type=LabelType.DISTANCE,
+              use_precomputed_features=True, random_seed=None, config_path="example_config.json"):
 
     with open(config_path, 'r') as file:
         config_dict = json.load(file)
@@ -166,9 +165,9 @@ def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_jso
         dataset_manager = DatasetManager(ConfigPaths.path_to_raw_data, label_type, random_seed=random_seed)
 
     generalization_dataset_manager = None
-    if generalization_set_folder != "":
-        generalization_dataset_manager = DatasetManager(generalization_set_folder, label_type, random_seed=random_seed,
-                                                        full_folder=True)
+    if ConfigPaths.generalization_set_folder:
+        generalization_dataset_manager = DatasetManager(ConfigPaths.generalization_set_folder, label_type,
+                                                        random_seed=random_seed, full_folder=True)
 
     store = FeatureStore(ConfigPaths.feature_store_path)
 
@@ -176,8 +175,8 @@ def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_jso
 
     train_loader, test_loader, score_loader, generalization_loader, num_classes = \
         make_dataloaders(
-            train_json,
-            test_json,
+            ConfigPaths.train_json,
+            ConfigPaths.test_json,
             label_type,
             preprocessor,
             feature_store=(store if use_precomputed_features else None),
diff --git a/vla/benchmarks/class_distribution_using_segmentation/config.py b/vla/benchmarks/class_distribution_using_segmentation/config.py
index 1d54b36..3133202 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/config.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/config.py
@@ -4,6 +4,9 @@
 class ConfigPaths:
     path_to_raw_data = None
     feature_store_path = "./precomputed"
+    generalization_set_folder = ""
+    train_json = "train_los_dataset.json"
+    test_json = "test_los_dataset.json"
 
 class TrainConfig:
     learning_rate = 1e-3
@@ -14,6 +17,7 @@ class TrainConfig:
     layers_sizes = [512]
 
 def construct_configs(path_to_raw_data=None, feature_store_path=None,
+        generalization_set_folder="", train_json="train_los_dataset.json", test_json="test_los_dataset.json",
         learning_rate=1e-3, momentum=0.9, epochs=10, output_activation="none", num_layers=1, layers_sizes=[512]):
     TrainConfig.learning_rate = learning_rate
     TrainConfig.momentum = momentum
@@ -29,4 +33,7 @@ def construct_configs(path_to_raw_data=None, feature_store_path=None,
         warnings.warn("Feature store path is not set by config. If benchmark will be ran with use_precomputed_features "
                       "default folder ./precomputed will be used to store features")
     else:
-        ConfigPaths.feature_store_path = feature_store_path
\ No newline at end of file
+        ConfigPaths.feature_store_path = feature_store_path
+    ConfigPaths.generalization_set_folder = generalization_set_folder
+    ConfigPaths.train_json = train_json
+    ConfigPaths.test_json = test_json
\ No newline at end of file
diff --git a/vla/benchmarks/class_distribution_using_segmentation/example_config.json b/vla/benchmarks/class_distribution_using_segmentation/example_config.json
index f310443..12f2c3c 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/example_config.json
+++ b/vla/benchmarks/class_distribution_using_segmentation/example_config.json
@@ -1,6 +1,9 @@
 {
   "path_to_raw_data" : "./2026_LOS_SEGM",
   "feature_store_path" : "./precomputed",
+  "generalization_set_folder" : "./2026_LOS_SEGM/Night_clear/Mountain_Range",
+  "train_json" : "train_los_dataset.json",
+  "test_json" : "test_los_dataset.json",
   "learning_rate" : 1e-3,
   "momentum" : 0.9,
   "epochs" : 20,
diff --git a/vla/benchmarks/class_distribution_using_segmentation/images_to_distribution.py b/vla/benchmarks/class_distribution_using_segmentation/images_to_distribution.py
index 7a701d0..968b045 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/images_to_distribution.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/images_to_distribution.py
@@ -1,11 +1,15 @@
 # This file is used to convert segmented images to npy files which contains class distribution
+import json
 import numpy as np
 from pathlib import Path
 from PIL import Image
 
+with open("example_config.json", 'r') as f:
+    config = json.load(f)
+
 paths = []
 unique_colors_list = []
-basepath = Path("./2026_LOS_SEGM")
+basepath = Path(config["path_to_raw_data"])
 segmentation_subfolders = list(basepath.rglob("segmentation"))
 
 stub = 0
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_dino_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_dino_example.py
index fe8d8c8..4df066c 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/test_dino_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_dino_example.py
@@ -8,7 +8,6 @@
         Dinov2Features_full(),
         preprocessor,
         random_seed=10,
-        generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
         config_path="example_config.json",
         use_precomputed_features=True
     )
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_eagle_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_eagle_example.py
index 9d1b913..ebb86fd 100755
--- a/vla/benchmarks/class_distribution_using_segmentation/test_eagle_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_eagle_example.py
@@ -34,7 +34,6 @@ def encode(self, x):
     ModelWrapper(),
     preprocess,
     random_seed=10,
-    generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
     config_path="example_config.json",
     use_precomputed_features=False
 )
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_llava_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_llava_example.py
index 949bec4..0ed13cf 100755
--- a/vla/benchmarks/class_distribution_using_segmentation/test_llava_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_llava_example.py
@@ -38,7 +38,6 @@ def encode(self, x):
     ModelWrapper(),
     preprocess,
     random_seed=10,
-    generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
     config_path="example_config.json",
     use_precomputed_features=True
 )
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_resnet_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_resnet_example.py
index 9dda7dc..14ad826 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/test_resnet_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_resnet_example.py
@@ -31,7 +31,6 @@ def encode(self, x):
         ResnetWrapper(),
         preprocess,
         random_seed=10,
-        generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
         config_path="example_config.json",
         use_precomputed_features=True
     )
\ No newline at end of file
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_siglip_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_siglip_example.py
index d955338..9063fcc 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/test_siglip_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_siglip_example.py
@@ -29,7 +29,6 @@ def encode(self, x):
         SiglipWrapper(),
         preprocess,
         random_seed=10,
-        generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
         config_path="example_config.json",
         use_precomputed_features=True
     )
\ No newline at end of file
diff --git a/vla/benchmarks/class_distribution_using_segmentation/test_yolo_example.py b/vla/benchmarks/class_distribution_using_segmentation/test_yolo_example.py
index 77d47db..ebc4d33 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/test_yolo_example.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/test_yolo_example.py
@@ -28,7 +28,6 @@ def encode(self, x):
         Yolo11Wrapper(),
         preprocess,
         random_seed=10,
-        generalization_set_folder="./2026_LOS_SEGM/Night_clear/Mountain_Range",
         config_path="example_config.json",
         use_precomputed_features=True
     )
\ No newline at end of file
diff --git a/vla/benchmarks/class_distribution_using_segmentation/train.py b/vla/benchmarks/class_distribution_using_segmentation/train.py
index fd70659..d1e189a 100644
--- a/vla/benchmarks/class_distribution_using_segmentation/train.py
+++ b/vla/benchmarks/class_distribution_using_segmentation/train.py
@@ -113,9 +113,8 @@ def train_classifier(
 
     return classifier
 
-def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_json="test_los_dataset.json",
-              use_precomputed_features=True, random_seed=None,
-              generalization_set_folder="", config_path="example_config.json"):
+def benchmark(model, preprocessor, use_precomputed_features=True, random_seed=None,
+              config_path="example_config.json"):
 
     with open(config_path, 'r') as file:
         config_dict = json.load(file)
@@ -127,8 +126,8 @@ def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_jso
         dataset_manager = DatasetManager(ConfigPaths.path_to_raw_data, random_seed=random_seed)
 
     generalization_dataset_manager = None
-    if generalization_set_folder != "":
-        generalization_dataset_manager = DatasetManager(generalization_set_folder, random_seed=random_seed,
+    if ConfigPaths.generalization_set_folder:
+        generalization_dataset_manager = DatasetManager(ConfigPaths.generalization_set_folder, random_seed=random_seed,
                                                         full_folder=True)
 
     store = FeatureStore(os.path.join(ConfigPaths.feature_store_path, type(model).__name__))
@@ -137,8 +136,8 @@ def benchmark(model, preprocessor, train_json="train_los_dataset.json", test_jso
 
     train_loader, test_loader, score_loader, generalization_loader, num_classes = \
         make_dataloaders(
-            train_json,
-            test_json,
+            ConfigPaths.train_json,
+            ConfigPaths.test_json,
             preprocessor,
             feature_store=(store if use_precomputed_features else None),
             workers=8,