Repository: MathisFederico/LearnRL (language: python)

Issue:
agents_order in Playground
We need an attribute agents_order and methods to set and shuffle the order of agents
This will come handy if we want to change agents rôles in a environement during training or testing


Discussion / hints:

Fix patch:
diff --git a/learnrl/playground.py b/learnrl/playground.py
index 2429c67..1bd0007 100644
--- a/learnrl/playground.py
+++ b/learnrl/playground.py
@@ -110,7 +110,7 @@ def __call__(self, *args) -> float:
 
 class Playground():
 
-    """A playground is used to run agent(s) on an environement
+    """A playground is used to run interactions between an environement and agent(s)
 
     Attributes:
         env (gym.Env):  Environement in which the agent(s) will play.
@@ -118,7 +118,7 @@ class Playground():
 
     """
 
-    def __init__(self, environement:Env, agents:Union[Agent, List[Agent]]):
+    def __init__(self, environement:Env, agents:Union[Agent, List[Agent]], agents_order=None):
         """A playground is used to run agent(s) on an environement
 
         Args:
@@ -136,6 +136,8 @@ def __init__(self, environement:Env, agents:Union[Agent, List[Agent]]):
 
         self.env = environement
         self.agents = agents
+        self.agents_order = None
+        self.set_agents_order(agents_order)
 
     @staticmethod
     def _get_episodes_cycle_len(episodes_cycle_len, episodes):
@@ -277,10 +279,13 @@ def _run_step(self, step, observation, previous,
             self.env.render()
 
         # Get playing agent (TurnEnv)
-        agent_id = self.env.turn(observation) if isinstance(self.env, TurnEnv) else 0
-        if agent_id >= len(previous):
-            raise ValueError(f'Not enough agents to play environement {self.env}')
-        agent = self.agents[agent_id]
+        turn_id = self.env.turn(observation) if isinstance(self.env, TurnEnv) else 0
+        agent_id = self.agents_order[turn_id]
+        try:
+            agent = self.agents[agent_id]
+        except IndexError as error:
+            error_msg = f'Not enough agents to play environement {self.env}'
+            raise ValueError(error_msg) from error
 
         # If the agent has played before, perform a learning step
         prev = previous[agent_id]
@@ -366,3 +371,38 @@ def test(self, episodes, **kwargs):
                 UserWarning
             )
         self.run(episodes, render=render, learn=learn, verbose=verbose, **kwargs)
+
+    def set_agents_order(self, agents_order: list) -> list:
+        """Change the agents_order.
+
+        This will update the agents order.
+
+        Args:
+            agents_order: New agents indices order.
+                Default is range(n_agents).
+
+        Returns:
+            The updated agents ordered indices list.
+
+        """
+        if agents_order is None:
+            self.agents_order = list(range(len(self.agents)))
+        else:
+            if len(agents_order) != len(self.agents):
+                raise ValueError(
+                    f"Not every agents have an order number.\n"
+                    f"Custom order: {agents_order} for {len(self.agents)} agents\n"
+                )
+
+            valid_order = True
+            for place in range(len(self.agents)):
+                if not place in agents_order:
+                    valid_order = False
+
+            if not valid_order:
+                raise ValueError(
+                    f"Custom order is not taking every index in [0, n_agents-1].\n"
+                    f"Custom order: {agents_order}"
+                )
+            self.agents_order = agents_order
+        return self.agents_order


Test patch:
diff --git a/tests/unit/test_playground.py b/tests/unit/test_playground.py
new file mode 100644
index 0000000..dfc8673
--- /dev/null
+++ b/tests/unit/test_playground.py
@@ -0,0 +1,46 @@
+import pytest
+
+from gym import Env
+from learnrl import Playground, Agent
+
+
+def test_single_agent():
+    env = Env()
+    agents = Agent()
+    pg = Playground(env, agents)
+    assert pg.agents == [agents], \
+        "A single agent should be transformed in a list containing itself"
+
+def test_agent_order():
+    env = Env()
+    n_agents = 5
+    agents = [Agent() for _ in range(n_agents)]
+
+    # Default order
+    pg = Playground(env, agents)
+    assert pg.agents_order == list(range(n_agents)), \
+        f"Default agents_order shoud be {list(range(n_agents))} but was {pg.agents_order}"
+
+    # Custom order at init
+    custom_order = [4, 3, 1, 2, 0]
+    pg = Playground(env, agents, agents_order=custom_order)
+    assert pg.agents_order == custom_order, \
+        f"Custom agents_order shoud be {custom_order} but was {pg.agents_order}"
+
+    # Not enough indexes in custom order
+    with pytest.raises(ValueError, match=r"Not every agents*"):
+        pg = Playground(env, agents, agents_order=[4, 3, 1, 2])
+
+    # Missing indexes in custom order
+    with pytest.raises(ValueError, match=r".*not taking every index*"):
+        pg = Playground(env, agents, agents_order=[4, 6, 1, 2, 0])
+
+    # Custom order after init
+    custom_order = [4, 3, 1, 2, 0]
+    pg = Playground(env, agents, agents_order=custom_order)
+
+    new_custom_order = [3, 4, 1, 2, 0]
+    pg.set_agents_order(new_custom_order)
+
+    assert pg.agents_order == new_custom_order, \
+        f"Custom agents_order shoud be {new_custom_order} but was {pg.agents_order}"