def __init__(self,
config,
vae_decoder,
mem_model,
dataset,
dir_name='plots'):
self.config = config
self.decoder = vae_decoder
self.model = mem_model
self.sequence_len = self.config.rnn['sequence_len']
self.latent_dim = self.config.vae['latent_space_dim']
# Check if destination dir exists
self.plots_dir = os.path.join(self.config.rnn['logs_dir'], dir_name)
create_directory(self.plots_dir)
# Prepare data
(states, actions), _ = dataset[0]
self.n_episodes = min(self.config.rnn['rend_n_episodes'], len(dataset))
self.eval_states = torch.zeros(
(self.n_episodes, self.sequence_len, states.shape[1]),
device=next(self.model.parameters()).device,
dtype=states.dtype)
self.eval_actions = torch.zeros(
(self.n_episodes, self.sequence_len, actions.shape[1]),
device=next(self.model.parameters()).device,
dtype=actions.dtype)
for i in range(self.n_episodes):
(states, actions), _ = dataset[i]
self.eval_states[i] = states
self.eval_actions[i] = actions
def train_mem(ctx, path, vae_path):
"""Train MDN-RNN model as specified in .json config with data at `PATH`."""
from third_party.torchtrainer import EarlyStopping, LambdaCallback, ModelCheckpoint, CSVLogger, RandomBatchSampler, evaluate
from torch.utils.data import DataLoader
config = obtain_config(ctx)
env = hrl.create_gym(config.general['game_name'])
# Create checkpoint directory, if it doesn't exist
create_directory(os.path.dirname(config.rnn['ckpt_path']))
# Create training DataLoader
dataset = MDNDataset(path, config.rnn['sequence_len'],
config.rnn['terminal_prob'],
config.rnn['dataset_fraction'])
data_loader = DataLoader(dataset,
batch_sampler=RandomBatchSampler(
dataset, config.rnn['batch_size']),
pin_memory=True)
# Build model
rnn = build_rnn_model(config.rnn, config.vae['latent_space_dim'],
env.action_space)
# Create callbacks
callbacks = [
EarlyStopping(metric='loss',
patience=config.rnn['patience'],
verbose=1),
LambdaCallback(on_batch_begin=lambda _, batch_size: rnn.model.
init_hidden(batch_size)),
ModelCheckpoint(config.rnn['ckpt_path'], metric='loss',
save_best=True),
CSVLogger(
filename=os.path.join(config.rnn['logs_dir'], 'train_mem.csv'))
]
# Evaluate and visualize memory progress if render allowed
if config.allow_render:
if vae_path is None:
raise ValueError("To render provide valid path to VAE checkpoint!")
# Build VAE model and load checkpoint
_, _, decoder = build_vae_model(config.vae,
config.general['state_shape'],
vae_path)
callbacks += [
MemoryVisualization(config, decoder, rnn.model, dataset,
'mdn_plots')
]
# Fit MDN-RNN model!
rnn.fit_loader(data_loader,
epochs=config.rnn['epochs'],
callbacks=callbacks)
dataset.close()
def save_es_ckpt_and_mind_weights(self, ckpt_path, mind_path, score):
# Create CMA-ES checkpoint dir if doesn't exist
create_directory(os.path.dirname(ckpt_path))
# Create Mind weights checkpoint dir if doesn't exist
mind_dir = os.path.dirname(mind_path)
create_directory(mind_dir)
# Create paths for best and mean Mind weights checkpoints
mind_name = os.path.basename(mind_path).split('.')[0]
best_path = os.path.join(mind_dir, mind_name + "_best.ckpt")
mean_path = os.path.join(mind_dir, mind_name + "_mean.ckpt")
with open(os.path.abspath(ckpt_path), 'wb') as f:
pickle.dump(self, f)
log.debug("Saved CMA-ES checkpoint in path: %s", ckpt_path)
if self.check_if_better(score):
log.info("New best score: %f", score)
with open(os.path.abspath(best_path), 'wb') as f:
pickle.dump(self.best_param(), f)
with open(os.path.abspath(mean_path), 'wb') as f:
pickle.dump(self.current_param(), f)
log.debug("Saved Mind weights in path: %s", mind_path)
def save_results(names, pred, y_true, path):
common_utils.create_directory(os.path.dirname(path))
with open(path, 'w') as f:
f.write("stay,prediction,y_true\n")
for (name, x, y) in zip(names, pred, y_true):
f.write("{},{:.6f},{}\n".format(name, x, y))
def train_vae(ctx, path):
"""Train VAE model as specified in .json config with data at `PATH`."""
from keras.callbacks import EarlyStopping, LambdaCallback, ModelCheckpoint, CSVLogger
config = obtain_config(ctx)
# Get dataset length and eight examples to evaluate VAE on
with h5.File(path, 'r') as hfile:
n_transitions = hfile.attrs['N_TRANSITIONS']
X_eval = hfile['states'][:8] / 255.
# Get training data
train_gen = HDF5DataGenerator(path,
'states',
'states',
batch_size=config.vae['batch_size'],
end=int(n_transitions * 0.8),
preprocess_fn=lambda X, y:
(X / 255., y / 255.))
val_gen = HDF5DataGenerator(path,
'states',
'states',
batch_size=config.vae['batch_size'],
start=int(n_transitions * 0.8),
preprocess_fn=lambda X, y:
(X / 255., y / 255.))
# Build VAE model
vae, _, _ = build_vae_model(config.vae, config.general['state_shape'])
# If render features enabled...
if config.allow_render:
# ...plot first eight training examples with VAE reconstructions
# at the beginning of every epoch
import matplotlib
matplotlib.use("Agg")
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
# Check if destination dir exists
plots_dir = os.path.join(config.vae['logs_dir'], "plots_vae")
if not os.path.exists(plots_dir):
os.makedirs(plots_dir)
# Evaluate VAE at the end of epoch
def plot_samples(epoch, logs):
pred = vae.predict(X_eval)
samples = np.empty_like(np.concatenate((X_eval, pred)))
samples[0::2] = X_eval
samples[1::2] = pred
_ = plt.figure(figsize=(4, 4))
gs = gridspec.GridSpec(4, 4)
gs.update(wspace=0.05, hspace=0.05)
for i, sample in enumerate(samples):
ax = plt.subplot(gs[i])
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_aspect('equal')
plt.imshow(sample.reshape(*config.general['state_shape']))
# Save figure to logs dir
plt.savefig(
os.path.join(
plots_dir, "interpreter_sample_{}".format(
dt.datetime.now().strftime("%d-%mT%H:%M"))))
plt.close()
else:
def plot_samples(epoch, logs):
pass
# Create checkpoint and logging directory, if it doesn't exist
create_directory(os.path.dirname(config.vae['ckpt_path']))
create_directory(os.path.dirname(config.vae['logs_dir']))
# Initialize callbacks
callbacks = [
EarlyStopping(patience=config.vae['patience']),
LambdaCallback(on_epoch_begin=plot_samples),
ModelCheckpoint(config.vae['ckpt_path'],
verbose=1,
save_best_only=True,
save_weights_only=True),
CSVLogger(filename=os.path.join(config.vae['logs_dir'],
'train_vae.csv'),
append=True)
]
# Fit VAE model!
vae.fit_generator(
generator=train_gen,
validation_data=val_gen,
epochs=config.vae['epochs'],
use_multiprocessing=False,
# NOTE: There is no need for more then one workers, we are disk IO bound (I suppose ...)
# NOTE2: h5py from conda should be threadsafe... but it apparently isn't and raises
# `OSError: Can't read data (wrong B-tree signature)` sporadically if `workers` = 1
# and always if `workers` > 1. That's why this generator needs to run in main thread
# (`workers` = 0).
workers=3,
max_queue_size=100,
shuffle=True, # It shuffles whole batches, not items in batches
callbacks=callbacks)