update

codes/QLearning/agent.py

@@ -5,8 +5,8 @@ Author: John
 Email: johnjim0816@gmail.com
 Date: 2020-09-11 23:03:00
 LastEditor: John
-LastEditTime: 2021-03-26 16:51:01
-Discription: 
+LastEditTime: 2021-04-29 16:59:41
+Discription: use defaultdict to define Q table
 Environment: 
 '''
 import numpy as np
@@ -15,7 +15,7 @@ import torch
 from collections import defaultdict
 
 class QLearning(object):
-    def __init__(self,
+    def __init__(self,state_dim,
                  action_dim,cfg):
         self.action_dim = action_dim  # dimension of acgtion
         self.lr = cfg.lr  # learning rate
@@ -26,17 +26,20 @@ class QLearning(object):
         self.epsilon_end = cfg.epsilon_end
         self.epsilon_decay = cfg.epsilon_decay
         self.Q_table  = defaultdict(lambda: np.zeros(action_dim)) # A nested dictionary that maps state -> (action -> action-value)
+        
     def choose_action(self, state):
         self.sample_count += 1
         self.epsilon = self.epsilon_end + (self.epsilon_start - self.epsilon_end) * \
             math.exp(-1. * self.sample_count / self.epsilon_decay)
         # e-greedy policy
         if np.random.uniform(0, 1) > self.epsilon:
-            action = np.argmax(self.Q_table[str(state)])
+            action = self.predict(state)
         else:
             action = np.random.choice(self.action_dim) 
         return action
-            
+    def predict(self,state):
+        action = np.argmax(self.Q_table[str(state)])
+        return action
     def update(self, state, action, reward, next_state, done):
         Q_predict = self.Q_table[str(state)][action]
         if done: