update src

src/QD_MARL/agents/agents.py

@@ -60,7 +60,6 @@ def has_policy(self):
     def __str__(self):
         return f"Name: {self._name}; Squad: {self._squad}; Set: {self._set}; Optimize: {str(self._to_optimize)}"
 
-
 class CoachAgent:
 
     # Coach agent which select the team of agent from the pool produced by the initial population
@@ -78,17 +77,24 @@ class CoachAgent:
     # set_evaluator: set the evaluator of the algorithm, evaluates the candidates
     # get_final_pop: get the final population of the algorithm
 
-    def __init__(self, config, me = None):
+    def __init__(self, config, me = None, n_teams = 1):
         self._config = config
         self._me = me
         self.random = random.Random()
         self.random.seed(self._config["seed"])
         self._pop_size = self._config["pop_size"]
         self._batch_size = self._config["batch_size"]
         self._algorithm = self.set_algorithm()
-        self._pop_fitnesses = None
-        self._pop_desc = None
-
+        self._pop_descs = []
+        self._pop_fitnesses = []
+        self._n_teams = n_teams
+
+    def get_n_teams(self):
+        return self._n_teams
+
+    def set_n_teams(self, n_teams):
+        self._n_teams = n_teams
+
     def set_algorithm(self):
         # Type of avilable algorithms:
         # ec.GA, ec.EvolutionaryComputation
@@ -97,9 +103,6 @@ def set_algorithm(self):
         return getattr(ec, name)(self.random)
 
     def init_algogrithm(self):
-        args = {
-            "setdefault"
-        }
         self._algorithm.terminator = ec.terminators.evaluation_termination
         self._algorithm.replacer = ec.replacers.generational_replacement
         self._algorithm.variator = [
@@ -111,7 +114,7 @@ def init_algogrithm(self):
     def set_generator(self, random, args):
         # generate candidates
         # return list of lists of indices in population
-        return [random.randint(0, len(self._pop_desc)-1) for _ in range(self._batch_size)]
+        return [random.randint(0, len(self._pop_descs)-1) for _ in range(self._batch_size)]
 
     def set_evaluator(self, candidates, args):
         # evaluate the candidates
@@ -121,21 +124,21 @@ def set_evaluator(self, candidates, args):
             team = []
             index = []
             for c in cs:
-                team.append(self._pop_fitnesses[c])       
-            res.append(np.mean(team))
+                team.append(self._pop_fitnesses[c])
+            res.append(getattr(np, self._config["statistics"]["team"]["type"])(a=team, **self._config["statistics"]["team"]["params"]))
         return res
 
-    def get_descriptions(self, index):
+    def get_descriptors(self, index):
         descriptors = []
         for i in index:
-            descriptors.append(self._pop_desc[i])
+            descriptors.append(self._pop_descs[i])
         return descriptors
 
-    def get_squad(self, n_squad):
-        solutions = []
-        me_pop = self._me._archive.data()
-        self._pop_desc = me_pop["solution"]
-        self._pop_fitnesses = me_pop["objective"]
+    def ask(self):
+        teams = []
+        me_data = self._me._archive.data()
+        self._pop_descs = me_data["solution"]
+        self._pop_fitnesses = me_data["objective"]
 
         final_pop = self._algorithm.evolve(
             generator=self.set_generator,
@@ -149,9 +152,35 @@ def get_squad(self, n_squad):
         sort_indexes = sorted(range(len(final_pop_fitnesses)), key=final_pop_fitnesses.__getitem__, reverse=True)
         final_pop_fitnesses = final_pop_fitnesses[sort_indexes]
         final_pop_candidates = final_pop_candidates[sort_indexes]
-        for i in range(n_squad):
-            solutions.append(self.get_descriptions(final_pop_candidates[i]))
-        return solutions
+        solution_fitness = final_pop_fitnesses[:self._n_teams]
+        solution_pop = final_pop_candidates[:self._n_teams]
+        for j in range(self._n_teams):
+            solution = []
+            for i in range(self._batch_size):
+                solution.append(self._pop_descs[solution_pop[j][i]])
+            me_pop = self._me.ask(solution)
+            teams.append(me_pop)
+        return teams
+
+    def tell(self, fitnesses, trees):
+        self._pop_trees = trees
+        self._pop_fitnesses = fitnesses
+        self.set_best_team(trees, fitnesses)
+
+    def set_best_team(self, teams, teams_fitnesses):
+        # Set the best squad to the coach agent
+        # squad: list of squads, each squad is a list of agents
+
+        for index, team in enumerate(teams):
+            team_fitness = np.mean(teams_fitnesses[index])
+            if team_fitness > self._best_fitness:
+                self._best_team = team
+                self._best_fitness = team_fitness
+                print_info(f"New best team with fitness: {team_fitness}")
+
+    def get_best_squad(self):
+        # Get the best squad
+        return self._best_team
 
     def __str__(self):
         return f"Coach config: {self._config}"