def test_add(self):
buffer_size = 500
obs_shape = (84, 84, 3)
act_dim = 10
rb = PrioritizedReplayBuffer(buffer_size, {
"obs": {
"shape": obs_shape
},
"act": {
"shape": act_dim
},
"rew": {},
"done": {}
},
next_of=("obs"))
obs = np.zeros(obs_shape)
act = np.ones(act_dim)
rew = 1
done = 0
rb.add(obs=obs, act=act, rew=rew, next_obs=obs, done=done)
ps = 1.5
rb.add(obs=obs,
act=act,
rew=rew,
next_obs=obs,
done=done,
priorities=ps)
self.assertAlmostEqual(rb.get_max_priority(), 1.5)
obs = np.stack((obs, obs))
act = np.stack((act, act))
rew = (1, 0)
done = (0.0, 1.0)
rb.add(obs=obs, act=act, rew=rew, next_obs=obs, done=done)
ps = (0.2, 0.4)
rb.add(obs=obs,
act=act,
rew=rew,
next_obs=obs,
done=done,
priorities=ps)
rb.clear()
self.assertEqual(rb.get_next_index(), 0)
self.assertEqual(rb.get_stored_size(), 0)
def test_cpdef_super(self):
buffer_size = 256
obs_dim = 15
act_dim = 3
prb = PrioritizedReplayBuffer(
buffer_size, {
"obs": {
"shape": obs_dim
},
"act": {
"shape": act_dim
},
"rew": {},
"next_obs": {
"shape": obs_dim
},
"done": {}
})
prb.clear()