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))
