<reponame>k4ntz/XmodRL
# deep neuroevolution as a genetic rl algo

import gym
import numpy as np
import os
import copy
import multiprocessing
import random

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from joblib import Parallel, delayed

from itertools import count
from PIL import Image
from atariari.benchmark.wrapper import AtariARIWrapper
from captum.attr import IntegratedGradients

from rtpt import RTPT

from xrl.features import feature_processing


PATH_TO_OUTPUTS = os.getcwd() + "/xrl/checkpoints/"
if not os.path.exists(PATH_TO_OUTPUTS):
    os.makedirs(PATH_TO_OUTPUTS)

model_name = lambda training_name : PATH_TO_OUTPUTS + training_name + "_model.pth"

# TODO: Fix serialization problem and decomment after,
# look at other genetic_rl.py file for details!
from xrl.genetic_rl import policy_net
# define policy network for genetic algo
#class policy_net(nn.Module):
#    def __init__(self, input, hidden, actions, make_hidden = True): 
#        super(policy_net, self).__init__()
#        # should make one hidden layer
#        self.make_hidden = make_hidden
#
#        if self.make_hidden:
#            self.h = nn.Linear(input, hidden)
#            self.out = nn.Linear(hidden, actions)
#        else:
#            self.out = nn.Linear(input, actions)
#
#
#    def forward(self, x):
#        if self.make_hidden:
#            x = F.relu(self.h(x))
#        return F.softmax(self.out(x), dim=1)


def init_weights(m):
    # nn.Conv2d weights are of shape [16, 1, 3, 3] i.e. # number of filters, 1, stride, stride
    # nn.Conv2d bias is of shape [16] i.e. # number of filters
    
    # nn.Linear weights are of shape [32, 24336] i.e. # number of input features, number of output features
    # nn.Linear bias is of shape [32] i.e. # number of output features
    
    if ((type(m) == nn.Linear) | (type(m) == nn.Conv2d)):
        torch.nn.init.xavier_uniform(m.weight)
        m.bias.data.fill_(0.00)


# function to create random agents of given count
def return_random_agents(n_inputs, num_agents, n_actions, cfg):
    agents = []
    if cfg.train.make_hidden:
        print("Agents have", n_inputs, "input nodes,", cfg.train.hidden_layer_size, "hidden nodes and", n_actions, "output nodes")
    else:
        print("Linear model, no hidden layer! Policy net has", n_inputs, "input nodes and", n_actions, "output nodes")
    for _ in range(num_agents):
        agent = policy_net(n_inputs, cfg.train.hidden_layer_size, n_actions, cfg.train.make_hidden)
        for param in agent.parameters():
            param.requires_grad = False
        init_weights(agent)
        agents.append(agent) 
    return agents


# function to select action by given features
def select_action(features, policy, random_tr = -1):
    sample = random.random()
    if sample > random_tr:
        # calculate probabilities of taking each action
        probs = policy(torch.tensor(features).unsqueeze(0).float())
        # sample an action from that set of probs
        sampler = Categorical(probs)
        action = sampler.sample()
    else:
        action = random.randint(0, 5)
    # return action
    return action


# function to run list of agents in given env
def run_agents(env, agents, cfg):
    reward_agents = []
    _ = env.reset()
    _, _, done, info = env.step(1)
    for agent in agents:
        agent.eval()
        raw_features, features, _, _ = feature_processing.do_step(env)
        r = 0
        t = 0
        while t < cfg.train.max_steps:
            if cfg.train.use_raw_features:
                features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
            action = select_action(features, agent)
            raw_features, features, reward, done = feature_processing.do_step(env, action, raw_features)
            r = r + reward
            if(done):
                break
            t += 1
        if t == cfg.train.max_steps:
            r = -25
        reward_agents.append(r)
    return reward_agents


# returns average score of given agent when it runs n times
def return_average_score(agent, runs, cfg):
    score = 0.
    env = AtariARIWrapper(gym.make(cfg.env_name))
    rtpt = RTPT(name_initials='DV', experiment_name=cfg.exp_name,
                    max_iterations=runs)
    rtpt.start()
    for i in range(runs):
        score += run_agents(env, [agent], cfg)[0]
        rtpt.step()
    avg_score = score/runs
    return avg_score


# gets avg score of every agent running n runs 
def run_agents_n_times(agents, runs, cfg):
    avg_score = []
    agents = tqdm(agents)
    cpu_cores = min(multiprocessing.cpu_count(), 100)
    avg_score = Parallel(n_jobs=cpu_cores)(delayed(return_average_score)(agent, runs, cfg) for agent in agents)
    return avg_score


