MuhammadOsamaNusrat/rl-gridworld

1import numpy as np
2
3class GridWorld:
4    def __init__(self, max_steps=50):
5        self.size = 4
6        self.start = (0, 0)
7        self.goal = (3, 3)
8        self.holes = [(1, 1), (1, 3), (3, 0)]
9        self.actions = ['up', 'down', 'left', 'right']
10        self.max_steps = max_steps
11        self.reset()
12
13    def reset(self):
14        self.state = self.start
15        self.steps = 0
16        return self.state
17
18    def step(self, action):
19        x, y = self.state
20        self.steps += 1
21
22        if action == 'up': x = max(0, x - 1)
23        elif action == 'down': x = min(self.size - 1, x + 1)
24        elif action == 'left': y = max(0, y - 1)
25        elif action == 'right': y = min(self.size - 1, y + 1)
26
27        self.state = (x, y)
28
29        if self.state in self.holes:
30            return self.state, -50, True
31        elif self.state == self.goal:
32            return self.state, 50, True
33        elif self.steps >= self.max_steps:
34            return self.state, -10, True  # episode timeout
35        else:
36            return self.state, -1, False
37
38