diff --git a/rl_lab.py b/rl_lab.py
index b93d72a..f1d7c52 100644
--- a/rl_lab.py
+++ b/rl_lab.py
@@ -71,9 +71,9 @@ def run(pos, N_trials, learning_rate, penalty_factor=0.99, forgetting_factor=0.9
 
 #		forget(forgetting_factor)
 
-		P = P_door[pos]
-		Pn = P/np.sum(P)
+		Pn = P_door[pos]/np.sum(P_door[pos])
 
+		P_door[pos,:] *= forgetting_factor
 		while True:
 			fail = False
 			door = choose_door(Pn)
@@ -100,8 +100,7 @@ def run(pos, N_trials, learning_rate, penalty_factor=0.99, forgetting_factor=0.9
 				P_door[pos][door] *= penalty_factor
 				continue
 			else:
-				P[door] = P[door] + learning_rate
-				P_door[pos] = P
+				P_door[pos][door] = P_door[pos][door] + learning_rate
 			break
 
 
@@ -112,7 +111,7 @@ def run(pos, N_trials, learning_rate, penalty_factor=0.99, forgetting_factor=0.9
 
 	return moves_needed, lab_visited
 
-for i in range(0, 1000):
+for i in range(0, 100):
 	moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.5)
 	print ('Visited map after {} moves:'.format(moves_needed))
 #	print (lab_visited)