Track utilities for value/policy iteration

jtrussell · jtrussell · commit 3f0e3f45b619 · 2017-11-22T22:10:55.000-05:00
diff --git a/mdp.py b/mdp.py
@@ -6,6 +6,7 @@
 dictionary of {state:number} pairs. We then define the value_iteration
 and policy_iteration algorithms."""
 
+from time import clock as now
 from utils import argmax, vector_add, orientations, turn_right, turn_left
 
 import random
@@ -118,22 +119,23 @@ def to_arrows(self, policy):
 # ______________________________________________________________________________
 
 
-def value_iteration(mdp, epsilon=0.001):
+def value_iteration(mdp, epsilon=0.001, max_iters=5000):
     """Solving an MDP by value iteration. [Figure 17.4]"""
     U1 = {s: 0 for s in mdp.states}
     R, T, gamma = mdp.R, mdp.T, mdp.gamma
     mdp.iters_ = 0
-    while True:
+    mdp.U_ = U = None
+    while mdp.iters_ < max_iters:
         mdp.iters_  = mdp.iters_ + 1
-        U = U1.copy()
+        mdp.U_ = U = U1.copy()
         delta = 0
         for s in mdp.states:
             U1[s] = R(s) + gamma * max([sum([p * U[s1] for (p, s1) in T(s, a)])
                                         for a in mdp.actions(s)])
             delta = max(delta, abs(U1[s] - U[s]))
         if delta < epsilon * (1 - gamma) / gamma:
             return U
-
+    return U
 
 def best_policy(mdp, U):
     """Given an MDP and a utility function U, determine the best policy,
@@ -151,14 +153,15 @@ def expected_utility(a, s, U, mdp):
 # ______________________________________________________________________________
 
 
-def policy_iteration(mdp):
+def policy_iteration(mdp, max_iters=2000):
     """Solve an MDP by policy iteration [Figure 17.7]"""
     U = {s: 0 for s in mdp.states}
     pi = {s: random.choice(mdp.actions(s)) for s in mdp.states}
     mdp.iters_ = 0
-    while True:
+    mdp.U_ = U
+    while mdp.iters_ < max_iters:
         mdp.iters_ = mdp.iters_ + 1
-        U = policy_evaluation(pi, U, mdp)
+        mdp.U_ = U = policy_evaluation(pi, U, mdp)
         unchanged = True
         for s in mdp.states:
             a = argmax(mdp.actions(s), key=lambda a: expected_utility(a, s, U, mdp))
@@ -167,6 +170,7 @@ def policy_iteration(mdp):
                 unchanged = False
         if unchanged:
             return pi
+    return pi
 
 
 def policy_evaluation(pi, U, mdp, k=20):
