diff --git a/agent.py b/agent.py
index b01c28c..1474bed 100644
--- a/agent.py
+++ b/agent.py
@@ -54,6 +54,9 @@ class Agent:
 		self.closed = []
 		self.path = []
 		self.open = PriorityQueue()
+		packages_to_go = [p for p in self.warehouse.packages if p.status != PackStatus.STORED]
+		if not packages_to_go and not self.transported_package:
+			return
 		if self.is_loaded:
 			rack = self.find_nearest_rack_for(self.transported_package)
 			self.dest = Node(rack.x_position, rack.y_position, is_rack=True)
diff --git a/main.py b/main.py
index 7d5d581..638dac3 100644
--- a/main.py
+++ b/main.py
@@ -26,7 +26,7 @@ class MainGameFrame:
         agent_radius = int(TILE_WIDTH/2)
         self.agent_tex = pygame.image.load('forklift.png')
         self.font = pygame.font.Font('freesansbold.ttf', 16)
-        self.warehouse_map = warehouse.Warehouse(20, 20, 180, 50)
+        self.warehouse_map = warehouse.Warehouse(20, 20, 150, 10)
         starting_x, starting_y = self.set_starting_agent_position()
         self.agent = agent.Agent(starting_x, starting_y, self.warehouse_map, agent_radius)
         self.clock = pygame.time.Clock()
@@ -42,8 +42,12 @@ class MainGameFrame:
             self.draw_agent()
             self.draw_nums()
             self.agent.move()
+            packages_to_go = [p for p in self.warehouse_map.packages if p.status != PackStatus.STORED]
+            if not packages_to_go and not self.agent.transported_package:
+                pygame.quit()
+                sys.exit()
             pygame.display.update()
-            self.clock.tick(12)
+            self.clock.tick(6)
 
     def draw_floor(self):
         for x in range(self.warehouse_map.width):
diff --git a/package_location_classifier/classifier.py b/package_location_classifier/classifier.py
index 39313dc..5f8e13c 100644
--- a/package_location_classifier/classifier.py
+++ b/package_location_classifier/classifier.py
@@ -33,13 +33,13 @@ class PackageLocationClassifier():
 		test_X = pd.get_dummies(testset[feature_cols])
 		test_y = testset.chance_of_survive
 		products = products.sample(frac=1)
-		X = pd.get_dummies(products[feature_cols])
-		y = products.chance_of_survive
-		dummies_names = X.columns.tolist()
+		X_train = pd.get_dummies(products[feature_cols])
+		y_train = products.chance_of_survive
+		dummies_names = X_train.columns.tolist()
 
-		X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.01, random_state=1, shuffle=True)
+		# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.01, random_state=1, shuffle=True)
 
-		clf = DecisionTreeRegressor(ccp_alpha=0.02, min_samples_leaf=5, max_depth=6)
+		clf = DecisionTreeRegressor(ccp_alpha=0.02, min_samples_leaf=5, max_depth=5)
 		self.predictor = clf.fit(X_train, y_train)
 
 
@@ -49,10 +49,10 @@ class PackageLocationClassifier():
 		evaluation = pd.DataFrame({'category': testset.category, 'temperature': testset.temperature , 'humid': testset.humidity ,'Actual': test_y, 'Predicted': y_pred})
 		evaluation['Prediction_diff'] = abs(evaluation['Actual'] - evaluation['Predicted'])
 		print("Prediction differs from actual value by average {}".format(round(evaluation['Prediction_diff'].mean(), 2)))
-		export_graphviz(clf, out_file=data, filled=True, rounded=True, special_characters=True, feature_names=dummies_names)
-		graph = pydotplus.graph_from_dot_data(data.getvalue())
-		graph.write_png('Ułożenie.png')
-		Image(graph.create_png())
+		# export_graphviz(clf, out_file=data, filled=True, rounded=True, special_characters=True, feature_names=dummies_names)
+		# graph = pydotplus.graph_from_dot_data(data.getvalue())
+		# graph.write_png('Drzewo.png')
+		# Image(graph.create_png())
 
 	def check_if_can_place(self, package, tile):
 		category = package.category
diff --git a/warehouse.py b/warehouse.py
index a92ed65..ef75911 100644
--- a/warehouse.py
+++ b/warehouse.py
@@ -221,13 +221,16 @@ class Warehouse:
         ]
         return diagonals
 
-    def get_all_racks(self):
+    def get_all_racks(self, all_storages=False):
         """:return list of Tile objects"""
         racks = []
-        for x in  self.tiles:
+        for x in self.tiles:
             # row_racks = [t for t in self.tiles[x] if t.category.name == 'Rack']
             for y in x:
-                if y.category.name == 'Rack':
+                if all_storages:
+                    if y.category.name in self.storage_types:
+                        racks.append(y)
+                elif y.category.name == 'Rack':
                     racks.append(y)
         return racks