| 12 | |
| 13 | |
| 14 | class Grid: # Environment |
| 15 | def __init__(self, rows, cols, start): |
| 16 | self.rows = rows |
| 17 | self.cols = cols |
| 18 | self.i = start[0] |
| 19 | self.j = start[1] |
| 20 | |
| 21 | def set(self, rewards, actions): |
| 22 | # rewards should be a dict of: (i, j): r (row, col): reward |
| 23 | # actions should be a dict of: (i, j): A (row, col): list of possible actions |
| 24 | self.rewards = rewards |
| 25 | self.actions = actions |
| 26 | |
| 27 | def set_state(self, s): |
| 28 | self.i = s[0] |
| 29 | self.j = s[1] |
| 30 | |
| 31 | def current_state(self): |
| 32 | return (self.i, self.j) |
| 33 | |
| 34 | def is_terminal(self, s): |
| 35 | return s not in self.actions |
| 36 | |
| 37 | def reset(self): |
| 38 | # put agent back in start position |
| 39 | self.i = 2 |
| 40 | self.j = 0 |
| 41 | return (self.i, self.j) |
| 42 | |
| 43 | def get_next_state(self, s, a): |
| 44 | # this answers: where would I end up if I perform action 'a' in state 's'? |
| 45 | i, j = s[0], s[1] |
| 46 | |
| 47 | # if this action moves you somewhere else, then it will be in this dictionary |
| 48 | if a in self.actions[(i, j)]: |
| 49 | if a == 'U': |
| 50 | i -= 1 |
| 51 | elif a == 'D': |
| 52 | i += 1 |
| 53 | elif a == 'R': |
| 54 | j += 1 |
| 55 | elif a == 'L': |
| 56 | j -= 1 |
| 57 | return i, j |
| 58 | |
| 59 | def move(self, action): |
| 60 | # check if legal move first |
| 61 | if action in self.actions[(self.i, self.j)]: |
| 62 | if action == 'U': |
| 63 | self.i -= 1 |
| 64 | elif action == 'D': |
| 65 | self.i += 1 |
| 66 | elif action == 'R': |
| 67 | self.j += 1 |
| 68 | elif action == 'L': |
| 69 | self.j -= 1 |
| 70 | # return a reward (if any) |
| 71 | return self.rewards.get((self.i, self.j), 0) |
no outgoing calls
no test coverage detected