| 29 | return self.count # Return the number of elements in the buffer |
| 30 | |
| 31 | def sample_batch(self, batch_size, keys_num): |
| 32 | if self.count < batch_size * keys_num / 2: # If the number of elements in buffer is less than batch_size * keys_num / 2 |
| 33 | batch_size = self.count // keys_num |
| 34 | |
| 35 | s_batch, a_batch, r_batch, s2_batch, d_batch = [], [], [], [], [] |
| 36 | |
| 37 | for _ in range(batch_size): |
| 38 | while True: |
| 39 | end_index = random.randint(keys_num, self.count) |
| 40 | sequence = list(islice(self.buffer, end_index - keys_num, end_index)) |
| 41 | # sequence = self.buffer[end_index - keys_num : end_index] |
| 42 | s_seq, a_seq, r_seq, s2_seq, d_seq, end_seq = zip(*sequence) |
| 43 | |
| 44 | # Check that the 'end' flag of the first keys_num - 1 entries is not equal to 1 |
| 45 | if all(end != 1 for end in end_seq[:-1]): |
| 46 | s_batch.append(s_seq) |
| 47 | a_batch.append(a_seq) |
| 48 | r_batch.append(r_seq) |
| 49 | s2_batch.append(s2_seq) |
| 50 | d_batch.append(d_seq) |
| 51 | break # Break out of the while loop to process the next batch |
| 52 | |
| 53 | s_batch = np.array(s_batch) |
| 54 | a_batch = np.array(a_batch).reshape(batch_size, keys_num, -1) |
| 55 | r_batch = np.array(r_batch).reshape(batch_size, keys_num, -1) |
| 56 | s2_batch = np.array(s2_batch) |
| 57 | d_batch = np.array(d_batch).reshape(batch_size, keys_num, -1) |
| 58 | |
| 59 | return s_batch, a_batch, r_batch, s2_batch, d_batch |
| 60 | |
| 61 | def clear(self): |
| 62 | self.buffer.clear() # Clear the buffer |