add state labels

ShangtongZhang · Jul 5, 2020 · d1d9cd5 · d1d9cd5
1 parent 1b1723e
commit d1d9cd5
Showing 1 changed file with 23 additions and 0 deletions.
diff --git a/chapter03/grid_world.py b/chapter03/grid_world.py
@@ -57,8 +57,20 @@ def draw_image(image):
 
     # Add cells
     for (i, j), val in np.ndenumerate(image):
+
+        # add state labels
+        if [i, j] == A_POS:
+            val = str(val) + " (A)"
+        if [i, j] == A_PRIME_POS:
+            val = str(val) + " (A')"
+        if [i, j] == B_POS:
+            val = str(val) + " (B)"
+        if [i, j] == B_PRIME_POS:
+            val = str(val) + " (B')"
+
         tb.add_cell(i, j, width, height, text=val,
                     loc='center', facecolor='white')
+
 
     # Row and column labels...
     for i in range(len(image)):
@@ -88,6 +100,17 @@ def draw_policy(optimal_values):
         val=''
         for ba in best_actions:
             val+=ACTIONS_FIGS[ba]
+
+        # add state labels
+        if [i, j] == A_POS:
+            val = str(val) + " (A)"
+        if [i, j] == A_PRIME_POS:
+            val = str(val) + " (A')"
+        if [i, j] == B_POS:
+            val = str(val) + " (B)"
+        if [i, j] == B_PRIME_POS:
+            val = str(val) + " (B')"
+
         tb.add_cell(i, j, width, height, text=val,
                 loc='center', facecolor='white')