/
test_frame_stack.py
50 lines (41 loc) · 1.24 KB
/
test_frame_stack.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import pytest
pytest.importorskip("gym.envs.atari")
import numpy as np
import gym
from gym.wrappers import FrameStack
try:
import lz4
except ImportError:
lz4 = None
@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v1", "Pong-v0"])
@pytest.mark.parametrize("num_stack", [2, 3, 4])
@pytest.mark.parametrize(
"lz4_compress",
[
pytest.param(
True,
marks=pytest.mark.skipif(
lz4 is None, reason="Need lz4 to run tests with compression"
),
),
False,
],
)
def test_frame_stack(env_id, num_stack, lz4_compress):
env = gym.make(env_id)
env.seed(0)
shape = env.observation_space.shape
env = FrameStack(env, num_stack, lz4_compress)
assert env.observation_space.shape == (num_stack,) + shape
assert env.observation_space.dtype == env.env.observation_space.dtype
dup = gym.make(env_id)
dup.seed(0)
obs = env.reset()
dup_obs = dup.reset()
assert np.allclose(obs[-1], dup_obs)
for _ in range(num_stack ** 2):
action = env.action_space.sample()
dup_obs, _, _, _ = dup.step(action)
obs, _, _, _ = env.step(action)
assert np.allclose(obs[-1], dup_obs)
assert len(obs) == num_stack