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