def test_ffmodel_for_atari() -> None:
atari = gym.make("BreakoutNoFrameskip-v0")
acvp_netfn = models.prepare_ff()
d = Device()
acvp_net = acvp_netfn((3, 210, 160), 4, d)
states, actions = [], []
atari.reset()
for _ in range(10):
s, _, _, _ = atari.step(0)
states.append(s.transpose(2, 0, 1))
actions.append(0)
states = d.tensor(states)
actions = d.tensor(actions, dtype=torch.long)
s_decoded = acvp_net(states, actions)
assert tuple(s_decoded.shape) == (10, 3, 210, 160)
def test_dim(net_gen: callable, input_dim: tuple) -> None:
device = Device()
vae_net = net_gen(input_dim, ACTION_DIM, device=device)
batch = torch.randn(BATCH_SIZE, *input_dim)
with torch.no_grad():
vae, policy, value = vae_net(device.tensor(batch))
assert vae.x.shape == torch.Size((BATCH_SIZE, *input_dim))
assert policy.dist.probs.shape == torch.Size((BATCH_SIZE, ACTION_DIM))
assert value.shape == torch.Size((BATCH_SIZE, ))
print(vae_net)
