| 57 | |
| 58 | ### The experience replay memory ### |
| 59 | class ReplayBuffer: |
| 60 | def __init__(self, obs_dim, act_dim, size): |
| 61 | self.obs1_buf = np.zeros([size, obs_dim], dtype=np.float32) |
| 62 | self.obs2_buf = np.zeros([size, obs_dim], dtype=np.float32) |
| 63 | self.acts_buf = np.zeros([size, act_dim], dtype=np.float32) |
| 64 | self.rews_buf = np.zeros(size, dtype=np.float32) |
| 65 | self.done_buf = np.zeros(size, dtype=np.float32) |
| 66 | self.ptr, self.size, self.max_size = 0, 0, size |
| 67 | |
| 68 | def store(self, obs, act, rew, next_obs, done): |
| 69 | self.obs1_buf[self.ptr] = obs |
| 70 | self.obs2_buf[self.ptr] = next_obs |
| 71 | self.acts_buf[self.ptr] = act |
| 72 | self.rews_buf[self.ptr] = rew |
| 73 | self.done_buf[self.ptr] = done |
| 74 | self.ptr = (self.ptr+1) % self.max_size |
| 75 | self.size = min(self.size+1, self.max_size) |
| 76 | |
| 77 | def sample_batch(self, batch_size=32): |
| 78 | idxs = np.random.randint(0, self.size, size=batch_size) |
| 79 | return dict(s=self.obs1_buf[idxs], |
| 80 | s2=self.obs2_buf[idxs], |
| 81 | a=self.acts_buf[idxs], |
| 82 | r=self.rews_buf[idxs], |
| 83 | d=self.done_buf[idxs]) |
| 84 | |
| 85 | |
| 86 | ### Implement the DDPG algorithm ### |