# function to mutate given agent to child agent
def mutate(agent):
    child_agent = copy.deepcopy(agent)
    mutation_power = 0.02 #hyper-parameter, set from https://arxiv.org/pdf/1712.06567.pdf
    for param in child_agent.parameters():
        if(len(param.shape)==4): #weights of Conv2D
            for i0 in range(param.shape[0]):
                for i1 in range(param.shape[1]):
                    for i2 in range(param.shape[2]):
                        for i3 in range(param.shape[3]):
                            param[i0][i1][i2][i3]+= mutation_power * np.random.randn()
        elif(len(param.shape)==2): #weights of linear layer
            for i0 in range(param.shape[0]):
                for i1 in range(param.shape[1]):
                    param[i0][i1]+= mutation_power * np.random.randn()
        elif(len(param.shape)==1): #biases of linear layer or conv layer
            for i0 in range(param.shape[0]):
                param[i0]+=mutation_power * np.random.randn()
    return child_agent


# function to add elite to childs 
def add_elite(agents, sorted_parent_indexes, cfg, elite_index=None, only_consider_top_n=10):
    candidate_elite_index = sorted_parent_indexes[:only_consider_top_n]
    if(elite_index is not None):
        candidate_elite_index = np.append(candidate_elite_index,[elite_index]) 
    top_score = None
    top_elite_index = None
    tqdmcandidate_elite_index = tqdm(candidate_elite_index)
    cpu_cores = min(100, max(multiprocessing.cpu_count(), only_consider_top_n))
    scores = Parallel(n_jobs=cpu_cores)(delayed(return_average_score)(agents[i], runs=5, cfg=cfg) for i in tqdmcandidate_elite_index)
    for i, score in enumerate(scores):
        i = candidate_elite_index[i]
        print("Score for elite i ", i, " is ", score)
        if(top_score is None):
            top_score = score
            top_elite_index = i
        elif(score > top_score):
            top_score = score
            top_elite_index = i
    print("Elite selected with index ",top_elite_index, " and score", top_score)
    child_agent = copy.deepcopy(agents[top_elite_index])
    return child_agent


# function to create and return children from given parent agents
def return_children(agents, sorted_parent_indexes, elite_index, cfg):  
    children_agents = []
    #first take selected parents from sorted_parent_indexes and generate N-1 children
    for i in range(len(agents) - 1):
        selected_agent_index = sorted_parent_indexes[np.random.randint(len(sorted_parent_indexes))]
        children_agents.append(mutate(agents[selected_agent_index]))
    #now add one elite
    elite_child = add_elite(agents, sorted_parent_indexes, cfg, elite_index)
    children_agents.append(elite_child)
    elite_index = len(children_agents) - 1 #it is the last one
    return children_agents, elite_index


# save model helper function
def save_agents(training_name, agents, generation, elite_index):
    if not os.path.exists(PATH_TO_OUTPUTS):
        os.makedirs(PATH_TO_OUTPUTS)
    model_path = model_name(training_name)
    print("Saving {}".format(model_path))
    torch.save({
            'agents': agents,
            'generation': generation,
            'elite_index': elite_index
            }, model_path)


# load agents if exists
def load_agents(model_path):
    print("{} does exist, loading ... ".format(model_path))
    checkpoint = torch.load(model_path)
    agents = checkpoint['agents']
    generation = checkpoint['generation']
    elite_index = None
    try:
        elite_index = checkpoint['elite_index']
    except:
        print("No elite index in save available...")
    return agents, generation, elite_index


