def worker(solutions, env):
fitness_solutions = []
if not isinstance(solutions, list):
solutions = [solutions]
for weights in solutions:
wm = World_Model(dirname + "\\vae.pt", dirname + "\\mdn_rnn.pt",
actions, device)
w = weights[0:actions * (hidden_size + latent_size)]
b = weights[actions * (hidden_size + latent_size)::]
w = nn.Parameter(
torch.tensor(np.reshape(
w, (actions, hidden_size +
latent_size))).type('torch.FloatTensor').to(device))
b = nn.Parameter(torch.tensor(b).type('torch.FloatTensor').to(device))
wm.set_controller(w, b)
fitness = []
for i in range(num_rollouts):
runner = Env_Runner(device)
wm.reset_rnn()
_, _, rewards = runner.run(env, wm, img_resize=(64, 64))
# append negative return, because ES will try to minimize it
fitness.append(-np.sum(np.array(rewards)))
env.close()
fitness_solutions.append(np.mean(np.array(fitness)))
return fitness_solutions
