def validate_step(self, x, y): """Perform a validation step on an ensemble of models without using bootstrapping weights Args: x: tf.Tensor a batch of validation inputs shaped like [batch_size, channels] y: tf.Tensor a batch of validation labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ # corrupt the inputs with noise x = cont_noise(x, self.continuous_noise_std) statistics = dict() # calculate the prediction error and accuracy of the model d = self.forward_model(x, training=False) nll = tf.keras.losses.mean_squared_error(y, d) # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d[:, 0]) statistics[f'{self.logger_prefix}/validate/nll'] = nll statistics[f'{self.logger_prefix}/validate/rank_corr'] = rank_correlation return statistics
def validate_step(self, x, y): """Perform a validation step on an ensemble of models without using bootstrapping weights Args: x: tf.Tensor a batch of validation inputs shaped like [batch_size, channels] y: tf.Tensor a batch of validation labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ statistics = dict() for i in range(self.bootstraps): fm = self.forward_models[i] # calculate the prediction error and accuracy of the model d = fm.get_distribution(x, training=False) nll = -d.log_prob(y)[:, 0] # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) statistics[f'oracle_{i}/validate/nll'] = nll statistics[f'oracle_{i}/validate/rank_corr'] = rank_correlation return statistics
def train_step(self, x, y): """Perform a training step of gradient descent on an ensemble using bootstrap weights for each model in the ensemble Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] b: tf.Tensor bootstrap indicators shaped like [batch_size, num_oracles] Returns: statistics: dict a dictionary that contains logging information """ # corrupt the inputs with noise x = cont_noise(x, self.continuous_noise_std) statistics = dict() with tf.GradientTape() as tape: # calculate the prediction error and accuracy of the model d = self.forward_model(x, training=True) nll = tf.keras.losses.mean_squared_error(y, d) # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d[:, 0]) multiplier_loss = 0.0 last_weight = self.forward_model.trainable_variables[-1] if tf.shape(tf.reshape(last_weight, [-1]))[0] == 1: statistics[f'{self.logger_prefix}/train/tanh_multipier'] = \ self.forward_model.trainable_variables[-1] # build the total loss and weight by the bootstrap total_loss = tf.reduce_mean(nll) + multiplier_loss grads = tape.gradient(total_loss, self.forward_model.trainable_variables) self.forward_model_optim.apply_gradients( zip(grads, self.forward_model.trainable_variables)) statistics[f'{self.logger_prefix}/train/nll'] = nll statistics[f'{self.logger_prefix}/train/rank_corr'] = rank_correlation return statistics
def validate_step(self, x, y): """Perform a validation step on the loss function of a conservative objective model Args: x: tf.Tensor a batch of validation inputs shaped like [batch_size, channels] y: tf.Tensor a batch of validation labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ statistics = dict() batch_dim = tf.shape(y)[0] # corrupt the inputs with noise x = cont_noise(x, self.continuous_noise_std) # calculate the prediction error and accuracy of the model d = self.forward_model(x, training=False) mse = tf.keras.losses.mean_squared_error(y, d) statistics[f'validate/mse'] = mse # evaluate how correct the rank fo the model predictions are rank_corr = spearman(y[:, 0], d[:, 0]) statistics[f'validate/rank_corr'] = rank_corr # calculate negative samples starting from the dataset x_pos = x x_pos = tf.where(tf.random.uniform([batch_dim] + [1 for _ in x.shape[1:]]) < self.negatives_fraction, x_pos, self.solution[:batch_dim]) x_neg = self.lookahead(x_pos, self.lookahead_steps, training=False) if not self.lookahead_backprop: x_neg = tf.stop_gradient(x_neg) # calculate the prediction error and accuracy of the model d_pos = self.forward_model( {"dataset": x, "mix": x_pos, "solution": self.solution[:batch_dim]} [self.constraint_type], training=False) d_neg = self.forward_model(x_neg, training=False) conservatism = d_neg[:, 0] - d_pos[:, 0] statistics[f'validate/conservatism'] = conservatism return statistics
def train_step(self, x, y, b): """Perform a training step of gradient descent on an ensemble using bootstrap weights for each model in the ensemble Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] b: tf.Tensor bootstrap indicators shaped like [batch_size, num_oracles] Returns: statistics: dict a dictionary that contains logging information """ statistics = dict() for i in range(self.bootstraps): model = self.forward_models[i] optim = self.forward_model_optims[i] # corrupt the inputs with noise if self.is_discrete: x0 = disc_noise(x, keep=self.keep, temp=self.temp) else: x0 = cont_noise(x, self.noise_std) with tf.GradientTape(persistent=True) as tape: # calculate the prediction error and accuracy of the model d = model.get_distribution(x0, training=True) nll = -d.log_prob(y) statistics[f'oracle_{i}/train/nll'] = nll # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) statistics[f'oracle_{i}/train/rank_corr'] = rank_correlation # build the total loss total_loss = tf.math.divide_no_nan( tf.reduce_sum(b[:, i] * nll), tf.reduce_sum(b[:, i])) var_list = model.trainable_variables optim.apply_gradients( zip(tape.gradient(total_loss, var_list), var_list)) return statistics
def train_step(self, x, y, b): """Perform a training step of gradient descent on an ensemble using bootstrap weights for each model in the ensemble Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] b: tf.Tensor bootstrap indicators shaped like [batch_size, num_oracles] Returns: statistics: dict a dictionary that contains logging information """ # corrupt the inputs with noise x0 = cont_noise(x, self.noise_std) statistics = dict() with tf.GradientTape(persistent=True) as tape: # calculate the prediction error and accuracy of the model d = self.fm.get_distribution(x0, training=True) nll = -d.log_prob(y) # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) # model loss that combines maximum likelihood model_loss = nll # build the total and lagrangian losses denom = tf.reduce_sum(b) total_loss = tf.math.divide_no_nan(tf.reduce_sum(b * model_loss), denom) grads = tape.gradient(total_loss, self.fm.trainable_variables) self.optim.apply_gradients(zip(grads, self.fm.trainable_variables)) statistics[f'train/nll'] = nll statistics[f'train/rank_corr'] = rank_correlation return statistics
def train_step(self, x, y, b, w): """Perform a training step of gradient descent on an ensemble using bootstrap weights for each model in the ensemble Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] b: tf.Tensor bootstrap indicators shaped like [batch_size, num_oracles] Returns: statistics: dict a dictionary that contains logging information """ statistics = dict() for i in range(self.bootstraps): fm = self.forward_models[i] fm_optim = self.forward_model_optims[i] with tf.GradientTape(persistent=True) as tape: # calculate the prediction error and accuracy of the model d = fm.get_distribution(x, training=True) nll = -d.log_prob(y)[:, 0] # evaluate how correct the rank fo the model predictions are rank_correlation = spearman(y[:, 0], d.mean()[:, 0]) # build the total loss and weight by the bootstrap total_loss = tf.math.divide_no_nan( tf.reduce_sum(w[:, 0] * b[:, i] * nll), tf.reduce_sum(b[:, i])) grads = tape.gradient(total_loss, fm.trainable_variables) fm_optim.apply_gradients(zip(grads, fm.trainable_variables)) statistics[f'oracle_{i}/train/nll'] = nll statistics[f'oracle_{i}/train/rank_corr'] = rank_correlation return statistics
def validate_step(self, x, y): """Perform a validation step on an ensemble of models without using bootstrapping weights Args: x: tf.Tensor a batch of validation inputs shaped like [batch_size, channels] y: tf.Tensor a batch of validation labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ statistics = dict() # calculate the prediction error and accuracy of the model d_pos = self.forward_model(x, training=False) mse = tf.keras.losses.mean_squared_error(y, d_pos) statistics[f'validate/mse'] = mse # evaluate how correct the rank fo the model predictions are rank_corr = spearman(y[:, 0], d_pos[:, 0]) statistics[f'validate/rank_corr'] = rank_corr # calculate negative samples starting from the dataset x_neg = self.outer_optimize( x, self.beta, self.outer_gradient_steps, training=False) # calculate the prediction error and accuracy of the model d_neg = self.forward_model(x_neg, training=False) conservatism = d_pos[:, 0] - d_neg[:, 0] statistics[f'validate/conservatism'] = conservatism return statistics
def evaluate_solution(xt): nonlocal evaluations, score # evaluate the design using the oracle and the forward model with tf.GradientTape() as tape: tape.watch(xt) model = forward_model(xt) # evaluate the predictions and gradient norm evaluations += 1 grads = tape.gradient(model, xt) model = model * st_y + mu_y for i, val in enumerate(validation_models): prediction = val(xt) logger.record(f"validation_model_{i}/prediction", prediction * st_y + mu_y, evaluations) # record the prediction and score to the logger logger.record("distance/travelled", tf.linalg.norm(xt - initial_x), evaluations) logger.record(f"train/prediction", model, evaluations) logger.record( f"train/grad_norm", tf.linalg.norm(tf.reshape(grads, [grads.shape[0], -1]), axis=-1), evaluations) if evaluations in config['evaluate_steps'] \ or len(config['evaluate_steps']) == 0 or score is None: solution = xt * st_x + mu_x if config['is_discrete']: solution = tf.math.softmax( tf.pad(solution, [[0, 0], [0, 0], [1, 0]]) / 0.001) score = task.score(solution) logger.record("score", score, evaluations, percentile=True) logger.record(f"rank_corr/model_to_real", spearman(model[:, 0], score[:, 0]), evaluations) return score, model
def train_step(self, x, y): """Perform a training step of gradient descent on an ensemble using bootstrap weights for each model in the ensemble Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ # corrupt the inputs with noise x = cont_noise(x, self.continuous_noise_std) statistics = dict() with tf.GradientTape(persistent=True) as tape: # calculate the prediction error and accuracy of the model d_pos = self.forward_model(x, training=True) mse = tf.keras.losses.mean_squared_error(y, d_pos) statistics[f'train/mse'] = mse # evaluate how correct the rank fo the model predictions are rank_corr = spearman(y[:, 0], d_pos[:, 0]) statistics[f'train/rank_corr'] = rank_corr # calculate negative samples starting from the dataset x_neg = self.outer_optimize( x, self.beta, self.outer_gradient_steps, training=False) x_neg = tf.stop_gradient(x_neg) # calculate the prediction error and accuracy of the model d_neg = self.forward_model(x_neg, training=False) conservatism = d_pos[:, 0] - d_neg[:, 0] statistics[f'train/conservatism'] = conservatism # build a lagrangian for dual descent alpha_loss = -(self.alpha * self.target_conservatism - self.alpha * conservatism) statistics[f'train/alpha'] = self.alpha # loss that combines maximum likelihood with a constraint model_loss = mse - self.alpha * conservatism total_loss = tf.reduce_mean(model_loss) alpha_loss = tf.reduce_mean(alpha_loss) # calculate gradients using the model alpha_grads = tape.gradient(alpha_loss, self.log_alpha) model_grads = tape.gradient( total_loss, self.forward_model.trainable_variables) # take gradient steps on the model self.alpha_opt.apply_gradients([[alpha_grads, self.log_alpha]]) self.forward_model_opt.apply_gradients(zip( model_grads, self.forward_model.trainable_variables)) return statistics
def train_step(self, x, y): """Perform a training step of gradient descent on the loss function of a conservative objective model Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] y: tf.Tensor a batch of training labels shaped like [batch_size, 1] Returns: statistics: dict a dictionary that contains logging information """ self.step.assign_add(1) statistics = dict() batch_dim = tf.shape(y)[0] # corrupt the inputs with noise x = cont_noise(x, self.continuous_noise_std) with tf.GradientTape(persistent=True) as tape: # calculate the prediction error and accuracy of the model d = self.forward_model(x, training=True) mse = tf.keras.losses.mean_squared_error(y, d) statistics[f'train/mse'] = mse # evaluate how correct the rank fo the model predictions are rank_corr = spearman(y[:, 0], d[:, 0]) statistics[f'train/rank_corr'] = rank_corr # calculate negative samples starting from the dataset x_pos = x x_pos = tf.where(tf.random.uniform([batch_dim] + [1 for _ in x.shape[1:]]) < self.negatives_fraction, x_pos, self.solution[:batch_dim]) x_neg = self.lookahead(x_pos, self.lookahead_steps, training=False) if not self.lookahead_backprop: x_neg = tf.stop_gradient(x_neg) # calculate the prediction error and accuracy of the model d_pos = self.forward_model( {"dataset": x, "mix": x_pos, "solution": self.solution[:batch_dim]} [self.constraint_type], training=False) d_neg = self.forward_model(x_neg, training=False) conservatism = d_neg[:, 0] - d_pos[:, 0] statistics[f'train/conservatism'] = conservatism # build a lagrangian for dual descent alpha_loss = (self.alpha * self.target_conservatism - self.alpha * conservatism) statistics[f'train/alpha'] = self.alpha multiplier_loss = 0.0 last_weight = self.forward_model.trainable_variables[-1] if tf.shape(tf.reshape(last_weight, [-1]))[0] == 1: statistics[f'train/tanh_multipier'] = \ self.forward_model.trainable_variables[-1] # loss that combines maximum likelihood with a constraint model_loss = mse + self.alpha * conservatism + multiplier_loss total_loss = tf.reduce_mean(model_loss) alpha_loss = tf.reduce_mean(alpha_loss) # initialize stateful variables at the first iteration if self.particle_loss is None: initialization = tf.zeros_like(conservatism) self.particle_loss = tf.Variable(initialization) self.particle_constraint = tf.Variable(initialization) # calculate gradients using the model alpha_grads = tape.gradient(alpha_loss, self.log_alpha) model_grads = tape.gradient( total_loss, self.forward_model.trainable_variables) # occasionally take gradient ascent steps on the solution if tf.logical_and( tf.equal(tf.math.mod(self.step, self.solver_interval), 0), tf.math.greater_equal(self.step, self.solver_warmup)): with tf.GradientTape() as tape: # take gradient steps on the model self.alpha_opt.apply_gradients([[alpha_grads, self.log_alpha]]) self.forward_model_opt.apply_gradients( zip(model_grads, self.forward_model.trainable_variables)) # calculate the predicted score of the current solution current_score_new_model = self.forward_model( self.solution, training=False)[:, 0] # look into the future and evaluate future solutions future_new_model = self.lookahead( self.solution, self.solver_steps, training=False) future_score_new_model = self.forward_model( future_new_model, training=False)[:, 0] # evaluate the conservatism of the current solution particle_loss = (self.solver_beta * future_score_new_model - current_score_new_model) update = (self.solution - self.solver_lr * tape.gradient(particle_loss, self.solution)) # if optimizer conservatism passes threshold stop optimizing self.solution.assign(tf.where(self.done, self.solution, update)) self.particle_loss.assign(particle_loss) self.particle_constraint.assign( future_score_new_model - current_score_new_model) else: # take gradient steps on the model self.alpha_opt.apply_gradients([[alpha_grads, self.log_alpha]]) self.forward_model_opt.apply_gradients( zip(model_grads, self.forward_model.trainable_variables)) statistics[f'train/done'] = tf.cast(self.done, tf.float32) statistics[f'train/particle_loss'] = self.particle_loss statistics[f'train/particle_constraint'] = self.particle_constraint return statistics
def gradient_ascent(config): """Train a Score Function to solve a Model-Based Optimization using gradient ascent on the input design Args: config: dict a dictionary of hyper parameters such as the learning rate """ # create the training task and logger logger = Logger(config['logging_dir']) task = StaticGraphTask(config['task'], **config['task_kwargs']) if config['normalize_ys']: task.map_normalize_y() if task.is_discrete and not config["use_vae"]: task.map_to_logits() if config['normalize_xs']: task.map_normalize_x() x = task.x y = task.y if task.is_discrete and config["use_vae"]: vae_model = SequentialVAE(task, hidden_size=config['vae_hidden_size'], latent_size=config['vae_latent_size'], activation=config['vae_activation'], kernel_size=config['vae_kernel_size'], num_blocks=config['vae_num_blocks']) vae_trainer = VAETrainer(vae_model, vae_optim=tf.keras.optimizers.Adam, vae_lr=config['vae_lr'], beta=config['vae_beta']) # create the training task and logger train_data, val_data = build_pipeline( x=x, y=y, batch_size=config['vae_batch_size'], val_size=config['val_size']) # estimate the number of training steps per epoch vae_trainer.launch(train_data, val_data, logger, config['vae_epochs']) # map the x values to latent space x = vae_model.encoder_cnn.predict(x)[0] mean = np.mean(x, axis=0, keepdims=True) standard_dev = np.std(x - mean, axis=0, keepdims=True) x = (x - mean) / standard_dev input_shape = x.shape[1:] input_size = np.prod(input_shape) # make several keras neural networks with different architectures forward_models = [ ForwardModel(input_shape, activations=activations, hidden_size=config['hidden_size'], initial_max_std=config['initial_max_std'], initial_min_std=config['initial_min_std']) for activations in config['activations'] ] # scale the learning rate based on the number of channels in x config['solver_lr'] *= np.sqrt(np.prod(x.shape[1:])) trs = [] for i, fm in enumerate(forward_models): # create a bootstrapped data set train_data, validate_data = build_pipeline( x=x, y=y, batch_size=config['batch_size'], val_size=config['val_size'], bootstraps=1) # create a trainer for a forward model with a conservative objective trainer = MaximumLikelihood( fm, forward_model_optim=tf.keras.optimizers.Adam, forward_model_lr=config['forward_model_lr'], noise_std=config.get('model_noise_std', 0.0)) # train the model for an additional number of epochs trs.append(trainer) trainer.launch(train_data, validate_data, logger, config['epochs'], header=f'oracle_{i}/') # select the top k initial designs from the dataset mean_x = tf.reduce_mean(x, axis=0, keepdims=True) indices = tf.math.top_k(y[:, 0], k=config['solver_samples'])[1] initial_x = tf.gather(x, indices, axis=0) x = initial_x # evaluate the starting point solution = x if task.is_normalized_y: preds = [ task.denormalize_y(fm.get_distribution(solution).mean()) for fm in forward_models ] else: preds = [fm.get_distribution(solution).mean() for fm in forward_models] # record the prediction and score to the logger logger.record("distance/travelled", tf.linalg.norm(solution - initial_x), 0) logger.record("distance/from_mean", tf.linalg.norm(solution - mean_x), 0) for n, prediction_i in enumerate(preds): logger.record(f"oracle_{n}/prediction", prediction_i, 0) if n > 0: logger.record(f"rank_corr/0_to_{n}", spearman(preds[0][:, 0], prediction_i[:, 0]), 0) # perform gradient ascent on the score through the forward model for i in range(1, config['solver_steps'] + 1): # back propagate through the forward model with tf.GradientTape() as tape: tape.watch(x) predictions = [] for fm in forward_models: solution = x predictions.append(fm.get_distribution(solution).mean()) if config['aggregation_method'] == 'mean': score = tf.reduce_min(predictions, axis=0) if config['aggregation_method'] == 'min': score = tf.reduce_min(predictions, axis=0) if config['aggregation_method'] == 'random': score = predictions[np.random.randint(len(predictions))] grads = tape.gradient(score, x) # use the conservative optimizer to update the solution x = x + config['solver_lr'] * grads solution = x # evaluate the design using the oracle and the forward model if task.is_normalized_y: preds = [ task.denormalize_y(fm.get_distribution(solution).mean()) for fm in forward_models ] else: preds = [ fm.get_distribution(solution).mean() for fm in forward_models ] # record the prediction and score to the logger logger.record("distance/travelled", tf.linalg.norm(solution - initial_x), i) logger.record("distance/from_mean", tf.linalg.norm(solution - mean_x), i) for n, prediction_i in enumerate(preds): logger.record(f"oracle_{n}/prediction", prediction_i, i) logger.record( f"oracle_{n}/grad_norm", tf.linalg.norm(tf.reshape(grads[n], [-1, input_size]), axis=-1), i) if n > 0: logger.record(f"rank_corr/0_to_{n}", spearman(preds[0][:, 0], prediction_i[:, 0]), i) logger.record( f"grad_corr/0_to_{n}", tfp.stats.correlation(grads[0], grads[n], sample_axis=0, event_axis=None), i) if task.is_discrete and config["use_vae"]: solution = solution * standard_dev + mean logits = vae_model.decoder_cnn.predict(solution) solution = tf.argmax(logits, axis=2, output_type=tf.int32) # save the current solution to the disk np.save(os.path.join(config["logging_dir"], f"solution.npy"), solution.numpy()) # evaluate the found solution and record a video score = task.predict(solution) if task.is_normalized_y: score = task.denormalize_y(score) logger.record("score", score, config['solver_steps'], percentile=True)
def coms_cleaned(config): """Train a forward model and perform model based optimization using a conservative objective function Args: config: dict a dictionary of hyper parameters such as the learning rate """ # create the training task and logger logger = Logger(config['logging_dir']) task = StaticGraphTask(config['task'], **config['task_kwargs']) # save the initial dataset statistics for safe keeping x = task.x y = task.y if config['is_discrete']: # clip the distribution probabilities to a max of discrete_clip p = np.full_like(x, 1 / float(x.shape[-1])) discrete_clip = config.get('discrete_clip', 5.0) x = discrete_clip * x + (1.0 - discrete_clip) * p # map the distribution probabilities to logits x = np.log(x) x = x[:, :, 1:] - x[:, :, :1] if config['normalize_ys']: # remove the mean from the score values mu_y = np.mean(y, axis=0, keepdims=True).astype(np.float32) y = y - mu_y # standardize the variance of the score values st_y = np.std(y, axis=0, keepdims=True).astype(np.float32).clip(1e-6, 1e9) y = y / st_y else: # create placeholder normalization statistics mu_y = 0.0 st_y = 1.0 if config['normalize_xs']: # remove the mean from the data vectors mu_x = np.mean(x, axis=0, keepdims=True).astype(np.float32) x = x - mu_x # standardize the variance of the data vectors st_x = np.std(x, axis=0, keepdims=True).astype(np.float32).clip(1e-6, 1e9) x = x / st_x else: # create placeholder normalization statistics mu_x = 0.0 st_x = 1.0 # record the inputs shape of the forward model input_shape = list(task.input_shape) if config['is_discrete']: input_shape[-1] = input_shape[-1] - 1 # compute the normalized learning rate of the model inner_lr = config['inner_lr'] * np.sqrt(np.prod(input_shape)) outer_lr = config['outer_lr'] * np.sqrt(np.prod(input_shape)) # make a neural network to predict scores forward_model = ForwardModel( input_shape, activations=config['activations'], hidden=config['hidden'], final_tanh=config['final_tanh']) # make a trainer for the forward model trainer = ConservativeObjectiveModel( forward_model, forward_model_opt=tf.keras.optimizers.Adam, forward_model_lr=config['forward_model_lr'], initial_alpha=config['initial_alpha'], alpha_opt=tf.keras.optimizers.Adam, alpha_lr=config['alpha_lr'], target_conservatism=config['target_conservatism'], inner_lr=inner_lr, outer_lr=outer_lr, inner_gradient_steps=config['inner_gradient_steps'], outer_gradient_steps=config['outer_gradient_steps'], beta=config['train_beta'], entropy_coefficient=config['entropy_coefficient'], continuous_noise_std=config['continuous_noise_std']) # create a data set train_data, validate_data = task.build( x=x, y=y, batch_size=config['batch_size'], val_size=config['val_size']) # train the forward model trainer.launch(train_data, validate_data, logger, config["epochs"]) # select the top k initial designs from the dataset indices = tf.math.top_k(y[:, 0], k=config['batch_size'])[1] initial_x = tf.gather(x, indices, axis=0) xt = initial_x scores = [] predictions = [] eval_beta = config['eval_beta'] for step in range(config['outer_gradient_steps']): xt = trainer.outer_optimize(xt, eval_beta, 1, training=False) prediction = forward_model( xt, training=False).numpy() * st_y + mu_y next_xt = trainer.inner_optimize(xt, training=False) next_prediction = forward_model( next_xt, training=False).numpy() * st_y + mu_y final_xt = trainer.outer_optimize( xt, eval_beta, config['outer_gradient_steps'], training=False) final_prediction = forward_model( final_xt, training=False).numpy() * st_y + mu_y solution = xt * st_x + mu_x if config['is_discrete']: solution = tf.math.softmax(tf.pad( solution, [[0, 0], [0, 0], [1, 0]]) / 0.001) score = task.score(solution) # record the prediction and score to the logger logger.record(f"score", score, step, percentile=True) logger.record(f"solver/model_to_real", spearman(prediction[:, 0], score[:, 0]), step) logger.record(f"solver/distance", tf.linalg.norm(xt - initial_x), step) logger.record(f"solver/prediction", prediction, step) logger.record(f"solver/beta_conservatism", prediction - eval_beta * next_prediction, step) logger.record(f"solver/conservatism", prediction - final_prediction, step) logger.record(f"solver/overestimation", prediction - score, step) scores.append(score) predictions.append(prediction) # save the model predictions and scores to be aggregated later np.save(os.path.join(config['logging_dir'], "scores.npy"), np.concatenate(scores, axis=1)) np.save(os.path.join(config['logging_dir'], "predictions.npy"), np.stack(predictions, axis=1))