# train main function
def train(cfg):
    print('Experiment name:', cfg.exp_name)

    writer = SummaryWriter(os.getcwd() + cfg.logdir + cfg.exp_name)

    generations = cfg.train.num_episodes
    print('Generations:', generations)
    print('Max Steps per Episode:', cfg.train.max_steps)

    # disable gradients as we will not use them
    torch.set_grad_enabled(False)

    # init env to get actions count and features space
    env = AtariARIWrapper(gym.make(cfg.env_name))
    n_actions = env.action_space.n
    _ = env.reset()
    raw_features, features, _, _ = feature_processing.do_step(env)
    if cfg.train.use_raw_features:
        features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
    # initialize N number of agents
    num_agents = 500
    print('Number of agents:', num_agents)
    agents = return_random_agents(len(features), num_agents, n_actions, cfg)
    generation = 0

    # load if exists
    model_path = model_name(cfg.exp_name)
    if os.path.isfile(model_path):
        agents, generation, _ = load_agents(model_path)

    # How many top agents to consider as parents
    top_limit = 20
    print('Number of top agents:', top_limit)

    # runs per generation
    n_gen_runs = 3
    print('Number of runs per generation:', n_gen_runs)

    elite_index = None

    rtpt = RTPT(name_initials='DV', experiment_name=cfg.exp_name,
                    max_iterations=generations)
    rtpt.start()
    while generation < generations:
        print("Starting generation", generation)
        # return rewards of agents
        rewards = run_agents_n_times(agents, n_gen_runs, cfg) #return average of 3 runs
 
        # sort by rewards
        # reverses and gives top values (argsort sorts by ascending by default) https://stackoverflow.com/questions/16486252/is-it-possible-to-use-argsort-in-descending-order
        sorted_parent_indexes = np.argsort(rewards)[::-1][:top_limit] 
        print("")
        print("")
        
        top_rewards = []
        for best_parent in sorted_parent_indexes:
            top_rewards.append(rewards[best_parent])
       
        print("Generation ", generation, " | Mean rewards: ", np.mean(rewards), " | Mean of top 5: ",np.mean(top_rewards[:5]))
        #print(rewards)
        print("Top ",top_limit," scores", sorted_parent_indexes)
        print("Rewards for top: ",top_rewards)
        
        # setup an empty list for containing children agents
        children_agents, elite_index = return_children(agents, sorted_parent_indexes, elite_index, cfg)
 
        # kill all agents, and replace them with their children
        agents = children_agents

        #log stuff
        writer.add_scalar('Train/Mean rewards', np.mean(rewards), generation)
        writer.add_scalar('Train/Mean of top 5', np.mean(top_rewards[:5]), generation)
        # save generation
        generation += 1
        save_agents(cfg.exp_name, agents, generation, elite_index)
        # make rtpt step
        rtpt.step()


# function to eval best agent of last generation
def eval_load(cfg):
    print('Experiment name:', cfg.exp_name)
    print('Evaluating Mode')
    # disable gradients as we will not use them
    torch.set_grad_enabled(False)
    # init env
    env = AtariARIWrapper(gym.make(cfg.env_name))
    n_actions = env.action_space.n
    _ = env.reset()
    raw_features, features, _, _ = feature_processing.do_step(env)
    if cfg.train.use_raw_features:
        features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
    # initialize N number of agents
    num_agents = 500
    print('Number of agents:', num_agents)
    agents = return_random_agents(len(features), num_agents, n_actions, cfg)
    generation = 0
    elite_index = 7 # will be overwritten when agents are loaded

    # load if exists
    model_path = model_name(cfg.exp_name)
    if os.path.isfile(model_path):
        agents, generation, t_elite_index = load_agents(model_path)
        if t_elite_index is not None:
            elite_index = t_elite_index
    print('Selected elite agent:', elite_index)
    elite_agent = agents[elite_index]
    # print nn structure
    dummy = policy_net(len(features), cfg.train.hidden_layer_size, n_actions, cfg.train.make_hidden)
    # because old trained runs does not have make_hidden param
    dummy.load_state_dict(elite_agent.state_dict())
    elite_agent = dummy
    return elite_agent


# deep neuroevolution as a genetic rl algo

import gym
import numpy as np
import os
import copy
import multiprocessing
import random

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from joblib import Parallel, delayed

from itertools import count
from PIL import Image
from atariari.benchmark.wrapper import AtariARIWrapper
from captum.attr import IntegratedGradients

from rtpt import RTPT

from xrl.features import feature_processing


PATH_TO_OUTPUTS = os.getcwd() + "/xrl/checkpoints/"
if not os.path.exists(PATH_TO_OUTPUTS):
    os.makedirs(PATH_TO_OUTPUTS)

model_name = lambda training_name : PATH_TO_OUTPUTS + training_name + "_model.pth"

# TODO: Fix serialization problem and decomment after,
# look at other genetic_rl.py file for details!
from xrl.genetic_rl import policy_net
# define policy network for genetic algo
#class policy_net(nn.Module):
#    def __init__(self, input, hidden, actions, make_hidden = True): 
#        super(policy_net, self).__init__()
#        # should make one hidden layer
#        self.make_hidden = make_hidden
#
#        if self.make_hidden:
#            self.h = nn.Linear(input, hidden)
#            self.out = nn.Linear(hidden, actions)
#        else:
#            self.out = nn.Linear(input, actions)
#
#
#    def forward(self, x):
#        if self.make_hidden:
#            x = F.relu(self.h(x))
#        return F.softmax(self.out(x), dim=1)


