From c96cb1e48eff067700c26ef06eeb7b3581c21ad0 Mon Sep 17 00:00:00 2001
From: Nicolai Czempin <nicolai.czempin@gmail.com>
Date: Sat, 1 Apr 2017 01:54:06 +0200
Subject: [PATCH] Clean up; do the terminal states

---
 gym_random_walk/envs/random_walk_env.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/gym_random_walk/envs/random_walk_env.py b/gym_random_walk/envs/random_walk_env.py
index 48017f9..bc9d22a 100644
--- a/gym_random_walk/envs/random_walk_env.py
+++ b/gym_random_walk/envs/random_walk_env.py
@@ -8,14 +8,24 @@ class RandomWalkEnv(gym.Env):
 
   def __init__(self):
     self.action_space = spaces.Discrete(2)
-    print("init")
+    #print("init")
   def _step(self, action):
-    print("step")
+    #print("step")
     reward = 0
     done = False
+    if (action == 0):
+       self.state -= 1
+    if (action == 1):
+        self.state += 1
+    if (self.state >= 6):
+        reward = 1
+        done = True
+    if (self.state <= 0):
+        done = True
     return np.array(self.state), reward, done, {}
   def _reset(self):
-    print("reset")
-    self.state = 0 # TODO start in a random position
+    #print("reset")
+    self.state = 1 # TODO start in a random position
   def _render(self, mode='human', close=False):
-    print("render")
+    #print("render")
+    print(self.state)