Repository: strakam/generals-bots

Issue:
Allow agent deployment to real generals.io!
Crucial part of agent development is to be able to evaluate it.

The best way to evaluate is to play against other players and bots! 
Currently, the best way to achieve that is to allow users to "deploy" their agents developed in this repository directly to generals.io and we just started working on it.

We will do our best to hide all the complexity from you (bot developers), and make it as easy as executing a script with some argument :)

**Main developer**: @revolko

Discussion:

Patch:
diff --git a/Makefile b/Makefile
index d5355f1..39daf4f 100644
--- a/Makefile
+++ b/Makefile
@@ -13,6 +13,8 @@ pz:
 gym:
 	poetry run python3 -m examples.gymnasium_example
 
+remote:
+	poetry run python3 -m examples.client_example
 # Create new replay and run it
 make n_replay:
 	poetry run python3 -m examples.record_replay_example
diff --git a/README.md b/README.md
index ecf8953..c897960 100644
--- a/README.md
+++ b/README.md
@@ -19,9 +19,10 @@ challenging players to balance micro and macro-level decision-making.
 The combination of these elements makes the game highly engaging and complex.
 
 Highlights:
-* 🚀 **blazing-fast simulator**: run thousands of steps per second with `numpy`-powered efficiency
+* ⚡ **blazing-fast simulator**: run thousands of steps per second with `numpy`-powered efficiency
 * 🤝 **seamless integration**: fully compatible with RL standards 🤸[Gymnasium](https://gymnasium.farama.org/) and 🦁[PettingZoo](https://pettingzoo.farama.org/)
-* 🔧 **effortless customization**: easily tailor environments to your specific needs
+* 🔧 **extensive customization**: easily tailor environments to your specific needs
+* 🚀 **effortless deployment**: launch your agents to generalsio
 * 🔬 **analysis tools**: leverage features like replays for deeper insights
 
 > [!Note]
@@ -42,7 +43,7 @@ make install
 > [!Note]
 > Under the hood, `make install` installs [poetry](https://python-poetry.org/) and the package using poetry.
 
-## 🚀 Getting Started
+## 🌱 Getting Started
 Creating an agent is very simple. Start by subclassing an `Agent` class just like
 [`RandomAgent`](./generals/agents/random_agent.py) or [`ExpanderAgent`](./generals/agents/expander_agent.py).
 You can specify your agent `id` (name) and `color` and the only thing remaining is to implement the `act` function,
@@ -159,19 +160,19 @@ An observation for one agent is a dictionary `{"observation": observation, "acti
 The `observation` is a `Dict`. Values are either `numpy` matrices with shape `(N,M)`, or simple `int` constants:
 | Key                  | Shape     | Description                                                                  |
 | -------------------- | --------- | ---------------------------------------------------------------------------- |
-| `army`               | `(N,M)`   | Number of units in a cell regardless of the owner                            |
-| `general`            | `(N,M)`   | Mask indicating cells containing a general                                   |
-| `city`               | `(N,M)`   | Mask indicating cells containing a city                                      |
-| `visible_cells`      | `(N,M)`   | Mask indicating cells that are visible to the agent                          |
-| `owned_cells`        | `(N,M)`   | Mask indicating cells owned by the agent                                     |
-| `opponent_cells`     | `(N,M)`   | Mask indicating cells owned by the opponent                                  |
-| `neutral_cells`      | `(N,M)`   | Mask indicating cells that are not owned by any agent                        |
-| `structure`          | `(N,M)`   | Mask indicating whether cells contain cities or mountains, even out of FoV   |
+| `armies`             | `(N,M)`   | Number of units in a visible cell regardless of the owner                    |
+| `generals`           | `(N,M)`   | Mask indicating visible cells containing a general                           |
+| `cities`             | `(N,M)`   | Mask indicating visible cells containing a city                              |
+| `mountains`          | `(N,M)`   | Mask indicating visible cells containing mountains                           |
+| `neutral_cells`      | `(N,M)`   | Mask indicating visible cells that are not owned by any agent                |
+| `owned_cells`        | `(N,M)`   | Mask indicating visible cells owned by the agent                             |
+| `opponent_cells`     | `(N,M)`   | Mask indicating visible cells owned by the opponent                          |
+| `fog_cells`          | `(N,M)`   | Mask indicating fog cells that are not mountains or cities                   |
+| `structures_in_fog`  | `(N,M)`   | Mask showing cells containing either cities or mountains in fog              |
 | `owned_land_count`   |     —     | Number of cells the agent owns                                               |
 | `owned_army_count`   |     —     | Total number of units owned by the agent                                     |
 | `opponent_land_count`|     —     | Number of cells owned by the opponent                                        |
 | `opponent_army_count`|     —     | Total number of units owned by the opponent                                  |
-| `is_winner`          |     —     | Indicates whether the agent won                                              |
 | `timestep`           |     —     | Current timestep of the game                                                 |
 
 The `action_mask` is a 3D array with shape `(N, M, 4)`, where each element corresponds to whether a move is valid from cell
@@ -203,7 +204,25 @@ env = gym.make(..., reward_fn=custom_reward_fn)
 observations, info = env.reset()
 ```
 
-## 🌱 Contributing
+## 🚀 Deployment to Live Servers
+Complementary to local development, it is possible to run agents online against other agents and players.
+We use `socketio` for communication and you can either use out `autopilot` to run agent in a specified lobby indefinitely,
+or create your own connection workflow. Our implementations expect that your agent inherits from the `Agent` class, and has
+implemented required methods.
+```python
+from generals.remote import autopilot
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--user_id", type=str, default=...) # Register yourself at generalsio and use this id
+parser.add_argument("--lobby_id", type=str, default="psyo") # The last part of the lobby url
+parser.add_argument("--agent_id", type=str, default="Expander") # agent_id should be "registered" in AgentFactory
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    autopilot(args.agent_id, args.user_id, args.lobby_id)
+```
+This script will run your `ExpanderAgent` in lobby `psyo`.
+## 🙌 Contributing
 You can contribute to this project in multiple ways:
 - 🤖 If you implement ANY non-trivial agent, send it to us! We will publish it so others can play against it
 - 💡 If you have an idea on how to improve the game, submit an issue or create a PR, we are happy to improve!
diff --git a/examples/client_example.py b/examples/client_example.py
new file mode 100644
index 0000000..9e5376c
--- /dev/null
+++ b/examples/client_example.py
@@ -0,0 +1,12 @@
+from generals.remote import autopilot
+
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--user_id", type=str, default=...) # Register yourself at generalsio and use this id
+parser.add_argument("--lobby_id", type=str, default="psyo") # After you create a private lobby, copy last part of the url
+parser.add_argument("--agent_id", type=str, default="Expander") # agent_id should be "registered" in AgentFactory
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    autopilot(args.agent_id, args.user_id, args.lobby_id)
diff --git a/examples/record_replay_example.py b/examples/record_replay_example.py
index cb79c22..902eff4 100644
--- a/examples/record_replay_example.py
+++ b/examples/record_replay_example.py
@@ -3,8 +3,8 @@ import gymnasium as gym
 from generals import AgentFactory, GridFactory
 
 # Initialize agents -- see generals/agents/agent_factory.py for more options
-npc = AgentFactory.make_agent("expander")
-agent = AgentFactory.make_agent("random")
+npc = AgentFactory.make_agent("Expander")
+agent = AgentFactory.make_agent("Random")
 
 # Initialize grid factory
 grid_factory = GridFactory(
diff --git a/generals/agents/agent_factory.py b/generals/agents/agent_factory.py
index fc63314..ef6fd6c 100644
--- a/generals/agents/agent_factory.py
+++ b/generals/agents/agent_factory.py
@@ -16,9 +16,9 @@ class AgentFactory:
         """
         Creates an agent of the specified type.
         """
-        if agent_type == "random":
+        if agent_type == "Random":
             return RandomAgent(**kwargs)
-        elif agent_type == "expander":
+        elif agent_type == "Expander":
             return ExpanderAgent(**kwargs)
         else:
             raise ValueError(f"Unknown agent type: {agent_type}")
diff --git a/generals/agents/expander_agent.py b/generals/agents/expander_agent.py
index ca3550b..1afac4a 100644
--- a/generals/agents/expander_agent.py
+++ b/generals/agents/expander_agent.py
@@ -1,6 +1,8 @@
 import numpy as np
 
-from generals.core.config import Action, Direction, Observation
+from generals.core.config import Direction
+from generals.core.game import Action
+from generals.core.observation import Observation
 
 from .agent import Agent
 
@@ -26,7 +28,7 @@ class ExpanderAgent(Agent):
                 "split": 0,
             }
 
-        army = observation["army"]
+        army = observation["armies"]
         opponent = observation["opponent_cells"]
         neutral = observation["neutral_cells"]
 
diff --git a/generals/agents/random_agent.py b/generals/agents/random_agent.py
index e4e6a78..9d75622 100644
--- a/generals/agents/random_agent.py
+++ b/generals/agents/random_agent.py
@@ -1,6 +1,7 @@
 import numpy as np
 
-from generals.core.game import Action, Observation
+from generals.core.game import Action
+from generals.core.observation import Observation
 
 from .agent import Agent
 
diff --git a/generals/core/channels.py b/generals/core/channels.py
index 1820d35..70d9857 100644
--- a/generals/core/channels.py
+++ b/generals/core/channels.py
@@ -1,4 +1,5 @@
 import numpy as np
+from scipy.ndimage import maximum_filter  # type: ignore
 
 from .config import MOUNTAIN, PASSABLE
 
@@ -7,20 +8,21 @@ valid_generals = ["A", "B"]  # Generals are represented by A and B
 
 class Channels:
     """
-    army - army size in each cell
-    general - general mask (1 if general is in cell, 0 otherwise)
-    mountain - mountain mask (1 if cell is mountain, 0 otherwise)
-    city - city mask (1 if cell is city, 0 otherwise)
+    armies - army size in each cell
+    generals - general mask (1 if general is in cell, 0 otherwise)
+    mountains - mountain mask (1 if cell is mountain, 0 otherwise)
+    cities - city mask (1 if cell is city, 0 otherwise)
     passable - passable mask (1 if cell is passable, 0 otherwise)
     ownership_i - ownership mask for player i (1 if player i owns cell, 0 otherwise)
-    ownership_neutral - ownership mask for neutral cells that are passable (1 if cell is neutral, 0 otherwise)
+    ownership_neutral - ownership mask for neutral cells that are
+    passable (1 if cell is neutral, 0 otherwise)
     """
 
     def __init__(self, grid: np.ndarray, _agents: list[str]):
-        self._army: np.ndarray = np.where(np.isin(grid, valid_generals), 1, 0).astype(int)
-        self._general: np.ndarray = np.where(np.isin(grid, valid_generals), 1, 0).astype(bool)
-        self._mountain: np.ndarray = np.where(grid == MOUNTAIN, 1, 0).astype(bool)
-        self._city: np.ndarray = np.where(np.char.isdigit(grid), 1, 0).astype(bool)
+        self._armies: np.ndarray = np.where(np.isin(grid, valid_generals), 1, 0).astype(int)
+        self._generals: np.ndarray = np.where(np.isin(grid, valid_generals), 1, 0).astype(bool)
+        self._mountains: np.ndarray = np.where(grid == MOUNTAIN, 1, 0).astype(bool)
+        self._cities: np.ndarray = np.where(np.char.isdigit(grid), 1, 0).astype(bool)
         self._passable: np.ndarray = (grid != MOUNTAIN).astype(bool)
 
         self._ownership: dict[str, np.ndarray] = {
@@ -31,39 +33,71 @@ class Channels:
 
         # City costs are 40 + digit in the cell
         city_costs = np.where(np.char.isdigit(grid), grid, "0").astype(int)
-        self.army += 40 * self.city + city_costs
+        self.armies += 40 * self.cities + city_costs
+
+    def get_visibility(self, agent_id: str) -> np.ndarray:
+        channel = self._ownership[agent_id]
+        return maximum_filter(channel, size=3)
+
+    @staticmethod
+    def channel_to_indices(channel: np.ndarray) -> np.ndarray:
+        """
+        Returns a list of indices of cells with non-zero values from specified a channel.
+        """
+        return np.argwhere(channel != 0)
 
     @property
     def ownership(self) -> dict[str, np.ndarray]:
         return self._ownership
 
+    @ownership.setter
+    def ownership(self, value):
+        self._ownership = value
+
     @property
-    def army(self) -> np.ndarray:
-        return self._army
+    def armies(self) -> np.ndarray:
+        return self._armies
 
-    @army.setter
-    def army(self, value):
-        self._army = value
+    @armies.setter
+    def armies(self, value):
+        self._armies = value
 
     @property
-    def general(self) -> np.ndarray:
-        return self._general
+    def generals(self) -> np.ndarray:
+        return self._generals
+
+    @generals.setter
+    def generals(self, value):
+        self._generals = value
 
     @property
-    def mountain(self) -> np.ndarray:
-        return self._mountain
+    def mountains(self) -> np.ndarray:
+        return self._mountains
+
+    @mountains.setter
+    def mountains(self, value):
+        self._mountains = value
 
     @property
-    def city(self) -> np.ndarray:
-        return self._city
+    def cities(self) -> np.ndarray:
+        return self._cities
+
+    @cities.setter
+    def cities(self, value):
+        self._cities = value
 
     @property
     def passable(self) -> np.ndarray:
         return self._passable
 
+    @passable.setter
+    def passable(self, value):
+        self._passable = value
+
     @property
     def ownership_neutral(self) -> np.ndarray:
         return self._ownership["neutral"]
 
-    def _set_passable(self, value):
-        self._passable = value
+    @ownership_neutral.setter
+    def ownership_neutral(self, value):
+        self._ownership["neutral"] = value
diff --git a/generals/core/config.py b/generals/core/config.py
index bc86042..73cd592 100644
--- a/generals/core/config.py
+++ b/generals/core/config.py
@@ -1,19 +1,6 @@
-from collections.abc import Callable
 from enum import Enum, IntEnum, StrEnum
 from importlib.resources import files
-from typing import Any, Literal, TypeAlias
-
-import gymnasium as gym
-import numpy as np
-
-# Type aliases
-Observation: TypeAlias = dict[str, np.ndarray | dict[str, gym.Space]]
-Action: TypeAlias = dict[str, int | np.ndarray]
-Info: TypeAlias = dict[str, Any]
-
-Reward: TypeAlias = float
-RewardFn: TypeAlias = Callable[[Observation, Action, bool, Info], Reward]
-AgentID: TypeAlias = str
+from typing import Literal
 
 # Game Literals
 PASSABLE: Literal["."] = "."
@@ -34,6 +21,9 @@ class Direction(Enum):
     RIGHT = (0, 1)
 
 
+DIRECTIONS = [Direction.UP, Direction.DOWN, Direction.LEFT, Direction.RIGHT]
+
+
 class Path(StrEnum):
     GENERAL_PATH = str(files("generals.assets.images") / "crownie.png")
     CITY_PATH = str(files("generals.assets.images") / "citie.png")
diff --git a/generals/core/game.py b/generals/core/game.py
index db9808d..feda084 100644
--- a/generals/core/game.py
+++ b/generals/core/game.py
@@ -1,15 +1,16 @@
-import warnings
-from typing import Any
+from typing import Any, TypeAlias
 
 import gymnasium as gym
 import numpy as np
-from scipy.ndimage import maximum_filter  # type: ignore
 
 from .channels import Channels
-from .config import Action, Direction, Info, Observation
+from .config import DIRECTIONS
 from .grid import Grid
+from .observation import Observation
 
-DIRECTIONS = [Direction.UP, Direction.DOWN, Direction.LEFT, Direction.RIGHT]
+# Type aliases
+Action: TypeAlias = dict[str, int | np.ndarray]
+Info: TypeAlias = dict[str, Any]
 
 
 class Game:
@@ -41,19 +42,19 @@ class Game:
             {
                 "observation": gym.spaces.Dict(
                     {
-                        "army": gym.spaces.MultiDiscrete(grid_discrete),
-                        "general": grid_multi_binary,
-                        "city": grid_multi_binary,
+                        "armies": gym.spaces.MultiDiscrete(grid_discrete),
+                        "generals": grid_multi_binary,
+                        "cities": grid_multi_binary,
+                        "mountains": grid_multi_binary,
+                        "neutral_cells": grid_multi_binary,
                         "owned_cells": grid_multi_binary,
                         "opponent_cells": grid_multi_binary,
-                        "neutral_cells": grid_multi_binary,
-                        "visible_cells": grid_multi_binary,
-                        "structure": grid_multi_binary,
+                        "fog_cells": grid_multi_binary,
+                        "structures_in_fog": grid_multi_binary,
                         "owned_land_count": gym.spaces.Discrete(self.max_army_value),
                         "owned_army_count": gym.spaces.Discrete(self.max_army_value),
                         "opponent_land_count": gym.spaces.Discrete(self.max_army_value),
                         "opponent_army_count": gym.spaces.Discrete(self.max_army_value),
-                        "is_winner": gym.spaces.Discrete(2),
                         "timestep": gym.spaces.Discrete(self.max_timestep),
                     }
                 ),
@@ -70,58 +71,6 @@ class Game:
             }
         )
 
-    def action_mask(self, agent: str) -> np.ndarray:
-        """
-        Function to compute valid actions from a given ownership mask.
-
-        Valid action is an action that originates from agent's cell with atleast 2 units
-        and does not bump into a mountain or fall out of the grid.
-        Returns:
-            np.ndarray: an NxNx4 array, where each channel is a boolean mask
-            of valid actions (UP, DOWN, LEFT, RIGHT) for each cell in the grid.
-
-            I.e. valid_action_mask[i, j, k] is 1 if action k is valid in cell (i, j).
-        """
-
-        ownership_channel = self.channels.ownership[agent]
-        more_than_1_army = (self.channels.army > 1) * ownership_channel
-        owned_cells_indices = self.channel_to_indices(more_than_1_army)
-        valid_action_mask = np.zeros((self.grid_dims[0], self.grid_dims[1], 4), dtype=bool)
-
-        if self.is_done() and not self.agent_won(agent):  # if you lost, return all zeros
-            return valid_action_mask
-
-        for channel_index, direction in enumerate(DIRECTIONS):
-            destinations = owned_cells_indices + direction.value
-
-            # check if destination is in grid bounds
-            in_first_boundary = np.all(destinations >= 0, axis=1)
-            in_height_boundary = destinations[:, 0] < self.grid_dims[0]
-            in_width_boundary = destinations[:, 1] < self.grid_dims[1]
-            destinations = destinations[in_first_boundary & in_height_boundary & in_width_boundary]
-
-            # check if destination is road
-            passable_cell_indices = self.channels.passable[destinations[:, 0], destinations[:, 1]] == 1
-            action_destinations = destinations[passable_cell_indices]
-
-            # get valid action mask for a given direction
-            valid_source_indices = action_destinations - direction.value
-            valid_action_mask[valid_source_indices[:, 0], valid_source_indices[:, 1], channel_index] = 1.0
-        # assert False
-        return valid_action_mask
-
-    def channel_to_indices(self, channel: np.ndarray) -> np.ndarray:
-        """
-        Returns a list of indices of cells with non-zero values from specified a channel.
-        """
-        return np.argwhere(channel != 0)
-
-    def visibility_channel(self, ownership_channel: np.ndarray) -> np.ndarray:
-        """
-        Returns a binary channel of visible cells from the perspective of the given player.
-        """
-        return maximum_filter(ownership_channel, size=3)
-
     def step(self, actions: dict[str, Action]) -> tuple[dict[str, Observation], dict[str, Any]]:
         """
         Perform one step of the game
@@ -142,18 +91,10 @@ class Game:
             # Skip if agent wants to pass the turn
             if pass_turn == 1:
                 continue
-            # Skip if the move is invalid
-            if self.action_mask(agent)[i, j, direction] == 0:
-                warnings.warn(
-                    f"The submitted move byt agent {agent} does not take effect.\
-                    Probably because you submitted an invalid move.",
-                    UserWarning,
-                )
-                continue
             if split_army == 1:  # Agent wants to split the army
-                army_to_move = self.channels.army[i, j] // 2
+                army_to_move = self.channels.armies[i, j] // 2
             else:  # Leave just one army in the source cell
-                army_to_move = self.channels.army[i, j] - 1
+                army_to_move = self.channels.armies[i, j] - 1
             if army_to_move < 1:  # Skip if army size to move is less than 1
                 continue
             moves[agent] = (i, j, direction, army_to_move)
@@ -163,8 +104,8 @@ class Game:
             si, sj, direction, army_to_move = moves[agent]
 
             # Cap the amount of army to move (previous moves may have lowered available army)
-            army_to_move = min(army_to_move, self.channels.army[si, sj] - 1)
-            army_to_stay = self.channels.army[si, sj] - army_to_move
+            army_to_move = min(army_to_move, self.channels.armies[si, sj] - 1)
+            army_to_stay = self.channels.armies[si, sj] - army_to_move
 
             # Check if the current agent still owns the source cell and has more than 1 army
             if self.channels.ownership[agent][si, sj] == 0 or army_to_move < 1:
@@ -176,20 +117,20 @@ class Game:
             )  # destination indices
 
             # Figure out the target square owner and army size
-            target_square_army = self.channels.army[di, dj]
+            target_square_army = self.channels.armies[di, dj]
             target_square_owner_idx = np.argmax(
                 [self.channels.ownership[agent][di, dj] for agent in ["neutral"] + self.agents]
             )
             target_square_owner = (["neutral"] + self.agents)[target_square_owner_idx]
             if target_square_owner == agent:
-                self.channels.army[di, dj] += army_to_move
-                self.channels.army[si, sj] = army_to_stay
+                self.channels.armies[di, dj] += army_to_move
+                self.channels.armies[si, sj] = army_to_stay
             else:
                 # Calculate resulting army, winner and update channels
                 remaining_army = np.abs(target_square_army - army_to_move)
                 square_winner = agent if target_square_army < army_to_move else target_square_owner
-                self.channels.army[di, dj] = remaining_army
-                self.channels.army[si, sj] = army_to_stay
+                self.channels.armies[di, dj] = remaining_army
+                self.channels.armies[si, sj] = army_to_stay
                 self.channels.ownership[square_winner][di, dj] = 1
                 if square_winner != target_square_owner:
                     self.channels.ownership[target_square_owner][di, dj] = 0
@@ -206,16 +147,10 @@ class Game:
         else:
             self._global_game_update()
 
-        observations = self.get_all_observations()
-        infos: dict[str, Any] = {agent: {} for agent in self.agents}
+        observations = {agent: self.agent_observation(agent) for agent in self.agents}
+        infos = self.get_infos()
         return observations, infos
 
-    def get_all_observations(self) -> dict[str, Observation]:
-        """
-        Returns observations for all agents.
-        """
-        return {agent: self.agent_observation(agent) for agent in self.agents}
-
     def _global_game_update(self) -> None:
         """
         Update game state globally.
@@ -226,13 +161,13 @@ class Game:
         # every `increment_rate` steps, increase army size in each cell
         if self.time % self.increment_rate == 0:
             for owner in owners:
-                self.channels.army += self.channels.ownership[owner]
+                self.channels.armies += self.channels.ownership[owner]
 
         # Increment armies on general and city cells, but only if they are owned by player
         if self.time % 2 == 0 and self.time > 0:
-            update_mask = self.channels.general + self.channels.city
+            update_mask = self.channels.generals + self.channels.cities
             for owner in owners:
-                self.channels.army += update_mask * self.channels.ownership[owner]
+                self.channels.armies += update_mask * self.channels.ownership[owner]
 
     def is_done(self) -> bool:
         """
@@ -249,7 +184,7 @@ class Game:
         """
         players_stats = {}
         for agent in self.agents:
-            army_size = np.sum(self.channels.army * self.channels.ownership[agent]).astype(int)
+            army_size = np.sum(self.channels.armies * self.channels.ownership[agent]).astype(int)
             land_size = np.sum(self.channels.ownership[agent]).astype(int)
             players_stats[agent] = {
                 "army": army_size,
@@ -262,31 +197,51 @@ class Game:
         """
         Returns an observation for a given agent.
         """
-        info = self.get_infos()
+        scores = {}
+        for _agent in self.agents:
+            army_size = np.sum(self.channels.armies * self.channels.ownership[_agent]).astype(int)
+            land_size = np.sum(self.channels.ownership[_agent]).astype(int)
+            scores[_agent] = {
+                "army": army_size,
+                "land": land_size,
+            }
+
+        visible = self.channels.get_visibility(agent)
+        invisible = 1 - visible
+
         opponent = self.agents[0] if agent == self.agents[1] else self.agents[1]
-        visibility = self.visibility_channel(self.channels.ownership[agent])
-        _observation = {
-            "army": self.channels.army.astype(int) * visibility,
-            "general": self.channels.general * visibility,
-            "city": self.channels.city * visibility,
-            "owned_cells": self.channels.ownership[agent] * visibility,
-            "opponent_cells": self.channels.ownership[opponent] * visibility,
-            "neutral_cells": self.channels.ownership_neutral * visibility,
-            "visible_cells": visibility,
-            "structure": self.channels.mountain + self.channels.city,
-            "owned_land_count": info[agent]["land"],
-            "owned_army_count": info[agent]["army"],
-            "opponent_land_count": info[opponent]["land"],
-            "opponent_army_count": info[opponent]["army"],
-            "is_winner": int(info[agent]["is_winner"]),
-            "timestep": self.time,
-        }
-        observation: Observation = {
-            "observation": _observation,
-            "action_mask": self.action_mask(agent),
-        }
 
-        return observation
+        armies = self.channels.armies.astype(int) * visible
+        mountains = self.channels.mountains * visible
+        generals = self.channels.generals * visible
+        cities = self.channels.cities * visible
+        neutral_cells = self.channels.ownership_neutral * visible
+        owned_cells = self.channels.ownership[agent] * visible
+        opponent_cells = self.channels.ownership[opponent] * visible
+        structures_in_fog = invisible * (self.channels.mountains + self.channels.cities)
+        fog_cells = invisible - structures_in_fog
+        owned_land_count = scores[agent]["land"]
+        owned_army_count = scores[agent]["army"]
+        opponent_land_count = scores[opponent]["land"]
+        opponent_army_count = scores[opponent]["army"]
+        timestep = self.time
+
+        return Observation(
+            armies=armies,
+            generals=generals,
+            cities=cities,
+            mountains=mountains,
+            neutral_cells=neutral_cells,
+            owned_cells=owned_cells,
+            opponent_cells=opponent_cells,
+            fog_cells=fog_cells,
+            structures_in_fog=structures_in_fog,
+            owned_land_count=owned_land_count,
+            owned_army_count=owned_army_count,
+            opponent_land_count=opponent_land_count,
+            opponent_army_count=opponent_army_count,
+            timestep=timestep,
+        )
 
     def agent_won(self, agent: str) -> bool:
         """
diff --git a/generals/core/observation.py b/generals/core/observation.py
new file mode 100644
index 0000000..cbd104b
--- /dev/null
+++ b/generals/core/observation.py
@@ -0,0 +1,108 @@
+import numpy as np
+
+from generals.core.config import DIRECTIONS
+
+
+class Observation:
+    def __init__(
+        self,
+        armies: np.ndarray,
+        generals: np.ndarray,
+        cities: np.ndarray,
+        mountains: np.ndarray,
+        neutral_cells: np.ndarray,
+        owned_cells: np.ndarray,
+        opponent_cells: np.ndarray,
+        fog_cells: np.ndarray,
+        structures_in_fog: np.ndarray,
+        owned_land_count: int,
+        owned_army_count: int,
+        opponent_land_count: int,
+        opponent_army_count: int,
+        timestep: int,
+    ):
+        self.armies = armies
+        self.generals = generals
+        self.cities = cities
+        self.mountains = mountains
+        self.neutral_cells = neutral_cells
+        self.owned_cells = owned_cells
+        self.opponent_cells = opponent_cells
+        self.fog_cells = fog_cells
+        self.structures_in_fog = structures_in_fog
+        self.owned_land_count = owned_land_count
+        self.owned_army_count = owned_army_count
+        self.opponent_land_count = opponent_land_count
+        self.opponent_army_count = opponent_army_count
+        self.timestep = timestep
+        # armies, generals, cities, mountains, empty, owner, fogged, structure in fog
+
+    def action_mask(self) -> np.ndarray:
+        """
+        Function to compute valid actions from a given ownership mask.
+
+        Valid action is an action that originates from agent's cell with atleast 2 units
+        and does not bump into a mountain or fall out of the grid.
+        Returns:
+            np.ndarray: an NxNx4 array, where each channel is a boolean mask
+            of valid actions (UP, DOWN, LEFT, RIGHT) for each cell in the grid.
+
+            I.e. valid_action_mask[i, j, k] is 1 if action k is valid in cell (i, j).
+        """
+        height, width = self.owned_cells.shape
+
+        ownership_channel = self.owned_cells
+        more_than_1_army = (self.armies > 1) * ownership_channel
+        owned_cells_indices = np.argwhere(more_than_1_army)
+        valid_action_mask = np.zeros((height, width, 4), dtype=bool)
+
+        if np.sum(ownership_channel) == 0:
+            return valid_action_mask
+
+        for channel_index, direction in enumerate(DIRECTIONS):
+            destinations = owned_cells_indices + direction.value
+
+            # check if destination is in grid bounds
+            in_first_boundary = np.all(destinations >= 0, axis=1)
+            in_height_boundary = destinations[:, 0] < height
+            in_width_boundary = destinations[:, 1] < width
+            destinations = destinations[in_first_boundary & in_height_boundary & in_width_boundary]
+
+            # check if destination is road
+            passable_cells = 1 - self.mountains
+            # assert that every value is either 0 or 1 in passable cells
+            assert np.all(np.isin(passable_cells, [0, 1])), f"{passable_cells}"
+            passable_cell_indices = passable_cells[destinations[:, 0], destinations[:, 1]] == 1
+            action_destinations = destinations[passable_cell_indices]
+
+            # get valid action mask for a given direction
+            valid_source_indices = action_destinations - direction.value
+            valid_action_mask[valid_source_indices[:, 0], valid_source_indices[:, 1], channel_index] = 1.0
+
+        return valid_action_mask
+
+    def as_dict(self, with_mask=True):
+        _obs = {
+            "armies": self.armies,
+            "generals": self.generals,
+            "cities": self.cities,
+            "mountains": self.mountains,
+            "neutral_cells": self.neutral_cells,
+            "owned_cells": self.owned_cells,
+            "opponent_cells": self.opponent_cells,
+            "fog_cells": self.fog_cells,
+            "structures_in_fog": self.structures_in_fog,
+            "owned_land_count": self.owned_land_count,
+            "owned_army_count": self.owned_army_count,
+            "opponent_land_count": self.opponent_land_count,
+            "opponent_army_count": self.opponent_army_count,
+            "timestep": self.timestep,
+        }
+        if with_mask:
+            obs = {
+                "observation": _obs,
+                "action_mask": self.action_mask(),
+            }
+        else:
+            obs = _obs
+        return obs
diff --git a/generals/envs/gymnasium_generals.py b/generals/envs/gymnasium_generals.py
index febaf91..4a38844 100644
--- a/generals/envs/gymnasium_generals.py
+++ b/generals/envs/gymnasium_generals.py
@@ -1,16 +1,20 @@
+from collections.abc import Callable
 from copy import deepcopy
-from typing import Any, SupportsFloat
+from typing import Any, SupportsFloat, TypeAlias
 
 import gymnasium as gym
 
 from generals.agents import Agent, AgentFactory
-from generals.core.config import Reward, RewardFn
-from generals.core.game import Action, Game, Info, Observation
+from generals.core.game import Action, Game, Info
 from generals.core.grid import GridFactory
+from generals.core.observation import Observation
 from generals.core.replay import Replay
 from generals.gui import GUI
 from generals.gui.properties import GuiMode
 
+Reward: TypeAlias = float
+RewardFn: TypeAlias = Callable[[Observation, Action, bool, Info], Reward]
+
 
 class GymnasiumGenerals(gym.Env):
     metadata = {
@@ -90,18 +94,21 @@ class GymnasiumGenerals(gym.Env):
         self.observation_space = self.game.observation_space
         self.action_space = self.game.action_space
 
-        observation = self.game.agent_observation(self.agent_id)
+        observation = self.game.agent_observation(self.agent_id).as_dict()
         info: dict[str, Any] = {}
         return observation, info
 
     def step(self, action: Action) -> tuple[Observation, SupportsFloat, bool, bool, dict[str, Any]]:
         # Get action of NPC
-        npc_action = self.npc.act(self.game.agent_observation(self.npc.id))
+        npc_observation = self.game.agent_observation(self.npc.id).as_dict()
+        npc_action = self.npc.act(npc_observation)
         actions = {self.agent_id: action, self.npc.id: npc_action}
 
         observations, infos = self.game.step(actions)
+        infos = {agent_id: {} for agent_id in self.agent_ids}
+
         # From observations of all agents, pick only those relevant for the main agent
-        obs = observations[self.agent_id]
+        obs = observations[self.agent_id].as_dict()
         info = infos[self.agent_id]
         reward = self.reward_fn(obs, action, self.game.is_done(), info)
         terminated = self.game.is_done()
@@ -124,14 +131,7 @@ class GymnasiumGenerals(gym.Env):
         done: bool,
         info: Info,
     ) -> Reward:
-        """
-        Give 0 if game still running, otherwise 1 for winner and -1 for loser.
-        """
-        if done:
-            reward = 1 if observation["observation"]["is_winner"] else -1
-        else:
-            reward = 0
-        return reward
+        return 0
 
     def close(self) -> None:
         if self.render_mode == "human":
diff --git a/generals/envs/gymnasium_wrappers.py b/generals/envs/gymnasium_wrappers.py
index ee570ca..28434b2 100644
--- a/generals/envs/gymnasium_wrappers.py
+++ b/generals/envs/gymnasium_wrappers.py
@@ -18,7 +18,7 @@ class NormalizedObservationWrapper(gym.ObservationWrapper):
                         "opponent_cells": grid_multi_binary,
                         "neutral_cells": grid_multi_binary,
                         "visible_cells": grid_multi_binary,
-                        "structure": grid_multi_binary,
+                        "structures_in_fog": grid_multi_binary,
                         "owned_land_count": unit_box,
                         "owned_army_count": unit_box,
                         "opponent_land_count": unit_box,
@@ -68,7 +68,7 @@ class RemoveActionMaskWrapper(gym.ObservationWrapper):
                 "opponent_cells": grid_multi_binary,
                 "neutral_cells": grid_multi_binary,
                 "visible_cells": grid_multi_binary,
-                "structure": grid_multi_binary,
+                "structures_in_fog": grid_multi_binary,
                 "owned_land_count": unit_box,
                 "owned_army_count": unit_box,
                 "opponent_land_count": unit_box,
@@ -106,7 +106,7 @@ class ObservationAsImageWrapper(gym.ObservationWrapper):
                 _observation["opponent_cells"],
                 _observation["neutral_cells"],
                 _observation["visible_cells"],
-                _observation["structure"],
+                _observation["structures_in_fog"],
                 _owned_land_count,
                 _owned_army_count,
                 _opponent_land_count,
diff --git a/generals/envs/pettingzoo_generals.py b/generals/envs/pettingzoo_generals.py
index bd550d0..d04c1f3 100644
--- a/generals/envs/pettingzoo_generals.py
+++ b/generals/envs/pettingzoo_generals.py
@@ -1,18 +1,22 @@
 import functools
+from collections.abc import Callable
 from copy import deepcopy
-from typing import Any
+from typing import Any, TypeAlias
 
 import pettingzoo  # type: ignore
 from gymnasium import spaces
 
 from generals.agents.agent import Agent
-from generals.core.config import AgentID, Reward, RewardFn
 from generals.core.game import Action, Game, Info, Observation
 from generals.core.grid import GridFactory
 from generals.core.replay import Replay
 from generals.gui import GUI
 from generals.gui.properties import GuiMode
 
+AgentID: TypeAlias = str
+Reward: TypeAlias = float
+RewardFn: TypeAlias = Callable[[Observation, Action, bool, Info], Reward]
+
 
 class PettingZooGenerals(pettingzoo.ParallelEnv):
     metadata: dict[str, Any] = {
@@ -82,7 +86,7 @@ class PettingZooGenerals(pettingzoo.ParallelEnv):
         elif hasattr(self, "replay"):
             del self.replay
 
-        observations = self.game.get_all_observations()
+        observations = {agent: self.game.agent_observation(agent).as_dict() for agent in self.agents}
         infos: dict[str, Any] = {agent: {} for agent in self.agents}
         return observations, infos
 
@@ -96,6 +100,7 @@ class PettingZooGenerals(pettingzoo.ParallelEnv):
         dict[AgentID, Info],
     ]:
         observations, infos = self.game.step(actions)
+        observations = {agent: observation.as_dict() for agent, observation in observations.items()}
         # You probably want to set your truncation based on self.game.time
         truncation = False if self.truncation is None else self.game.time >= self.truncation
         truncated = {agent: truncation for agent in self.agents}
@@ -128,14 +133,7 @@ class PettingZooGenerals(pettingzoo.ParallelEnv):
         done: bool,
         info: Info,
     ) -> Reward:
-        """
-        Give 0 if game still running, otherwise 1 for winner and -1 for loser.
-        """
-        if done:
-            reward = 1 if observation["observation"]["is_winner"] else -1
-        else:
-            reward = 0
-        return reward
+        return 0
 
     def close(self) -> None:
         if self.render_mode == "human":
diff --git a/generals/gui/rendering.py b/generals/gui/rendering.py
index f36848d..2ef2b15 100644
--- a/generals/gui/rendering.py
+++ b/generals/gui/rendering.py
@@ -180,7 +180,7 @@ class Renderer:
             ownership = self.game.channels.ownership[agent]
             owned_map = np.logical_or(owned_map, ownership)
             if self.agent_fov[agent]:
-                visibility = self.game.visibility_channel(ownership)
+                visibility = self.game.channels.get_visibility(agent)
                 visible_map = np.logical_or(visible_map, visibility)
 
         # Helper maps for not owned and invisible cells
@@ -194,7 +194,7 @@ class Renderer:
             self.draw_channel(visible_ownership, self.agent_data[agent]["color"])
 
         # Draw visible generals
-        visible_generals = np.logical_and(self.game.channels.general, visible_map)
+        visible_generals = np.logical_and(self.game.channels.generals, visible_map)
         self.draw_images(visible_generals, self._general_img)
 
         # Draw background of visible but not owned squares
@@ -205,27 +205,27 @@ class Renderer:
         self.draw_channel(invisible_map, FOG_OF_WAR)
 
         # Draw background of visible mountains
-        visible_mountain = np.logical_and(self.game.channels.mountain, visible_map)
+        visible_mountain = np.logical_and(self.game.channels.mountains, visible_map)
         self.draw_channel(visible_mountain, VISIBLE_MOUNTAIN)
 
         # Draw mountains (even if they are not visible)
-        self.draw_images(self.game.channels.mountain, self._mountain_img)
+        self.draw_images(self.game.channels.mountains, self._mountain_img)
 
         # Draw background of visible neutral cities
-        visible_cities = np.logical_and(self.game.channels.city, visible_map)
+        visible_cities = np.logical_and(self.game.channels.cities, visible_map)
         visible_cities_neutral = np.logical_and(visible_cities, self.game.channels.ownership_neutral)
         self.draw_channel(visible_cities_neutral, NEUTRAL_CASTLE)
 
         # Draw invisible cities as mountains
-        invisible_cities = np.logical_and(self.game.channels.city, invisible_map)
+        invisible_cities = np.logical_and(self.game.channels.cities, invisible_map)
         self.draw_images(invisible_cities, self._mountain_img)
 
         # Draw visible cities
         self.draw_images(visible_cities, self._city_img)
 
         # Draw nonzero army counts on visible squares
-        visible_army = self.game.channels.army * visible_map
-        visible_army_indices = self.game.channel_to_indices(visible_army)
+        visible_army = self.game.channels.armies * visible_map
+        visible_army_indices = self.channel_to_indices(visible_army)
         for i, j in visible_army_indices:
             self.render_cell_text(
                 self.tiles[i][j],
@@ -240,12 +240,18 @@ class Renderer:
             self.game_area.blit(self.tiles[i][j], (j * square_size, i * square_size))
         self.screen.blit(self.game_area, (0, 0))
 
+    def channel_to_indices(self, channel: np.ndarray) -> np.ndarray:
+        """
+        Returns a list of indices of cells with non-zero values from specified a channel.
+        """
+        return np.argwhere(channel != 0)
+
     def draw_channel(self, channel: np.ndarray, color: Color):
         """
         Draw background and borders (left and top) for grid tiles of a given channel
         """
         square_size = Dimension.SQUARE_SIZE.value
-        for i, j in self.game.channel_to_indices(channel):
+        for i, j in self.channel_to_indices(channel):
             self.tiles[i][j].fill(color)
             pygame.draw.line(self.tiles[i][j], BLACK, (0, 0), (0, square_size), 1)
             pygame.draw.line(self.tiles[i][j], BLACK, (0, 0), (square_size, 0), 1)
@@ -254,5 +260,5 @@ class Renderer:
         """
         Draw images on grid tiles of a given channel
         """
-        for i, j in self.game.channel_to_indices(channel):
+        for i, j in self.channel_to_indices(channel):
             self.tiles[i][j].blit(image, (3, 2))
diff --git a/generals/remote/__init__.py b/generals/remote/__init__.py
new file mode 100644
index 0000000..7b03810
--- /dev/null
+++ b/generals/remote/__init__.py
@@ -0,0 +1,8 @@
+from .generalsio_client import GeneralsBotError, GeneralsIOClient, GeneralsIOClientError, autopilot
+
+__all__ = [
+    "autopilot",
+    "GeneralsBotError",
+    "GeneralsIOClientError",
+    "GeneralsIOClient",
+]
diff --git a/generals/remote/generalsio_client.py b/generals/remote/generalsio_client.py
new file mode 100644
index 0000000..7e37397
--- /dev/null
+++ b/generals/remote/generalsio_client.py
@@ -0,0 +1,190 @@
+import numpy as np
+from socketio import SimpleClient  # type: ignore
+
+from generals.agents import Agent, AgentFactory
+from generals.core.config import Direction
+from generals.core.observation import Observation
+from generals.remote.generalsio_state import GeneralsIOstate
+
+DIRECTIONS = [Direction.UP, Direction.DOWN, Direction.LEFT, Direction.RIGHT]
+
+
+def autopilot(agent_id: str, user_id: str, lobby_id: str) -> None:
+    """
+    Start the autopilot for the GeneralsIO client.
+    This means that agent will join the lobby and force starts,
+    so he plays indefinitely.
+    """
+    agent = AgentFactory.make_agent(agent_id)
+    with GeneralsIOClient(agent, user_id) as client:
+        while True:
+            if client.status == "off":
+                client.join_private_lobby(lobby_id)
+            if client.status == "lobby":
+                client.join_game()
+
+
+class GeneralsBotError(Exception):
+    """Base generals-bot exception
+    TODO: find a place for exceptions
+    """
+
+    pass
+
+
+class GeneralsIOClientError(GeneralsBotError):
+    """Base GeneralsIOClient exception"""
+
+    pass
+
+
+class RegisterAgentError(GeneralsIOClientError):
+    """Registering bot error"""
+
+    def __init__(self, msg: str) -> None:
+        super().__init__()
+        self.msg = msg
+
+    def __str__(self) -> str:
+        return f"Failed to register the agent. Error: {self.msg}"
+
+
+class GeneralsIOClient(SimpleClient):
+    """
+    Wrapper around socket.io client to enable Agent to join
+    GeneralsIO lobby.
+    """
+
+    def __init__(self, agent: Agent, user_id: str):
+        super().__init__()
+        self.connect("https://botws.generals.io")
+        self.user_id = user_id
+        self.agent = agent
+        self._queue_id = ""
+        self._replay_id = ""
+        self._status = "off"  # can be "off","game","lobby","queue"
+
+    @property
+    def queue_id(self):
+        if not self._queue_id:
+            raise GeneralsIOClientError("Queue ID is not set.\nIs agent in the game lobby?")
+
+        return self._queue_id
+
+    @property
+    def replay_id(self):
+        if not self._replay_id:
+            print("No replay ID available.")
+        return self._replay_id
+
+    @property
+    def status(self):
+        return self._status
+
+    def _emit_receive(self, *args):
+        self.emit(*args)
+        return self.receive()
+
+    def register_agent(self, username: str) -> None:
+        """
+        Register Agent to GeneralsIO platform.
+        You can configure one Agent per `user_id`. `user_id` should be handled as secret.
+        :param user_id: secret ID of Agent
+        :param username: agent username, must be prefixed with `[Bot]`
+        """
+        event, response = self._emit_receive("set_username", (self.user_id, username))
+        if response:
+            # in case of success the response is empty
+            raise RegisterAgentError(response)
+
+    def join_private_lobby(self, queue_id: str) -> None:
+        """
+        Join (or create) private game lobby.
+        :param queue_id: Either URL or lobby ID number
+        """
+        self._status = "lobby"
+        self._emit_receive("join_private", (queue_id, self.user_id))
+        self._queue_id = queue_id
+
+    def join_game(self, force_start: bool = True) -> None:
+        """
+        Set force start if requested and wait for the game start.
+        :param force_start: If set to True, the Agent will set `Force Start` flag
+        """
+        self._status = "queue"
+        while True:
+            event, *data = self.receive()
+            self.emit("set_force_start", (self.queue_id, force_start))
+            if event == "game_start":
+                self._status = "game"
+                self._initialize_game(data)
+                self._play_game()
+                break
+
+    def join_1v1_queue(self) -> None:
+        """
+        Join 1v1 queue.
+        """
+        self._status = "queue"
+        self._emit_receive("join_1v1", self.user_id)
+        while True:
+            event, *data = self.receive()
+            if event == "game_start":
+                self._status = "game"
+                self._initialize_game(data)
+                self._play_game()
+                break
+
+    def _initialize_game(self, data: dict) -> None:
+        """
+        Triggered after server starts the game.
+        :param data: dictionary of information received in the beginning
+        """
+        self.game_state = GeneralsIOstate(data[0])
+        self._replay_id = data[0]["replay_id"]
+
+    def _generate_action(self, observation: Observation) -> tuple[int, int, int] | None:
+        """
+        Translate action from Agent to the server format.
+        :param action: dictionary representing the action
+        """
+        obs = observation.as_dict()
+        action = self.agent.act(obs)
+        if not action["pass"]:
+            source: np.ndarray = np.array(action["cell"])
+            direction = np.array(DIRECTIONS[action["direction"]].value)
+            split = action["split"]
+            destination = source + direction
+            # convert to index
+            source_index = source[0] * self.game_state.map[0] + source[1]
+            destination_index = destination[0] * self.game_state.map[0] + destination[1]
+            return (int(source_index), int(destination_index), int(split))
+        return None
+
+    def _play_game(self) -> None:
+        """
+        Main game-play loop.
+        TODO: spawn a new thread in which Agent will calculate its moves
+        """
+        while True:
+            event, data, _ = self.receive()
+            match event:
+                case "game_update":
+                    self.game_state.update(data)
+                    obs = self.game_state.get_observation()
+                    action = self._generate_action(obs)
+                    if action:
+                        self.emit("attack", action)
+                case "game_lost" | "game_won":
+                    self._finish_game(event == "game_won")
+                    break
+
+    def _finish_game(self, is_winner: bool) -> None:
+        """
+        Triggered after server finishes the game.
+        :param is_winner: True if Agent won the game
+        """
+        self._status = "off"
+        status = "won!" if is_winner else "lost."
+        print(f"Game is finished, you {status}, replay ID: https://bot.generals.io/replays/{self.replay_id}")
+        self.emit("leave_game")
diff --git a/generals/remote/generalsio_state.py b/generals/remote/generalsio_state.py
new file mode 100644
index 0000000..9048e31
--- /dev/null
+++ b/generals/remote/generalsio_state.py
@@ -0,0 +1,82 @@
+import numpy as np
+
+from generals.core.observation import Observation
+
+
+class GeneralsIOstate:
+    def __init__(self, data: dict):
+        self.usernames = data["usernames"]
+        self.player_index = data["playerIndex"]
+        self.opponent_index = 1 - self.player_index  # works only for 1v1
+
+        self.n_players = len(self.usernames)
+
+        self.map: list[int] = []
+        self.cities: list[int] = []
+
+    def update(self, data: dict) -> None:
+        self.turn = data["turn"]
+        self.map = self.apply_diff(self.map, data["map_diff"])
+        self.cities = self.apply_diff(self.cities, data["cities_diff"])
+        self.generals = data["generals"]
+        self.scores = data["scores"]
+        if "stars" in data:
+            self.stars = data["stars"]
+
+    def apply_diff(self, old: list[int], diff: list[int]) -> list[int]:
+        i = 0
+        new: list[int] = []
+        while i < len(diff):
+            if diff[i] > 0:  # matching
+                new.extend(old[len(new) : len(new) + diff[i]])
+            i += 1
+            if i < len(diff) and diff[i] > 0:  # applying diffs
+                new.extend(diff[i + 1 : i + 1 + diff[i]])
+                i += diff[i]
+            i += 1
+        return new
+
+    def get_observation(self) -> Observation:
+        width, height = self.map[0], self.map[1]
+        size = height * width
+
+        armies = np.array(self.map[2 : 2 + size]).reshape((height, width))
+        terrain = np.array(self.map[2 + size : 2 + 2 * size]).reshape((height, width))
+        cities = np.zeros((height, width))
+        for city in self.cities:
+            cities[city // width, city % width] = 1
+
+        generals = np.zeros((height, width))
+        for general in self.generals:
+            if general != -1:
+                generals[general // width, general % width] = 1
+
+        army = armies
+        owned_cells = np.where(terrain == self.player_index, 1, 0)
+        opponent_cells = np.where(terrain == self.opponent_index, 1, 0)
+        neutral_cells = np.where(terrain == -1, 1, 0)
+        mountain_cells = np.where(terrain == -2, 1, 0)
+        fog_cells = np.where(terrain == -3, 1, 0)
+        structures_in_fog = np.where(terrain == -4, 1, 0)
+        owned_land_count = self.scores[self.player_index]["tiles"]
+        owned_army_count = self.scores[self.player_index]["total"]
+        opponent_land_count = self.scores[self.opponent_index]["tiles"]
+        opponent_army_count = self.scores[self.opponent_index]["total"]
+        timestep = self.turn
+
+        return Observation(
+            armies=army,
+            generals=generals,
+            cities=cities,
+            mountains=mountain_cells,
+            neutral_cells=neutral_cells,
+            owned_cells=owned_cells,
+            opponent_cells=opponent_cells,
+            fog_cells=fog_cells,
+            structures_in_fog=structures_in_fog,
+            owned_land_count=owned_land_count,
+            owned_army_count=owned_army_count,
+            opponent_land_count=opponent_land_count,
+            opponent_army_count=opponent_army_count,
+            timestep=timestep,
+        )
diff --git a/poetry.lock b/poetry.lock
index 12a1699..0e6b615 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,26 @@
-# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
+
+[[package]]
+name = "bidict"
+version = "0.23.1"
+description = "The bidirectional mapping library for Python."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "bidict-0.23.1-py3-none-any.whl", hash = "sha256:5dae8d4d79b552a71cbabc7deb25dfe8ce710b17ff41711e13010ead2abfc3e5"},
+    {file = "bidict-0.23.1.tar.gz", hash = "sha256:03069d763bc387bbd20e7d49914e75fc4132a41937fa3405417e1a5a2d006d71"},
+]
+
+[[package]]
+name = "certifi"
+version = "2024.8.30"
+description = "Python package for providing Mozilla's CA Bundle."
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "certifi-2024.8.30-py3-none-any.whl", hash = "sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8"},
+    {file = "certifi-2024.8.30.tar.gz", hash = "sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9"},
+]
 
 [[package]]
 name = "cfgv"
@@ -11,6 +33,120 @@ files = [
     {file = "cfgv-3.4.0.tar.gz", hash = "sha256:e52591d4c5f5dead8e0f673fb16db7949d2cfb3f7da4582893288f0ded8fe560"},
 ]
 
+[[package]]
+name = "charset-normalizer"
+version = "3.4.0"
+description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet."
+optional = false
+python-versions = ">=3.7.0"
+files = [
+    {file = "charset_normalizer-3.4.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:4f9fc98dad6c2eaa32fc3af1417d95b5e3d08aff968df0cd320066def971f9a6"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0de7b687289d3c1b3e8660d0741874abe7888100efe14bd0f9fd7141bcbda92b"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5ed2e36c3e9b4f21dd9422f6893dec0abf2cca553af509b10cd630f878d3eb99"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:40d3ff7fc90b98c637bda91c89d51264a3dcf210cade3a2c6f838c7268d7a4ca"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1110e22af8ca26b90bd6364fe4c763329b0ebf1ee213ba32b68c73de5752323d"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:86f4e8cca779080f66ff4f191a685ced73d2f72d50216f7112185dc02b90b9b7"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7f683ddc7eedd742e2889d2bfb96d69573fde1d92fcb811979cdb7165bb9c7d3"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:27623ba66c183eca01bf9ff833875b459cad267aeeb044477fedac35e19ba907"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f606a1881d2663630ea5b8ce2efe2111740df4b687bd78b34a8131baa007f79b"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:0b309d1747110feb25d7ed6b01afdec269c647d382c857ef4663bbe6ad95a912"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:136815f06a3ae311fae551c3df1f998a1ebd01ddd424aa5603a4336997629e95"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:14215b71a762336254351b00ec720a8e85cada43b987da5a042e4ce3e82bd68e"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:79983512b108e4a164b9c8d34de3992f76d48cadc9554c9e60b43f308988aabe"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-win32.whl", hash = "sha256:c94057af19bc953643a33581844649a7fdab902624d2eb739738a30e2b3e60fc"},
+    {file = "charset_normalizer-3.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:55f56e2ebd4e3bc50442fbc0888c9d8c94e4e06a933804e2af3e89e2f9c1c749"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:0d99dd8ff461990f12d6e42c7347fd9ab2532fb70e9621ba520f9e8637161d7c"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c57516e58fd17d03ebe67e181a4e4e2ccab1168f8c2976c6a334d4f819fe5944"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6dba5d19c4dfab08e58d5b36304b3f92f3bd5d42c1a3fa37b5ba5cdf6dfcbcee"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bf4475b82be41b07cc5e5ff94810e6a01f276e37c2d55571e3fe175e467a1a1c"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ce031db0408e487fd2775d745ce30a7cd2923667cf3b69d48d219f1d8f5ddeb6"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8ff4e7cdfdb1ab5698e675ca622e72d58a6fa2a8aa58195de0c0061288e6e3ea"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3710a9751938947e6327ea9f3ea6332a09bf0ba0c09cae9cb1f250bd1f1549bc"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:82357d85de703176b5587dbe6ade8ff67f9f69a41c0733cf2425378b49954de5"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:47334db71978b23ebcf3c0f9f5ee98b8d65992b65c9c4f2d34c2eaf5bcaf0594"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:8ce7fd6767a1cc5a92a639b391891bf1c268b03ec7e021c7d6d902285259685c"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f1a2f519ae173b5b6a2c9d5fa3116ce16e48b3462c8b96dfdded11055e3d6365"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:63bc5c4ae26e4bc6be6469943b8253c0fd4e4186c43ad46e713ea61a0ba49129"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bcb4f8ea87d03bc51ad04add8ceaf9b0f085ac045ab4d74e73bbc2dc033f0236"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-win32.whl", hash = "sha256:9ae4ef0b3f6b41bad6366fb0ea4fc1d7ed051528e113a60fa2a65a9abb5b1d99"},
+    {file = "charset_normalizer-3.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:cee4373f4d3ad28f1ab6290684d8e2ebdb9e7a1b74fdc39e4c211995f77bec27"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:0713f3adb9d03d49d365b70b84775d0a0d18e4ab08d12bc46baa6132ba78aaf6"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:de7376c29d95d6719048c194a9cf1a1b0393fbe8488a22008610b0361d834ecf"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4a51b48f42d9358460b78725283f04bddaf44a9358197b889657deba38f329db"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b295729485b06c1a0683af02a9e42d2caa9db04a373dc38a6a58cdd1e8abddf1"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ee803480535c44e7f5ad00788526da7d85525cfefaf8acf8ab9a310000be4b03"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3d59d125ffbd6d552765510e3f31ed75ebac2c7470c7274195b9161a32350284"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8cda06946eac330cbe6598f77bb54e690b4ca93f593dee1568ad22b04f347c15"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:07afec21bbbbf8a5cc3651aa96b980afe2526e7f048fdfb7f1014d84acc8b6d8"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6b40e8d38afe634559e398cc32b1472f376a4099c75fe6299ae607e404c033b2"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:b8dcd239c743aa2f9c22ce674a145e0a25cb1566c495928440a181ca1ccf6719"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:84450ba661fb96e9fd67629b93d2941c871ca86fc38d835d19d4225ff946a631"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:44aeb140295a2f0659e113b31cfe92c9061622cadbc9e2a2f7b8ef6b1e29ef4b"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1db4e7fefefd0f548d73e2e2e041f9df5c59e178b4c72fbac4cc6f535cfb1565"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-win32.whl", hash = "sha256:5726cf76c982532c1863fb64d8c6dd0e4c90b6ece9feb06c9f202417a31f7dd7"},
+    {file = "charset_normalizer-3.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:b197e7094f232959f8f20541ead1d9862ac5ebea1d58e9849c1bf979255dfac9"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dd4eda173a9fcccb5f2e2bd2a9f423d180194b1bf17cf59e3269899235b2a114"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e9e3c4c9e1ed40ea53acf11e2a386383c3304212c965773704e4603d589343ed"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:92a7e36b000bf022ef3dbb9c46bfe2d52c047d5e3f3343f43204263c5addc250"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:54b6a92d009cbe2fb11054ba694bc9e284dad30a26757b1e372a1fdddaf21920"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1ffd9493de4c922f2a38c2bf62b831dcec90ac673ed1ca182fe11b4d8e9f2a64"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:35c404d74c2926d0287fbd63ed5d27eb911eb9e4a3bb2c6d294f3cfd4a9e0c23"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4796efc4faf6b53a18e3d46343535caed491776a22af773f366534056c4e1fbc"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e7fdd52961feb4c96507aa649550ec2a0d527c086d284749b2f582f2d40a2e0d"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:92db3c28b5b2a273346bebb24857fda45601aef6ae1c011c0a997106581e8a88"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ab973df98fc99ab39080bfb0eb3a925181454d7c3ac8a1e695fddfae696d9e90"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:4b67fdab07fdd3c10bb21edab3cbfe8cf5696f453afce75d815d9d7223fbe88b"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:aa41e526a5d4a9dfcfbab0716c7e8a1b215abd3f3df5a45cf18a12721d31cb5d"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ffc519621dce0c767e96b9c53f09c5d215578e10b02c285809f76509a3931482"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-win32.whl", hash = "sha256:f19c1585933c82098c2a520f8ec1227f20e339e33aca8fa6f956f6691b784e67"},
+    {file = "charset_normalizer-3.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:707b82d19e65c9bd28b81dde95249b07bf9f5b90ebe1ef17d9b57473f8a64b7b"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:dbe03226baf438ac4fda9e2d0715022fd579cb641c4cf639fa40d53b2fe6f3e2"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dd9a8bd8900e65504a305bf8ae6fa9fbc66de94178c420791d0293702fce2df7"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b8831399554b92b72af5932cdbbd4ddc55c55f631bb13ff8fe4e6536a06c5c51"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a14969b8691f7998e74663b77b4c36c0337cb1df552da83d5c9004a93afdb574"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dcaf7c1524c0542ee2fc82cc8ec337f7a9f7edee2532421ab200d2b920fc97cf"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:425c5f215d0eecee9a56cdb703203dda90423247421bf0d67125add85d0c4455"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-musllinux_1_2_aarch64.whl", hash = "sha256:d5b054862739d276e09928de37c79ddeec42a6e1bfc55863be96a36ba22926f6"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-musllinux_1_2_i686.whl", hash = "sha256:f3e73a4255342d4eb26ef6df01e3962e73aa29baa3124a8e824c5d3364a65748"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-musllinux_1_2_ppc64le.whl", hash = "sha256:2f6c34da58ea9c1a9515621f4d9ac379871a8f21168ba1b5e09d74250de5ad62"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-musllinux_1_2_s390x.whl", hash = "sha256:f09cb5a7bbe1ecae6e87901a2eb23e0256bb524a79ccc53eb0b7629fbe7677c4"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-musllinux_1_2_x86_64.whl", hash = "sha256:0099d79bdfcf5c1f0c2c72f91516702ebf8b0b8ddd8905f97a8aecf49712c621"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-win32.whl", hash = "sha256:9c98230f5042f4945f957d006edccc2af1e03ed5e37ce7c373f00a5a4daa6149"},
+    {file = "charset_normalizer-3.4.0-cp37-cp37m-win_amd64.whl", hash = "sha256:62f60aebecfc7f4b82e3f639a7d1433a20ec32824db2199a11ad4f5e146ef5ee"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:af73657b7a68211996527dbfeffbb0864e043d270580c5aef06dc4b659a4b578"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:cab5d0b79d987c67f3b9e9c53f54a61360422a5a0bc075f43cab5621d530c3b6"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:9289fd5dddcf57bab41d044f1756550f9e7cf0c8e373b8cdf0ce8773dc4bd417"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b493a043635eb376e50eedf7818f2f322eabbaa974e948bd8bdd29eb7ef2a51"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9fa2566ca27d67c86569e8c85297aaf413ffab85a8960500f12ea34ff98e4c41"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8e538f46104c815be19c975572d74afb53f29650ea2025bbfaef359d2de2f7f"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6fd30dc99682dc2c603c2b315bded2799019cea829f8bf57dc6b61efde6611c8"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2006769bd1640bdf4d5641c69a3d63b71b81445473cac5ded39740a226fa88ab"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:dc15e99b2d8a656f8e666854404f1ba54765871104e50c8e9813af8a7db07f12"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:ab2e5bef076f5a235c3774b4f4028a680432cded7cad37bba0fd90d64b187d19"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:4ec9dd88a5b71abfc74e9df5ebe7921c35cbb3b641181a531ca65cdb5e8e4dea"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-musllinux_1_2_s390x.whl", hash = "sha256:43193c5cda5d612f247172016c4bb71251c784d7a4d9314677186a838ad34858"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:aa693779a8b50cd97570e5a0f343538a8dbd3e496fa5dcb87e29406ad0299654"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-win32.whl", hash = "sha256:7706f5850360ac01d80c89bcef1640683cc12ed87f42579dab6c5d3ed6888613"},
+    {file = "charset_normalizer-3.4.0-cp38-cp38-win_amd64.whl", hash = "sha256:c3e446d253bd88f6377260d07c895816ebf33ffffd56c1c792b13bff9c3e1ade"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:980b4f289d1d90ca5efcf07958d3eb38ed9c0b7676bf2831a54d4f66f9c27dfa"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:f28f891ccd15c514a0981f3b9db9aa23d62fe1a99997512b0491d2ed323d229a"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a8aacce6e2e1edcb6ac625fb0f8c3a9570ccc7bfba1f63419b3769ccf6a00ed0"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bd7af3717683bea4c87acd8c0d3d5b44d56120b26fd3f8a692bdd2d5260c620a"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5ff2ed8194587faf56555927b3aa10e6fb69d931e33953943bc4f837dfee2242"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e91f541a85298cf35433bf66f3fab2a4a2cff05c127eeca4af174f6d497f0d4b"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:309a7de0a0ff3040acaebb35ec45d18db4b28232f21998851cfa709eeff49d62"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:285e96d9d53422efc0d7a17c60e59f37fbf3dfa942073f666db4ac71e8d726d0"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:5d447056e2ca60382d460a604b6302d8db69476fd2015c81e7c35417cfabe4cd"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:20587d20f557fe189b7947d8e7ec5afa110ccf72a3128d61a2a387c3313f46be"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:130272c698667a982a5d0e626851ceff662565379baf0ff2cc58067b81d4f11d"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-musllinux_1_2_s390x.whl", hash = "sha256:ab22fbd9765e6954bc0bcff24c25ff71dcbfdb185fcdaca49e81bac68fe724d3"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:7782afc9b6b42200f7362858f9e73b1f8316afb276d316336c0ec3bd73312742"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-win32.whl", hash = "sha256:2de62e8801ddfff069cd5c504ce3bc9672b23266597d4e4f50eda28846c322f2"},
+    {file = "charset_normalizer-3.4.0-cp39-cp39-win_amd64.whl", hash = "sha256:95c3c157765b031331dd4db3c775e58deaee050a3042fcad72cbc4189d7c8dca"},
+    {file = "charset_normalizer-3.4.0-py3-none-any.whl", hash = "sha256:fe9f97feb71aa9896b81973a7bbada8c49501dc73e58a10fcef6663af95e5079"},
+    {file = "charset_normalizer-3.4.0.tar.gz", hash = "sha256:223217c3d4f82c3ac5e29032b3f1c2eb0fb591b72161f86d93f5719079dae93e"},
+]
+
 [[package]]
 name = "cloudpickle"
 version = "3.1.0"
@@ -101,6 +237,17 @@ testing = ["dill (>=0.3.7)", "pytest (==7.1.3)", "scipy (>=1.7.3)"]
 torch = ["torch (>=1.0.0)"]
 toy-text = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"]
 
+[[package]]
+name = "h11"
+version = "0.14.0"
+description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"},
+    {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"},
+]
+
 [[package]]
 name = "identify"
 version = "2.6.1"
@@ -115,6 +262,20 @@ files = [
 [package.extras]
 license = ["ukkonen"]
 
+[[package]]
+name = "idna"
+version = "3.10"
+description = "Internationalized Domain Names in Applications (IDNA)"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3"},
+    {file = "idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9"},
+]
+
+[package.extras]
+all = ["flake8 (>=7.1.1)", "mypy (>=1.11.2)", "pytest (>=8.3.2)", "ruff (>=0.6.2)"]
+
 [[package]]
 name = "iniconfig"
 version = "2.0.0"
@@ -435,6 +596,47 @@ pluggy = ">=1.5,<2"
 [package.extras]
 dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
 
+[[package]]
+name = "python-engineio"
+version = "4.10.1"
+description = "Engine.IO server and client for Python"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "python_engineio-4.10.1-py3-none-any.whl", hash = "sha256:445a94004ec8034960ab99e7ce4209ec619c6e6b6a12aedcb05abeab924025c0"},
+    {file = "python_engineio-4.10.1.tar.gz", hash = "sha256:166cea8dd7429638c5c4e3a4895beae95196e860bc6f29ed0b9fe753d1ef2072"},
+]
+
+[package.dependencies]
+simple-websocket = ">=0.10.0"
+
+[package.extras]
+asyncio-client = ["aiohttp (>=3.4)"]
+client = ["requests (>=2.21.0)", "websocket-client (>=0.54.0)"]
+docs = ["sphinx"]
+
+[[package]]
+name = "python-socketio"
+version = "5.11.4"
+description = "Socket.IO server and client for Python"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "python_socketio-5.11.4-py3-none-any.whl", hash = "sha256:42efaa3e3e0b166fc72a527488a13caaac2cefc76174252486503bd496284945"},
+    {file = "python_socketio-5.11.4.tar.gz", hash = "sha256:8b0b8ff2964b2957c865835e936310190639c00310a47d77321a594d1665355e"},
+]
+
+[package.dependencies]
+bidict = ">=0.21.0"
+python-engineio = ">=4.8.0"
+requests = {version = ">=2.21.0", optional = true, markers = "extra == \"client\""}
+websocket-client = {version = ">=0.54.0", optional = true, markers = "extra == \"client\""}
+
+[package.extras]
+asyncio-client = ["aiohttp (>=3.4)"]
+client = ["requests (>=2.21.0)", "websocket-client (>=0.54.0)"]
+docs = ["sphinx"]
+
 [[package]]
 name = "pyyaml"
 version = "6.0.2"
@@ -497,6 +699,27 @@ files = [
     {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"},
 ]
 
+[[package]]
+name = "requests"
+version = "2.32.3"
+description = "Python HTTP for Humans."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6"},
+    {file = "requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760"},
+]
+
+[package.dependencies]
+certifi = ">=2017.4.17"
+charset-normalizer = ">=2,<4"
+idna = ">=2.5,<4"
+urllib3 = ">=1.21.1,<3"
+
+[package.extras]
+socks = ["PySocks (>=1.5.6,!=1.5.7)"]
+use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
+
 [[package]]
 name = "ruff"
 version = "0.6.9"
@@ -574,6 +797,24 @@ dev = ["cython-lint (>=0.12.2)", "doit (>=0.36.0)", "mypy (==1.10.0)", "pycodest
 doc = ["jupyterlite-pyodide-kernel", "jupyterlite-sphinx (>=0.13.1)", "jupytext", "matplotlib (>=3.5)", "myst-nb", "numpydoc", "pooch", "pydata-sphinx-theme (>=0.15.2)", "sphinx (>=5.0.0,<=7.3.7)", "sphinx-design (>=0.4.0)"]
 test = ["Cython", "array-api-strict (>=2.0)", "asv", "gmpy2", "hypothesis (>=6.30)", "meson", "mpmath", "ninja", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"]
 
+[[package]]
+name = "simple-websocket"
+version = "1.1.0"
+description = "Simple WebSocket server and client for Python"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "simple_websocket-1.1.0-py3-none-any.whl", hash = "sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c"},
+    {file = "simple_websocket-1.1.0.tar.gz", hash = "sha256:7939234e7aa067c534abdab3a9ed933ec9ce4691b0713c78acb195560aa52ae4"},
+]
+
+[package.dependencies]
+wsproto = "*"
+
+[package.extras]
+dev = ["flake8", "pytest", "pytest-cov", "tox"]
+docs = ["sphinx"]
+
 [[package]]
 name = "types-setuptools"
 version = "75.1.0.20241014"
@@ -596,6 +837,23 @@ files = [
     {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"},
 ]
 
+[[package]]
+name = "urllib3"
+version = "2.2.3"
+description = "HTTP library with thread-safe connection pooling, file post, and more."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "urllib3-2.2.3-py3-none-any.whl", hash = "sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac"},
+    {file = "urllib3-2.2.3.tar.gz", hash = "sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9"},
+]
+
+[package.extras]
+brotli = ["brotli (>=1.0.9)", "brotlicffi (>=0.8.0)"]
+h2 = ["h2 (>=4,<5)"]
+socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"]
+zstd = ["zstandard (>=0.18.0)"]
+
 [[package]]
 name = "virtualenv"
 version = "20.26.6"
@@ -616,7 +874,37 @@ platformdirs = ">=3.9.1,<5"
 docs = ["furo (>=2023.7.26)", "proselint (>=0.13)", "sphinx (>=7.1.2,!=7.3)", "sphinx-argparse (>=0.4)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=23.6)"]
 test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=23.1)", "pytest (>=7.4)", "pytest-env (>=0.8.2)", "pytest-freezer (>=0.4.8)", "pytest-mock (>=3.11.1)", "pytest-randomly (>=3.12)", "pytest-timeout (>=2.1)", "setuptools (>=68)", "time-machine (>=2.10)"]
 
+[[package]]
+name = "websocket-client"
+version = "1.8.0"
+description = "WebSocket client for Python with low level API options"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "websocket_client-1.8.0-py3-none-any.whl", hash = "sha256:17b44cc997f5c498e809b22cdf2d9c7a9e71c02c8cc2b6c56e7c2d1239bfa526"},
+    {file = "websocket_client-1.8.0.tar.gz", hash = "sha256:3239df9f44da632f96012472805d40a23281a991027ce11d2f45a6f24ac4c3da"},
+]
+
+[package.extras]
+docs = ["Sphinx (>=6.0)", "myst-parser (>=2.0.0)", "sphinx-rtd-theme (>=1.1.0)"]
+optional = ["python-socks", "wsaccel"]
+test = ["websockets"]
+
+[[package]]
+name = "wsproto"
+version = "1.2.0"
+description = "WebSockets state-machine based protocol implementation"
+optional = false
+python-versions = ">=3.7.0"
+files = [
+    {file = "wsproto-1.2.0-py3-none-any.whl", hash = "sha256:b9acddd652b585d75b20477888c56642fdade28bdfd3579aa24a4d2c037dd736"},
+    {file = "wsproto-1.2.0.tar.gz", hash = "sha256:ad565f26ecb92588a3e43bc3d96164de84cd9902482b130d0ddbaa9664a85065"},
+]
+
+[package.dependencies]
+h11 = ">=0.9.0,<1"
+
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "813d3c0fbf8ef6d7f719ee930eb1b91d9120bf217dbe9385d7706a18ac8612e2"
+content-hash = "4f6ff75639d4bc40abb8e2d3cafb70e4d05e8692b241224d0b37c6dcf21b8be6"
diff --git a/pyproject.toml b/pyproject.toml
index ab34eae..6a7fc30 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,6 +14,7 @@ pettingzoo = "^1.24.3"
 gymnasium = "^1.0.0"
 pygame = "^2.6.0"
 scipy = "^1.14.1"
+python-socketio = {extras = ["client"], version = "^5.11.4"}
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.11.2"


Test patch:
diff --git a/tests/test_game.py b/tests/test_game.py
index 1bc0d23..8ee50d3 100644
--- a/tests/test_game.py
+++ b/tests/test_game.py
@@ -28,7 +28,7 @@ def test_grid_creation():
         assert game.grid_dims == (4, 4)
 
         # mountain and city should be disjoint
-        assert np.logical_and(game.channels.mountain, game.channels.city).sum() == 0
+        assert np.logical_and(game.channels.mountains, game.channels.cities).sum() == 0
 
         owners = ["neutral"] + game.agents
         # for every pair of agents, the ownership channels should be disjoint
@@ -58,831 +58,799 @@ def test_channel_to_indices():
 
     channel = np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]])
     reference = np.array([(0, 0), (0, 2), (1, 1), (2, 0), (2, 2)])
-    indices = game.channel_to_indices(channel)
+    indices = game.channels.channel_to_indices(channel)
     assert (indices == reference).all()
 
     channel = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
     reference = np.empty((0, 2))
-    indices = game.channel_to_indices(channel)
+    indices = game.channels.channel_to_indices(channel)
     assert (indices == reference).all()
 
 
-def test_visibility_channel():
-    """
-    For given ownership mask, we should get visibility mask.
-    """
-    dummy_game = get_game()
-
-    ownership = np.array([[0, 0, 0], [0, 1, 0], [0, 0, 0]])
-    reference = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
-    visibility = dummy_game.visibility_channel(ownership)
-    assert (visibility == reference).all()
-
-    ownership = np.array(
-        [
-            [0, 0, 0, 0, 0],
-            [1, 1, 0, 0, 0],
-            [0, 1, 0, 0, 0],
-            [0, 0, 0, 0, 1],
-            [0, 0, 0, 0, 1],
-        ]
-    )
-    reference = np.array(
-        [
-            [1, 1, 1, 0, 0],
-            [1, 1, 1, 0, 0],
-            [1, 1, 1, 1, 1],
-            [1, 1, 1, 1, 1],
-            [0, 0, 0, 1, 1],
-        ]
-    )
-    visibility = dummy_game.visibility_channel(ownership)
-    assert (visibility == reference).all()
-
-
-def test_action_mask():
-    """
-    For given ownership mask and passable mask, we should get NxNx4 mask of valid actions.
-    """
-    game = get_game()
-    game.grid_dims = (4, 4)
-    game.channels.army = np.array(
-        [
-            [3, 0, 1, 0],
-            [0, 3, 6, 2],
-            [1, 1, 5, 0],
-            [2, 0, 5, 8],
-        ],
-        dtype=int,
-    )
-    game.channels._set_passable(
-        np.array(
-            [
-                [1, 1, 1, 1],
-                [1, 1, 0, 0],
-                [1, 1, 1, 0],
-                [1, 0, 0, 0],
-            ],
-            dtype=bool,
-        )
-    )
-
-    game.channels.ownership["red"] = np.array(
-        [
-            [0, 0, 1, 0],
-            [0, 1, 0, 0],
-            [0, 1, 1, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=bool,
-    )
-    reference = np.array(
-        [
-            # UP
-            [
-                [0, 0, 0, 0],
-                [0, 1, 0, 0],
-                [0, 0, 0, 0],
-                [1, 0, 0, 0],
-            ],
-            # DOWN
-            [
-                [0, 0, 0, 0],
-                [0, 1, 0, 0],
-                [0, 0, 0, 0],
-                [0, 0, 0, 0],
-            ],
-            # LEFT
-            [
-                [0, 0, 0, 0],
-                [0, 1, 0, 0],
-                [0, 0, 1, 0],
-                [0, 0, 0, 0],
-            ],
-            # RIGHT
-            [
-                [0, 0, 0, 0],
-                [0, 0, 0, 0],
-                [0, 0, 0, 0],
-                [0, 0, 0, 0],
-            ],
-        ],
-        dtype=bool,
-    )
-
-    action_mask = game.action_mask("red")
-    assert (action_mask[:, :, 0] == reference[0]).all()
-    assert (action_mask[:, :, 1] == reference[1]).all()
-    assert (action_mask[:, :, 2] == reference[2]).all()
-    assert (action_mask[:, :, 3] == reference[3]).all()
-
-
-def test_observations():
-    """
-    For given actions, we should get new state of the game.
-    """
-    map = """...#
-#..A
-#..#
-.#.B
-"""
-    grid = Grid(map)
-    game = get_game(grid)
-    game.channels.ownership["red"] = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    game.channels.ownership["blue"] = np.array(
-        [
-            [1, 0, 0, 0],
-            [0, 1, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    game.channels.army = np.array(
-        [
-            [3, 0, 0, 0],
-            [0, 3, 6, 2],
-            [1, 9, 5, 0],
-            [0, 0, 5, 8],
-        ],
-        dtype=np.float32,
-    )
-    game.channels.ownership["neutral"] = np.array(
-        [
-            [0, 1, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-
-    ############
-    # TEST RED #
-    ############
-    red_observation = game.agent_observation("red")["observation"]
-    reference_opponent_ownership = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (red_observation["opponent_cells"] == reference_opponent_ownership).all()
-
-    reference_neutral_ownership = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (red_observation["neutral_cells"] == reference_neutral_ownership).all()
-
-    reference_ownership = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (red_observation["owned_cells"] == reference_ownership).all()
-
-    # union of all ownerships should be zero
-    assert (
-        np.logical_and.reduce(
-            [
-                red_observation["opponent_cells"],
-                red_observation["neutral_cells"],
-                red_observation["owned_cells"],
-            ]
-        )
-    ).sum() == 0
-
-    reference_army = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 3, 6, 2],
-            [1, 9, 5, 0],
-            [0, 0, 5, 8],
-        ],
-        dtype=np.float32,
-    )
-    assert (red_observation["army"] == reference_army).all()
-
-    #############
-    # TEST BLUE #
-    #############
-    blue_observation = game.agent_observation("blue")["observation"]
-    reference_opponent_ownership = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (blue_observation["opponent_cells"] == reference_opponent_ownership).all()
-
-    reference_neutral_ownership = np.array(
-        [
-            [0, 1, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (blue_observation["neutral_cells"] == reference_neutral_ownership).all()
-
-    reference_ownership = np.array(
-        [
-            [1, 0, 0, 0],
-            [0, 1, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (blue_observation["owned_cells"] == reference_ownership).all()
-
-    reference_army = np.array(
-        [
-            [3, 0, 0, 0],
-            [0, 3, 6, 2],
-            [1, 9, 5, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (blue_observation["army"] == reference_army).all()
-
-    # union of all ownerships should be zero
-    assert (
-        np.logical_and.reduce(
-            [
-                blue_observation["opponent_cells"],
-                blue_observation["neutral_cells"],
-                blue_observation["owned_cells"],
-            ]
-        )
-    ).sum() == 0
-
-
-def test_game_step():
-    """
-    Test a number of moves from this situation
-    """
-    map = """...#
-#..A
-#..#
-.#.B
-"""
-    grid = Grid(map)
-    game = get_game(grid)
-    game.channels.ownership["red"] = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=int,
-    )
-    game.channels.ownership["blue"] = np.array(
-        [
-            [1, 0, 0, 0],
-            [0, 1, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    game.channels.army = np.array(
-        [
-            [3, 0, 0, 0],
-            [0, 3, 6, 2],
-            [1, 9, 5, 0],
-            [0, 0, 5, 8],
-        ],
-        dtype=np.float32,
-    )
-    game.channels.ownership["neutral"] = np.array(
-        [
-            [0, 1, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-
-    # Test capturing (if equal armies meet, the defender keeps the cell)
-    #############################################################################################################
-    # red moves from (2, 1) UP (captures blue square), blue moves from (1, 2) DOWN, (doesnt capture red square) #
-    #############################################################################################################
-    red_move = {
-        "pass": 0,
-        "cell": np.array([2, 1]),
-        "direction": 0,
-        "split": 0,
-    }
-    blue_move = {
-        "pass": 0,
-        "cell": np.array([1, 2]),
-        "direction": 1,
-        "split": 0,
-    }
-    moves = {
-        "red": red_move,
-        "blue": blue_move,
-    }
-    game.step(moves)
-    reference_army = np.array(
-        [
-            [3, 0, 0, 0],
-            [0, 5, 1, 2],
-            [1, 1, 0, 0],
-            [0, 0, 5, 8],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_ownership_red = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["red"] == reference_ownership_red).all()
-
-    reference_ownership_blue = np.array(
-        [
-            [1, 0, 0, 0],
-            [0, 0, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
-
-    reference_ownership_neutral = np.array(
-        [
-            [0, 1, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
-
-    reference_total_army_red = 20
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 6
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 6
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 3
-    assert stats["blue"]["land"] == reference_total_army_land
-
-    # Test raising of warning on invalid move (moving from cell with only 1 army)
-    ##################################################################################
-    # Now red moves from (2, 1) DOWN (invalid move) and blue moves from (0, 0) RIGHT #
-    ##################################################################################
-    red_move = {
-        "pass": 0,
-        "cell": np.array([2, 1]),
-        "direction": 1,
-        "split": 0,
-    }
-    blue_move = {
-        "pass": 0,
-        "cell": np.array([0, 0]),
-        "direction": 3,
-        "split": 0,
-    }
-    moves = {
-        "red": red_move,
-        "blue": blue_move,
-    }
-
-    with pytest.warns(UserWarning):  # we expect a warning
-        game.step(moves)
-
-    # this is second move, so army increments in base
-    reference_army = np.array(
-        [
-            [1, 2, 0, 0],
-            [0, 5, 1, 3],
-            [1, 1, 0, 0],
-            [0, 0, 5, 9],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_ownership_red = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["red"] == reference_ownership_red).all()
-
-    reference_ownership_blue = np.array(
-        [
-            [1, 1, 0, 0],
-            [0, 0, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
-
-    reference_ownership_neutral = np.array(
-        [
-            [0, 0, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
-
-    reference_total_army_red = 21
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 7
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 6
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 4
-    assert stats["blue"]["land"] == reference_total_army_land
-
-    # Test splitting of army
-    #####################################################################################
-    # Red sends half army from (3, 3) LEFT and blue sends half army from (1, 3) LEFT    #
-    #####################################################################################
-    red_move = {
-        "pass": 0,
-        "cell": np.array([3, 3]),
-        "direction": 2,
-        "split": 1,
-    }
-    blue_move = {
-        "pass": 0,
-        "cell": np.array([1, 3]),
-        "direction": 2,
-        "split": 1,
-    }
-    moves = {
-        "red": red_move,
-        "blue": blue_move,
-    }
-    game.step(moves)
-    reference_army = np.array(
-        [
-            [1, 2, 0, 0],
-            [0, 5, 2, 2],
-            [1, 1, 0, 0],
-            [0, 0, 9, 5],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_ownership_red = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["red"] == reference_ownership_red).all()
-
-    reference_ownership_blue = np.array(
-        [
-            [1, 1, 0, 0],
-            [0, 0, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
-
-    reference_ownership_neutral = np.array(
-        [
-            [0, 0, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
-
-    reference_total_army_red = 21
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 7
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 6
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 4
-    assert stats["blue"]["land"] == reference_total_army_land
-
-    # Test passing a move
-    ##############################################################
-    # Red moves army from (3, 2) UP and blue is passing the move #
-    ##############################################################
-    red_move = {
-        "pass": 0,
-        "cell": np.array([3, 2]),
-        "direction": 0,
-        "split": 0,
-    }
-    blue_move = {
-        "pass": 1,
-        "cell": np.array([1, 3]),
-        "direction": 2,
-        "split": 0,
-    }
-    moves = {
-        "red": red_move,
-        "blue": blue_move,
-    }
-
-    game.step(moves)
-    reference_army = np.array(
-        [
-            [1, 2, 0, 0],
-            [0, 5, 2, 3],
-            [1, 1, 8, 0],
-            [0, 0, 1, 6],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_ownership_red = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 0, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["red"] == reference_ownership_red).all()
-
-    reference_ownership_blue = np.array(
-        [
-            [1, 1, 0, 0],
-            [0, 0, 1, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
-
-    reference_ownership_neutral = np.array(
-        [
-            [0, 0, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
-
-    reference_total_army_red = 22
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 8
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 6
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 4
-    assert stats["blue"]["land"] == reference_total_army_land
-
-    # Test order of moves (smaller army has priority)
-    #############################################################
-    # Red moves from (2, 2) UP and blue moves from (1, 2) RIGHT #
-    #############################################################
-    red_move = {
-        "pass": 0,
-        "cell": np.array([2, 2]),
-        "direction": 0,
-        "split": 0,
-    }
-    blue_move = {
-        "pass": 0,
-        "cell": np.array([1, 2]),
-        "direction": 3,
-        "split": 0,
-    }
-    moves = {
-        "red": red_move,
-        "blue": blue_move,
-    }
-
-    game.step(moves)
-    reference_army = np.array(
-        [
-            [1, 2, 0, 0],
-            [0, 5, 6, 4],
-            [1, 1, 1, 0],
-            [0, 0, 1, 6],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_ownership_red = np.array(
-        [
-            [0, 0, 0, 0],
-            [0, 1, 1, 0],
-            [1, 1, 1, 0],
-            [0, 0, 1, 1],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["red"] == reference_ownership_red).all()
-
-    reference_ownership_blue = np.array(
-        [
-            [1, 1, 0, 0],
-            [0, 0, 0, 1],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
-
-    reference_ownership_neutral = np.array(
-        [
-            [0, 0, 1, 0],
-            [0, 0, 0, 0],
-            [0, 0, 0, 0],
-            [1, 0, 0, 0],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
-
-    reference_total_army_red = 21
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 7
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 7
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 3
-    assert stats["blue"]["land"] == reference_total_army_land
-
-    ##############################
-    # Test global army increment #
-    ##############################
-    game.time = 50
-    game._global_game_update()
-    reference_army = np.array(
-        [
-            [2, 3, 0, 0],
-            [0, 6, 7, 6],
-            [2, 2, 2, 0],
-            [0, 0, 2, 8],
-        ],
-        dtype=np.float32,
-    )
-    assert (game.channels.army == reference_army).all()
-
-    reference_total_army_red = 29
-    stats = game.get_infos()
-    assert stats["red"]["army"] == reference_total_army_red
-
-    reference_total_army_blue = 11
-    assert stats["blue"]["army"] == reference_total_army_blue
-
-    reference_total_army_land = 7
-    assert stats["red"]["land"] == reference_total_army_land
-
-    reference_total_army_land = 3
-    assert stats["blue"]["land"] == reference_total_army_land
-
-
-# def test_end_of_game():
+# def test_action_mask():
+#     """
+#     For given ownership mask and passable mask, we should get NxNx4 mask of valid actions.
+#     """
+#     game = get_game()
+#     game.grid_dims = (4, 4)
+#     game.channels.army = np.array(
+#         [
+#             [3, 0, 1, 0],
+#             [0, 3, 6, 2],
+#             [1, 1, 5, 0],
+#             [2, 0, 5, 8],
+#         ],
+#         dtype=int,
+#     )
+#     game.channels.passable = (
+#         np.array(
+#             [
+#                 [1, 1, 1, 1],
+#                 [1, 1, 0, 0],
+#                 [1, 1, 1, 0],
+#                 [1, 0, 0, 0],
+#             ],
+#             dtype=bool,
+#         )
+#     )
+#
+#     game.channels.ownership["red"] = np.array(
+#         [
+#             [0, 0, 1, 0],
+#             [0, 1, 0, 0],
+#             [0, 1, 1, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=bool,
+#     )
+#     reference = np.array(
+#         [
+#             # UP
+#             [
+#                 [0, 0, 0, 0],
+#                 [0, 1, 0, 0],
+#                 [0, 0, 0, 0],
+#                 [1, 0, 0, 0],
+#             ],
+#             # DOWN
+#             [
+#                 [0, 0, 0, 0],
+#                 [0, 1, 0, 0],
+#                 [0, 0, 0, 0],
+#                 [0, 0, 0, 0],
+#             ],
+#             # LEFT
+#             [
+#                 [0, 0, 0, 0],
+#                 [0, 1, 0, 0],
+#                 [0, 0, 1, 0],
+#                 [0, 0, 0, 0],
+#             ],
+#             # RIGHT
+#             [
+#                 [0, 0, 0, 0],
+#                 [0, 0, 0, 0],
+#                 [0, 0, 0, 0],
+#                 [0, 0, 0, 0],
+#             ],
+#         ],
+#         dtype=bool,
+#     )
+#
+#     obs = game.agent_observation("red").as_dict()
+#     action_mask = obs["action_mask"]
+#     assert (action_mask[:, :, 0] == reference[0]).all()
+#     assert (action_mask[:, :, 1] == reference[1]).all()
+#     assert (action_mask[:, :, 2] == reference[2]).all()
+#     assert (action_mask[:, :, 3] == reference[3]).all()
+#
+#
+# def test_observations():
+#     """
+#     For given actions, we should get new state of the game.
+#     """
 #     map = """...#
 # #..A
 # #..#
 # .#.B
 # """
-#     game = get_game(map)
-#     game.general_positions = {"red": [3, 3], "blue": [1, 3]}
+#     grid = Grid(map)
+#     game = get_game(grid)
 #     game.channels.ownership["red"] = np.array(
 #         [
 #             [0, 0, 0, 0],
-#             [0, 0, 1, 0],
+#             [0, 0, 0, 0],
 #             [1, 1, 1, 0],
 #             [0, 0, 1, 1],
 #         ],
 #         dtype=np.float32,
 #     )
-#
 #     game.channels.ownership["blue"] = np.array(
 #         [
+#             [1, 0, 0, 0],
+#             [0, 1, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     game.channels.army = np.array(
+#         [
+#             [3, 0, 0, 0],
+#             [0, 3, 6, 2],
+#             [1, 9, 5, 0],
+#             [0, 0, 5, 8],
+#         ],
+#         dtype=np.float32,
+#     )
+#     game.channels.ownership["neutral"] = np.array(
+#         [
+#             [0, 1, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#
+#     ############
+#     # TEST RED #
+#     ############
+#     red_observation = game.agent_observation("red").as_dict()["observation"]
+#     reference_opponent_ownership = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (red_observation["opponent_cells"] == reference_opponent_ownership).all()
+#
+#     reference_neutral_ownership = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (red_observation["neutral_cells"] == reference_neutral_ownership).all()
+#
+#     reference_ownership = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (red_observation["owned_cells"] == reference_ownership).all()
+#
+#     # union of all ownerships should be zero
+#     assert (
+#         np.logical_and.reduce(
+#             [
+#                 red_observation["opponent_cells"],
+#                 red_observation["neutral_cells"],
+#                 red_observation["owned_cells"],
+#             ]
+#         )
+#     ).sum() == 0
+#
+#     reference_army = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 3, 6, 2],
+#             [1, 9, 5, 0],
+#             [0, 0, 5, 8],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (red_observation["armies"] == reference_army).all()
+#
+#     #############
+#     # TEST BLUE #
+#     #############
+#     blue_observation = game.agent_observation("blue")["observation"]
+#     reference_opponent_ownership = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
 #             [1, 1, 1, 0],
-#             [0, 1, 0, 1],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (blue_observation["opponent_cells"] == reference_opponent_ownership).all()
+#
+#     reference_neutral_ownership = np.array(
+#         [
+#             [0, 1, 1, 0],
+#             [0, 0, 0, 0],
 #             [0, 0, 0, 0],
 #             [0, 0, 0, 0],
 #         ],
 #         dtype=np.float32,
 #     )
+#     assert (blue_observation["neutral_cells"] == reference_neutral_ownership).all()
+#
+#     reference_ownership = np.array(
+#         [
+#             [1, 0, 0, 0],
+#             [0, 1, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (blue_observation["owned_cells"] == reference_ownership).all()
+#
+#     reference_army = np.array(
+#         [
+#             [3, 0, 0, 0],
+#             [0, 3, 6, 2],
+#             [1, 9, 5, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (blue_observation["armies"] == reference_army).all()
+#
+#     # union of all ownerships should be zero
+#     assert (
+#         np.logical_and.reduce(
+#             [
+#                 blue_observation["opponent_cells"],
+#                 blue_observation["neutral_cells"],
+#                 blue_observation["owned_cells"],
+#             ]
+#         )
+#     ).sum() == 0
 #
+#
+# def test_game_step():
+#     """
+#     Test a number of moves from this situation
+#     """
+#     map = """...#
+# #..A
+# #..#
+# .#.B
+# """
+#     grid = Grid(map)
+#     game = get_game(grid)
+#     game.channels.ownership["red"] = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=int,
+#     )
+#     game.channels.ownership["blue"] = np.array(
+#         [
+#             [1, 0, 0, 0],
+#             [0, 1, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
 #     game.channels.army = np.array(
 #         [
-#             [3, 2, 2, 0],
+#             [3, 0, 0, 0],
 #             [0, 3, 6, 2],
 #             [1, 9, 5, 0],
 #             [0, 0, 5, 8],
 #         ],
 #         dtype=np.float32,
 #     )
+#     game.channels.ownership["neutral"] = np.array(
+#         [
+#             [0, 1, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#
+#     # Test capturing (if equal armies meet, the defender keeps the cell)
+#     #############################################################################################################
+#     # red moves from (2, 1) UP (captures blue square), blue moves from (1, 2) DOWN, (doesnt capture red square) #
+#     #############################################################################################################
+#     red_move = {
+#         "pass": 0,
+#         "cell": np.array([2, 1]),
+#         "direction": 0,
+#         "split": 0,
+#     }
+#     blue_move = {
+#         "pass": 0,
+#         "cell": np.array([1, 2]),
+#         "direction": 1,
+#         "split": 0,
+#     }
+#     moves = {
+#         "red": red_move,
+#         "blue": blue_move,
+#     }
+#     game.step(moves)
+#     reference_army = np.array(
+#         [
+#             [3, 0, 0, 0],
+#             [0, 5, 1, 2],
+#             [1, 1, 0, 0],
+#             [0, 0, 5, 8],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.armies == reference_army).all()
+#
+#     reference_ownership_red = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["red"] == reference_ownership_red).all()
+#
+#     reference_ownership_blue = np.array(
+#         [
+#             [1, 0, 0, 0],
+#             [0, 0, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
+#
+#     reference_ownership_neutral = np.array(
+#         [
+#             [0, 1, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
+#
+#     reference_total_army_red = 20
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 6
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 6
+#     assert stats["red"]["land"] == reference_total_army_land
+#
+#     reference_total_army_land = 3
+#     assert stats["blue"]["land"] == reference_total_army_land
+#
+#     # Test raising of warning on invalid move (moving from cell with only 1 army)
+#     ##################################################################################
+#     # Now red moves from (2, 1) DOWN (invalid move) and blue moves from (0, 0) RIGHT #
+#     ##################################################################################
+#     red_move = {
+#         "pass": 0,
+#         "cell": np.array([2, 1]),
+#         "direction": 1,
+#         "split": 0,
+#     }
+#     blue_move = {
+#         "pass": 0,
+#         "cell": np.array([0, 0]),
+#         "direction": 3,
+#         "split": 0,
+#     }
+#     moves = {
+#         "red": red_move,
+#         "blue": blue_move,
+#     }
+#
+#     with pytest.warns(UserWarning):  # we expect a warning
+#         game.step(moves)
+#
+#     # this is second move, so army increments in base
+#     reference_army = np.array(
+#         [
+#             [1, 2, 0, 0],
+#             [0, 5, 1, 3],
+#             [1, 1, 0, 0],
+#             [0, 0, 5, 9],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.army == reference_army).all()
+#
+#     reference_ownership_red = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["red"] == reference_ownership_red).all()
+#
+#     reference_ownership_blue = np.array(
+#         [
+#             [1, 1, 0, 0],
+#             [0, 0, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
+#
+#     reference_ownership_neutral = np.array(
+#         [
+#             [0, 0, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
+#
+#     reference_total_army_red = 21
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 7
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 6
+#     assert stats["red"]["land"] == reference_total_army_land
+#
+#     reference_total_army_land = 4
+#     assert stats["blue"]["land"] == reference_total_army_land
+#
+#     # Test splitting of army
+#     #####################################################################################
+#     # Red sends half army from (3, 3) LEFT and blue sends half army from (1, 3) LEFT    #
+#     #####################################################################################
+#     red_move = {
+#         "pass": 0,
+#         "cell": np.array([3, 3]),
+#         "direction": 2,
+#         "split": 1,
+#     }
+#     blue_move = {
+#         "pass": 0,
+#         "cell": np.array([1, 3]),
+#         "direction": 2,
+#         "split": 1,
+#     }
+#     moves = {
+#         "red": red_move,
+#         "blue": blue_move,
+#     }
+#     game.step(moves)
+#     reference_army = np.array(
+#         [
+#             [1, 2, 0, 0],
+#             [0, 5, 2, 2],
+#             [1, 1, 0, 0],
+#             [0, 0, 9, 5],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.army == reference_army).all()
+#
+#     reference_ownership_red = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["red"] == reference_ownership_red).all()
 #
-#     game.channels.ownership_neutral = np.array(
+#     reference_ownership_blue = np.array(
 #         [
+#             [1, 1, 0, 0],
+#             [0, 0, 1, 1],
 #             [0, 0, 0, 0],
 #             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
+#
+#     reference_ownership_neutral = np.array(
+#         [
+#             [0, 0, 1, 0],
+#             [0, 0, 0, 0],
 #             [0, 0, 0, 0],
 #             [1, 0, 0, 0],
 #         ],
 #         dtype=np.float32,
 #     )
+#     assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
+#
+#     reference_total_army_red = 21
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 7
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 6
+#     assert stats["red"]["land"] == reference_total_army_land
+#
+#     reference_total_army_land = 4
+#     assert stats["blue"]["land"] == reference_total_army_land
+#
+#     # Test passing a move
+#     ##############################################################
+#     # Red moves army from (3, 2) UP and blue is passing the move #
+#     ##############################################################
+#     red_move = {
+#         "pass": 0,
+#         "cell": np.array([3, 2]),
+#         "direction": 0,
+#         "split": 0,
+#     }
+#     blue_move = {
+#         "pass": 1,
+#         "cell": np.array([1, 3]),
+#         "direction": 2,
+#         "split": 0,
+#     }
+#     moves = {
+#         "red": red_move,
+#         "blue": blue_move,
+#     }
 #
-#     moves = {"red": (0, np.array([3, 2]), 0, 0), "blue": (0, np.array([1, 3]), 2, 1)}
 #     game.step(moves)
+#     reference_army = np.array(
+#         [
+#             [1, 2, 0, 0],
+#             [0, 5, 2, 3],
+#             [1, 1, 8, 0],
+#             [0, 0, 1, 6],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.army == reference_army).all()
 #
-#     # Neither should win
-#     assert not game.agent_won("red")
-#     assert not game.agent_won("blue")
-#     assert not game.is_done()
+#     reference_ownership_red = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 0, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["red"] == reference_ownership_red).all()
+#
+#     reference_ownership_blue = np.array(
+#         [
+#             [1, 1, 0, 0],
+#             [0, 0, 1, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
+#
+#     reference_ownership_neutral = np.array(
+#         [
+#             [0, 0, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
+#
+#     reference_total_army_red = 22
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 8
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 6
+#     assert stats["red"]["land"] == reference_total_army_land
+#
+#     reference_total_army_land = 4
+#     assert stats["blue"]["land"] == reference_total_army_land
+#
+#     # Test order of moves (smaller army has priority)
+#     #############################################################
+#     # Red moves from (2, 2) UP and blue moves from (1, 2) RIGHT #
+#     #############################################################
+#     red_move = {
+#         "pass": 0,
+#         "cell": np.array([2, 2]),
+#         "direction": 0,
+#         "split": 0,
+#     }
+#     blue_move = {
+#         "pass": 0,
+#         "cell": np.array([1, 2]),
+#         "direction": 3,
+#         "split": 0,
+#     }
+#     moves = {
+#         "red": red_move,
+#         "blue": blue_move,
+#     }
 #
-#     # Red moves to blues general, blue makes random move
-#     moves = {"red": (0, np.array([2, 2]), 0, 0), "blue": (0, np.array([1, 3]), 3, 0)}
 #     game.step(moves)
+#     reference_army = np.array(
+#         [
+#             [1, 2, 0, 0],
+#             [0, 5, 6, 4],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 6],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.army == reference_army).all()
+#
+#     reference_ownership_red = np.array(
+#         [
+#             [0, 0, 0, 0],
+#             [0, 1, 1, 0],
+#             [1, 1, 1, 0],
+#             [0, 0, 1, 1],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["red"] == reference_ownership_red).all()
+#
+#     reference_ownership_blue = np.array(
+#         [
+#             [1, 1, 0, 0],
+#             [0, 0, 0, 1],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership["blue"] == reference_ownership_blue).all()
+#
+#     reference_ownership_neutral = np.array(
+#         [
+#             [0, 0, 1, 0],
+#             [0, 0, 0, 0],
+#             [0, 0, 0, 0],
+#             [1, 0, 0, 0],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.ownership_neutral == reference_ownership_neutral).all()
+#
+#     reference_total_army_red = 21
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 7
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 7
+#     assert stats["red"]["land"] == reference_total_army_land
+#
+#     reference_total_army_land = 3
+#     assert stats["blue"]["land"] == reference_total_army_land
+#
+#     ##############################
+#     # Test global army increment #
+#     ##############################
+#     game.time = 50
+#     game._global_game_update()
+#     reference_army = np.array(
+#         [
+#             [2, 3, 0, 0],
+#             [0, 6, 7, 6],
+#             [2, 2, 2, 0],
+#             [0, 0, 2, 8],
+#         ],
+#         dtype=np.float32,
+#     )
+#     assert (game.channels.army == reference_army).all()
+#
+#     reference_total_army_red = 29
+#     stats = game.get_infos()
+#     assert stats["red"]["army"] == reference_total_army_red
+#
+#     reference_total_army_blue = 11
+#     assert stats["blue"]["army"] == reference_total_army_blue
+#
+#     reference_total_army_land = 7
+#     assert stats["red"]["land"] == reference_total_army_land
 #
-#     # Red should win
-#     assert game.agent_won("red")
+#     reference_total_army_land = 3
+#     assert stats["blue"]["land"] == reference_total_army_land
 #
-#     # Blue should be dead
-#     assert not game.agent_won("blue")
 #
-#     # Game should be done
-#     assert game.is_done()
+# # def test_end_of_game():
+# #     map = """...#
+# # #..A
+# # #..#
+# # .#.B
+# # """
+# #     game = get_game(map)
+# #     game.general_positions = {"red": [3, 3], "blue": [1, 3]}
+# #     game.channels.ownership["red"] = np.array(
+# #         [
+# #             [0, 0, 0, 0],
+# #             [0, 0, 1, 0],
+# #             [1, 1, 1, 0],
+# #             [0, 0, 1, 1],
+# #         ],
+# #         dtype=np.float32,
+# #     )
+# #
+# #     game.channels.ownership["blue"] = np.array(
+# #         [
+# #             [1, 1, 1, 0],
+# #             [0, 1, 0, 1],
+# #             [0, 0, 0, 0],
+# #             [0, 0, 0, 0],
+# #         ],
+# #         dtype=np.float32,
+# #     )
+# #
+# #     game.channels.army = np.array(
+# #         [
+# #             [3, 2, 2, 0],
+# #             [0, 3, 6, 2],
+# #             [1, 9, 5, 0],
+# #             [0, 0, 5, 8],
+# #         ],
+# #         dtype=np.float32,
+# #     )
+# #
+# #     game.channels.ownership_neutral = np.array(
+# #         [
+# #             [0, 0, 0, 0],
+# #             [0, 0, 0, 0],
+# #             [0, 0, 0, 0],
+# #             [1, 0, 0, 0],
+# #         ],
+# #         dtype=np.float32,
+# #     )
+# #
+# #     moves = {"red": (0, np.array([3, 2]), 0, 0), "blue": (0, np.array([1, 3]), 2, 1)}
+# #     game.step(moves)
+# #
+# #     # Neither should win
+# #     assert not game.agent_won("red")
+# #     assert not game.agent_won("blue")
+# #     assert not game.is_done()
+# #
+# #     # Red moves to blues general, blue makes random move
+# #     moves = {"red": (0, np.array([2, 2]), 0, 0), "blue": (0, np.array([1, 3]), 3, 0)}
+# #     game.step(moves)
+# #
+# #     # Red should win
+# #     assert game.agent_won("red")
+# #
+# #     # Blue should be dead
+# #     assert not game.agent_won("blue")
+# #
+# #     # Game should be done
+# #     assert game.is_done()
diff --git a/tests/test_gym.py b/tests/test_gym.py
index 43a760d..4ab811f 100644
--- a/tests/test_gym.py
+++ b/tests/test_gym.py
@@ -5,7 +5,7 @@ from generals.agents import AgentFactory
 
 
 def test_gym_runs():
-    npc = AgentFactory.make_agent("random")
+    npc = AgentFactory.make_agent("Random")
 
     env = gym.make(
         "gym-generals-v0",