def train_epoch_joint(motion_train_loader, X_motion_test, y_motion_test, conv_encoder, predictor, forward_steps, aux_coeff):
for batch_id, (X_motion_batch, y_motion_batch) in enumerate(motion_train_loader):
X_motion_batch = Variable(X_motion_batch)
y_motion_batch = Variable(y_motion_batch)
optimizer_motion.zero_grad()
loss_auxiliary, loss_pred_recons, pred_recons_batch = get_losses(conv_encoder, predictor, X_motion_batch, y_motion_batch, forward_steps = forward_steps)
reg_conv = conv_encoder.get_regularization(source = ["weight", "bias"]) * reg_amp_conv * reg_multiplier[epoch]
reg_predictor = predictor.get_regularization(source = ["weight", "bias"]) * reg_amp_predictor * reg_multiplier[epoch]
loss_train = loss_auxiliary * aux_coeff + loss_pred_recons + reg_conv + reg_predictor
loss_train.backward()
optimizer_motion.step()
loss_auxiliary_test, loss_pred_recons_test, pred_recons_test = get_losses(conv_encoder, predictor, X_motion_test, y_motion_test, forward_steps = forward_steps)
loss_test = loss_auxiliary_test * aux_coeff + loss_pred_recons_test + reg_conv + reg_predictor
to_stop = early_stopping_motion.monitor(loss_test.data[0])
print("epoch {0}\tloss_train: {1:.6f}\tloss_test: {2:.6f}\tloss_aux: {3:.6f}\tloss_pred: {4:.6f}\treg_conv: {5:.6f}\treg_predictor: {6:.6f}".format(
epoch, loss_train.data[0], loss_test.data[0], loss_auxiliary_test.data[0] * aux_coeff, loss_pred_recons_test.data[0], reg_conv.data[0], reg_predictor.data[0]))
if epoch % 10 == 0:
print("epoch {0}:".format(epoch))
plot_matrices(np.concatenate((X_motion_batch[0].cpu().data.numpy(), y_motion_batch[:, torch.LongTensor(np.array(forward_steps) - 1).cuda()][0].cpu().data.numpy())))
plot_matrices(np.concatenate((forward(conv_encoder, X_motion_batch)[0].cpu().data.numpy(), pred_recons_batch[0].cpu().data.numpy())))
print("encoding:")
plot_encoding(X_motion_batch[:,:1].contiguous(), conv_encoder, target = "encoding")
print("axis:")
plot_encoding(X_motion_batch[:,:1].contiguous(), conv_encoder, target = "axis")
print("\n\n")
return to_stop
def get_weights_bias(self,
W_source=None,
b_source=None,
layer_ids=None,
isplot=False,
raise_error=True):
layer_ids = range(len(
self.struct_param)) if layer_ids is None else layer_ids
W_list = []
b_list = []
if W_source is not None:
for k in range(len(self.struct_param)):
if k in layer_ids:
if W_source == "core":
try:
W, _ = getattr(
self,
"layer_{0}".format(k)).get_weights_bias()
except Exception as e:
if raise_error:
raise
else:
print(e)
W = np.array([np.NaN])
else:
raise Exception(
"W_source '{0}' not recognized!".format(W_source))
W_list.append(W)
if b_source is not None:
for k in range(len(self.struct_param)):
if k in layer_ids:
if b_source == "core":
try:
_, b = getattr(
self,
"layer_{0}".format(k)).get_weights_bias()
except Exception as e:
if raise_error:
raise
else:
print(e)
b = np.array([np.NaN])
else:
raise Exception(
"b_source '{0}' not recognized!".format(b_source))
b_list.append(b)
if isplot:
if W_source is not None:
print("weight {0}:".format(W_source))
plot_matrices(W_list)
if b_source is not None:
print("bias {0}:".format(b_source))
plot_matrices(b_list)
return W_list, b_list
def train_epoch_pretrain(train_loader, X_test_all, autoencoder, optimizer_pre):
for batch_id, X_batch in enumerate(train_loader):
X_batch = Variable(X_batch)
optimizer_pre.zero_grad()
reconstruct = forward(autoencoder, X_batch)
reg = autoencoder.get_regularization(source = ["weight", "bias"]) * reg_amp_autoencoder
reg_latent = forward(autoencoder, X_batch).mean() * reg_amp_latent
loss_train = nn.MSELoss()(reconstruct, X_batch) + reg + reg_latent
loss_train.backward()
optimizer_pre.step()
reconstruct_test = forward(autoencoder, X_test_all)
loss_test = nn.MSELoss()(reconstruct_test, X_test_all)
to_stop = early_stopping_pre.monitor(loss_test.data[0])
print("epoch {0} \tloss_train: {1:.6f}\tloss_test: {2:.6f}\treg: {3:.6f}\treg_latent: {4:.6f}".format(epoch, loss_train.data[0], loss_test.data[0], reg.data[0], reg_latent.data[0]))
if epoch % 10 == 0:
plot_matrices(X_batch[0].cpu().data.numpy(), images_per_row = 5)
latent = forward(autoencoder.encode, X_batch)
print("latent: {0}".format(latent.cpu().data.numpy()[0]))
plot_matrices(reconstruct[0].cpu().data.numpy(), images_per_row = 5)
plot_encoding(X_train_all[:,:1], autoencoder, "axis")
return to_stop
def plot_tasks(tasks, autoencoder, forward_steps, num_tasks = 3):
task_keys = np.random.choice(list(tasks.keys()), num_tasks, replace = False)
for i, task_key in enumerate(task_keys):
task = tasks[task_key]
((X_train_obs, y_train_obs), (X_test_obs, y_test_obs)), _ = task
X_train = forward(autoencoder.encode, X_train_obs)
y_train = forward(autoencoder.encode, y_train_obs[:, forward_steps_idx])
X_test = forward(autoencoder.encode, X_test_obs)
y_test = forward(autoencoder.encode, y_test_obs[:, forward_steps_idx])
# Plotting:
if X_train.is_cuda:
print("Task {0}:".format(task_key))
plot_matrices(np.concatenate((X_test_obs[0].cpu().data.numpy(), y_test_obs[:, torch.LongTensor(np.array(forward_steps) - 1).cuda()][0].cpu().data.numpy())))
latent_pred = get_forward_pred(generative_Net, X_test, forward_steps)
pred_recons = forward(autoencoder.decode, latent_pred)
plot_matrices(np.concatenate((forward(autoencoder.decode, X_test.view(X_test.size(0), -1, 2))[0].cpu().data.numpy(), pred_recons[0].cpu().data.numpy())))
else:
print("Task {0}:".format(task_key))
plot_matrices(np.concatenate((X_test_obs[0].data.numpy(), y_test_obs[:, torch.LongTensor(np.array(forward_steps) - 1)][0].data.numpy())))
latent_pred = get_forward_pred(generative_Net, X_test, forward_steps)
pred_recons = forward(autoencoder.decode, latent_pred)
plot_matrices(np.concatenate((forward(autoencoder.decode, X_test.view(X_test.size(0), -1, 2))[0].data.numpy(), pred_recons[0].data.numpy())))
# ### Regression of statistics vs. z:
# In[ ]:
from sklearn import linear_model
reg = linear_model.LinearRegression()
# reg = linear_model.Ridge(alpha = .5)
# reg = linear_model.Lasso(alpha = 0.1)
# reg = linear_model.BayesianRidge()
# reg = linear_model.HuberRegressor()
coeff_list = []
for i in range(statistics_list_test.shape[1]):
reg.fit(z_list_test, statistics_list_test[:, i])
coeff_list.append(reg.coef_)
coeff_list = np.array(coeff_list)
plot_matrices([abs(coeff_list)])
print(coeff_list)
# ### Fitting neural network to statistics vs. z:
# In[ ]:
struct_param = [
[120, "Simple_Layer", {
"activation": "leakyRelu"
}],
[120, "Simple_Layer", {
"activation": "leakyRelu"
}],
[120, "Simple_Layer", {
"activation": "leakyRelu"
def get_env_data(
env_name,
data_format = "states",
input_dims = 1,
output_dims = None,
ball_idx = 0,
num_examples = 10000,
test_size = 0.2,
is_cuda = False,
isplot = 1,
verbose = True,
**kwargs
):
if env_name[:9] == "piecewise":
env_name_split = env_name.split("-")
input_size = int(env_name_split[1][:-1])
num_pieces = int(env_name_split[2][:-1])
num_boundaries = int(env_name_split[3][:-1])
func_types = []
for letter in env_name_split[4]:
if letter == "l":
func_types.append("linear")
elif letter == "q":
func_types.append("quadratic")
else:
raise Exception("letter {0} is not a valid function type!".format(letter))
((X_train, y_train), (X_test, y_test), (reflected_train, reflected_test)), info = get_piecewise_dataset(input_size = input_size,
num_pieces = num_pieces,
num_boundaries = num_boundaries,
func_types = func_types,
x_range = kwargs["x_range"] if "x_range" in kwargs else (0, 20),
num_examples = num_examples,
is_cuda = is_cuda,
isplot = isplot,
)
else:
from mela.settings.a2c_env_settings import ENV_SETTINGS_CHOICE
from mela.variational.envs import make_env
from mela.util import plot_matrices
import random
# Obtain settings from kwargs:
time_steps = kwargs["time_steps"] if "time_steps" in kwargs else 3
forward_steps = kwargs["forward_steps"] if "forward_steps" in kwargs else [1]
episode_length = kwargs["episode_length"] if "episode_length" in kwargs else 30
is_flatten = kwargs["is_flatten"] if "is_flatten" in kwargs else False
bounce_focus = kwargs["bounce_focus"] if "bounce_focus" in kwargs else True
normalize = kwargs["normalize"] if "normalize" in kwargs else True
translation = kwargs["translation"] if "translation" in kwargs else None
render = kwargs["render"] if "render" in kwargs else False
env_name_split = env_name.split("-")
if "nobounce" in env_name_split:
env_name_core = "-".join(env_name_split[:-1])
else:
env_name_core = env_name
env_settings = {key: random.choice(value) if isinstance(value, list) else value for key, value in ENV_SETTINGS_CHOICE[env_name_core].items()}
env_settings["info_contents"] = ["coordinates"]
max_distance = env_settings["max_distance"] if "max_distance" in env_settings else None
max_range = env_settings["max_range"] if "max_range" in env_settings else None
if max_range is not None:
value_min, value_max = max_range
input_dims = env_settings["input_dims"] if "input_dims" in env_settings else input_dims
# Reset certain aspects of the environment:
if "screen_width" in kwargs:
print("corrected screen_width: {0}".format(kwargs["screen_width"]))
env_settings["screen_width"] = kwargs["screen_width"]
if "screen_height" in kwargs:
print("corrected screen_height: {0}".format(kwargs["screen_height"]))
env_settings["screen_height"] = kwargs["screen_height"]
if "physics" in kwargs:
print("corrected physics: {0}".format(kwargs["physics"]))
env_settings["physics"] = kwargs["physics"]
if "boundaries" in kwargs:
print("corrected boundaries: {0}".format(kwargs["boundaries"]))
env_settings["boundaries"] = kwargs["boundaries"]
if "ball_vmax" in kwargs:
print("corrected ball_vmax: {0}".format(kwargs["ball_vmax"]))
env_settings["ball_vmax"] = kwargs["ball_vmax"]
if "step_dt" in kwargs:
print("corrected step_dt: {0}".format(kwargs["step_dt"]))
env_settings["step_dt"] = kwargs["step_dt"]
env = make_env("Breakout_Custom-v0", 1, 0, "", clip_rewards = False, env_settings = env_settings)()
env.allow_early_resets = True
obs_var = []
info_list = []
bouncing_list = []
k = 0
num_episodes_candidate = max(1, 1.5 * int(num_examples / (episode_length - time_steps - max(forward_steps))))
if bounce_focus:
num_episodes_candidate * 2
while k < num_episodes_candidate:
obs = env.reset()
obs_var_candidate = []
info_list_candidate = []
bouncing_list_candidate = []
ball_x = None
ball_y = None
is_break = False
# Obtain the frames:
for i in range(episode_length):
obs, _, _, info = env.step(1)
obs_var_candidate.append(obs)
coordinates = info["coordinates"]
info_list_candidate.append(coordinates)
bouncing_list_candidate.append(info["ball_bouncing_info"])
if max_distance is not None:
last_ball_x = ball_x
last_ball_y = ball_y
ball_x, ball_y = coordinates["ball"][ball_idx]
if last_ball_x is not None:
if abs(ball_x - last_ball_x) > max_distance:
is_break = True
if verbose:
print("{0} break for too large velocity.".format(k))
break
if max_range is not None:
ball_x, ball_y = coordinates["ball"][ball_idx]
if ball_x < value_min or ball_x > value_max or ball_y < value_min or ball_y > value_max:
is_break = True
if verbose:
print("{0} break for going outsize the max_range".format(k))
break
if render:
time.sleep(0.1)
env.render('human')
# Only add the episode if it is does not break:
if not is_break:
obs_var = obs_var + obs_var_candidate
info_list = info_list + info_list_candidate
bouncing_list = bouncing_list + bouncing_list_candidate
obs_var.append({})
info_list.append({})
bouncing_list.append({})
k += 1
if isplot > 0:
plot_matrices(np.array(obs_var[:30]).squeeze())
# Process the info_list into numpy format:
perception_dict = process_object_info(info_list, chosen_dim = input_dims)
bouncing_list = [len(element[ball_idx]) if len(element) > 0 else np.NaN for element in bouncing_list]
if data_format == "images":
obs_array = np.array([element if len(element) > 0 else np.full(obs_var[0].shape, np.nan) for element in obs_var])
else:
obs_array = None
trajectory0 = perception_dict["ball_{0}".format(ball_idx)]
width = env_settings["screen_width"] if input_dims == 0 else env_settings["screen_height"]
((X_train, y_train), (X_test, y_test), (reflected_train, reflected_test)), info = get_task(trajectory0,
num_examples = num_examples,
bouncing_list = bouncing_list,
obs_array = obs_array,
time_steps = time_steps,
forward_steps = forward_steps,
is_flatten = is_flatten,
output_dims = output_dims,
is_cuda = is_cuda,
width = width,
test_size = test_size,
bounce_focus = bounce_focus,
normalize = normalize,
translation = translation,
)
if "nobounce" in env_name_split:
if obs_array is None:
X_train = X_train[reflected_train.unsqueeze(1) == 0].view(-1, X_train.size(1))
y_train = y_train[reflected_train.unsqueeze(1) == 0].view(-1, y_train.size(1))
X_test = X_test[reflected_test.unsqueeze(1) == 0].view(-1, X_test.size(1))
y_test = y_test[reflected_test.unsqueeze(1) == 0].view(-1, y_test.size(1))
else:
X_train = X_train[reflected_train.view(reflected_train.size(0), 1,1,1) == 0].view(-1, *X_train.size()[1:])
y_train = y_train[reflected_train.view(reflected_train.size(0), 1,1,1) == 0].view(-1, *y_train.size()[1:])
X_test = X_test[reflected_test.view(reflected_test.size(0), 1,1,1) == 0].view(-1, *X_test.size()[1:])
y_test = y_test[reflected_test.view(reflected_test.size(0), 1,1,1) == 0].view(-1, *y_test.size()[1:])
reflected_train = reflected_train[reflected_train == 0]
reflected_test = reflected_test[reflected_test == 0]
return ((X_train, y_train), (X_test, y_test), (reflected_train, reflected_test)), info