def test_controller():
print("\n\n\n\n\n")
# env_name = "SonicTheHedgehog-Genesis" # None
# env_name = "SonicTheHedgehog2-Genesis"
# env_name = "SonicAndKnuckles-Genesis"
# env_name = "SonicTheHedgehog3-Genesis"
env_name = "SonicAndKnuckles3-Genesis"
env = retro.make(env_name)
# print(env.observation_space) # Box(224, 320, 3)
# print(env.action_space) # MultiBinary(12)
# print(env.action_space.sample()) # [1 1 1 0 1 0 1 0 0 1 1 1]
conv_vae_filename = "weights/conv_vae_SonicAndKnuckles.pkl"
lstm_mdn_filename = "weights/lstm_mdn_SonicAndKnuckles.pkl"
solver_filename = "weights/solver_sigma_mu_weights_22_0.31758243.npz"
# only forward pass
conv_vae = ConvVAE((3, 128, 128), 4608)
if os.path.exists(conv_vae_filename):
print("loading cnn vae weights")
conv_vae.load_state_dict(torch.load(conv_vae_filename))
# only forward pass
lstm_mdn = LSTM(sequence_len=1)
if os.path.exists(lstm_mdn_filename):
print("loading lstm mdn weights")
lstm_mdn.load_state_dict(torch.load(lstm_mdn_filename))
controller = Controller(input_size=6656, action_size=12)
if os.path.exists(solver_filename):
print("loading controller weights")
data = np.load(solver_filename)
weights = data["weights"]
# solver.mu = data["mu"]
# solver.sigma = data["sigma"]
controller.set_weights(weights)
for episode in range(1, 2):
img = env.reset()
while True:
img = dumyshape(img)
img = torch.FloatTensor(img)
# latent vector
z = conv_vae(img, encode=True)
# hidden state, cell state
h, c = lstm_mdn.encode(z.unsqueeze(0)) # [2, 1, 512] [2, 1, 512]
h = h.view(1, -1) # [1, 1024]
c = c.view(1, -1) # [1, 1024]
input = torch.cat((z, h, c), dim=1) # [1, 6656]
actions = controller(
input) # [[1., 1., 1., 0., 1., 0., 1., 1., 0., 1., 1., 0.]]
actions = actions.squeeze(0).cpu().data.numpy()
actions = actions.astype("int") # [1 0 0 0 1 1 0 0 0 1 1 1]
# decode latent vector
deconv_img = conv_vae.decode(z)
deconv_img = deconv_img.squeeze(0).cpu().data.numpy()
deconv_img = np.transpose(deconv_img, (1, 2, 0))
deconv_img *= 255
deconv_img = cv2.cvtColor(deconv_img, cv2.COLOR_BGR2RGB)
env.render()
cv2.imshow("conv vae", deconv_img)
cv2.waitKey(1)
img, reward, done, info = env.step(actions)
if done:
break
def test_lstm_mdn():
print("\n\n\n\n\n")
# env_name = "SonicTheHedgehog-Genesis" # None
env_name = "SonicTheHedgehog2-Genesis"
# env_name = "SonicAndKnuckles-Genesis"
# env_name = "SonicTheHedgehog3-Genesis"
# env_name = "SonicAndKnuckles3-Genesis"
env = retro.make(env_name)
# print(env.observation_space) # Box(224, 320, 3)
# print(env.action_space) # MultiBinary(12)
# print(env.action_space.sample()) # [1 1 1 0 1 0 1 0 0 1 1 1]
# conv_vae_filename = "weights/conv_vae_SonicAndKnuckles.pkl" # 4608
# lstm_mdn_filename = "weights/lstm_mdn_SonicAndKnuckles.pkl"
# conv_vae_filename = "weights/conv_vae_gray_edges.pkl"
# lstm_mdn_filename = "weights/lstm_mdn_gray_edges.pkl"
conv_vae_filename = "weights/conv_vae_gray.pkl"
lstm_mdn_filename = "weights/lstm_mdn_gray.pkl"
# only forward pass
conv_vae = ConvVAE((1, 128, 128), 1024)
if os.path.exists(conv_vae_filename):
print("loading cnn vae weights")
conv_vae.load_state_dict(torch.load(conv_vae_filename))
# only forward pass
lstm_mdn = LSTM(vector_size=1024)
if os.path.exists(lstm_mdn_filename):
print("loading lstm mdn weights")
lstm_mdn.load_state_dict(torch.load(lstm_mdn_filename))
for episode in range(1, 2):
img = env.reset()
while True:
# img = dumyshape(img)
# img = dumyshape_gray_edges(img)
img = dumyshape_gray(img)
img = torch.FloatTensor(img) # [3, 128, 128]
img = img.unsqueeze(0) # [1, 3, 128, 128]
z = conv_vae(img, encode=True)
z = z.unsqueeze(0) # [1, 1, 4608]
z_t = lstm_mdn.predict(z) # [1, 1, 1024]
z_t = z_t.squeeze(0) # [1, 1024]
deconv_img = conv_vae.decode(z_t) # [1, 1, 128, 128]
deconv_img = deconv_img.squeeze(0).cpu().data.numpy()
deconv_img = np.transpose(deconv_img, (1, 2, 0))
# deconv_img *= 255
env.render()
cv2.imshow("lstm_mdn reconstruct", deconv_img)
# cv2.imshow("lstm_mdn reconstruct", cv2.cvtColor(deconv_img, cv2.COLOR_BGR2RGB))
cv2.waitKey(1)
action = env.action_space.sample()
action[7] = 1
img, reward, done, info = env.step(action)
if done:
break
print()