reset_graph()
# Creating the VAE model as an object of the ConvVAE class
vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
learning_rate=learning_rate,
kl_tolerance=kl_tolerance,
is_training=False,
reuse=False,
gpu_mode=True)
# Loading the weights of the VAE model
vae.load_json(os.path.join(model_path_name, 'vae.json'))
# Running the main code that generates the data from the VAE model for the MDN-RNN model
mu_dataset = []
logvar_dataset = []
for i in range(len(dataset)):
data_batch = dataset[i]
mu, logvar, z = encode_batch(data_batch)
mu_dataset.append(mu.astype(np.float16))
logvar_dataset.append(logvar.astype(np.float16))
if ((i + 1) % 100 == 0):
print(i + 1)
action_dataset = np.array(action_dataset)
mu_dataset = np.array(mu_dataset)
logvar_dataset = np.array(logvar_dataset)
class Model:
''' simple one layer model for car racing '''
def __init__(self, load_model=True):
self.env_name = "carracing"
self.vae = ConvVAE(batch_size=1,
gpu_mode=False,
is_training=False,
reuse=True)
self.rnn = MDNRNN(hps_sample, gpu_mode=False, reuse=True)
if load_model:
self.vae.load_json('vae/vae.json')
self.rnn.load_json('rnn/rnn.json')
self.state = rnn_init_state(self.rnn)
self.rnn_mode = True
self.input_size = rnn_output_size(EXP_MODE)
self.z_size = 32
if EXP_MODE == MODE_Z_HIDDEN: # one hidden layer
self.hidden_size = 40
self.weight_hidden = np.random.randn(self.input_size,
self.hidden_size)
self.bias_hidden = np.random.randn(self.hidden_size)
self.weight_output = np.random.randn(self.hidden_size, 3)
self.bias_output = np.random.randn(3)
self.param_count = ((self.input_size + 1) *
self.hidden_size) + (self.hidden_size * 3 + 3)
else:
self.weight = np.random.randn(self.input_size, 3)
self.bias = np.random.randn(3)
self.param_count = (self.input_size) * 3 + 3
self.render_mode = False
def make_env(self, seed=-1, render_mode=False, full_episode=False):
self.render_mode = render_mode
self.env = make_env(self.env_name,
seed=seed,
render_mode=render_mode,
full_episode=full_episode)
def reset(self):
self.state = rnn_init_state(self.rnn)
def encode_obs(self, obs):
# convert raw obs to z, mu, logvar
result = np.copy(obs).astype(np.float) / 255.0
result = result.reshape(1, 96, 96, 3)
mu, logvar = self.vae.encode_mu_logvar(result)
mu = mu[0]
logvar = logvar[0]
s = logvar.shape
z = mu + np.exp(logvar / 2.0) * np.random.randn(*s)
return z, mu, logvar
def get_action(self, z):
h = rnn_output(self.state, z, EXP_MODE)
'''
action = np.dot(h, self.weight) + self.bias
action[0] = np.tanh(action[0])
action[1] = sigmoid(action[1])
action[2] = clip(np.tanh(action[2]))
'''
if EXP_MODE == MODE_Z_HIDDEN: # one hidden layer
h = np.tanh(np.dot(h, self.weight_hidden) + self.bias_hidden)
action = np.tanh(np.dot(h, self.weight_output) + self.bias_output)
else:
action = np.tanh(np.dot(h, self.weight) + self.bias)
action[1] = (action[1] + 1.0) / 2.0
action[2] = clip(action[2])
self.state = rnn_next_state(self.rnn, z, action, self.state)
return action
def set_model_params(self, model_params):
if EXP_MODE == MODE_Z_HIDDEN: # one hidden layer
params = np.array(model_params)
cut_off = (self.input_size + 1) * self.hidden_size
params_1 = params[:cut_off]
params_2 = params[cut_off:]
self.bias_hidden = params_1[:self.hidden_size]
self.weight_hidden = params_1[self.hidden_size:].reshape(
self.input_size, self.hidden_size)
self.bias_output = params_2[:3]
self.weight_output = params_2[3:].reshape(self.hidden_size, 3)
else:
self.bias = np.array(model_params[:3])
self.weight = np.array(model_params[3:]).reshape(
self.input_size, 3)
def load_model(self, filename):
with open(filename) as f:
data = json.load(f)
print('loading file %s' % (filename))
self.data = data
model_params = np.array(data[0]) # assuming other stuff is in data
self.set_model_params(model_params)
def get_random_model_params(self, stdev=0.1):
#return np.random.randn(self.param_count)*stdev
return np.random.standard_cauchy(
self.param_count) * stdev # spice things up
def init_random_model_params(self, stdev=0.1):
params = self.get_random_model_params(stdev=stdev)
self.set_model_params(params)
vae_params = self.vae.get_random_model_params(stdev=stdev)
self.vae.set_model_params(vae_params)
rnn_params = self.rnn.get_random_model_params(stdev=stdev)
self.rnn.set_model_params(rnn_params)
class Model:
# Initializing all the parameters and variables of the Model class
def __init__(self, load_model=True):
self.env_name = "carracing"
self.vae = ConvVAE(batch_size=1,
gpu_mode=False,
is_training=False,
reuse=True)
self.rnn = MDNRNN(hps_sample, gpu_mode=False, reuse=True)
if load_model:
self.vae.load_json('Weights/vae_weights.json')
self.rnn.load_json('Weights/rnn_weights.json')
self.state = rnn_init_state(self.rnn)
self.rnn_mode = True
self.input_size = rnn_output_size(EXP_MODE)
self.z_size = 32
if EXP_MODE == MODE_Z_HIDDEN:
self.hidden_size = 40
self.weight_hidden = np.random.randn(self.input_size,
self.hidden_size)
self.bias_hidden = np.random.randn(self.hidden_size)
self.weight_output = np.random.randn(self.hidden_size, 3)
self.bias_output = np.random.randn(3)
self.param_count = ((self.input_size + 1) *
self.hidden_size) + (self.hidden_size * 3 + 3)
else:
self.weight = np.random.randn(self.input_size, 3)
self.bias = np.random.randn(3)
self.param_count = (self.input_size) * 3 + 3
self.render_mode = False
# Making a method that creates an environment (in our case the CarRacing game) inside which both the AI and HI will play
def make_env(self, seed=-1, render_mode=False, full_episode=False):
self.render_mode = render_mode
self.env = make_env(self.env_name,
seed=seed,
render_mode=render_mode,
full_episode=full_episode)
# Making a method that reinitiates the states for the RNN model
def reset(self):
self.state = rnn_init_state(self.rnn)
# Making a method that encodes the observations (input frames)
def encode_obs(self, obs):
result = np.copy(obs).astype(np.float) / 255.0
result = result.reshape(1, 64, 64, 3)
mu, logvar = self.vae.encode_mu_logvar(result)
mu = mu[0]
logvar = logvar[0]
s = logvar.shape
z = mu + np.exp(logvar / 2.0) * np.random.randn(*s)
return z, mu, logvar
# Making a method that samples an action based on the latent vector z
def get_action(self, z):
h = rnn_output(self.state, z, EXP_MODE)
if EXP_MODE == MODE_Z_HIDDEN:
h = np.tanh(np.dot(h, self.weight_hidden) + self.bias_hidden)
action = np.tanh(np.dot(h, self.weight_output) + self.bias_output)
else:
action = np.tanh(np.dot(h, self.weight) + self.bias)
action[1] = (action[1] + 1.0) / 2.0
action[2] = clip(action[2])
self.state = rnn_next_state(self.rnn, z, action, self.state)
return action
# Making a method that sets the initialized/loaded weights into the model
def set_model_params(self, model_params):
if EXP_MODE == MODE_Z_HIDDEN:
params = np.array(model_params)
cut_off = (self.input_size + 1) * self.hidden_size
params_1 = params[:cut_off]
params_2 = params[cut_off:]
self.bias_hidden = params_1[:self.hidden_size]
self.weight_hidden = params_1[self.hidden_size:].reshape(
self.input_size, self.hidden_size)
self.bias_output = params_2[:3]
self.weight_output = params_2[3:].reshape(self.hidden_size, 3)
else:
self.bias = np.array(model_params[:3])
self.weight = np.array(model_params[3:]).reshape(
self.input_size, 3)
# Making a method that loads the model weights
def load_model(self, filename):
with open(filename) as f:
data = json.load(f)
print('loading file %s' % (filename))
self.data = data
model_params = np.array(data[0])
self.set_model_params(model_params)
# Making a method that randomly initializes the weights
def get_random_model_params(self, stdev=0.1):
return np.random.standard_cauchy(self.param_count) * stdev
# Making a method that randomly initializes the weights for all 3 parts of the Full World Model (VAE, RNN, Controller)
def init_random_model_params(self, stdev=0.1):
params = self.get_random_model_params(stdev=stdev)
self.set_model_params(params)
vae_params = self.vae.get_random_model_params(stdev=stdev)
self.vae.set_model_params(vae_params)
rnn_params = self.rnn.get_random_model_params(stdev=stdev)
self.rnn.set_model_params(rnn_params)
print("batch_size {}".format(batch_size))
print("crop_size {}".format(crop_size))
cv2.imshow("crop",crops[0])
cv2.waitKey(0)
reset_graph()
test_vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
is_training=False,
reuse=False,
gpu_mode=True)
# show reconstruction example
test_vae.load_json("../../models/0/vae_{}.json".format(180000))
z = test_vae.encode(crops)
print(z.shape)
rec = test_vae.decode(z)
print(rec.shape)
np.save("../../output/z_{}.npy".format(batch_size), z)
np.save("../../output/rec_{}.npy".format(batch_size), rec)
#for img_idx, img in enumerate(test_batch):
vis = np.concatenate((crops[0], rec[0]), axis=1)
cv2.imshow("org vs. rec",cv2.resize(vis,(0,0),fx=4.0,fy=4.0))
key = cv2.waitKey(0) & 0xFF