109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
import time
|
|
import pytest
|
|
import numpy as np
|
|
from tianshou.env import FrameStack, VectorEnv, SubprocVectorEnv, RayVectorEnv
|
|
|
|
if __name__ == '__main__':
|
|
from env import MyTestEnv
|
|
else: # pytest
|
|
from test.base.env import MyTestEnv
|
|
|
|
|
|
def test_framestack(k=4, size=10):
|
|
env = MyTestEnv(size=size)
|
|
fsenv = FrameStack(env, k)
|
|
fsenv.seed()
|
|
obs = fsenv.reset()
|
|
assert abs(obs - np.array([0, 0, 0, 0])).sum() == 0
|
|
for i in range(5):
|
|
obs, rew, done, info = fsenv.step(1)
|
|
assert abs(obs - np.array([2, 3, 4, 5])).sum() == 0
|
|
for i in range(10):
|
|
obs, rew, done, info = fsenv.step(0)
|
|
assert abs(obs - np.array([0, 0, 0, 0])).sum() == 0
|
|
for i in range(9):
|
|
obs, rew, done, info = fsenv.step(1)
|
|
assert abs(obs - np.array([6, 7, 8, 9])).sum() == 0
|
|
assert (rew, done) == (0, False)
|
|
obs, rew, done, info = fsenv.step(1)
|
|
assert abs(obs - np.array([7, 8, 9, 10])).sum() == 0
|
|
assert (rew, done) == (1, True)
|
|
with pytest.raises(ValueError):
|
|
obs, rew, done, info = fsenv.step(0)
|
|
# assert abs(obs - np.array([8, 9, 10, 10])).sum() == 0
|
|
# assert (rew, done) == (0, True)
|
|
fsenv.close()
|
|
|
|
|
|
def test_vecenv(size=10, num=8, sleep=0.001):
|
|
verbose = __name__ == '__main__'
|
|
env_fns = [lambda: MyTestEnv(size=size, sleep=sleep) for _ in range(num)]
|
|
venv = [
|
|
VectorEnv(env_fns, reset_after_done=True),
|
|
SubprocVectorEnv(env_fns, reset_after_done=True),
|
|
]
|
|
if verbose:
|
|
venv.append(RayVectorEnv(env_fns, reset_after_done=True))
|
|
for v in venv:
|
|
v.seed()
|
|
action_list = [1] * 5 + [0] * 10 + [1] * 15
|
|
if not verbose:
|
|
o = [v.reset() for v in venv]
|
|
for i, a in enumerate(action_list):
|
|
o = [v.step([a] * num) for v in venv]
|
|
for i in zip(*o):
|
|
for j in range(1, len(i)):
|
|
assert (i[0] == i[j]).all()
|
|
else:
|
|
t = [0, 0, 0]
|
|
for i, e in enumerate(venv):
|
|
t[i] = time.time()
|
|
e.reset()
|
|
for a in action_list:
|
|
e.step([a] * num)
|
|
t[i] = time.time() - t[i]
|
|
print(f'VectorEnv: {t[0]:.6f}s')
|
|
print(f'SubprocVectorEnv: {t[1]:.6f}s')
|
|
print(f'RayVectorEnv: {t[2]:.6f}s')
|
|
for v in venv:
|
|
v.close()
|
|
|
|
|
|
def test_vecenv2():
|
|
verbose = __name__ == '__main__'
|
|
env_fns = [
|
|
lambda: MyTestEnv(size=1),
|
|
lambda: MyTestEnv(size=2),
|
|
lambda: MyTestEnv(size=3),
|
|
lambda: MyTestEnv(size=4),
|
|
]
|
|
num = len(env_fns)
|
|
venv = [
|
|
VectorEnv(env_fns, reset_after_done=False),
|
|
SubprocVectorEnv(env_fns, reset_after_done=False),
|
|
]
|
|
if verbose:
|
|
venv.append(RayVectorEnv(env_fns, reset_after_done=False))
|
|
for v in venv:
|
|
v.seed()
|
|
o = [v.reset() for v in venv]
|
|
action_list = [1] * 6
|
|
for i, a in enumerate(action_list):
|
|
o = [v.step([a] * num) for v in venv]
|
|
if verbose:
|
|
print(o[0])
|
|
print(o[1])
|
|
print(o[2])
|
|
print('---')
|
|
for i in zip(*o):
|
|
for j in range(1, len(i)):
|
|
assert (i[0] == i[j]).all()
|
|
for v in venv:
|
|
v.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_framestack()
|
|
test_vecenv()
|
|
test_vecenv2()
|