def test_env_manager_init(self):
stack_size = 3
env = FrameStackEnvManager('Breakout-v0',
preprocess=preprocess,
frame_stack_size=stack_size)
self.assertEqual(0, len(env.frames))
env.reset()
self.assertEqual(env.state()[0].shape, torch.Size(img_size))
self.assertEqual(env.state().shape[0], stack_size)
def test_env_manager_state(self):
frame_stack_size = 5
env = FrameStackEnvManager('Breakout-v0',
preprocess=preprocess,
frame_stack_size=frame_stack_size)
env.reset()
state = env.state()
self.assertEqual(frame_stack_size, state.shape[0])
screen = cv2.resize(screen, (64, 64), interpolation=cv2.INTER_NEAREST)
screen = np.ascontiguousarray(screen)
return screen
env = FrameStackEnvManager('Breakout-v0', preprocess=preprocess)
raw_screen = env.get_raw_screen()
plt.figure()
plt.title('Raw Screen')
plt.axis('off')
plt.imshow(raw_screen) #.transpose(1,2,0))
plt.show()
env.reset()
state = env.state()
plt.suptitle('Initial State')
plt.subplot(141)
plt.axis('off')
plt.imshow(state[0], cmap=plt.get_cmap('gray'))
plt.subplot(142)
plt.axis('off')
plt.imshow(state[1], cmap=plt.get_cmap('gray'))
plt.subplot(143)
plt.axis('off')
plt.imshow(state[2], cmap=plt.get_cmap('gray'))
plt.subplot(144)
plt.axis('off')
plt.imshow(state[3], cmap=plt.get_cmap('gray'))
plt.show()