improve toy_text printing

yichao-liang · Feb 12, 2017 · 622429c · 622429c
1 parent 4647fb1
commit 622429c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/gym/envs/toy_text/discrete.py b/gym/envs/toy_text/discrete.py
@@ -47,6 +47,7 @@ def _seed(self, seed=None):
 
     def _reset(self):
         self.s = categorical_sample(self.isd, self.np_random)
+        self.lastaction=None
         return self.s
 
     def _step(self, a):

diff --git a/gym/envs/toy_text/frozen_lake.py b/gym/envs/toy_text/frozen_lake.py
@@ -76,13 +76,13 @@ def __init__(self, desc=None, map_name="4x4",is_slippery=True):
         def to_s(row, col):
             return row*ncol + col
         def inc(row, col, a):
-            if a==0:
+            if a==0: # left
                 col = max(col-1,0)
-            elif a==1:
+            elif a==1: # down
                 row = min(row+1,nrow-1)
-            elif a==2:
+            elif a==2: # right
                 col = min(col+1,ncol-1)
-            elif a==3:
+            elif a==3: # up
                 row = max(row-1,0)
             return (row, col)
 
@@ -116,17 +116,16 @@ def inc(row, col, a):
     def _render(self, mode='human', close=False):
         if close:
             return
-
         outfile = StringIO() if mode == 'ansi' else sys.stdout
 
         row, col = self.s // self.ncol, self.s % self.ncol
         desc = self.desc.tolist()
         desc = [[c.decode('utf-8') for c in line] for line in desc]
         desc[row][col] = utils.colorize(desc[row][col], "red", highlight=True)
-        outfile.write("\n".join(''.join(line) for line in desc)+"\n")
         if self.lastaction is not None:
             outfile.write("  ({})\n".format(["Left","Down","Right","Up"][self.lastaction]))
         else:
             outfile.write("\n")
+        outfile.write("\n".join(''.join(line) for line in desc)+"\n")
 
         return outfile