MuhammadOsamaNusrat/rl-gridworld

1import random
2from collections import defaultdict
3
4class RLAgent:
5    def __init__(self, actions, alpha=0.1, gamma=1.0, epsilon=0.1, method='sarsa'):
6        self.Q = defaultdict(lambda: {a: 0.0 for a in actions})
7        self.visits = defaultdict(int)
8        self.alpha = alpha
9        self.gamma = gamma
10        self.epsilon = epsilon
11        self.actions = actions
12        self.method = method
13
14    def choose_action(self, state):
15        self.visits[state] += 1
16        if random.random() < self.epsilon:
17            return random.choice(self.actions)
18        return max(self.Q[state], key=self.Q[state].get)
19
20    def update(self, s, a, r, s_, a_):
21        predict = self.Q[s][a]
22        if self.method == 'sarsa':
23            target = r + self.gamma * self.Q[s_][a_]
24        else:
25            target = r + self.gamma * max(self.Q[s_].values())
26        self.Q[s][a] += self.alpha * (target - predict)