diff --git a/main.py b/main.py
index ffd3381..77dcb23 100644
--- a/main.py
+++ b/main.py
@@ -67,9 +67,8 @@ def init_global(maze_filename):
     maze_cost[MAZE == 'T'] = 50
     maze_cost[MAZE == 'G'] = -1
     G1_X = maze_cost.copy()[state_mask]
-    # maze_cost[(MAZE=='0') | (MAZE=='S') | (MAZE=='G')] += 1
-    # G2_X = maze_cost.copy()[state_mask]
-    G2_X = G1_X + 1
+    maze_cost[(MAZE=='0') | (MAZE=='S') | (MAZE=='G')] += 1
+    G2_X = maze_cost.copy()[state_mask]
 
     # Actual environment modelling
     U_OF_X = np.zeros((SN, len(A2)), dtype=np.bool)