logger.remove()
logger.add(sys.stdout,
colorize=True,
format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
"<level>{level}</level> | " +
"<light-black>{file.path}:{line}</light-black> | " + "{message}")
# Parse arguments
opt = get_arguments().parse_args()
opt = post_config(opt)
# Init game specific inputs
replace_tokens = {}
sprite_path = opt.game + '/sprites'
if opt.game == 'mario':
opt.ImgGen = MarioLevelGen(sprite_path)
replace_tokens = MARIO_REPLACE_TOKENS
downsample = special_mario_downsampling
elif opt.game == 'mariokart':
opt.ImgGen = MariokartLevelGen(sprite_path)
replace_tokens = MARIOKART_REPLACE_TOKENS
downsample = special_mariokart_downsampling
else:
NameError("name of --game not recognized. Supported: mario, mariokart")
real = read_level(opt, None, replace_tokens).to(opt.device)
################################################################################
# MarioVAE
N, C, H_real, W_real = real.shape
def main():
""" Main Training funtion. Parses inputs, inits logger, trains, and then generates some samples. """
# torch.autograd.set_detect_anomaly(True)
# Logger init
logger.remove()
logger.add(sys.stdout,
colorize=True,
format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | " +
"<level>{level}</level> | " +
"<light-black>{file.path}:{line}</light-black> | " +
"{message}")
# Parse arguments
opt = get_arguments().parse_args()
opt = post_config(opt)
# Init wandb
run = wandb.init(project="mario",
tags=get_tags(opt),
config=opt,
dir=opt.out)
opt.out_ = run.dir
# Init game specific inputs
replace_tokens = {}
sprite_path = opt.game + '/sprites'
if opt.game == 'mario':
opt.ImgGen = MarioLevelGen(sprite_path)
replace_tokens = MARIO_REPLACE_TOKENS
downsample = special_mario_downsampling
elif opt.game == 'mariokart':
opt.ImgGen = MariokartLevelGen(sprite_path)
replace_tokens = MARIOKART_REPLACE_TOKENS
downsample = special_mariokart_downsampling
else:
NameError("name of --game not recognized. Supported: mario, mariokart")
# Read level according to input arguments
real = read_level(opt, None, replace_tokens).to(opt.device)
# Train!
generators, noise_maps, reals, noise_amplitudes = train(real, opt)
# Generate Samples of same size as level
logger.info("Finished training! Generating random samples...")
in_s = None
generate_samples(generators,
noise_maps,
reals,
noise_amplitudes,
opt,
in_s=in_s)
# Generate samples of smaller size than level
logger.info("Generating arbitrary sized random samples...")
scale_v = 0.8 # Arbitrarily chosen scales
scale_h = 0.4
real_down = downsample(1, [[scale_v, scale_h]], real, opt.token_list)
real_down = real_down[0]
# necessary for correct input shape
in_s = torch.zeros(real_down.shape, device=opt.device)
generate_samples(generators,
noise_maps,
reals,
noise_amplitudes,
opt,
in_s=in_s,
scale_v=scale_v,
scale_h=scale_h,
save_dir="arbitrary_random_samples")