def encode(_input, weights, latent, reconstruct):
conv_vae = VAE()
conv_vae.make_vae(_input + ".npz", int(latent))
conv_vae.load_model(weights)
latent_vectors = []
# for i in range(len(conv_vae.x_test)):
# latent_vectors.append(conv_vae.encode_image(conv_vae.x_test[i:i+1]))
raw_data = np.load(_input + ".npz")
for f in sorted(raw_data.files):
images = raw_data[f]
latent_vectors.append(
[conv_vae.encode_image(np.array([image]))[0] for image in images])
np.savez_compressed(_input + "_latent.npz", *latent_vectors)
if reconstruct:
data = np.load(_input + "_latent.npz")
# files = data.files
# vectors = np.array(data[files[0]])
reconstructed_images = []
for f in sorted(data.files):
latents = data[f]
# a = a.reshape(-1, a.shape[-2], a.shape[-1])
reconstructed_images.append(
[conv_vae.decode_latent(np.array([l]))[0] for l in latents])
# print(np.shape(np.array(reconstructed_images)))
# plt.imshow(reconstructed_images[1])
# plt.show()
np.savez_compressed(_input + "_recon.npz", *reconstructed_images)
def vae_process_images(images,
weights,
latent_size,
latents=None,
encode=True,
decode=True,
image_size=None):
conv_vae = VAE()
sys.stdout = open(os.devnull, 'w')
if image_size:
conv_vae.make_vae_shape(image_size, image_size, latent_size)
else:
conv_vae.make_vae(images + ".npz", latent_size)
conv_vae.load_model(weights)
sys.stdout = sys.__stdout__
latent_vectors = []
raw_data = np.load(images + ".npz")
if encode:
print("Encoding images...")
for i in range(len(raw_data.files)):
f = 'arr_' + str(i)
file_images = raw_data[f]
latent_vectors.append([
conv_vae.encode_image(np.array([image]))[0]
for image in file_images
])
np.savez_compressed(images + "_latent.npz", *latent_vectors)
if decode:
print("Decoding latent vectors...")
data = ""
if latents:
data = np.load(latents + ".npz")
else:
data = np.load(images + "_latent.npz")
# files = data.files
# vectors = np.array(data[files[0]])
reconstructed_images = []
for f in sorted(data.files):
latents = data[f]
# a = a.reshape(-1, a.shape[-2], a.shape[-1])
reconstructed_images.append(
[conv_vae.decode_latent(np.array([l]))[0] for l in latents])
# print(np.shape(np.array(reconstructed_images)))
# plt.imshow(reconstructed_images[1])
# plt.show()
np.savez_compressed(images + "_recon.npz", *reconstructed_images)
class Simulation:
def __init__(self, path, controller_weights=None):
self.params = json.load(open(path))[0]
self.load_model(controller_weights=controller_weights)
self.env = gym.make(self.params['env_name'])
def load_model(self, controller_weights=None):
p = self.params
self.action_utils = ActionUtils(p['env_name'])
self.action_size = self.action_utils.action_size()
self.vae = VAE()
sys.stdout = open(os.devnull, 'w')
self.vae.make_vae_shape(p['img_size'], p['img_size'], p['latent_size'])
sys.stdout = sys.__stdout__
self.vae.load_model('../' + p['vae_hps']['weights_path'])
# TODO: Make MDN just take in all of params.
mdn_hps = p['mdn_hps']
mdn_hps['max_seq_len'] = p['max_seq_len']
mdn_hps['in_width'] = p['latent_size'] + self.action_size
mdn_hps['out_width'] = p['latent_size']
mdn_hps['action_size'] = self.action_size
mdn_hps['rnn_size'] = p['hidden_size']
mdn_hps['batch_size'] = 1
mdn_hps['max_seq_len'] = 1
mdn_hps['use_recurrent_dropout'] = 0
mdn_hps['training'] = 0
# self.mdn_rnn = MDNRNN(mdn_hps)
# hps_inf = MDNRNN.set_hps_to_inference(hps)
self.mdn_rnn = MDNRNN(mdn_hps)
self.mdn_rnn.load('../' + p['mdn_hps']['weights_path'])
self.controller = ControllerModel(
[p['latent_size'] + p['hidden_size'], self.action_size])
if controller_weights:
self.controller.load_weights(controller_weights)
def simulate(self, dreaming=False, render=False):
rewards = []
for i in range(1):
obs = self.env.reset()
# initialize hidden + action variables
state = self.mdn_rnn.rnn_init_state()
a = self.env.action_space.sample()
h = np.zeros((1, self.params['hidden_size']))
c = np.zeros((1, self.params['hidden_size']))
total_reward = 0
if dreaming:
img = self.env.render(mode='rgb_array')
img = compress_image(img, size=self.params['img_size'])
z = self.vae.encode_image(np.array([img]))[0]
for t in range(self.params['max_seq_len']):
z_current = z.copy()
z, state = self.mdn_rnn.sample_z(
z_current, self.action_utils.action_to_input(a), state)
z = z[0][0]
h, c = state[0], state[1]
out = self.controller.get_action(
np.concatenate((z_current, h[0])))
obs, reward, done, info = self.env.step(
self.action_utils.output_to_action(out))
total_reward += reward
if done:
print('Episode finished after {} timesteps'.format(t +
1))
break
else:
for t in range(self.params['max_seq_len']):
img = self.env.render(mode='rgb_array')
img = compress_image(img, size=self.params['img_size'])
# compute action
z = self.vae.encode_image(np.array([img]))[0]
print(a)
state = self.mdn_rnn.rnn_next_state(
z, self.action_utils.action_to_input(a), state)
h, c = state[0], state[1]
out = self.controller.get_action(np.concatenate((z, h[0])))
a = self.action_utils.output_to_action(out)
obs, reward, done, info = self.env.step(a)
total_reward += reward
if done:
print('Episode finished after {} timesteps'.format(t +
1))
break
rewards.append(total_reward)
return -np.mean(rewards)