def init_weights(m):
    # nn.Conv2d weights are of shape [16, 1, 3, 3] i.e. # number of filters, 1, stride, stride
    # nn.Conv2d bias is of shape [16] i.e. # number of filters
    
    # nn.Linear weights are of shape [32, 24336] i.e. # number of input features, number of output features
    # nn.Linear bias is of shape [32] i.e. # number of output features
    
    if ((type(m) == nn.Linear) | (type(m) == nn.Conv2d)):
        torch.nn.init.xavier_uniform(m.weight)
        m.bias.data.fill_(0.00)


# function to create random agents of given count
def return_random_agents(n_inputs, num_agents, n_actions, cfg):
    agents = []
    if cfg.train.make_hidden:
        print("Agents have", n_inputs, "input nodes,", cfg.train.hidden_layer_size, "hidden nodes and", n_actions, "output nodes")
    else:
        print("Linear model, no hidden layer! Policy net has", n_inputs, "input nodes and", n_actions, "output nodes")
    for _ in range(num_agents):
        agent = policy_net(n_inputs, cfg.train.hidden_layer_size, n_actions, cfg.train.make_hidden)
        for param in agent.parameters():
            param.requires_grad = False
        init_weights(agent)
        agents.append(agent) 
    return agents


# function to select action by given features
def select_action(features, policy, random_tr = -1):
    sample = random.random()
    if sample > random_tr:
        # calculate probabilities of taking each action
        probs = policy(torch.tensor(features).unsqueeze(0).float())
        # sample an action from that set of probs
        sampler = Categorical(probs)
        action = sampler.sample()
    else:
        action = random.randint(0, 5)
    # return action
    return action


# function to run list of agents in given env
def run_agents(env, agents, cfg):
    reward_agents = []
    _ = env.reset()
    _, _, done, info = env.step(1)
    for agent in agents:
        agent.eval()
        raw_features, features, _, _ = feature_processing.do_step(env)
        r = 0
        t = 0
        while t < cfg.train.max_steps:
            if cfg.train.use_raw_features:
                features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
            action = select_action(features, agent)
            raw_features, features, reward, done = feature_processing.do_step(env, action, raw_features)
            r = r + reward
            if(done):
                break
            t += 1
        if t == cfg.train.max_steps:
            r = -25
        reward_agents.append(r)
    return reward_agents


# returns average score of given agent when it runs n times
def return_average_score(agent, runs, cfg):
    score = 0.
    env = AtariARIWrapper(gym.make(cfg.env_name))
    rtpt = RTPT(name_initials='DV', experiment_name=cfg.exp_name,
                    max_iterations=runs)
    rtpt.start()
    for i in range(runs):
        score += run_agents(env, [agent], cfg)[0]
        rtpt.step()
    avg_score = score/runs
    return avg_score


# gets avg score of every agent running n runs 
def run_agents_n_times(agents, runs, cfg):
    avg_score = []
    agents = tqdm(agents)
    cpu_cores = min(multiprocessing.cpu_count(), 100)
    avg_score = Parallel(n_jobs=cpu_cores)(delayed(return_average_score)(agent, runs, cfg) for agent in agents)
    return avg_score


# function to mutate given agent to child agent
def mutate(agent):
    child_agent = copy.deepcopy(agent)
    mutation_power = 0.02 #hyper-parameter, set from https://arxiv.org/pdf/1712.06567.pdf
    for param in child_agent.parameters():
        if(len(param.shape)==4): #weights of Conv2D
            for i0 in range(param.shape[0]):
                for i1 in range(param.shape[1]):
                    for i2 in range(param.shape[2]):
                        for i3 in range(param.shape[3]):
                            param[i0][i1][i2][i3]+= mutation_power * np.random.randn()
        elif(len(param.shape)==2): #weights of linear layer
            for i0 in range(param.shape[0]):
                for i1 in range(param.shape[1]):
                    param[i0][i1]+= mutation_power * np.random.randn()
        elif(len(param.shape)==1): #biases of linear layer or conv layer
            for i0 in range(param.shape[0]):
                param[i0]+=mutation_power * np.random.randn()
    return child_agent


