Update environment.py

environment.py

@@ -13,14 +13,18 @@ DEFAULT_DATASET_PATH = BASE_DIR / "dataset.json"
 class SecurityEnv:
     def __init__(self, dataset_path=None):
         dataset_file = Path(dataset_path) if dataset_path else DEFAULT_DATASET_PATH
-        if not dataset_file.is_absolute():
+
+        # only join if relative
+        if not dataset_file.exists():
             dataset_file = BASE_DIR / dataset_file
 
+        if not dataset_file.exists():
+            raise FileNotFoundError(f"Dataset not found: {dataset_file}")
+
         with dataset_file.open("r", encoding="utf-8") as handle:
             self.data = json.load(handle)
 
         self.current_sample = None
-        self.step_count = 0
 
     def reset(self):
         self.current_sample = random.choice(self.data)