diff --git a/rl_lab.py b/rl_lab.py
index 16cf899..b93d72a 100644
--- a/rl_lab.py
+++ b/rl_lab.py
@@ -7,27 +7,31 @@ num_doors_per_room = 4
 
 
 P_door = np.ndarray((lab_size_x, lab_size_y, num_doors_per_room))
-P_door = np.random.uniform(0,1,P_door.shape)
+P_door = np.ones(P_door.shape)
 
 print('P_door.shape = ', P_door.shape)
 lab = np.rint(np.random.uniform(0, num_doors_per_room-1, size=(lab_size_x, lab_size_y)))
 
 
-#    1
+#    2
 #    |
-# 4 -0- 2
+# 3 -0- 1
 #    |
-#    3
+#    4
 
-move = np.array([(0,0), (-1,0), (0,1), (1,0), (0,-1)])
+move = np.array([(0,0), (0,1), (1,0), (0,-1), (-1,0)])
 lab = np.array([
-	[3, 0, 0, 0, -1],
-	[3, 0, 0, 2, 1],
-	[2, 2, 3, 1, 4],
-	[3, 4, 4, 2, 1],
-	[2, 2, 2, 1, 0]
+	[4, 0, 0, 0, -1],
+	[4, 0, 0, 0, 2],
+	[4, 0, 1, 1, 2],
+	[4, 0, 2, 0, 0],
+	[1, 1, 2, 0, 0]
 ])
 
+def to_this_door(door_last_room):
+	q, r = divmod(door_last_room + 2, 4)
+	return r
+
 def choose_door(Pn):
 	Pnorm = np.cumsum(Pn)
 	size = len(Pn)
@@ -41,49 +45,87 @@ def choose_door(Pn):
 
 	return result
 
-def run(pos, N_trials, learning_rate):
-	lab_visited = np.zeros(P_door.shape)
+def normalize(P):
+	for x in range(0, lab_size_x):
+		for y in range(0, lab_size_y):
+			pos = (x,y)
+			P0 = P[pos]
+			P[pos] = P0/np.sum(P0)
+
+	return P
+
+def forget(forgetting_factor):
+	for x in range(0, lab_size_x):
+		for y in range(0, lab_size_y):
+			pos = (x,y)
+			P_door[pos,:] *= forgetting_factor
+
+def run(pos, N_trials, learning_rate, penalty_factor=0.99, forgetting_factor=0.999):
+	lab_visited = np.zeros((lab_size_x, lab_size_y))
 	moves_needed = 0
+	door_last = -1
 	for n in range(0, N_trials):
 
-#		print ("Pos={}".format(pos))
 		if lab[pos] == -1:
 			break
 
+#		forget(forgetting_factor)
+
 		P = P_door[pos]
 		Pn = P/np.sum(P)
 
 		while True:
+			fail = False
 			door = choose_door(Pn)
+			if door_last >= 0:
+				this_door = to_this_door(door_last)
+				if this_door == door:
+					P_door[pos][door] *= (1.0 - learning_rate)
+
 			pos_new = tuple(pos + move[door + 1])
 
 			if pos_new[0] < 0 or pos_new[1] < 0:
-				continue
+				fail = True
 
 			if pos_new[0] >= lab_size_x or pos_new[1] >= lab_size_y:
-				continue
+				fail = True
 
-			if lab[pos_new] == 0:
-				continue
+			try:
+				if lab[pos_new] == 0:
+					fail = True
+			except:
+				pass
 
+			if fail:
+				P_door[pos][door] *= penalty_factor
+				continue
+			else:
+				P[door] = P[door] + learning_rate
+				P_door[pos] = P
 			break
 
-		moves_needed += 1
-		k = lab_visited[pos][door]
-		P[door] = max(0, P[door] - learning_rate)
-		P_door[pos] = P
 
 		pos = pos_new
-		lab_visited[pos][door] += 1
+		lab_visited[pos] += 1
+		moves_needed += 1
+		door_last = door
 
 	return moves_needed, lab_visited
 
-for i in range(0, 100):
-	moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.001)
-	print ('Visited map after {} trials: {}'.format(moves_needed, lab_visited))
-	print ('P_door after {} trials: {}'.format(moves_needed, P_door))
+for i in range(0, 1000):
+	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)
+#	print ('P_door after {} moves:'.format(moves_needed))
+#	print (P_door)
 
 moves_needed, lab_visited = run(pos=(0,0), N_trials=1000, learning_rate=0.0)
-print ('Finished after {} trials'.format(moves_needed))
+print ('Finished after {} moves'.format(moves_needed))
+print('Visited map after {} moves:'.format(moves_needed))
+print(lab_visited)
+print ('P_door after {} moves:'.format(moves_needed))
+
+P = normalize(P_door)
+for i in range(0, num_doors_per_room):
+	print (P[:,:,i])
 
-#	P_door[pos, door] += learning_rate