update codes

codes/HierarchicalDQN/agent.py

@@ -11,23 +11,62 @@ Environment:
 '''
 import torch
 import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
 import numpy as np
 import random,math
-import torch.optim as optim
-from common.model import MLP
-from common.memory import ReplayBuffer
 
+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 MLP(nn.Module):
+    def __init__(self, input_dim,output_dim,hidden_dim=128):
+        """ 初始化q网络，为全连接网络
+            input_dim: 输入的特征数即环境的状态数
+            output_dim: 输出的动作维度
+        """
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim) # 输入层
+        self.fc2 = nn.Linear(hidden_dim,hidden_dim) # 隐藏层
+        self.fc3 = nn.Linear(hidden_dim, output_dim) # 输出层
+        
+    def forward(self, x):
+        # 各层对应的激活函数
+        x = F.relu(self.fc1(x)) 
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+        
 class HierarchicalDQN:
-    def __init__(self,state_dim,action_dim,cfg):
-        self.state_dim = state_dim
-        self.action_dim = action_dim
+    def __init__(self,n_states,n_actions,cfg):
+        self.n_states = n_states
+        self.n_actions = n_actions
         self.gamma = cfg.gamma
         self.device = cfg.device
         self.batch_size = cfg.batch_size
-        self.frame_idx = 0 
+        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.policy_net = MLP(2*state_dim, action_dim,cfg.hidden_dim).to(self.device)
-        self.meta_policy_net = MLP(state_dim, state_dim,cfg.hidden_dim).to(self.device)
+        self.policy_net = MLP(2*n_states, n_actions,cfg.hidden_dim).to(self.device)
+        self.meta_policy_net = MLP(n_states, n_states,cfg.hidden_dim).to(self.device)
         self.optimizer = optim.Adam(self.policy_net.parameters(),lr=cfg.lr)
         self.meta_optimizer = optim.Adam(self.meta_policy_net.parameters(),lr=cfg.lr)
         self.memory = ReplayBuffer(cfg.memory_capacity)
@@ -37,7 +76,7 @@ class HierarchicalDQN:
         self.losses = []
         self.meta_losses = []
     def to_onehot(self,x):
-        oh = np.zeros(self.state_dim)
+        oh = np.zeros(self.n_states)
         oh[x - 1] = 1.
         return oh
     def set_goal(self,state):
@@ -46,7 +85,7 @@ class HierarchicalDQN:
                 state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(0)
                 goal = self.meta_policy_net(state).max(1)[1].item() 
         else:
-            goal = random.randrange(self.state_dim)
+            goal = random.randrange(self.n_states)
         return goal
     def choose_action(self,state):
         self.frame_idx += 1
@@ -56,7 +95,7 @@ class HierarchicalDQN:
                 q_value = self.policy_net(state)
                 action = q_value.max(1)[1].item()  
         else:
-            action = random.randrange(self.action_dim)
+            action = random.randrange(self.n_actions)
         return action
     def update(self):
         self.update_policy()