ප්රමුඛතාඅත්දැකීම් නැවත ධාවනය කිරීමේ බෆරය

16import random
17
18import numpy as np

21class ReplayBuffer:

91    def __init__(self, capacity, alpha):

96        self.capacity = capacity

98        self.alpha = alpha

101        self.priority_sum = [0 for _ in range(2 * self.capacity)]
102        self.priority_min = [float('inf') for _ in range(2 * self.capacity)]

105        self.max_priority = 1.

108        self.data = {
109            'obs': np.zeros(shape=(capacity, 4, 84, 84), dtype=np.uint8),
110            'action': np.zeros(shape=capacity, dtype=np.int32),
111            'reward': np.zeros(shape=capacity, dtype=np.float32),
112            'next_obs': np.zeros(shape=(capacity, 4, 84, 84), dtype=np.uint8),
113            'done': np.zeros(shape=capacity, dtype=np.bool)
114        }

117        self.next_idx = 0

120        self.size = 0

122    def add(self, obs, action, reward, next_obs, done):

128        idx = self.next_idx

131        self.data['obs'][idx] = obs
132        self.data['action'][idx] = action
133        self.data['reward'][idx] = reward
134        self.data['next_obs'][idx] = next_obs
135        self.data['done'][idx] = done

138        self.next_idx = (idx + 1) % self.capacity

140        self.size = min(self.capacity, self.size + 1)

143        priority_alpha = self.max_priority ** self.alpha

145        self._set_priority_min(idx, priority_alpha)
146        self._set_priority_sum(idx, priority_alpha)

148    def _set_priority_min(self, idx, priority_alpha):

154        idx += self.capacity
155        self.priority_min[idx] = priority_alpha

159        while idx >= 2:

161            idx //= 2

163            self.priority_min[idx] = min(self.priority_min[2 * idx], self.priority_min[2 * idx + 1])

165    def _set_priority_sum(self, idx, priority):

171        idx += self.capacity

173        self.priority_sum[idx] = priority

177        while idx >= 2:

179            idx //= 2

181            self.priority_sum[idx] = self.priority_sum[2 * idx] + self.priority_sum[2 * idx + 1]

183    def _sum(self):

189        return self.priority_sum[1]

191    def _min(self):

197        return self.priority_min[1]

199    def find_prefix_sum_idx(self, prefix_sum):

205        idx = 1
206        while idx < self.capacity:

208            if self.priority_sum[idx * 2] > prefix_sum:

210                idx = 2 * idx
211            else:

214                prefix_sum -= self.priority_sum[idx * 2]
215                idx = 2 * idx + 1

219        return idx - self.capacity

221    def sample(self, batch_size, beta):

227        samples = {
228            'weights': np.zeros(shape=batch_size, dtype=np.float32),
229            'indexes': np.zeros(shape=batch_size, dtype=np.int32)
230        }

233        for i in range(batch_size):
234            p = random.random() * self._sum()
235            idx = self.find_prefix_sum_idx(p)
236            samples['indexes'][i] = idx

239        prob_min = self._min() / self._sum()

241        max_weight = (prob_min * self.size) ** (-beta)
242
243        for i in range(batch_size):
244            idx = samples['indexes'][i]

246            prob = self.priority_sum[idx + self.capacity] / self._sum()

248            weight = (prob * self.size) ** (-beta)

251            samples['weights'][i] = weight / max_weight

254        for k, v in self.data.items():
255            samples[k] = v[samples['indexes']]
256
257        return samples

259    def update_priorities(self, indexes, priorities):

264        for idx, priority in zip(indexes, priorities):

266            self.max_priority = max(self.max_priority, priority)

269            priority_alpha = priority ** self.alpha

271            self._set_priority_min(idx, priority_alpha)
272            self._set_priority_sum(idx, priority_alpha)

274    def is_full(self):

278        return self.capacity == self.size