| 22 | |
| 23 | |
| 24 | class ReplayMemory(object): |
| 25 | def __init__(self, max_size, state_shape, history_len, dtype='uint8'): |
| 26 | """ |
| 27 | Args: |
| 28 | state_shape (tuple[int]): shape (without history) of state |
| 29 | dtype: numpy dtype for the state |
| 30 | """ |
| 31 | self.max_size = int(max_size) |
| 32 | self.state_shape = state_shape |
| 33 | assert len(state_shape) in [1, 2, 3], state_shape |
| 34 | # self._output_shape = self.state_shape + (history_len + 1, ) |
| 35 | self.history_len = int(history_len) |
| 36 | self.dtype = dtype |
| 37 | |
| 38 | all_state_shape = (self.max_size,) + state_shape |
| 39 | logger.info("Creating experience replay buffer of {:.1f} GB ... " |
| 40 | "use a smaller buffer if you don't have enough CPU memory.".format( |
| 41 | np.prod(all_state_shape) / 1024.0**3)) |
| 42 | self.state = np.zeros(all_state_shape, dtype=self.dtype) |
| 43 | self.action = np.zeros((self.max_size,), dtype='int32') |
| 44 | self.reward = np.zeros((self.max_size,), dtype='float32') |
| 45 | self.isOver = np.zeros((self.max_size,), dtype='bool') |
| 46 | |
| 47 | self._curr_size = 0 |
| 48 | self._curr_pos = 0 |
| 49 | |
| 50 | self.writer_lock = threading.Lock() # a lock to guard writing to the memory |
| 51 | |
| 52 | def append(self, exp): |
| 53 | """ |
| 54 | Args: |
| 55 | exp (Experience): |
| 56 | """ |
| 57 | if self._curr_size < self.max_size: |
| 58 | self._assign(self._curr_pos, exp) |
| 59 | self._curr_pos = (self._curr_pos + 1) % self.max_size |
| 60 | self._curr_size += 1 |
| 61 | else: |
| 62 | self._assign(self._curr_pos, exp) |
| 63 | self._curr_pos = (self._curr_pos + 1) % self.max_size |
| 64 | |
| 65 | def sample(self, idx): |
| 66 | """ return a tuple of (s,r,a,o), |
| 67 | where s is of shape self._output_shape, which is |
| 68 | [H, W, (hist_len+1) * channel] if input is (H, W, channel)""" |
| 69 | idx = (self._curr_pos + idx) % self._curr_size |
| 70 | k = self.history_len + 1 |
| 71 | if idx + k <= self._curr_size: |
| 72 | state = self.state[idx: idx + k] |
| 73 | reward = self.reward[idx: idx + k] |
| 74 | action = self.action[idx: idx + k] |
| 75 | isOver = self.isOver[idx: idx + k] |
| 76 | else: |
| 77 | end = idx + k - self._curr_size |
| 78 | state = self._slice(self.state, idx, end) |
| 79 | reward = self._slice(self.reward, idx, end) |
| 80 | action = self._slice(self.action, idx, end) |
| 81 | isOver = self._slice(self.isOver, idx, end) |
no outgoing calls
no test coverage detected
searching dependent graphs…