update projects

projects/codes/DQN/dqn.py

@@ -5,7 +5,7 @@
 @Email: johnjim0816@gmail.com
 @Date: 2020-06-12 00:50:49
 @LastEditor: John
-LastEditTime: 2022-07-20 23:57:16
+LastEditTime: 2022-08-11 09:52:23
 @Discription: 
 @Environment: python 3.7.7
 '''
@@ -14,77 +14,39 @@ LastEditTime: 2022-07-20 23:57:16
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import torch.optim as optim
 import random
 import math
 import numpy as np
 
-class MLP(nn.Module):
-    def __init__(self, n_states,n_actions,hidden_dim=128):
-        """ 初始化q网络，为全连接网络
-            n_states: 输入的特征数即环境的状态维度
-            n_actions: 输出的动作维度
-        """
-        super(MLP, self).__init__()
-        self.fc1 = nn.Linear(n_states, hidden_dim) # 输入层
-        self.fc2 = nn.Linear(hidden_dim,hidden_dim) # 隐藏层
-        self.fc3 = nn.Linear(hidden_dim, n_actions) # 输出层
-        
-    def forward(self, x):
-        # 各层对应的激活函数
-        x = F.relu(self.fc1(x)) 
-        x = F.relu(self.fc2(x))
-        return self.fc3(x)
-
-class ReplayBuffer:
-    def __init__(self, capacity):
-        self.capacity = capacity # 经验回放的容量
-        self.buffer = [] # 缓冲区
-        self.position = 0 
-    
-    def push(self, state, action, reward, next_state, done):
-        ''' 缓冲区是一个队列，容量超出时去掉开始存入的转移(transition)
-        '''
-        if len(self.buffer) < self.capacity:
-            self.buffer.append(None)
-        self.buffer[self.position] = (state, action, reward, next_state, done)
-        self.position = (self.position + 1) % self.capacity 
-    
-    def sample(self, batch_size):
-        batch = random.sample(self.buffer, batch_size) # 随机采出小批量转移
-        state, action, reward, next_state, done =  zip(*batch) # 解压成状态，动作等
-        return state, action, reward, next_state, done
-    
-    def __len__(self):
-        ''' 返回当前存储的量
-        '''
-        return len(self.buffer)
-
 class DQN:
-    def __init__(self, n_states,n_actions,cfg):
+    def __init__(self,n_actions,model,memory,cfg):
 
         self.n_actions = n_actions  
         self.device = torch.device(cfg.device)  # cpu or cuda
         self.gamma = cfg.gamma  # 奖励的折扣因子
         # e-greedy策略相关参数
-        self.frame_idx = 0  # 用于epsilon的衰减计数
-        self.epsilon = lambda frame_idx: cfg.epsilon_end + \
-            (cfg.epsilon_start - cfg.epsilon_end) * \
-            math.exp(-1. * frame_idx / cfg.epsilon_decay)
+        self.sample_count = 0  # 用于epsilon的衰减计数
+        self.epsilon = cfg.epsilon_start
+        self.sample_count = 0  
+        self.epsilon_start = cfg.epsilon_start
+        self.epsilon_end = cfg.epsilon_end
+        self.epsilon_decay = cfg.epsilon_decay
         self.batch_size = cfg.batch_size
-        self.policy_net = MLP(n_states,n_actions).to(self.device)
-        self.target_net = MLP(n_states,n_actions).to(self.device)
+        self.policy_net = model.to(self.device)
+        self.target_net = model.to(self.device)
         for target_param, param in zip(self.target_net.parameters(),self.policy_net.parameters()): # 复制参数到目标网路targe_net
             target_param.data.copy_(param.data)
         self.optimizer = optim.Adam(self.policy_net.parameters(), lr=cfg.lr) # 优化器
-        self.memory = ReplayBuffer(cfg.memory_capacity) # 经验回放
+        self.memory = memory # 经验回放
 
-    def choose_action(self, state):
+    def sample(self, state):
         ''' 选择动作
         '''
-        self.frame_idx += 1
-        if random.random() > self.epsilon(self.frame_idx):
+        self.sample_count += 1
+        self.epsilon = self.epsilon_end + (self.epsilon_start - self.epsilon_end) * \
+            math.exp(-1. * self.sample_count / self.epsilon_decay) # epsilon是会递减的，这里选择指数递减
+        if random.random() > self.epsilon:
             with torch.no_grad():
                 state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
                 q_values = self.policy_net(state)
@@ -92,11 +54,16 @@ class DQN:
         else:
             action = random.randrange(self.n_actions)
         return action
+    def predict(self,state):
+        with torch.no_grad():
+            state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+            q_values = self.policy_net(state)
+            action = q_values.max(1)[1].item() # 选择Q值最大的动作
+        return action
     def update(self):
         if len(self.memory) < self.batch_size: # 当memory中不满足一个批量时，不更新策略
             return
         # 从经验回放中(replay memory)中随机采样一个批量的转移(transition)
-        # print('updating')
         
         state_batch, action_batch, reward_batch, next_state_batch, done_batch = self.memory.sample(
             self.batch_size)
@@ -118,9 +85,11 @@ class DQN:
         self.optimizer.step() 
 
     def save(self, path):
-        torch.save(self.target_net.state_dict(), path+'dqn_checkpoint.pth')
+        from pathlib import Path
+        Path(path).mkdir(parents=True, exist_ok=True)
+        torch.save(self.target_net.state_dict(), path+'checkpoint.pth')
 
     def load(self, path):
-        self.target_net.load_state_dict(torch.load(path+'dqn_checkpoint.pth'))
+        self.target_net.load_state_dict(torch.load(path+'checkpoint.pth'))
         for target_param, param in zip(self.target_net.parameters(), self.policy_net.parameters()):
             param.data.copy_(target_param.data)