# function to add elite to childs 
def add_elite(agents, sorted_parent_indexes, cfg, elite_index=None, only_consider_top_n=10):
    candidate_elite_index = sorted_parent_indexes[:only_consider_top_n]
    if(elite_index is not None):
        candidate_elite_index = np.append(candidate_elite_index,[elite_index]) 
    top_score = None
    top_elite_index = None
    tqdmcandidate_elite_index = tqdm(candidate_elite_index)
    cpu_cores = min(100, max(multiprocessing.cpu_count(), only_consider_top_n))
    scores = Parallel(n_jobs=cpu_cores)(delayed(return_average_score)(agents[i], runs=5, cfg=cfg) for i in tqdmcandidate_elite_index)
    for i, score in enumerate(scores):
        i = candidate_elite_index[i]
        print("Score for elite i ", i, " is ", score)
        if(top_score is None):
            top_score = score
            top_elite_index = i
        elif(score > top_score):
            top_score = score
            top_elite_index = i
    print("Elite selected with index ",top_elite_index, " and score", top_score)
    child_agent = copy.deepcopy(agents[top_elite_index])
    return child_agent


# function to create and return children from given parent agents
def return_children(agents, sorted_parent_indexes, elite_index, cfg):  
    children_agents = []
    #first take selected parents from sorted_parent_indexes and generate N-1 children
    for i in range(len(agents) - 1):
        selected_agent_index = sorted_parent_indexes[np.random.randint(len(sorted_parent_indexes))]
        children_agents.append(mutate(agents[selected_agent_index]))
    #now add one elite
    elite_child = add_elite(agents, sorted_parent_indexes, cfg, elite_index)
    children_agents.append(elite_child)
    elite_index = len(children_agents) - 1 #it is the last one
    return children_agents, elite_index


# save model helper function
def save_agents(training_name, agents, generation, elite_index):
    if not os.path.exists(PATH_TO_OUTPUTS):
        os.makedirs(PATH_TO_OUTPUTS)
    model_path = model_name(training_name)
    print("Saving {}".format(model_path))
    torch.save({
            'agents': agents,
            'generation': generation,
            'elite_index': elite_index
            }, model_path)


# load agents if exists
def load_agents(model_path):
    print("{} does exist, loading ... ".format(model_path))
    checkpoint = torch.load(model_path)
    agents = checkpoint['agents']
    generation = checkpoint['generation']
    elite_index = None
    try:
        elite_index = checkpoint['elite_index']
    except:
        print("No elite index in save available...")
    return agents, generation, elite_index


# train main function
def train(cfg):
    print('Experiment name:', cfg.exp_name)

    writer = SummaryWriter(os.getcwd() + cfg.logdir + cfg.exp_name)

    generations = cfg.train.num_episodes
    print('Generations:', generations)
    print('Max Steps per Episode:', cfg.train.max_steps)

    # disable gradients as we will not use them
    torch.set_grad_enabled(False)

    # init env to get actions count and features space
    env = AtariARIWrapper(gym.make(cfg.env_name))
    n_actions = env.action_space.n
    _ = env.reset()
    raw_features, features, _, _ = feature_processing.do_step(env)
    if cfg.train.use_raw_features:
        features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
    # initialize N number of agents
    num_agents = 500
    print('Number of agents:', num_agents)
    agents = return_random_agents(len(features), num_agents, n_actions, cfg)
    generation = 0

    # load if exists
    model_path = model_name(cfg.exp_name)
    if os.path.isfile(model_path):
        agents, generation, _ = load_agents(model_path)

    # How many top agents to consider as parents
    top_limit = 20
    print('Number of top agents:', top_limit)

    # runs per generation
    n_gen_runs = 3
    print('Number of runs per generation:', n_gen_runs)

    elite_index = None

    rtpt = RTPT(name_initials='DV', experiment_name=cfg.exp_name,
                    max_iterations=generations)
    rtpt.start()
    while generation < generations:
        print("Starting generation", generation)
        # return rewards of agents
        rewards = run_agents_n_times(agents, n_gen_runs, cfg) #return average of 3 runs
 
        # sort by rewards
        # reverses and gives top values (argsort sorts by ascending by default) https://stackoverflow.com/questions/16486252/is-it-possible-to-use-argsort-in-descending-order
        sorted_parent_indexes = np.argsort(rewards)[::-1][:top_limit] 
        print("")
        print("")
        
        top_rewards = []
        for best_parent in sorted_parent_indexes:
            top_rewards.append(rewards[best_parent])
       
        print("Generation ", generation, " | Mean rewards: ", np.mean(rewards), " | Mean of top 5: ",np.mean(top_rewards[:5]))
        #print(rewards)
        print("Top ",top_limit," scores", sorted_parent_indexes)
        print("Rewards for top: ",top_rewards)
        
        # setup an empty list for containing children agents
        children_agents, elite_index = return_children(agents, sorted_parent_indexes, elite_index, cfg)
 
        # kill all agents, and replace them with their children
        agents = children_agents

        #log stuff
        writer.add_scalar('Train/Mean rewards', np.mean(rewards), generation)
        writer.add_scalar('Train/Mean of top 5', np.mean(top_rewards[:5]), generation)
        # save generation
        generation += 1
        save_agents(cfg.exp_name, agents, generation, elite_index)
        # make rtpt step
        rtpt.step()


