Update DQN_mountain_car.py

dqn/DQN_mountain_car.py

@@ -13,51 +13,105 @@ GAMMA = 0.9
 LR = 0.01
 MEMORY_CAPACITY = 200
 Q_NETWORK_ITERATION = 50
-BATCH_SIZE = 4
-
+BATCH_SIZE = 32
 
+EPISODES = 400
 env = gym.make('MountainCar-v0')
-
-NUM_STATES = env.state_space.n
-NUM_ACTIONS = 2
+env = env.unwrapped
+NUM_STATES = env.observation_space.shape[0] # 2
+NUM_ACTIONS = env.action_space.n
 
 
 class Net(nn.Module):
     def __init__(self):
         super(Net, self).__init__()
-        self.num_states = NUM_STATES
-        self.num_actions = NUM_ACTIONS
-        self.fc1 = nn.Linear(self.num_states, 10)
+
+        self.fc1 = nn.Linear(NUM_STATES, 50)
         self.fc1.weight.data.normal_(0, 0.1)
-        self.fc2 = nn.Linear(10, self.num_actions)
+        self.fc2 = nn.Linear(50, 30)
         self.fc2.weight.data.normal_(0, 0.1)
+        self.out = nn.Linear(30, NUM_STATES)
 
-    def forward(self, state):
-        state = self.fc1(state)
-        state = F.relu(state)
-        action = self.fc2(state)
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x)
+        x = F.relu(x)
+        x = self.out(x)
 
-        return action
+        return x
 
 class Dqn():
     def __init__(self):
         self.eval_net, self.target_net = Net(), Net()
-        self.memory = np.zeros((MEMORY_CAPACITY, 4))
+        self.memory = np.zeros((MEMORY_CAPACITY, NUM_STATES *2 +2))
+        # state, action ,reward and next state
         self.memory_counter = 0
         self.learn_counter = 0
-        # state, action ,reward and next state 4
+        self.optimizer = optim.Adam(self.eval_net.parameters(), LR)
+        self.loss = nn.MSELoss()
 
     def store_trans(self, state, action, reward, next_state):
         index = self.memory_counter % MEMORY_CAPACITY
-        trans = np.hstack((state, action, reward, next_state))
+        trans = np.hstack((state, [action], [reward], next_state))
         self.memory[index,] = trans
         self.memory_counter += 1
 
     def choose_action(self, state):
+        # notation that the function return the action's index nor the real action
         # EPSILON
+        state = torch.unsqueeze(torch.FloatTensor(state) ,0)
         if np.random.randn() <= EPSILON:
             action_value = self.eval_net.forward(state)
-            action = torch.max(action_value, 1)[0].data.numpy()
+            action = torch.max(action_value, 1)[1].data.numpy() # get action whose q is max
+            action = action[0] #get the action index
         else:
-            action = np.random.choice()
+            action = np.random.randint(0,NUM_ACTIONS)
 
+        return action
+
+    def learn(self):
+        # learn 100 times then the target network update
+        if self.learn_counter % Q_NETWORK_ITERATION ==0:
+            self.target_net.load_state_dict(self.eval_net.state_dict())
+        self.learn_counter+=1
+
+        sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)
+        batch_memory = self.memory[sample_index, :]
+        batch_state = torch.FloatTensor(batch_memory[:, :NUM_STATES])
+        #note that the action must be a int
+        batch_action = torch.LongTensor(batch_memory[:, NUM_STATES:NUM_STATES+1].astype(int))
+        batch_reward = torch.FloatTensor(batch_memory[:, NUM_STATES+1: NUM_STATES+2])
+        batch_next_state = torch.FloatTensor(batch_memory[:, -NUM_STATES:])
+
+        q_eval = self.eval_net(batch_state).gather(1, batch_action)
+        q_next = self.target_net(batch_next_state).detach()
+        q_target = batch_reward + GAMMA*q_next.max(1)[0].view(BATCH_SIZE, 1)
+
+        loss = self.loss(q_eval, q_target)
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+
+def main():
+    net = Dqn()
+    for episode in range(EPISODES):
+        state = env.reset()
+        while True:
+            env.render()
+            action = net.choose_action(state)
+            next_state, reward, done, info = env.step(action)
+            net.store_trans(state, action, reward, next_state)
+
+            if net.memory_counter >= MEMORY_CAPACITY:
+                net.learn()
+                if done:
+                    print("episode {}, the reward is {}", episode, round(reward), 3)
+            if done:
+                break
+
+            state = next_state
+
+if __name__ == '__main__':
+    main()