Views
No views yet
SpaceInvadersNoFrameskip-v4 environment.1import torch
2import gymnasium as gym
3
4# Load models
5vae = VAE(latent_dim=32)
6vae.load_state_dict(torch.load('vae_model.pt'))
7
8rnn = MDNRNN(latent_dim=32, action_dim=6)
9rnn.load_state_dict(torch.load('mdnrnn_model.pt'))
10
11controller = Controller(latent_dim=32, hidden_dim=256)
12controller.load_state_dict(torch.load('controller_model.pt'))
13
14# Run agent
15env = gym.make('SpaceInvadersNoFrameskip-v4')
16# ... (see repository for full inference code)1@article{ha2018worldmodels,
2 title={World Models},
3 author={Ha, David and Schmidhuber, J{\"u}rgen},
4 journal={arXiv preprint arXiv:1803.10122},
5 year={2018}
6}