# function to eval best agent of last generation
def eval_load(cfg):
    print('Experiment name:', cfg.exp_name)
    print('Evaluating Mode')
    # disable gradients as we will not use them
    torch.set_grad_enabled(False)
    # init env
    env = AtariARIWrapper(gym.make(cfg.env_name))
    n_actions = env.action_space.n
    _ = env.reset()
    raw_features, features, _, _ = feature_processing.do_step(env)
    if cfg.train.use_raw_features:
        features = np.array(np.array([[0,0] if x==None else x for x in raw_features]).tolist()).flatten()
    # initialize N number of agents
    num_agents = 500
    print('Number of agents:', num_agents)
    agents = return_random_agents(len(features), num_agents, n_actions, cfg)
    generation = 0
    elite_index = 7 # will be overwritten when agents are loaded

    # load if exists
    model_path = model_name(cfg.exp_name)
    if os.path.isfile(model_path):
        agents, generation, t_elite_index = load_agents(model_path)
        if t_elite_index is not None:
            elite_index = t_elite_index
    print('Selected elite agent:', elite_index)
    elite_agent = agents[elite_index]
    # print nn structure
    dummy = policy_net(len(features), cfg.train.hidden_layer_size, n_actions, cfg.train.make_hidden)
    # because old trained runs does not have make_hidden param
    dummy.load_state_dict(elite_agent.state_dict())
    elite_agent = dummy
    return elite_agent
# deep neuroevolution as a genetic rl algo # TODO: Fix serialization problem and decomment after, # look at other genetic_rl.py file for details! # define policy network for genetic algo #class policy_net(nn.Module): #    def __init__(self, input, hidden, actions, make_hidden = True): #        super(policy_net, self).__init__() #        # should make one hidden layer #        self.make_hidden = make_hidden # #        if self.make_hidden: #            self.h = nn.Linear(input, hidden) #            self.out = nn.Linear(hidden, actions) #        else: #            self.out = nn.Linear(input, actions) # # #    def forward(self, x): #        if self.make_hidden: #            x = F.relu(self.h(x)) #        return F.softmax(self.out(x), dim=1) # nn.Conv2d weights are of shape [16, 1, 3, 3] i.e. # number of filters, 1, stride, stride # nn.Conv2d bias is of shape [16] i.e. # number of filters # nn.Linear weights are of shape [32, 24336] i.e. # number of input features, number of output features # nn.Linear bias is of shape [32] i.e. # number of output features # function to create random agents of given count # function to select action by given features # calculate probabilities of taking each action # sample an action from that set of probs # return action # function to run list of agents in given env # returns average score of given agent when it runs n times # gets avg score of every agent running n runs # function to mutate given agent to child agent #hyper-parameter, set from https://arxiv.org/pdf/1712.06567.pdf #weights of Conv2D #weights of linear layer #biases of linear layer or conv layer # function to add elite to childs # function to create and return children from given parent agents #first take selected parents from sorted_parent_indexes and generate N-1 children #now add one elite #it is the last one # save model helper function # load agents if exists # train main function # disable gradients as we will not use them # init env to get actions count and features space # initialize N number of agents # load if exists # How many top agents to consider as parents # runs per generation # return rewards of agents #return average of 3 runs # sort by rewards # reverses and gives top values (argsort sorts by ascending by default) https://stackoverflow.com/questions/16486252/is-it-possible-to-use-argsort-in-descending-order #print(rewards) # setup an empty list for containing children agents # kill all agents, and replace them with their children #log stuff # save generation # make rtpt step # function to eval best agent of last generation # disable gradients as we will not use them # init env # initialize N number of agents # will be overwritten when agents are loaded # load if exists # print nn structure # because old trained runs does not have make_hidden param