WorldModels-VAE-train

VAE-train.py

def train():
	es = cma.CMAEvolutionStrategy(_NUM_PARAMS * [0], 0.1, {'popsize': 16})
	rewards_through_gens = []
	generation = 1
	try:
		while not es.stop():
			solutions = es.ask()
			with mp.Pool(mp.cpu_count()) as p:
				rewards = list(tqdm.tqdm(p.imap(play, list(solutions)), total=len(solutions)))

			es.tell(solutions, rewards)

			rewards = np.array(rewards) *(-1.)
		
			generation+=1
			rewards_through_gens.append(rewards)
			np.save('rewards', rewards_through_gens)

	except (KeyboardInterrupt, SystemExit):
		print("Manual Interrupt")
	except Exception as e:
		print("Exception: {}".format(e))
	return es