Tianshou/test/test_buffer.py

26 lines
782 B
Python
Raw Normal View History

2020-03-11 17:28:51 +08:00
from tianshou.data import ReplayBuffer
2020-03-11 18:02:19 +08:00
if __name__ == '__main__':
from test_env import MyTestEnv
2020-03-12 22:20:33 +08:00
else: # pytest
2020-03-11 18:02:19 +08:00
from test.test_env import MyTestEnv
2020-03-11 17:28:51 +08:00
2020-03-11 18:02:19 +08:00
def test_replaybuffer(size=10, bufsize=20):
env = MyTestEnv(size)
2020-03-11 17:28:51 +08:00
buf = ReplayBuffer(bufsize)
obs = env.reset()
2020-03-11 18:02:19 +08:00
action_list = [1] * 5 + [0] * 10 + [1] * 15
2020-03-11 17:28:51 +08:00
for i, a in enumerate(action_list):
obs_next, rew, done, info = env.step(a)
buf.add(obs, a, rew, done, obs_next, info)
assert len(buf) == min(bufsize, i + 1), print(len(buf), i)
indice = buf.sample_indice(4)
data = buf.sample(4)
2020-03-11 18:02:19 +08:00
assert (indice < len(buf)).all()
assert (data.obs < size).all()
assert (0 <= data.done).all() and (data.done <= 1).all()
2020-03-11 17:28:51 +08:00
if __name__ == '__main__':
2020-03-11 18:02:19 +08:00
test_replaybuffer()