From 5b5d78c309e6b8b9d33d111f1ad1963c56883dfc Mon Sep 17 00:00:00 2001
From: JohnJim0816 <johnjim0816@gmail.com>
Date: Mon, 15 Mar 2021 17:17:16 +0800
Subject: [PATCH] update common

---
 codes/common/memory.py | 34 ++++++++++++++++++++++++++++++++++
 codes/common/model.py  | 26 ++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 codes/common/memory.py

diff --git a/codes/common/memory.py b/codes/common/memory.py
new file mode 100644
index 0000000..a9906e9
--- /dev/null
+++ b/codes/common/memory.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# coding=utf-8
+'''
+@Author: John
+@Email: johnjim0816@gmail.com
+@Date: 2020-06-10 15:27:16
+@LastEditor: John
+LastEditTime: 2021-03-13 11:37:15
+@Discription: 
+@Environment: python 3.7.7
+'''
+import random
+
+class ReplayBuffer:
+    
+    def __init__(self, capacity):
+        self.capacity = capacity
+        self.buffer = []
+        self.position = 0
+    
+    def push(self, state, action, reward, next_state, done):
+        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)
+
diff --git a/codes/common/model.py b/codes/common/model.py
index 0f367f2..c99d4a3 100644
--- a/codes/common/model.py
+++ b/codes/common/model.py
@@ -5,19 +5,19 @@ Author: John
 Email: johnjim0816@gmail.com
 Date: 2021-03-12 21:14:12
 LastEditor: John
-LastEditTime: 2021-03-13 11:51:38
+LastEditTime: 2021-03-13 13:48:35
 Discription: 
 Environment: 
 '''
 import torch.nn as nn
 import torch.nn.functional as F
-class MLP(nn.Module):
+class MLP1(nn.Module):
     ''' 多层感知机
         输入：state维度
         输出：概率
     '''
     def __init__(self,n_states,hidden_dim = 36):
-        super(MLP, self).__init__()
+        super(MLP1, self).__init__()
         # 24和36为hidden layer的层数，可根据state_dim, n_actions的情况来改变
         self.fc1 = nn.Linear(n_states, hidden_dim)
         self.fc2 = nn.Linear(hidden_dim,hidden_dim)
@@ -27,4 +27,22 @@ class MLP(nn.Module):
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         x = F.sigmoid(self.fc3(x))
-        return x
\ No newline at end of file
+        return x
+
+
+class MLP2(nn.Module):
+    def __init__(self, n_states,n_actions,hidden_dim=128):
+        """ 初始化q网络，为全连接网络
+            n_states: 输入的feature即环境的state数目
+            n_actions: 输出的action总个数
+        """
+        super(MLP2, 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)
\ No newline at end of file