Repository: jonbinney/deep_rabbit_hole
Pull request #158: Use Quoridor class for game logic
Author: login: jonbinney
Created: 2025-04-09T00:55:33Z | Merged: 2025-04-30T16:05:58Z

... in the petting zoo environment.

Good things:
- Tests pass! `quoridor_test.py` now tests both the QuoridorEnv and the Quoridor class at the same time.
- Using the classes and enums from quoridor.py should be a bit more readable than using numbers
- We now have the combined wall and player positions grid from the Quoridor class that might be a good representation to give to a neural net (@muralx maybe interesting to try with a conv net that can use the 2d structure?)

Meh things:
- Currently `profile_observe.py` takes about 4x as long as the old implementation on my computer. I think I can close that gap and make the new implementation faster in the next couple of days, but I didn't want to spend more time optimizing before getting this PR merged. For now we're using pretty much the same DFS code, which isn't yet optimized for the representation used in the `Board ` class.
- Some (all?) of the agents are broken. Especially the greedy agent since it made heavy use of the environment's API. I Started fixing this; will continue tomorrow. Unless @alejandromarcu has some time :-P . `QuoridorEnv._get_observation()` tries to provide the same arrays (walls, a grid of player positions, etc) so that the agents are easy to keep working.




Diff:
diff --git a/deep_quoridor/src/agents/core/trainable_agent.py b/deep_quoridor/src/agents/core/trainable_agent.py
index 20835609..02fd4acf 100644
--- a/deep_quoridor/src/agents/core/trainable_agent.py
+++ b/deep_quoridor/src/agents/core/trainable_agent.py
@@ -9,6 +9,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+from quoridor import ActionEncoder
 from utils import SubargsBase, my_device, resolve_path
 
 import wandb
@@ -108,6 +109,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
 
         # Setup device
         self.device = my_device()
@@ -155,23 +157,56 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        opponent_observation_before_action,
+        my_observation_after_opponent_action,
+        opponent_observation_after_action,
+        opponent_reward,
+        opponent_action,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation_after_action,
+        observation_after_action,
+        reward,
+        action,
+        done=False,
+    ):
         self.steps += 1
         if not self.training_mode:
             return
-        reward = self.handle_step_outcome_all(observation_before_action, action, game, self.player_id)
+        reward = self.handle_step_outcome_all(
+            observation_before_action,
+            opponent_observation_after_action,
+            observation_after_action,
+            reward,
+            action,
+            self.player_id,
+            done,
+        )
         self.current_episode_reward += reward
 
-    def handle_step_outcome_all(self, observation_before_action, action, game, player_id):
-        reward = self.adjust_reward(game.rewards[player_id], game)
+    def handle_step_outcome_all(
+        self,
+        observation_before_action,
+        opponent_observation_after_action,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done,
+    ):
+        # reward = self.adjust_reward(game.rewards[player_id], game)
+        reward = self.adjust_reward(reward, done)
 
         # Handle end of episode
         if action is None:
@@ -185,24 +220,22 @@ def handle_step_outcome_all(self, observation_before_action, action, game, playe
                     return reward
             return 0
 
-        state_before_action = self.observation_to_tensor(observation_before_action, player_id)
-        state_after_action = self.observation_to_tensor(game.observe(player_id), player_id)
-        opponent_state = game.observe(self.get_opponent_player_id(player_id))
+        state_before_action = self.observation_to_tensor(observation_before_action["observation"], agent_id)
+        state_after_action = self.observation_to_tensor(observation_after_action["observation"], agent_id)
         next_state_mask = None
         if self.params.mask_targetq:
             # next action mask is stored with the same rotation of the next state
             # if we want to mask actions on next state
             next_state_mask = self.convert_action_mask_to_tensor_for_player(
-                opponent_state["action_mask"], self.get_opponent_player_id(player_id)
+                opponent_observation_after_action["action_mask"], self.get_opponent_player_id(agent_id)
             )
-        done = game.is_done()
 
         # If next_state_mask is None, we just add a zero tensor. It is not really used anyway
         # Ideally for off policy training we could collect all moves and all the information
         # store it and use it for training without running "matches" every time.
         self.replay_buffer.add(
             state_before_action.cpu().numpy(),
-            self.convert_to_tensor_index_from_action(action, player_id),
+            self.convert_to_tensor_index_from_action(action, agent_id),
             reward,
             state_after_action.cpu().numpy()
             if state_after_action is not None
@@ -259,19 +292,13 @@ def update_target_network(self):
         """Copy parameters from online network to target network."""
         self.target_network.load_state_dict(self.online_network.state_dict())
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
         """Select an action using epsilon-greedy policy."""
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
-
-        mask = observation["action_mask"]
-
         if random.random() < self.epsilon:
-            valid_actions = self._get_valid_actions(mask)
+            valid_actions = self._get_valid_actions(action_mask)
             return self._select_random_action(valid_actions)
 
-        return self._get_best_action(game, observation, mask)
+        return self._get_best_action(observation, action_mask)
 
     def _get_valid_actions(self, mask):
         """Get valid actions from the action mask."""
@@ -296,7 +323,7 @@ def convert_to_action_from_tensor_index_for_player(self, action_index_in_tensor,
     def convert_to_tensor_index_from_action(self, action, player_id):
         return action
 
-    def _log_action(self, game, q_values):
+    def _log_action(self, q_values):
         if not self.action_log.is_enabled():
             return
 
@@ -305,13 +332,13 @@ def _log_action(self, game, q_values):
         # Log the 5 best actions, as long as the value is > -100 (arbitrary value)
         top_values, top_indices = torch.topk(q_values, min(5, len(q_values)))
         scores = {
-            game.action_index_to_params(int(self.convert_to_action_from_tensor_index(i.item()))): v.item()
+            self.action_encoder.index_to_action(int(self.convert_to_action_from_tensor_index(i.item()))): v.item()
             for v, i in zip(top_values, top_indices)
             if v.item() >= -100
         }
         self.action_log.action_score_ranking(scores)
 
-    def _get_best_action(self, game, observation, mask):
+    def _get_best_action(self, observation, mask):
         """Get the best action based on Q-values."""
         state = self.observation_to_tensor(observation, self.player_id)
         with torch.no_grad():
@@ -319,7 +346,7 @@ def _get_best_action(self, game, observation, mask):
 
         mask_tensor = self.convert_action_mask_to_tensor(mask)
         q_values = q_values * mask_tensor - 1e9 * (1 - mask_tensor)
-        self._log_action(game, q_values)
+        self._log_action(q_values)
 
         if self.training_mode and self.params.softmax_exploration:
             # Apply softmax to the Q-values to get action probabilities
diff --git a/deep_quoridor/src/agents/dexp.py b/deep_quoridor/src/agents/dexp.py
index c91db5eb..05caabfc 100644
--- a/deep_quoridor/src/agents/dexp.py
+++ b/deep_quoridor/src/agents/dexp.py
@@ -127,18 +127,31 @@ def _create_network(self):
         """Create the neural network model."""
         return DExpNetwork(self.board_size, self.action_size, self.params.split, self.params.turn)
 
-    def handle_opponent_step_outcome(self, opponent_observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        self,
+        opponent_observation_before_action,
+        my_observation_after_opponent_action,
+        opponent_observation_after_action,
+        opponent_reward,
+        opponent_action,
+        done=False,
+    ):
         if not self.training_mode or not self.params.use_opponents_actions:
             return
 
         self.handle_step_outcome_all(
-            opponent_observation_before_action, action, game, self.get_opponent_player_id(self.player_id)
+            opponent_observation_before_action,
+            my_observation_after_opponent_action,
+            opponent_observation_after_action,
+            opponent_reward,
+            opponent_action,
+            self.get_opponent_player_id(self.agent_id),
+            done,
         )
 
     def observation_to_tensor(self, observation, obs_player_id):
         """Convert the observation dict to a flat tensor."""
-        obs = observation["observation"]
-        obs_player_turn = 1 if obs["my_turn"] else 0
+        obs_player_turn = 1 if observation["my_turn"] else 0
 
         should_rotate = False
         if self.params.rotate and not self.params.target_as_source_for_opponent:
@@ -148,16 +161,16 @@ def observation_to_tensor(self, observation, obs_player_id):
             # This ensures board always faces to the player that will act on it
             should_rotate = (obs_player_id == "player_1") ^ (not obs_player_turn)
 
-        board = rotation.rotate_board(obs["board"]) if should_rotate else obs["board"]
-        walls = rotation.rotate_walls(obs["walls"]) if should_rotate else obs["walls"]
+        board = rotation.rotate_board(observation["board"]) if should_rotate else observation["board"]
+        walls = rotation.rotate_walls(observation["walls"]) if should_rotate else observation["walls"]
 
         # Create position matrices for player and opponent
         player_board = (board == 1).astype(np.float32)
         opponent_board = (board == 2).astype(np.float32)
 
         # Get wall counts
-        player_walls = np.array([obs["my_walls_remaining"]])
-        opponent_walls = np.array([obs["opponent_walls_remaining"]])
+        player_walls = np.array([observation["my_walls_remaining"]])
+        opponent_walls = np.array([observation["opponent_walls_remaining"]])
 
         # Swap boards and walls if not player's turn. It means this is a target state
         # Target states are played by the opponents, so board and walls should be in the opponents POV
diff --git a/deep_quoridor/src/agents/greedy.py b/deep_quoridor/src/agents/greedy.py
index aa45bc50..24d7dfc9 100644
--- a/deep_quoridor/src/agents/greedy.py
+++ b/deep_quoridor/src/agents/greedy.py
@@ -4,6 +4,7 @@
 from typing import Optional, TypeAlias
 
 import numpy as np
+from quoridor import ActionEncoder, Quoridor, MoveAction, Player, WallAction, WallOrientation, construct_game_from_observation
 from utils import SubargsBase
 
 from agents.core import Agent
@@ -21,6 +22,10 @@ class GreedyAgent(Agent):
     def __init__(self, params=GreedyParams(), **kwargs):
         super().__init__()
         self.params = params
+        self.action_space = kwargs["action_space"]
+        self.board_size = kwargs["board_size"]
+        self.wall_size = self.board_size - 1
+        self.action_encoder = ActionEncoder(self.board_size)
 
     @classmethod
     def params_class(cls):
@@ -34,16 +39,16 @@ def name(self):
     def start_game(self, game, player_id):
         self.player_id = player_id
 
-    def _valid_pawn_movements(self, game, action_mask: np.ndarray) -> list[Position]:
+    def _valid_pawn_movements(self, action_mask: np.ndarray) -> list[Position]:
         """Returns the coordinates of the possible movements based on the action mask.
         This is useful for the first move, since there could be complicated jumps.  After the
         first move, you can't easily use this because you'd need to generate a new action_mask
         """
-        movement_mask = action_mask[: game.board_size**2]
+        movement_mask = action_mask[: self.board_size**2]
         valid_moves = np.argwhere(movement_mask == 1).reshape(-1)
-        return [game.action_index_to_params(i)[:2] for i in valid_moves]
+        return [self.action_encoder.index_to_action(i).destination for i in valid_moves]
 
-    def _next_moves(self, game, pos: Position) -> list[Position]:
+    def _next_moves(self, game: Quoridor, pos: Position) -> list[Position]:
         """Given a position in the board, it returns the possible next moves.  Notice that it ignores
         the other pawn, because this function is meant to be used for moves other than the first one,
         so the pawn may not be there anymore.
@@ -52,11 +57,11 @@ def _next_moves(self, game, pos: Position) -> list[Position]:
         for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
             row = pos[0] + dr
             col = pos[1] + dc
-            if game.is_in_board(row, col) and not game.is_wall_between(pos[0], pos[1], row, col):
+            if game.board.is_position_on_board((row, col)) and not game.board.is_wall_between(pos, (row, col)):
                 moves.append((row, col))
         return moves
 
-    def _shortest_path_from(self, game, pos: Position, target_row: int) -> list[Position]:
+    def _shortest_path_from(self, game: Quoridor, pos: Position, target_row: int) -> list[Position]:
         """
         Given a position in the board, it returns the shortest path to get to the target_row.
         """
@@ -64,7 +69,7 @@ def _shortest_path_from(self, game, pos: Position, target_row: int) -> list[Posi
             return []
 
         # How many steps are needed to get to that position, starting from pos (None means unreachable)
-        distances = [[None] * game.board_size for _ in range(game.board_size)]
+        distances = [[None] * game.board.board_size for _ in range(game.board.board_size)]
 
         # Queue for the BFS, containing the neighbords to process
         q = Queue()
@@ -102,10 +107,10 @@ def _shortest_path_from(self, game, pos: Position, target_row: int) -> list[Posi
 
         return path[::-1]
 
-    def _shortest_path_from_me(self, game, observation: dict, target_row: int) -> list[Position]:
+    def _shortest_path_from_me(self, game: Quoridor, action_mask: np.ndarray, target_row: int) -> list[Position]:
         """Return the shortest path from wherever the agent currently is."""
         # The first move is based on the action mask, that already calculated what positions are valid
-        first_moves = self._valid_pawn_movements(game, observation["action_mask"])
+        first_moves = self._valid_pawn_movements(action_mask)
         shortest_path = None
 
         # Find the shortest path for each of the first moves and keep the shortest one.
@@ -117,24 +122,24 @@ def _shortest_path_from_me(self, game, observation: dict, target_row: int) -> li
         assert shortest_path is not None, "No path found"
         return shortest_path
 
-    def _get_opponent_position(self, game, board: np.ndarray) -> Position:
+    def _get_opponent_position(self, game: Quoridor, board: np.ndarray) -> Position:
         coords = np.argwhere(board == 2)
         assert len(coords) == 1, "Expected exactly one opponent position"
         return (int(coords[0][0]), int(coords[0][1]))
 
-    def _get_block_action(self, game, opponent_shortest_path: list[Position], action_mask: np.ndarray):
+    def _get_block_action(self, game: Quoridor, opponent_shortest_path: list[Position], action_mask: np.ndarray):
         for p0, p1 in zip(opponent_shortest_path, opponent_shortest_path[1:]):
             # For this movement of the opponent, get the possible wall placements
             if p0[0] == p1[0]:
                 # Horizontal movement
                 rows = [p0[0], p0[0] - 1]
                 cols = [min(p0[1], p1[1])]
-                orientation = 1
+                orientation = WallOrientation.VERTICAL
             elif p0[1] == p1[1]:
                 # Vertical movement
                 rows = [min(p0[0], p1[0])]
                 cols = [p0[1], p0[1] - 1]
-                orientation = 2
+                orientation = WallOrientation.HORIZONTAL
             else:
                 assert "Expected horizontal or vertical movement"
 
@@ -143,17 +148,21 @@ def _get_block_action(self, game, opponent_shortest_path: list[Position], action
             # by looking into which makes the difference of the distances more favorable
             for r in rows:
                 for c in cols:
-                    if r < 0 or c < 0 or r >= game.wall_size or c >= game.wall_size:
+                    if r < 0 or c < 0 or r >= self.wall_size or c >= self.wall_size:
                         continue
 
-                    action = game.action_params_to_index(r, c, orientation)
+                    action = self.action_encoder.action_to_index(WallAction((r, c), orientation))
                     if action_mask[action] == 1:
                         return action
         # No actions found
         return None
 
     def _log_action(
-        self, game, observation: dict, my_shortest_path: list[Position], opponent_shortest_path: list[Position]
+        self,
+        observation: dict,
+        action_mask: np.ndarray,
+        my_shortest_path: list[Position],
+        opponent_shortest_path: list[Position],
     ):
         if not self.action_log.is_enabled():
             return
@@ -161,40 +170,40 @@ def _log_action(
         self.action_log.clear()
 
         # Log the possible next movements
-        movement_mask = observation["action_mask"][: game.board_size**2]
+        movement_mask = action_mask[: self.board_size**2]
         for action in np.argwhere(movement_mask == 1).reshape(-1):
-            self.action_log.action_text(game.action_index_to_params(int(action)), "")
+            # self.action_log.action_text(self.action_encoder.index_to_action(int(action)), "")
+            pass
 
-        my_coords = np.argwhere(observation["observation"]["board"] == 1)
+        my_coords = np.argwhere(observation["board"] == 1)
         path = my_shortest_path[:]
         path.insert(0, (int(my_coords[0][0]), int(my_coords[0][1])))
-        self.action_log.path(path)
+        # self.action_log.path(path)
 
-        opp_coords = np.argwhere(observation["observation"]["board"] == 2)
+        opp_coords = np.argwhere(observation["board"] == 2)
         path = opponent_shortest_path[:]
         path.insert(0, (int(opp_coords[0][0]), int(opp_coords[0][1])))
-        self.action_log.path(path)
+        # self.action_log.path(path)
 
-    def get_action(self, game):
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
+    def get_action(self, observation, action_mask):
+        # Reconstruct the game from the observation.
+        game, player, opponent = construct_game_from_observation(observation, self.player_id)
 
         if random.random() < self.params.p_random:
             if self.action_log.is_enabled():
                 self.action_log.clear()
                 # TO DO: when we have the functionality in the log to output a message,
                 # use it to say that the move will be random.
-            return game.action_space(game.agent_selection).sample(observation["action_mask"])
+            return self.action_space.sample(action_mask)
 
-        goal_row = game.board_size - 1 if self.player_id == "player_0" else 0
-        opponent_row = game.board_size - 1 - goal_row
+        goal_row = game.get_goal_row(player)
+        opponent_goal_row = game.get_goal_row(opponent)
 
-        my_shortest_path = self._shortest_path_from_me(game, observation, goal_row)
-        opponent_pos = self._get_opponent_position(game, observation["observation"]["board"])
-        opponent_shortest_path = self._shortest_path_from(game, opponent_pos, opponent_row)
+        my_shortest_path = self._shortest_path_from_me(game, action_mask, goal_row)
+        opponent_pos = game.board.get_player_position(opponent)
+        opponent_shortest_path = self._shortest_path_from(game, opponent_pos, opponent_goal_row)
 
-        self._log_action(game, observation, my_shortest_path, opponent_shortest_path)
+        self._log_action(observation, action_mask, my_shortest_path, opponent_shortest_path)
 
         # TODO: use a more elaborate logic to decide whether to block, could be probabilistic
         block = False
@@ -202,8 +211,8 @@ def get_action(self, game):
             if len(opponent_shortest_path) < 5:
                 block = True
         if block:
-            action = self._get_block_action(game, opponent_shortest_path, observation["action_mask"])
+            action = self._get_block_action(game, opponent_shortest_path, action_mask)
             if action is not None:
                 return action
 
-        return game.action_params_to_index(my_shortest_path[0][0], my_shortest_path[0][1], 0)
+        return self.action_encoder.action_to_index(MoveAction((my_shortest_path[0][0], my_shortest_path[0][1])))
diff --git a/deep_quoridor/src/agents/human.py b/deep_quoridor/src/agents/human.py
index 6d72b9bd..ceb08685 100644
--- a/deep_quoridor/src/agents/human.py
+++ b/deep_quoridor/src/agents/human.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
+from quoridor import ActionEncoder, MoveAction, WallAction, WallOrientation
 from utils import SubargsBase
 
 from agents.core import Agent
@@ -15,6 +16,7 @@ class HumanAgent(Agent):
     def __init__(self, params=HumanParams(), **kwargs):
         super().__init__()
         self.params = params
+        self.action_encoder = ActionEncoder(kwargs["board_size"])
 
     @classmethod
     def params_class(cls):
@@ -25,21 +27,22 @@ def name(self):
             return self.params.nick
         return "Human"
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
         # Import this here to avoid circular dependencies
         from renderers.pygame import PygameQuoridor
 
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
-
         # Transform the action mask into a set of valid moves
         valid_moves = set()
-        for action, value in enumerate(observation["action_mask"]):
+        for action, value in enumerate(action_mask):
             if value == 1:
-                r, c, type = game.action_index_to_params(action)
-                valid_moves.add((r, c, type))
+                valid_action = self.action_encoder.index_to_action(action)
+                if isinstance(valid_action, MoveAction):
+                    valid_moves.add((valid_action.destination[0], valid_action.destination[1], 0))
+                elif isinstance(valid_action, WallAction):
+                    valid_moves.add(
+                        (valid_action.position[0], valid_action.position[1], valid_action.orientation.value - 1)
+                    )
 
         result = PygameQuoridor.instance().get_human_input(valid_moves)
 
-        return None if result is None else game.action_params_to_index(*result)
+        return None if result is None else self.action_encoder.action_to_index(*result)
diff --git a/deep_quoridor/src/agents/random.py b/deep_quoridor/src/agents/random.py
index 339ccce1..dc21bb86 100644
--- a/deep_quoridor/src/agents/random.py
+++ b/deep_quoridor/src/agents/random.py
@@ -5,9 +5,7 @@ class RandomAgent(Agent):
     def __init__(self, **kwargs):
         super().__init__()
 
-    def get_action(self, game):
-        observation, _, termination, truncation, _ = game.last()
-        mask = observation["action_mask"]
-        if termination or truncation:
-            return None
-        return game.action_space(game.agent_selection).sample(mask)
+        self.action_space = kwargs["action_space"]
+
+    def get_action(self, observation, action_mask):
+        return self.action_space.sample(action_mask)
diff --git a/deep_quoridor/src/agents/sb3_ppo.py b/deep_quoridor/src/agents/sb3_ppo.py
index bb6aaf39..450166e9 100644
--- a/deep_quoridor/src/agents/sb3_ppo.py
+++ b/deep_quoridor/src/agents/sb3_ppo.py
@@ -155,7 +155,7 @@ def start_game(self, game, player_id):
                 print("No policy found. The agent will not work correctly.")
                 return
 
-    def get_action(self, game):
+    def get_action(self, _observation, _action_mask):
         """
         Get the action to take based on the current game state using the PPO model.
 
diff --git a/deep_quoridor/src/agents/simple.py b/deep_quoridor/src/agents/simple.py
index 6c852041..349ca91a 100644
--- a/deep_quoridor/src/agents/simple.py
+++ b/deep_quoridor/src/agents/simple.py
@@ -1,31 +1,36 @@
+import random
+
+from quoridor import ActionEncoder, Player, Quoridor, construct_game_from_observation
+from quoridor_env import StepRewardCalculator
+
 from agents.core import Agent
 
 
-def sample_random_action_sequence(game, max_path_length):
+def sample_random_action_sequence(game: Quoridor, player: Player, opponent: Player, max_path_length: int):
     """
     Sample a random sequence of actions for a given game. Stops early if the game terminates."""
-    agent_name = game.agent_selection
-
     action_sequence = []
     total_reward = 0.0
+    reward_calculator = StepRewardCalculator(game, player, opponent)
     while len(action_sequence) < max_path_length:
-        observation, reward, termination, truncation, _ = game.last()
-        mask = observation["action_mask"]
-
         # For now, assume the other agent takes random actions.
-        if game.agent_selection != agent_name:
-            action = game.action_space(game.agent_selection).sample(mask)
-            game.step(action)
-            continue
+        valid_actions = game.get_valid_actions()
+        assert len(valid_actions) > 0, "No valid actions available... this shouldn't be possible."
 
-        total_reward += reward
-        if termination or truncation:
-            break
+        action = random.choice(valid_actions)
+
+        if game.get_current_player() == player:
+            action_sequence.append(action)
+            reward_calculator.before_step()
 
-        action = game.action_space(game.agent_selection).sample(mask)
-        action_sequence.append(action)
         game.step(action)
 
+        if game.get_current_player() == player:
+            total_reward += reward_calculator.after_step()
+
+        if game.is_game_over():
+            break
+
     return action_sequence, total_reward
 
 
@@ -34,17 +39,23 @@ def __init__(self, sequence_length=3, num_sequences=10, **kwargs):
         super().__init__()
         self.sequence_length = sequence_length
         self.num_sequences = num_sequences
+        self.board_size = kwargs["board_size"]
+        self.action_encoder = ActionEncoder(self.board_size)
 
-    def get_action(self, game):
-        _, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
+    def start_game(self, game, player_id):
+        self.player_id = player_id
 
+    def get_action(self, observation, action_mask):
         possible_action_sequences = []
         for _ in range(self.num_sequences):
-            action_sequence, total_reward = sample_random_action_sequence(game.copy(), self.sequence_length)
+            game, player, opponent = construct_game_from_observation(observation, self.player_id)
+            action_sequence, total_reward = sample_random_action_sequence(game, player, opponent, self.sequence_length)
             possible_action_sequences.append((action_sequence, total_reward))
 
         # Choose the action sequence with the highest reward.
         best_sequence, _ = max(possible_action_sequences, key=lambda x: x[1])
-        return best_sequence[0]
+        game, player, opponent = construct_game_from_observation(observation, self.player_id)
+        assert game.is_action_valid(best_sequence[0]), "The best action is not valid."
+        action_id = self.action_encoder.action_to_index(best_sequence[0])
+        assert action_mask[action_id] == 1, "The action is not valid according to the action mask."
+        return action_id
diff --git a/deep_quoridor/src/arena.py b/deep_quoridor/src/arena.py
index 04f8e2f8..f6ef2df6 100644
--- a/deep_quoridor/src/arena.py
+++ b/deep_quoridor/src/arena.py
@@ -50,32 +50,58 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
-            observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+        for agent_id in self.game.agent_iter():
+            observation_before_action, _, termination, truncation, _ = self.game.last()
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observation_before_action,
+                        opponent_observation_after_action=self.game.observe(opponent_agent_id),
+                        observation_after_action=None,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation_before_action["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(
+                agent.get_action(observation_before_action["observation"], observation_before_action["action_mask"])
+            )
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observation_before_action,
+                    opponent_observation_after_action=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=self.game.rewards[agent_id],
+                    action=action,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    opponent_observation_before_action=observation_before_action,
+                    my_observation_after_opponent_action=self.game.observe(opponent_agent_id),
+                    opponent_observation_after_action=self.game.observe(agent_id),
+                    opponent_reward=self.game.rewards[opponent_agent_id],
+                    opponent_action=action,
+                    done=self.game.is_done(),
+                )
 
-            self.plugins.after_action(self.game, step, player_id, action)
+            self.plugins.after_action(self.game, step, agent_id, action)
             step += 1
             # TODO: Move max steps with proper truncation to the environment
             if step >= self.max_steps:
-                print(f"Game {game_id} reached max steps. Player {agent.name()} vs Player {opponent_agent.name()}.")
                 break
 
         end_time = time.time()
@@ -104,7 +130,14 @@ def _play_games(self, players: list[str | Agent], times: int, mode: PlayMode):
                 agents.append(p)
             else:
                 agents.append(
-                    AgentRegistry.create_from_encoded_name(p, board_size=self.board_size, max_walls=self.max_walls)
+                    AgentRegistry.create_from_encoded_name(
+                        p,
+                        board_size=self.board_size,
+                        max_walls=self.max_walls,
+                        action_space=self.game.action_space(
+                            None
+                        ),  # We might need to do something fancier with the action space if we add agent wrappers.
+                    )
                 )
 
         if mode == PlayMode.ALL_VS_ALL:
diff --git a/deep_quoridor/src/quoridor.py b/deep_quoridor/src/quoridor.py
index 127ecc2f..9d111e27 100644
--- a/deep_quoridor/src/quoridor.py
+++ b/deep_quoridor/src/quoridor.py
@@ -1,11 +1,9 @@
+import copy
 from dataclasses import dataclass
 from enum import IntEnum, unique
-from typing import TypeAlias
 
 import numpy as np
 
-Position: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.uint8]]  # (row, col)
-
 
 @unique
 class Player(IntEnum):
@@ -23,149 +21,199 @@ class Action:
     pass
 
 
-@dataclass
+@dataclass(frozen=True)  # Frozen to make it hashable.
 class MoveAction(Action):
-    def __init__(self, destination: Position):
-        assert is_valid_position_type(destination)
-        self.destination = destination
-
-    destination: Position
+    destination: tuple[int, int]  # Destination cell (row, col)
 
 
-@dataclass
+@dataclass(frozen=True)  # Frozen to make it hashable.
 class WallAction(Action):
-    def __init__(self, position: Position, orientation: WallOrientation):
-        assert is_valid_position_type(position)
-        assert isinstance(orientation, WallOrientation)
-        self.position = np.array((position))
-        self.orientation = orientation
-
-    position: Position  # Start of wall
+    position: tuple[int, int]  # Start position of the wall (row, col)
     orientation: WallOrientation
 
 
-def is_valid_position_type(position: Position) -> bool:
-    return isinstance(position, np.ndarray) and position.shape == (2,)
+class ActionEncoder:
+    def __init__(self, board_size: int):
+        self._board_size = board_size
+        self._wall_size = board_size - 1
 
+    def action_to_index(self, action) -> int:
+        """
+        Converts an action object to an action index
+        """
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self._board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self._board_size**2 + action.position[0] * self._wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self._board_size**2 + self._wall_size**2 + action.position[0] * self._wall_size + action.position[1]
+        else:
+            raise ValueError(f"Invalid action type: {action}")
 
-class Board:
-    """
-    Grid which has both walls and possible pawn positions.
+    def index_to_action(self, idx) -> Action:
+        """
+        Converts an action index to an action object.
+        """
+        action = None
+        if idx < self._board_size**2:  # Pawn movement
+            action = MoveAction(divmod(idx, self._board_size))
+        elif idx < self._board_size**2 + self._wall_size**2:
+            action = WallAction(divmod(idx - self._board_size**2, self._wall_size), WallOrientation.VERTICAL)
+        elif idx < self._board_size**2 + (self._wall_size**2) * 2:
+            action = WallAction(
+                divmod(idx - self._board_size**2 - self._wall_size**2, self._wall_size),
+                WallOrientation.HORIZONTAL,
+            )
+        else:
+            raise ValueError(f"Invalid action index: {idx}")
+
+        return action
 
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
 
+class Board:
     # Possible values for each cell in the grid.
-    OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
+
+    # When we check whether a new wall could potentially block all routes to the goal for a player,
+    # we first check to see whether it spans between two walls already on the board, or between
+    # one wall on the board and the border of the board. There are three possible places a new wall
+    # could touch an existing wall (or the border). At the start of the new wall, in the middle of the
+    # new wall, or at the end of the new wall. To simplify these checks, we create this list of wall
+    # grid cells that touch each of those three places on a new wall. THe positions here are relative
+    # to the starting position of the new wall.
+    _potential_wall_neighbors = {
+        WallOrientation.VERTICAL: [
+            np.array([(-1, -1), (-2, 0), (-1, 1)]),
+            np.array([(1, -1), (1, 1)]),
+            np.array([(3, -1), (4, 0), (3, 1)]),
+        ],
+        WallOrientation.HORIZONTAL: [
+            np.array([(-1, -1), (0, -2), (1, -1)]),
+            np.array([(-1, 1), (1, 1)]),
+            np.array([(-1, 3), (0, 4), (1, 3)]),
+        ],
+    }
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((self.board_size * 2 + 3, self.board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
-        self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
-        for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+        self._player_positions = [(0, self.board_size // 2), (self.board_size - 1, self.board_size // 2)]
         self._walls_remaining = [self.max_walls, self.max_walls]
+        for player, position in zip(self._players, self._player_positions):
+            self.set_player_cell(position, player)
 
-    def get_player_position(self, player: Player) -> Position:
+    def get_player_position(self, player: Player) -> tuple[int, int]:
         """
         Get the position of the player's pawn.
         """
-        assert isinstance(player, Player)
         return self._player_positions[player]
 
-    def move_player(self, player: Player, new_position: Position):
+    def move_player(self, player: Player, new_position: tuple[int, int]):
         """
         Move the player's pawn. Doesn't check if the move is valid according to game rules.
         """
-        assert isinstance(player, Player)
-        assert is_valid_position_type(new_position)
-
         if not self.is_position_on_board(new_position):
             raise ValueError(f"Position {new_position} is out of bounds")
 
         old_position = self._player_positions[player]
-        self._grid[*old_position * 2] = Board.FREE
-        self._grid[*new_position * 2] = player
+        self.set_player_cell(old_position, Board.FREE)
+        self.set_player_cell(new_position, player)
         self._player_positions[player] = new_position
 
-    def get_player_cell(self, position: Position) -> int:
+    def get_player_cell(self, position: tuple[int, int]) -> int:
         """
         Get the value of a "player cell" in the grid.
         """
-        assert is_valid_position_type(position)
+        assert self.is_position_on_board(position)
 
-        if not self.is_position_on_board(position):
-            raise ValueError(f"Position {position} is out of bounds")
+        return self._grid[position[0] * 2 + 2, position[1] * 2 + 2]
 
-        return self._grid[*position * 2]
+    def set_player_cell(self, position: tuple[int, int], value: int):
+        self._grid[position[0] * 2 + 2, position[1] * 2 + 2] = value
 
-    def add_wall(self, player: Player, position: Position, orientation: WallOrientation):
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        return self._walls_remaining[player]
+
+    def set_walls_remaining(self, player: Player, walls_remaining: int):
+        """
+        Set the number of walls remaining for the player.
+        """
+        self._walls_remaining[player] = walls_remaining
+
+    def add_wall(self, player: Player, position: tuple[int, int], orientation: WallOrientation, check_if_valid=True):
         """
         Mark the grid cells corresponding to the wall as occupied.
         """
-        assert isinstance(player, Player)
-        assert is_valid_position_type(position)
-        assert isinstance(orientation, WallOrientation)
+        assert (not check_if_valid) or self.can_place_wall(player, position, orientation)
 
-        if self._walls_remaining[player] < 1:
-            raise ValueError("Cannot place wall, no walls remaining for player {player}")
+        self._grid[self._get_wall_slice(position, orientation)] = Board.WALL
+        self._walls_remaining[player] -= 1
 
-        if self.can_place_wall(position, orientation):
-            self._grid[self._get_wall_slice(position, orientation)] = 10 + player
-            self._walls_remaining[player] -= 1
-        else:
-            raise ValueError("Cannot place wall at position {position} with orientation {orientation}")
+        # Update the old-style wall representation.
+        self._old_style_walls[*position, orientation] = 1
 
-    def can_place_wall(self, position: Position, orientation: WallOrientation) -> bool:
+    def remove_wall(self, player: Player, position: tuple[int, int], orientation: WallOrientation):
+        self._grid[self._get_wall_slice(position, orientation)] = Board.FREE
+        self._walls_remaining[player] += 1
+
+        # Update the old-style wall representation.
+        self._old_style_walls[*position, orientation] = 0
+
+    def can_place_wall(self, player: Player, position: tuple[int, int], orientation: WallOrientation) -> bool:
         """
         Returns True if the wall can be placed at the given position and orientation.
 
         Checks that the player has walls remaining, the wall is within the bounds of the board, and doesn't overlap with other walls.
         """
-        assert is_valid_position_type(position)
-        assert isinstance(orientation, WallOrientation)
-
-        try:
-            return (self._grid[self._get_wall_slice(position, orientation)] == Board.FREE).all()
-        except IndexError:
+        if self._walls_remaining[player] < 1:
             return False
 
-    def is_wall_between(self, position1: Position, position2: Position) -> bool:
+        return (self._grid[self._get_wall_slice(position, orientation)] == Board.FREE).all()
+
+    def is_wall_between(self, position1: tuple[int, int], position2: tuple[int, int]) -> bool:
         """
         Check if there is a wall between two positions.
         """
-        assert is_valid_position_type(position1)
-        assert is_valid_position_type(position2)
-        assert np.sum(np.abs(position1 - position2)) == 1, "Positions must be adjacent"
+        assert np.sum(np.abs(np.subtract(position1, position2))) == 1, "Positions must be adjacent"
 
         position1_on_board = self.is_position_on_board(position1)
         position2_on_board = self.is_position_on_board(position2)
         assert position1_on_board or position2_on_board, "At least one position must be on the board"
 
         if position1_on_board and position2_on_board:
-            wall_position = position1 + position2
-            return self._grid[*wall_position] == Board.PLAYER1_WALL
+            wall_position = (position1[0] + position2[0] + 2, position1[1] + position2[1] + 2)
+            return self._grid[*wall_position] == Board.WALL
         else:
             # By convention we treat the border as a "wall". This is makes checking jumps more convenient, since
             # players are allowed to jump diagonally if they are adjacent to another player and the border of the
             # board is on the other side of that player.
             return True
 
-    def is_position_on_board(self, position: Position) -> bool:
-        assert is_valid_position_type(position)
+    def get_old_style_walls(self):
+        return copy.copy(self._old_style_walls)
 
+    def is_position_on_board(self, position: tuple[int, int]) -> bool:
         return (
             (position[0] >= 0)
             and (position[0] < self.board_size)
@@ -173,41 +221,145 @@ def is_position_on_board(self, position: Position) -> bool:
             and (position[1] < self.board_size)
         )
 
-    def _get_wall_slice(self, position: Position, orientation: WallOrientation) -> tuple[slice, slice]:
+    def _wall_position_to_grid_index(self, position: tuple[int, int], orientation: WallOrientation) -> tuple[int, int]:
         """
-        Get a tuple of slices that correspond to the wall's cells in the grid.
+        Returns the grid index of the topmost/leftmost cell of the wall.
         """
-        assert is_valid_position_type(position)
-        assert isinstance(orientation, WallOrientation)
+        if orientation == WallOrientation.VERTICAL:
+            return (position[0] * 2 + 2, position[1] * 2 + 3)
+        elif orientation == WallOrientation.HORIZONTAL:
+            return (position[0] * 2 + 3, position[1] * 2 + 2)
 
+        raise ValueError("Invalid wall orientation")
+
+    def _get_wall_slice(self, position: tuple[int, int], orientation: WallOrientation) -> tuple[slice, slice]:
+        """
+        Get a tuple of slices that correspond to the wall's cells in the grid.
+        """
+        r, c = self._wall_position_to_grid_index(position, orientation)
         if orientation == WallOrientation.VERTICAL:
-            wall_slice = (slice(position[0] * 2, position[0] * 2 + 3), position[1] * 2 + 1)
+            wall_slice = (slice(r, r + 3), c)
         elif orientation == WallOrientation.HORIZONTAL:
-            wall_slice = (position[0] * 2 + 1, slice(position[1] * 2, position[1] * 2 + 3))
+            wall_slice = (r, slice(c, c + 3))
 
         return wall_slice
 
+    def _is_wall_potential_block(self, position, orientation):
+        wall_start_index = self._wall_position_to_grid_index(position, orientation)
+
+        touches = 0
+        for neighbor_offsets in Board._potential_wall_neighbors[orientation]:
+            neighbors = wall_start_index + neighbor_offsets
+            if (self._grid[neighbors[:, 0], neighbors[:, 1]] == Board.WALL).any():
+                touches += 1
+
+        return touches > 1
+
+    def _bfs(self, start_position, target_row, visited):
+        """
+        Performs a breadth-first search to find the shortest path to the target row.
+
+
+        Args:
+            row (int): The current row of the pawn
+            col (int): The current column of the pawn
+            target_row (int): The target row to reach
+            visited (numpy.array): A 2D boolean array with the same shape as the board,
+                indicating which positions have been visited
+
+        Returns:
+            int: Number of steps to reach the target or -1 if it's unreachable
+        """
+        row, col = start_position
+        if target_row == row:
+            return 0
+
+        queue = [(row, col, 0)]
+        visited[row, col] = True
+
+        while queue:
+            curr_row, curr_col, steps = queue.pop(0)
+
+            for new_row, new_col in [
+                (curr_row + 1, curr_col),
+                (curr_row - 1, curr_col),
+                (curr_row, curr_col - 1),
+                (curr_row, curr_col + 1),
+            ]:
+                if (
+                    self.is_position_on_board((new_row, new_col))
+                    and not self.is_wall_between((curr_row, curr_col), (new_row, new_col))
+                    and not visited[new_row, new_col]
+                ):
+                    visited[new_row, new_col] = True
+                    if target_row == new_row:
+                        return steps + 1
+                    else:
+                        queue.append((new_row, new_col, steps + 1))
+
+        return -1
+
+    def _dfs(self, start_position, target_row, visited):
+        """
+        Performs a depth-first search to find whether the pawn can reach the target row.
+
+        Args:
+            row (int): The current row of the pawn
+            col (int): The current column of the pawn
+            target_row (int): The target row to reach
+            visited (numpy.array): A 2D boolean array with the same shape as the board,
+                indicating which positions have been visited
+
+        Returns:
+            int: Number of steps to reach the target or -1 if it's unreachable
+        """
+        row, col = start_position
+        if row == target_row:
+            return 0
+
+        visited[row, col] = True
+
+        # Find out the forward direction to try it first and maybe get to the target faster
+        fwd = 1 if target_row > row else -1
+
+        moves = [(row + fwd, col), (row, col - 1), (row, col + 1), (row - fwd, col)]
+        for new_row, new_col in moves:
+            if (
+                self.is_position_on_board((new_row, new_col))
+                and not self.is_wall_between((row, col), (new_row, new_col))
+                and not visited[new_row, new_col]
+            ):
+                dfs = self._dfs((new_row, new_col), target_row, visited)
+                if dfs != -1:
+                    return dfs + 1
+
+        return -1
+
     def __str__(self):
         """
         Return a pretty-printed string representation of the grid.
         """
-        display_grid = np.full(self._grid.shape, " ", dtype=str)
+        grid_without_border = self._grid[2:-2, 2:-2]
+        display_grid = np.full(grid_without_border.shape, " ", dtype=str)
         display_grid[::2, ::2] = "."
-        display_grid[self._grid == Board.PLAYER1_WALL] = "w"
-        display_grid[self._grid == Board.PLAYER2_WALL] = "W"
-        for player, _ in enumerate(self._player_positions):
-            display_grid[*self._player_positions[player] * 2] = str(player + 1)
+        display_grid[grid_without_border == Board.WALL] = "#"
+        for player, (row, col) in enumerate(self._player_positions):
+            display_grid[row * 2, col * 2] = str(player + 1)
+        for idx, value in np.ndenumerate(grid_without_border):
+            if value >= 5 and value <= 9:
+                display_grid[idx] = str(value)  # Useful for debugging.
         return "\n".join([" ".join(row) for row in display_grid]) + "\n"
 
 
 class Quoridor:
-    def __init__(self, board: Board):
+    def __init__(self, board: Board, current_player: Player = Player.ONE):
         """
         If you want to start from the initial game state, pass in board=Board(board_size, max_walls).
         """
         self.board = board
-        self.current_player = Player.ONE
+        self.current_player = current_player
 
+        self._goal_rows = {Player.ONE: self.board.board_size - 1, Player.TWO: 0}
         self._jump_checks = create_jump_checks()
 
     def step(self, action: Action, validate: bool = True):
@@ -228,62 +380,135 @@ def step(self, action: Action, validate: bool = True):
         else:
             raise ValueError("Invalid action type")
 
-        self.current_player = Player(1 - self.current_player)
+        self.go_to_next_player()
 
-    def is_action_valid(self, action: Action):
+    def is_action_valid(self, action: Action, for_player: Player = None) -> bool:
         """
         Check whether the given action is valid given the current game state.
         """
-        if isinstance(action, MoveAction):
+        if for_player is None:
             player = self.get_current_player()
+        else:
+            player = for_player
+        opponent = Player(1 - player)
+
+        is_valid = True
+        if isinstance(action, MoveAction):
             current_position = self.board.get_player_position(player)
-            opponent = Player(1 - player)
             opponent_position = self.board.get_player_position(opponent)
-            opponent_offset = opponent_position - current_position
-            destination_position = np.asarray(action.destination, dtype=np.uint8)
-            position_delta = tuple(destination_position - current_position)  # Tuple so we can use it as a key.
+            opponent_offset = (opponent_position[0] - current_position[0], opponent_position[1] - current_position[1])
+            position_delta = (action.destination[0] - current_position[0], action.destination[1] - current_position[1])
 
             # Destination cell must be free
-            is_valid = True
-            if self.board.get_player_cell(destination_position) != Board.FREE:
+            if self.board.get_player_cell(action.destination) != Board.FREE:
                 is_valid = False
 
             elif np.sum(np.abs(position_delta)) == 1:
                 # Moving to an adjacent cell, just make sure no wall is in the way.
-                if self.board.is_wall_between(current_position, destination_position):
+                if self.board.is_wall_between(current_position, action.destination):
                     is_valid = False
 
-            elif (tuple(opponent_offset), tuple(position_delta)) in self._jump_checks:
-                jump_checks = self._jump_checks[(tuple(opponent_offset), tuple(position_delta))]
+            elif (opponent_offset, position_delta) in self._jump_checks:
+                jump_checks = self._jump_checks[(opponent_offset, position_delta)]
 
                 for check_delta_1, check_delta_2 in jump_checks["wall"]:
                     if not self.board.is_wall_between(
-                        current_position + check_delta_1, current_position + check_delta_2
+                        (current_position[0] + check_delta_1[0], current_position[1] + check_delta_1[1]),
+                        (current_position[0] + check_delta_2[0], current_position[1] + check_delta_2[1]),
                     ):
                         is_valid = False
                 for check_delta_1, check_delta_2 in jump_checks["nowall"]:
-                    if self.board.is_wall_between(current_position + check_delta_1, current_position + check_delta_2):
+                    if self.board.is_wall_between(
+                        (current_position[0] + check_delta_1[0], current_position[1] + check_delta_1[1]),
+                        (current_position[0] + check_delta_2[0], current_position[1] + check_delta_2[1]),
+                    ):
                         is_valid = False
             else:
                 # This isn't a move by 1, and it isn't a jump, so it's invalid.
                 is_valid = False
 
         elif isinstance(action, WallAction):
-            is_valid = self.board.can_place_wall(action.position, action.orientation)
-
+            if self.board.can_place_wall(player, action.position, action.orientation):
+                if self.board._is_wall_potential_block(action.position, action.orientation):
+                    # Temporarily place the wall so that we can check player paths to their goals.
+                    self.board.add_wall(self.current_player, action.position, action.orientation, check_if_valid=False)
+                    for p in [Player.ONE, Player.TWO]:
+                        if not self.can_reach(self.board.get_player_position(p), self.get_goal_row(p)):
+                            is_valid = False
+                            break
+                    # Restore the board to its previous state
+                    self.board.remove_wall(self.current_player, action.position, action.orientation)
+            else:
+                is_valid = False
         else:
             raise ValueError("Invalid action type")
 
         return is_valid
 
-    def set_current_player(self, player: int):
-        if player not in [0, 1]:
-            raise ValueError(f"Invalid player {player}")
-        self.current_player = player
+    def get_valid_move_actions(self, player: Player = None) -> list[MoveAction]:
+        if player is None:
+            player = self.get_current_player()
+
+        position = self.board.get_player_position(player)
+
+        valid_move_actions = list()
+        for delta_row in range(-2, 3):
+            for delta_col in range(-2, 3):
+                destination = (position[0] + delta_row, position[1] + delta_col)
+                if self.board.is_position_on_board(destination):
+                    move_action = MoveAction(destination)
+                    if self.is_action_valid(move_action, player):
+                        valid_move_actions.append(move_action)
+        return valid_move_actions
+
+    def get_valid_wall_actions(self, player: Player = None) -> list[WallAction]:
+        if player is None:
+            player = self.get_current_player()
+
+        valid_wall_actions = list()
+        for row in range(self.board.board_size - 1):
+            for col in range(self.board.board_size - 1):
+                for orientation in [WallOrientation.VERTICAL, WallOrientation.HORIZONTAL]:
+                    wall_action = WallAction((row, col), orientation)
+                    if self.is_action_valid(wall_action, player):
+                        valid_wall_actions.append(wall_action)
+        return valid_wall_actions
+
+    def get_valid_actions(self, player: Player = None) -> list[Action]:
+        return self.get_valid_move_actions(player) + self.get_valid_wall_actions(player)
+
+    def go_to_next_player(self):
+        self.current_player = Player(1 - self.current_player)
 
     def get_current_player(self) -> int:
         return self.current_player
 
+    def get_goal_row(self, player: Player) -> int:
+        return self._goal_rows[player]
+
+    def check_win(self, player):
+        row, _ = self.board.get_player_position(player)
+        return row == self.get_goal_row(player)
+
+    def can_reach(self, start_position: tuple[int, int], target_row: int):
+        visited = np.zeros((self.board.board_size, self.board.board_size), dtype="bool")
+        return self.board._dfs(start_position, target_row, visited) != -1
+
+    def distance_to_target(self, start_position, target_row):
+        """
+        Returns the minimum number of moves it takes to reach the target row, or -1 if it's not reachable.
+        """
+        visited = np.zeros((self.board.board_size, self.board.board_size), dtype="bool")
+        return self.board._bfs(start_position, target_row, visited)
+
+    def player_distance_to_target(self, player: Player):
+        start_position = self.board.get_player_position(player)
+        target_row = self.get_goal_row(player)
+        return self.distance_to_target(start_position, target_row)
+
+    def is_game_over(self):
+        return self.check_win(Player.ONE) or self.check_win(Player.TWO)
+
     def __str__(self):
         return str(self.board)
 
@@ -316,25 +541,71 @@ def create_jump_checks():
             "nowall": [((0, 0), (1, 0)), ((1, 0), (1, -1))],
         },
     }
-    # Rotate the checks for all 4 possible directions and convert to numpy arrays.
+    # Rotate the checks for all 4 possible directions
     checks = {}
     for opponent_offset in np.array([(1, 0), (0, 1), (-1, 0), (0, -1)]):
         rotation_matrix = np.array((opponent_offset, -opponent_offset[::-1])).T
         for move in checks_for_one_direction:
-            rotated_move = np.dot(rotation_matrix, move)  # Tuple so that we can use it in the key.
+            rotated_move = np.dot(rotation_matrix, move)
             lookup_key = (tuple(opponent_offset), tuple(rotated_move))
             checks[lookup_key] = {}
             for check_type in checks_for_one_direction[move]:
                 checks[lookup_key][check_type] = []
                 for wall_start, wall_end in checks_for_one_direction[move][check_type]:
-                    rotated_wall = (np.dot(rotation_matrix, wall_start), np.dot(rotation_matrix, wall_end))
+                    rotated_wall = (
+                        tuple(np.dot(rotation_matrix, wall_start)),
+                        tuple(np.dot(rotation_matrix, wall_end)),
+                    )
                     checks[lookup_key][check_type].append(rotated_wall)
     return checks
 
 
+def construct_game_from_observation(observation: dict, player_id: str) -> tuple[Quoridor, Player, Player]:
+    if player_id == "player_0":
+        player = Player.ONE
+        opponent = Player.TWO
+    elif player_id == "player_1":
+        player = Player.TWO
+        opponent = Player.ONE
+    else:
+        raise ValueError(f"Invalid player ID: {player_id}")
+
+    if observation["my_turn"]:
+        current_player = player
+    else:
+        current_player = opponent
+
+    # Hack: we set max walls very high because the observation doesn't actually tell us the maximum
+    # number of walls. After we add the walls from the observation to the board, we set the walls_remaining
+    # values to match those in the observation.
+    board = Board(board_size=observation["board"].shape[0], max_walls=99999)
+
+    player_positions = np.argwhere(observation["board"] > 0)
+    assert len(player_positions) == 2, "There should be exactly two players on the board."
+
+    for row, col in player_positions:
+        player_on_board = Player(
+            observation["board"][row, col] - 1
+        )  # Players are 1 and 2 on the board, but we use 0 and 1.
+        board.move_player(player_on_board, (row, col))
+
+    for row, col, orientation in np.argwhere(observation["walls"] == 1):
+        board.add_wall(
+            player,
+            (row, col),
+            WallOrientation(orientation),
+            check_if_valid=False,
+        )
+
+    board.set_walls_remaining(player, observation["my_walls_remaining"])
+    board.set_walls_remaining(opponent, observation["opponent_walls_remaining"])
+
+    return Quoridor(board, current_player), player, opponent
+
+
 if __name__ == "__main__":
     game = Quoridor(Board(board_size=9, max_walls=10))
-    game.step(MoveAction(np.array((1, 4))))
-    game.step(WallAction(np.array((0, 4)), WallOrientation.HORIZONTAL))
-    game.step(WallAction(np.array((4, 2)), WallOrientation.VERTICAL))
+    game.step(MoveAction((1, 4)))
+    game.step(WallAction((0, 4), WallOrientation.HORIZONTAL))
+    game.step(WallAction((4, 2), WallOrientation.VERTICAL))
     print(str(game))
diff --git a/deep_quoridor/src/quoridor_env.py b/deep_quoridor/src/quoridor_env.py
index 44207aca..d3225e1c 100644
--- a/deep_quoridor/src/quoridor_env.py
+++ b/deep_quoridor/src/quoridor_env.py
@@ -27,20 +27,36 @@
 import copy
 import functools
 import random
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
 from gymnasium import spaces
 from pettingzoo import AECEnv
 from pettingzoo.utils import wrappers
+from quoridor import ActionEncoder, Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
 
 
 class QuoridorEnv(AECEnv):
     metadata = {"render_modes": ["ansi"], "name": "quoridor_v0"}
 
     def __init__(
-        self, board_size: int = 9, max_walls: int = 10, step_rewards: bool = False, render_mode: str = "human", **kwargs
+        self,
+        board_size: int = 9,
+        max_walls: int = 10,
+        step_rewards: bool = False,
+        render_mode: str = "human",
+        game_start_state: Optional[Quoridor] = None,
+        **kwargs,
     ):
+        """
+        Constructs a Quoridor environment.
+
+        Args:
+            board_size (int): Size of the square board (e.g., 9 for a 9x9 board).
+            max_walls (int): Maximum number of walls each player can place (e.g., 10).
+            step_rewards (bool): Whether to provide (heuristic) incremental rewards.
+            game_start_state (Optional[Quoridor]): An optional starting state of the game to initialize the environment. (mainly for testing)
+        """
         super(AECEnv, self).__init__()
 
         self.render_mode = render_mode
@@ -50,15 +66,22 @@ def __init__(
         self.wall_size = self.board_size - 1  # grid for walls
         self.max_walls = max_walls  # Each player gets 10 walls
 
+        self._action_encoder = ActionEncoder(board_size)
+
+        self._game_start_state = game_start_state
+        if self._game_start_state is not None:
+            assert self._game_start_state.board.board_size == self.board_size
+            assert self._game_start_state.board.max_walls == self.max_walls
+
         self.possible_agents = ["player_0", "player_1"]
 
+        # The Quoridor class uses an enum for the players, whereas petting zoo uses strings.
+        self.agent_to_player = {"player_0": Player.ONE, "player_1": Player.TWO}
+        self.player_to_agent = {Player.ONE: "player_0", Player.TWO: "player_1"}
+
         self.reset()
 
     def copy(self):
-        """
-        A deep copy is not very efficient; eventually we should break out the game logic and state
-        into a separate class that can be used by agents during planning
-        """
         return copy.deepcopy(self)
 
     def reset(self, seed=None, options=None):
@@ -66,13 +89,10 @@ def reset(self, seed=None, options=None):
         self.agent_order = self.agents.copy()
         self.last_action_mask = {a: [] for a in self.agents}
 
-        self.walls = np.zeros((self.wall_size, self.wall_size, 2), dtype=np.int8)
-        self.walls_remaining = {agent: self.max_walls for agent in self.agents}
-        # Positions are (row, col)
-        self.positions = {
-            "player_0": (0, self.board_size // 2),
-            "player_1": (self.board_size - 1, self.board_size // 2),
-        }
+        if self._game_start_state is None:
+            self.game = Quoridor(Board(self.board_size, self.max_walls))
+        else:
+            self.game = copy.deepcopy(self._game_start_state)
 
         self.rewards = {agent: 0 for agent in self.agents}
         self.terminations = {agent: False for agent in self.agents}
@@ -86,32 +106,39 @@ def reset(self, seed=None, options=None):
 
         return None
 
-    def step(self, action):
+    def step(self, action_index):
         """
-        Players move by selecting an index from 0-80 (9×9 board).
-        Wall placement is mapped to 81-208 (8×8×2).
+        Players move by selecting an index from 0-80 (9x9 board).
+        Wall placement is mapped to 81-208 (8x8x2).
         """
         agent = self.agent_selection
+        player = self.agent_to_player[agent]
+        opponent = self.agent_to_player[self.get_opponent(agent)]
+
         if self.terminations[agent]:
             self._next_player()
+            self.game.go_to_next_player()
             return
 
-        if self.last_action_mask[agent][action] != 1:
-            raise RuntimeError(f"Action not allowed by mask {action}")
+        if self.last_action_mask[agent][action_index] != 1:
+            raise RuntimeError(f"Action not allowed by mask {action_index}")
 
-        step_reward_calculator = StepRewardCalculator(self)
+        # If the environment and game get out of step, weird things happen.
+        assert player == self.game.get_current_player()
+
+        step_reward_calculator = StepRewardCalculator(self.game, player, opponent)
         if self.step_rewards:
             step_reward_calculator.before_step()
-        (row, col, action_type) = self.action_index_to_params(action)
-        if action_type == 0:
-            self._move(agent, (row, col))
-        else:
-            self.place_wall(agent, (row, col), action_type - 1)
 
-        if self._check_win(agent):
+        action = self._action_encoder.index_to_action(action_index)
+        self.game.step(action)
+
+        if self.game.check_win(player):
             self.terminations = {a: True for a in self.agents}
             self.rewards[agent] = 1
-            self.rewards[self.get_opponent(agent)] = -1
+            for a in self.agents:
+                if a != agent:
+                    self.rewards[a] = -1
         elif self.step_rewards:
             # Assign rewards as the difference in distance to the goal divided by
             # three times the board size.
@@ -143,223 +170,47 @@ def observe(self, agent_id):
         }
 
     def _get_observation(self, agent_id):
+        player = self.agent_to_player[agent_id]
+        opponent = self.agent_to_player[self.get_opponent(agent_id)]
         # TODO: Do we need to make copies of the state or can we return references directly?
 
         # Calculate board from self.positions
-        # NOTE: The board uses 1 to indicate where the agent is and 2 to indicate where the opponent is
-        # Obviously, this will be different for each player
-        board = np.zeros((self.board_size, self.board_size), dtype=np.int8)
-        board[self.positions[agent_id]] = 1
-        board[self.positions[self.get_opponent(agent_id)]] = 2
+        board = np.zeros((self.game.board.board_size, self.game.board.board_size), dtype=np.int8)
+        player_one_position = self.game.board.get_player_position(Player.ONE)
+        player_two_position = self.game.board.get_player_position(Player.TWO)
+        board[player_one_position] = 1
+        board[player_two_position] = 2
 
         # Make a copy of walls
-        walls = self.walls.copy()
+        walls = self.game.board.get_old_style_walls()
 
         return {
             "my_turn": self.agent_selection == agent_id,
             "board": board,
             "walls": walls,
-            "my_walls_remaining": self.walls_remaining[agent_id],
-            "opponent_walls_remaining": self.walls_remaining[self.get_opponent(agent_id)],
+            "my_walls_remaining": self.game.board.get_walls_remaining(player),
+            "opponent_walls_remaining": self.game.board.get_walls_remaining(opponent),
         }
 
-    def is_wall_between(self, row_a, col_a, row_b, col_b):
-        """
-        Returns True if there is a wall between pos_a and pos_b
-        NOTE: The min and max tricks there are just a lazy way to avoid overindexing
-        Since the wall grid is one smaller than the board (because boards take two spaces)
-        """
-        if row_a == row_b:  # Horizontal movement - check vertical walls
-            wall_col = min(col_a, col_b)
-            return (
-                self.walls[min(row_a, self.wall_size - 1), wall_col, 0] == 1
-                or self.walls[max(row_a - 1, 0), wall_col, 0] == 1
-            )
-        elif col_a == col_b:  # Vertical movement - check horizontal walls
-            wall_row = min(row_a, row_b)
-            return (
-                self.walls[wall_row, min(col_a, self.wall_size - 1), 1] == 1
-                or self.walls[wall_row, max(col_a - 1, 0), 1] == 1
-            )
-        else:
-            raise ValueError(f"Invalid movement from {row_a, col_a} to {row_b, col_b}")
-
-    def is_in_board(self, row, col):
-        return 0 <= row < self.board_size and 0 <= col < self.board_size
-
     def _get_action_mask(self, agent_id):
         # Start with an empty mask (nothing possible)
-        mask = np.zeros((self.board_size**2 + (self.wall_size**2) * 2,), dtype=np.int8)
-
-        # Calculate valid "moves" (as opposed to wall placements)
-        # Start with the four basic directions
-        for row, col in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
-            old_row, old_col = self.positions[agent_id]
-            new_row, new_col = old_row + row, old_col + col
-
-            # Check if the new position is still inside the board
-            if not self.is_in_board(new_row, new_col):
-                continue
-
-            # Check if that direction is blocked by a wall
-            if self.is_wall_between(old_row, old_col, new_row, new_col):
-                continue
-
-            # Check if the opponent is in the target position
-            if self.positions[self.get_opponent(agent_id)] == (new_row, new_col):
-                # Can we jump straight over the opponent?
-                straight_row = new_row + row
-                straight_col = new_col + col
-
-                if (
-                    # That new position falls within the board
-                    self.is_in_board(straight_row, straight_col)
-                    # That new position is not blocked by a wall behind them
-                    and not self.is_wall_between(new_row, new_col, straight_row, straight_col)
-                ):
-                    mask[self.action_params_to_index(straight_row, straight_col, 0)] = 1
-                    continue
-                else:
-                    # Try the diagonals, to the sides of the oponent
-                    for side in [1, -1]:
-                        if row == 0:
-                            diag_row = new_row + side
-                            diag_col = new_col
-                        else:
-                            diag_row = new_row
-                            diag_col = new_col + side
-
-                        if (
-                            # That new position falls within the board
-                            self.is_in_board(diag_row, diag_col)
-                            # That new position is not blocked by a wall
-                            and not self.is_wall_between(new_row, new_col, diag_row, diag_col)
-                        ):
-                            mask[self.action_params_to_index(diag_row, diag_col, 0)] = 1
-                    continue
-
-            # If we get to this point it's because the target is within the board,
-            # not blocked by a wall and the opponent is not on the way
-            mask[self.action_params_to_index(new_row, new_col, 0)] = 1
-
-        # Calculate wall placements
-        # Iterate over all possible wall positions, only of there are any left to place
-        if self.walls_remaining[agent_id] > 0:
-            for row, col, orientation in np.ndindex(self.walls.shape):
-                if self._is_wall_overlap(row, col, orientation):
-                    continue
-
-                if not self._is_wall_potential_block(row, col, orientation) or self._can_place_wall_without_blocking(
-                    row, col, orientation
-                ):
-                    mask[self.action_params_to_index(row, col, orientation + 1)] = 1
-
-        self.last_action_mask[agent_id] = mask
-        return mask
-
-    def _is_wall_potential_block(self, row, col, orientation):
-        # Area to find other walls that may touch this wall
-        top = max(0, row - 1)
-        bottom = min(self.wall_size - 1, row + 1)
-        left = max(0, col - 1)
-        right = min(self.wall_size - 1, col + 1)
-
-        # Whether the wall was touched in the left/top, middle, right/bottom
-        touches = [False, False, False]
-
-        if orientation == 1:  # Horizontal
-            # On the left border or touching another horizontal wall on the left side
-            if (col == 0) or (col >= 2 and self.walls[row, col - 2, 1] == 1):
-                touches[0] = True
-
-            # On the right border or touching another horizontal wall on the right sinde
-            if (col == self.wall_size - 1) or (col < self.wall_size - 2 and self.walls[row, col + 2, 1] == 1):
-                touches[2] = True
-
-            # Check for vertical walls touching it
-            for r in range(top, bottom + 1):
-                for c in range(left, right + 1):
-                    if self.walls[r, c, 0]:
-                        touches[c - col + 1] = True
-        else:  # Vertical
-            # On the top border or touching another verticall wall on top
-            if (row == 0) or (row >= 2 and self.walls[row - 2, col, 0] == 1):
-                touches[0] = True
-
-            # On the bottom border or touching another verticall wall on the bottom
-            if (row == self.wall_size - 1) or (row < self.wall_size - 2 and self.walls[row + 2, col, 0] == 1):
-                touches[2] = True
-
-            # Check for horizontal walls touching it
-            for r in range(top, bottom + 1):
-                for c in range(left, right + 1):
-                    if self.walls[r, c, 1]:
-                        touches[r - row + 1] = True
-
-        return sum(touches) > 1
-
-    def _is_wall_overlap(self, row, col, orientation):
-        """
-        Checks whether there is any wall overlaping where this wall would be placed
-        """
-        # Check if it doesn't collide with a crossing wall
-        if self.walls[row, col, (orientation + 1) % 2] == 1:
-            return True
-
-        # Check that in the 2 wall segments there's no other segment already
-        if orientation == 0:
-            if self.is_wall_between(row, col, row, col + 1) or self.is_wall_between(row + 1, col, row + 1, col + 1):
-                return True
-        else:
-            if self.is_wall_between(row, col, row + 1, col) or self.is_wall_between(row, col + 1, row + 1, col + 1):
-                return True
+        player = self.agent_to_player[agent_id]
+        action_mask = np.zeros((self.board_size**2 + (self.wall_size**2) * 2,), dtype=np.int8)
 
-        return False
+        # Mark valid moves
+        for action in self.game.get_valid_actions(player):
+            action_i = self._action_encoder.action_to_index(action)
+            action_mask[action_i] = 1
+
+        self.last_action_mask[agent_id] = action_mask
+        return action_mask
 
     def _get_info(self):
         # This is for now unused, returning empty dict
         return {}
 
     def render(self):
-        board_str = ""
-
-        for row in range(self.board_size):
-            # Render board row with player positions and vertical walls
-            for col in range(self.board_size):
-                if self.positions["player_0"] == (row, col):
-                    board_str += " P0 "
-                elif self.positions["player_1"] == (row, col):
-                    board_str += " P1 "
-                else:
-                    board_str += " .  "
-
-                # Render vertical walls (|)
-                if col < self.wall_size:
-                    if row < self.wall_size and self.walls[row, col, 0]:
-                        board_str += "|"
-                    elif row > 0 and self.walls[row - 1, col, 0]:  # Continuation of a vertical wall
-                        board_str += "|"
-                    else:
-                        board_str += " "
-                else:
-                    board_str += " "
-
-            board_str += "\n"
-
-            # Render horizontal walls (───)
-            if row < self.wall_size:
-                for col in range(self.board_size):
-                    if col < self.wall_size and self.walls[row, col, 1]:
-                        board_str += "────+"  # Wall segment
-                    elif col > 0 and self.walls[row, col - 1, 1]:
-                        board_str += "──── "  # Continuation of a wall from the left
-                    elif col < self.wall_size and self.walls[row, col, 0]:
-                        board_str += "    +"  # Vertical wall going through this intersection
-                    else:
-                        board_str += "     "
-                board_str += "\n"
-
-        print(board_str)
+        print(str(self.game))
 
     @functools.lru_cache(maxsize=None)
     def observation_space(self, agent):
@@ -385,76 +236,13 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
-        """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
-        """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
-        else:
-            raise ValueError(f"Invalid action type: {action_type}")
-
-    def action_index_to_params(self, idx) -> Tuple[int, int, int]:
-        """
-        Takes an action index to action parameters (row, col, movement_type)
-        movement_type = 0 for moving, 1 for vertical wall placement, 2 for horizontal wall placement
-        """
-        if idx < self.board_size**2:  # Pawn movement
-            action_type = 0
-            row, col = divmod(idx, self.board_size)
-        elif idx >= self.board_size**2 and idx < self.board_size**2 + self.wall_size**2:
-            # Vertical wall placement
-            action_type = 1
-            row, col = divmod(idx - self.board_size**2, self.wall_size)
-        elif idx >= self.board_size**2 + self.wall_size**2 and idx < self.board_size**2 + (self.wall_size**2) * 2:
-            # Horizontal wall placement
-            action_type = 2
-            row, col = divmod(idx - self.board_size**2 - self.wall_size**2, self.wall_size)
-        else:
-            raise ValueError(f"Invalid action index: {idx}")
-        return (row, col, action_type)
-
-    def _move(self, agent, position):
-        """
-        TODO: Only allow valid moves
-        """
-        row, col = position
-        if (row, col) in self.positions.values():
-            print("WTF: Invalid action: Occupied")
-            return  # Invalid move (occupied)
-        self.positions[agent] = (row, col)
-
-    def place_wall(self, agent, position, orientation):
-        """
-        Place a wall at the specified position and orientation.
-        Since each wall is of length 2, the wall will coneinue through two spaces.
-
-        TODO: Prevent placing walls conflicting with existing ones
-        TODO: Prevent completely blocking the path from one end to the other
-        """
-        if self.walls_remaining[agent] == 0:
-            print("WTF: Invalid action: No walls left")
-            return  # No walls left
-        (row, col) = position
-        if self.walls[row, col, orientation] == 0:
-            self.walls[row, col, orientation] = 1  # Mark as wall
-            self.walls_remaining[agent] -= 1
-
-    def _check_win(self, agent):
-        row, _ = self.positions[agent]
-        return row == self.get_goal_row(agent)
-
     def winner(self) -> Optional[int]:
         """
         Return the index of the winner (0 for player 1, 1 for player 2) or None if there's no winner
         """
         for idx, agent in enumerate(self.agent_order):
-            if self._check_win(agent):
+            player = self.agent_to_player[agent]
+            if self.game.check_win(player):
                 return idx
 
         return None
@@ -466,120 +254,6 @@ def _next_player(self):
     def get_opponent(self, agent):
         return "player_1" if agent == "player_0" else "player_0"
 
-    def get_goal_row(self, agent):
-        return 0 if agent == "player_1" else self.board_size - 1
-
-    def _bfs(self, row, col, target_row, visited):
-        """
-        Performs a breadth-first search to find the shortest path to the target row.
-
-
-        Args:
-            row (int): The current row of the pawn
-            col (int): The current column of the pawn
-            target_row (int): The target row to reach
-            visited (numpy.array): A 2D boolean array with the same shape as the board,
-                indicating which positions have been visited
-
-        Returns:
-            int: Number of steps to reach the target or -1 if it's unreachable
-        """
-        if target_row == row:
-            return 0
-
-        queue = [(row, col, 0)]
-        visited[row, col] = True
-
-        while queue:
-            curr_row, curr_col, steps = queue.pop(0)
-
-            for new_row, new_col in [
-                (curr_row + 1, curr_col),
-                (curr_row - 1, curr_col),
-                (curr_row, curr_col - 1),
-                (curr_row, curr_col + 1),
-            ]:
-                if (
-                    self.is_in_board(new_row, new_col)
-                    and not self.is_wall_between(curr_row, curr_col, new_row, new_col)
-                    and not visited[new_row, new_col]
-                ):
-                    visited[new_row, new_col] = True
-                    if target_row == new_row:
-                        return steps + 1
-                    else:
-                        queue.append((new_row, new_col, steps + 1))
-
-        return -1
-
-    def _dfs(self, row, col, target_row, visited):
-        """
-        Performs a depth-first search to find whether the pawn can reach the target row.
-
-        Args:
-            row (int): The current row of the pawn
-            col (int): The current column of the pawn
-            target_row (int): The target row to reach
-            visited (numpy.array): A 2D boolean array with the same shape as the board,
-                indicating which positions have been visited
-            If any_path is set to true, the first path to the target row will be returned (faster).
-            Otherwise, the shortest path will be returned (potentially slower)
-
-        Returns:
-            int: Number of steps to reach the target or -1 if it's unreachable
-        """
-        if row == target_row:
-            return 0
-
-        visited[row, col] = True
-
-        # Find out the forward direction to try it first and maybe get to the target faster
-        fwd = 1 if target_row > row else -1
-
-        moves = [(row + fwd, col), (row, col - 1), (row, col + 1), (row - fwd, col)]
-        for new_row, new_col in moves:
-            if (
-                self.is_in_board(new_row, new_col)
-                and not self.is_wall_between(row, col, new_row, new_col)
-                and not visited[new_row, new_col]
-            ):
-                dfs = self._dfs(new_row, new_col, target_row, visited)
-                if dfs != -1:
-                    return dfs + 1
-
-        return -1
-
-    def can_reach(self, row, col, target_row):
-        visited = np.zeros((self.board_size, self.board_size), dtype="bool")
-        return self._dfs(row, col, target_row, visited) != -1
-
-    def distance_to_target(self, row, col, target_row):
-        """
-        Returns the approximate number of moves it takes to reach the target row, or -1 if it's not reachable.
-        If any_path is set to true, the first path to the target row will be returned (faster).
-        Otherwise, the shortest path will be returned (potentially slower)
-        """
-        visited = np.zeros((self.board_size, self.board_size), dtype="bool")
-        return self._bfs(row, col, target_row, visited)
-
-    def _can_place_wall_without_blocking(self, row, col, orientation):
-        """
-        Returns whether a wall can be placed in the specified coordinates an orientations such that
-        both pawns can still reach the target row
-        """
-
-        # Temporarily place the wall so that we can easily check for walls
-        previous = self.walls[row, col, orientation]
-        self.walls[row, col, orientation] = 1
-        result = all(self.can_reach(*self.positions[agent], self.get_goal_row(agent)) for agent in self.agents)
-        self.walls[row, col, orientation] = previous
-
-        return result
-
-    def get_distance_to_target(self, agent):
-        agent_row, agent_col = self.positions[agent]
-        return self.distance_to_target(agent_row, agent_col, self.get_goal_row(agent))
-
 
 # Wrapping the environment for PettingZoo compatibility
 def env(**kwargs):
@@ -593,30 +267,21 @@ def env(**kwargs):
 
 
 class StepRewardCalculator:
-    def __init__(self, env: QuoridorEnv = None):
-        self.env = env
-
-    def _get_distances(self):
-        """
-        Gets the current distances to target for both agent and opponent.
-        Returns tuple of (agent_distance, opponent_distance)
-        """
-        agent = self.env.agent_selection
-        agent_distance = self.env.get_distance_to_target(agent)
-
-        opponent = self.env.get_opponent(agent)
-        opponent_distance = self.env.get_distance_to_target(opponent)
-
-        return agent_distance, opponent_distance
+    def __init__(self, game: Quoridor, player: Player, opponent: Player):
+        self.game = game
+        self.player = player
+        self.opponent = opponent
 
     def before_step(self):
-        self.orig_agent_distance, self.orig_opponent_distance = self._get_distances()
+        self.orig_player_distance = self.game.player_distance_to_target(self.player)
+        self.orig_opponent_distance = self.game.player_distance_to_target(self.opponent)
 
     def after_step(self):
-        agent_distance, opponent_distance = self._get_distances()
+        player_distance = self.game.player_distance_to_target(self.player)
+        opponent_distance = self.game.player_distance_to_target(self.opponent)
 
         # Calculate the reward based on the distance to the goal
-        distance_change = (self.orig_agent_distance - agent_distance) - (
+        distance_change = (self.orig_player_distance - player_distance) - (
             self.orig_opponent_distance - opponent_distance
         )
-        return distance_change / (3 * self.env.board_size)
+        return distance_change / (3 * self.game.board.board_size)
diff --git a/deep_quoridor/src/renderers/pygame.py b/deep_quoridor/src/renderers/pygame.py
index be51a0aa..9e73e8db 100644
--- a/deep_quoridor/src/renderers/pygame.py
+++ b/deep_quoridor/src/renderers/pygame.py
@@ -9,6 +9,7 @@
 from agents import ActionLog, Agent
 from agents.human import HumanAgent
 from arena import GameResult
+from quoridor import Player
 
 from renderers import Renderer
 
@@ -489,7 +490,7 @@ def _wait_not_in_states(self, states: List[GameState] | GameState):
 
             time.sleep(0.01)
 
-    def start_game(self, game, agent1: Agent, agent2: Agent):
+    def start_game(self, env, agent1: Agent, agent2: Agent):
         gui = PygameQuoridor.instance()
 
         # We always enable the log from the agents and decide later if we want to show it, since it's
@@ -499,12 +500,12 @@ def start_game(self, game, agent1: Agent, agent2: Agent):
         self.has_human_player = isinstance(agent1, HumanAgent) or isinstance(agent2, HumanAgent)
 
         initial_state = BoardState(
-            p1_position=game.positions["player_0"],
-            p1_walls_remaining=game.walls_remaining["player_0"],
-            p2_position=game.positions["player_1"],
-            p2_walls_remaining=game.walls_remaining["player_1"],
+            p1_position=env.game.board.get_player_position(Player.ONE),
+            p1_walls_remaining=env.game.board.get_walls_remaining(Player.ONE),
+            p2_position=env.game.board.get_player_position(Player.TWO),
+            p2_walls_remaining=env.game.board.get_walls_remaining(Player.TWO),
             is_p1_turn=True,
-            walls=game.walls,
+            walls=env.game.board.get_old_style_walls(),
         )
         gui.start_game(agent1.name(), agent2.name(), initial_state)
 
@@ -543,17 +544,17 @@ def before_action(self, game, agent):
                     time.sleep(0.01)
                 self.last_action_time = time.time()
 
-    def after_action(self, game, step, agent_id, action):
+    def after_action(self, env, step, agent_id, action):
         gui = PygameQuoridor.instance()
         gui.update_log(None)
-        gui.update_turn(game.agent_selection == "player_0")
+        gui.update_turn(env.agent_selection == "player_0")
 
         gui.update_players_state(
-            p1_position=game.positions["player_0"],
-            p1_walls_remaining=game.walls_remaining["player_0"],
-            p2_position=game.positions["player_1"],
-            p2_walls_remaining=game.walls_remaining["player_1"],
-            walls=game.walls,
+            p1_position=env.game.board.get_player_position(Player.ONE),
+            p1_walls_remaining=env.game.board.get_walls_remaining(Player.ONE),
+            p2_position=env.game.board.get_player_position(Player.TWO),
+            p2_walls_remaining=env.game.board.get_walls_remaining(Player.TWO),
+            walls=env.game.board.get_old_style_walls(),
         )
 
     def main_thread(self):
diff --git a/deep_quoridor/src/train_sb3.py b/deep_quoridor/src/train_sb3.py
index 312f8383..557a3463 100644
--- a/deep_quoridor/src/train_sb3.py
+++ b/deep_quoridor/src/train_sb3.py
@@ -18,14 +18,13 @@
 import time
 
 import quoridor_env
-from agents.sb3_ppo import DictFlattenExtractor, make_env_fn
+from agents.sb3_ppo import DictFlattenExtractor, SB3PPOAgent, make_env_fn
 from sb3_contrib import MaskablePPO
 from sb3_contrib.common.maskable.policies import MaskableActorCriticPolicy
 from sb3_contrib.common.wrappers import ActionMasker
 from utils import set_deterministic
 
 import wandb
-from deep_quoridor.src.agents.sb3_ppo import SB3PPOAgent
 from wandb.integration.sb3 import WandbCallback
 
 
diff --git a/deep_quoridor/test/agents/dexp_test.py b/deep_quoridor/test/agents/dexp_test.py
index badc520f..589f57c5 100644
--- a/deep_quoridor/test/agents/dexp_test.py
+++ b/deep_quoridor/test/agents/dexp_test.py
@@ -7,13 +7,11 @@
 @pytest.fixture
 def mock_observation():
     return {
-        "observation": {
-            "board": np.array([[0, 0, 0], [1, 0, 0], [0, 0, 2]]),
-            "walls": np.array([[[0, 0], [0, 0]], [[0, 0], [0, 0]]]),
-            "my_walls_remaining": 5,
-            "opponent_walls_remaining": 4,
-            "my_turn": True,
-        }
+        "board": np.array([[0, 0, 0], [1, 0, 0], [0, 0, 2]]),
+        "walls": np.array([[[0, 0], [0, 0]], [[0, 0], [0, 0]]]),
+        "my_walls_remaining": 5,
+        "opponent_walls_remaining": 4,
+        "my_turn": True,
     }
 
 
@@ -60,7 +58,7 @@ def test_base_observation_conversion(base_agent, mock_observation):
 
     # Verify board state conversion
     board_part = result[:9].numpy()
-    expected_board = mock_observation["observation"]["board"].flatten()
+    expected_board = mock_observation["board"].flatten()
     np.testing.assert_array_almost_equal(board_part, expected_board)
 
 
@@ -69,7 +67,7 @@ def test_rotation_for_player1(rotate_agent, mock_observation):
 
     # For player_1 with rotation, the board should be rotated 180 degrees
     board_part = result[:9].numpy()
-    expected_board = np.rot90(mock_observation["observation"]["board"], k=2)
+    expected_board = np.rot90(mock_observation["board"], k=2)
     np.testing.assert_array_almost_equal(board_part, expected_board.flatten())
 
 
@@ -83,8 +81,8 @@ def test_split_board_representation(split_agent, mock_observation):
     player_channel = result[:9].numpy()
     opponent_channel = result[9:18].numpy()
 
-    expected_player = (mock_observation["observation"]["board"] == 1).astype(np.float32).flatten()
-    expected_opponent = (mock_observation["observation"]["board"] == 2).astype(np.float32).flatten()
+    expected_player = (mock_observation["board"] == 1).astype(np.float32).flatten()
+    expected_opponent = (mock_observation["board"] == 2).astype(np.float32).flatten()
 
     np.testing.assert_array_almost_equal(player_channel, expected_player)
     np.testing.assert_array_almost_equal(opponent_channel, expected_opponent)
@@ -92,7 +90,7 @@ def test_split_board_representation(split_agent, mock_observation):
 
 def test_state_conversion_with_target_as_source(target_as_source_agent, mock_observation):
     # Test when it's not my turn (target state)
-    mock_observation["observation"]["my_turn"] = False
+    mock_observation["my_turn"] = False
 
     # When it's not my turn and target_as_source_for_opponent is True,
     # the board should be from opponent's perspective
@@ -101,5 +99,5 @@ def test_state_conversion_with_target_as_source(target_as_source_agent, mock_obs
     result = target_as_source_agent.observation_to_tensor(mock_observation, obs_player_id="player_0")
     wall_counts = result[-2:].numpy()
     # Wall counts should be swapped because this is a target state
-    assert wall_counts[0] == mock_observation["observation"]["opponent_walls_remaining"]
-    assert wall_counts[1] == mock_observation["observation"]["my_walls_remaining"]
+    assert wall_counts[0] == mock_observation["opponent_walls_remaining"]
+    assert wall_counts[1] == mock_observation["my_walls_remaining"]
diff --git a/deep_quoridor/test/profile_observe.py b/deep_quoridor/test/profile_observe.py
index 4bc2c14b..02e8ee36 100644
--- a/deep_quoridor/test/profile_observe.py
+++ b/deep_quoridor/test/profile_observe.py
@@ -1,26 +1,29 @@
 import cProfile
 import time
 
+import numpy as np
+from quoridor import Player, WallOrientation
 from quoridor_env import env
 
-e = env(9)
+e = env(board_size=9)
+
 walls = [
-    ((3, 0), 1),
-    ((3, 2), 1),
-    ((3, 4), 1),
-    ((3, 6), 1),
-    ((2, 7), 0),
-    ((5, 7), 1),
-    ((4, 0), 0),
-    ((4, 2), 0),
-    ((4, 4), 0),
+    ((3, 0), WallOrientation.HORIZONTAL),
+    ((3, 2), WallOrientation.HORIZONTAL),
+    ((3, 4), WallOrientation.HORIZONTAL),
+    ((3, 6), WallOrientation.HORIZONTAL),
+    ((2, 7), WallOrientation.VERTICAL),
+    ((5, 7), WallOrientation.HORIZONTAL),
+    ((4, 0), WallOrientation.VERTICAL),
+    ((4, 2), WallOrientation.VERTICAL),
+    ((4, 4), WallOrientation.VERTICAL),
 ]
 
 
 def run():
     st = time.time()
-    for w in walls:
-        e.place_wall("player_0", w[0], w[1])
+    for position, orientation in walls:
+        e.game.board.add_wall(Player.ONE, position, orientation)
         e.observe("player_0")
         e.observe("player_1")
 
@@ -37,5 +40,5 @@ def run():
 # Using _is_wall_potential_block optimization: 25,000 usec
 
 # Switch the comments to do a profiling
-run()
-# cProfile.run("run()", sort="tottime")
+# run()
+cProfile.run("run()", sort="tottime")
diff --git a/deep_quoridor/test/quoridor_env_test.py b/deep_quoridor/test/quoridor_env_test.py
deleted file mode 100644
index e7faec82..00000000
--- a/deep_quoridor/test/quoridor_env_test.py
+++ /dev/null
@@ -1,362 +0,0 @@
-from quoridor_env import env as quoridor_env
-
-
-def parse_board(board):
-    """
-    Parse a string representation of a Quoridor board (of any size) and convert it into a QuoridorEnv instance.
-
-    The board should contain N rows for the cells, which are composed by:
-    - '1': Position of player 1
-    - '2': Position of player 2
-    - '.': Empty cell
-    - '*': Potential move positions for player 1
-    - '|': Vertical walls
-    - 'v': A vertical wall cannot be placed here (for testing wall placement)
-    - ' ': Empty spaces (just for formatting)
-    - '+': Vertical wall continuation (just for formatting)
-
-    Additionally, between 2 rows for the cells, a row with horizontal walls can be included.
-    The position of the wall will be with respect to the row above.  The row is composed by:
-    - '-': Horizontal walls
-    - '>': A horizontal wall cannot be placed here (for testing wall placement)
-    - ' ': Empty spaces (just for formatting)
-    - '+': Horizontal wall continuation (just for formatting)
-
-    Returns:
-        tuple: A tuple containing three elements:
-            - QuoridorEnv: An environment instance representing the parsed board state
-            - list: A list of tuples (row, col) representing potential move positions
-            - list: A list of tuples (row, col, orientation) represeting where the walls are forbidden to be placed
-
-    """
-    rows = [r for r in board.split("\n") if r.strip()]
-    size = len([ch for ch in rows[0] if ch in "12.*"])
-    env = quoridor_env(board_size=size)
-    potential_moves = []
-    forbidden_walls = []
-
-    row_n = 0
-    col_positions = {}  # map from position in the string to column number
-
-    for row in rows:
-        col_n = 0
-
-        is_cell_row = set(row) <= set("12.*| +v")
-        is_h_wall_row = set(row) <= set("- +>")
-        assert is_cell_row or is_h_wall_row, f"Row contains invalid characters {row}"
-
-        if is_h_wall_row:
-            for i, col_n in col_positions.items():
-                if i >= len(row):
-                    break
-
-                ch = row[i]
-                if ch == "-" and not env.is_wall_between(row_n - 1, col_n, row_n, col_n):
-                    env.place_wall("player_0", (row_n - 1, col_n), 1)
-
-                if ch == ">" or ch == "-":
-                    if col_n < size - 1:
-                        forbidden_walls.append((row_n - 1, col_n, 1))
-
-            col_positions = {}
-
-        else:
-            for i, ch in enumerate(row):
-                if ch == " " or ch == "+":
-                    continue
-
-                if ch == "v" or ch == "|":
-                    if row_n < size - 1:
-                        forbidden_walls.append((row_n, col_n - 1, 0))
-                    if ch == "|" and not env.is_wall_between(row_n, col_n - 1, row_n, col_n):
-                        env.place_wall("player_0", (row_n, col_n - 1), 0)
-                    continue
-
-                if ch == "*":
-                    potential_moves.append((row_n, col_n))
-                elif ch == "1":
-                    env.positions["player_0"] = row_n, col_n
-                elif ch == "2":
-                    env.positions["player_1"] = row_n, col_n
-
-                col_positions[i] = col_n
-                col_n += 1
-
-            assert len(col_positions) == size, f"The row in the board doesn't contain {size} elements: {row}"
-            row_n += 1
-
-    assert row_n == size, f"Was expecting {size} rows, but found {row_n}"
-    return env, potential_moves, forbidden_walls
-
-
-class TestQuoridorEnv:
-    def _test_pawn_movements(self, s):
-        env, potential_moves, _ = parse_board(s)
-        N = env.board_size
-        action_mask = env.observe("player_0")["action_mask"]
-
-        env_moves = []
-        for i, value in enumerate(action_mask[: N**2]):
-            if value == 1:
-                env_moves.append(divmod(i, N))
-
-        assert env_moves == potential_moves
-
-    def _test_wall_placements(self, s, just_highlighted=False):
-        env, _, forbidden_walls = parse_board(s)
-        N = env.board_size
-        action_mask = env.observe("player_0")["action_mask"]
-
-        env_walls = []
-        for i in range(N**2, len(action_mask)):
-            if action_mask[i] == 0:
-                row, col, action_type = env.action_index_to_params(i)
-
-                env_walls.append((row, col, action_type - 1))
-
-        if just_highlighted:
-            diff = set(forbidden_walls).difference(set(env_walls))
-            assert not diff
-        else:
-            assert set(env_walls) == set(forbidden_walls)
-
-    def _test_distance_to_target(self, s, moves_p1, moves_p2):
-        env, _, _ = parse_board(s)
-        N = env.board_size
-        (row, col) = env.positions["player_0"]
-        assert moves_p1 == env.distance_to_target(row, col, N - 1)
-        (row, col) = env.positions["player_1"]
-        assert moves_p2 == env.distance_to_target(row, col, 0)
-
-    def test_distance_to_goal(self):
-        self._test_distance_to_target(
-            """
-            1 . .
-            . . .
-            . . 2
-        """,
-            2,
-            2,
-        )
-
-        self._test_distance_to_target(
-            """
-            . * .
-            . . .
-            2 . 1
-        """,
-            0,
-            2,
-        )
-
-        self._test_distance_to_target(
-            """
-            1 . .
-            - -
-            . . .
-            . . 2
-        """,
-            4,
-            2,
-        )
-
-        self._test_distance_to_target(
-            """
-            1 .|.
-            - -+
-            . .|.
-            . . 2
-        """,
-            -1,
-            2,
-        )
-
-        self._test_distance_to_target(
-            """
-            1 . .
-            - -
-            . .|.
-            . .|2
-        """,
-            4,
-            2,
-        )
-
-    def test_corner_movements(self):
-        self._test_pawn_movements("""
-            1 * .
-            * . .
-            . . 2
-        """)
-
-        self._test_pawn_movements("""
-            . * 1
-            . . *
-            . . 2
-        """)
-
-        self._test_pawn_movements("""
-            2 . .
-            . . *
-            . * 1
-        """)
-
-        self._test_pawn_movements("""
-            2 . .
-            * . .
-            1 * .
-        """)
-
-    def test_edge_movements(self):
-        self._test_pawn_movements("""
-            * . .
-            1 * .
-            * . 2
-        """)
-
-        self._test_pawn_movements("""
-            * 1 *
-            . * .
-            . . 2
-        """)
-
-        self._test_pawn_movements("""
-            . . *
-            2 * 1
-            . . *
-        """)
-
-        self._test_pawn_movements("""
-            . . .
-            2 * 1
-            * 1 *
-        """)
-
-    def test_center_movements(self):
-        self._test_pawn_movements("""
-            . * .
-            * 1 *
-            . * 2
-        """)
-
-    def test_simple_jumps(self):
-        self._test_pawn_movements("""
-            * . .
-            1 2 *
-            * . .
-        """)
-
-        self._test_pawn_movements("""
-            * 1 *
-            . 2 .
-            . * .
-        """)
-
-        self._test_pawn_movements("""
-            . . *
-            * 2 1
-            . . *
-        """)
-
-        self._test_pawn_movements("""
-            . * .
-            . 2 .
-            * 1 *
-        """)
-
-    def test_jumps_with_walls(self):
-        self._test_pawn_movements("""
-            * *|.
-            1 2|.
-            * * .
-        """)
-
-        self._test_pawn_movements("""
-           . * *|.
-           * 1 2|.
-             - -
-           . . . .
-           . . . .
-        """)
-
-        self._test_pawn_movements("""
-           . . *|.
-           -+-  +
-           * 1 2|.
-             -+-
-           . . . .
-           . . . .
-        """)
-
-    def test_forbidden_walls(self):
-        self._test_wall_placements("""
-            . . 1 . .
-            . . .v. .
-            . . .|. .
-                >
-            . . .|. .
-            . . 2 . .
-        """)
-
-        self._test_wall_placements("""
-            . . 1 . .
-            . . .v. .
-              > - -
-            . . . . .
-            . . . . .
-            . . 2 . .
-        """)
-
-        self._test_wall_placements("""
-            . .v1 .v.
-            . .|.v.|.
-              >+-+-+
-            . .|. .|.
-            . . . . .
-            . . 2 . .
-        """)
-
-    def test_forbidden_walls_due_to_blocking(self):
-        self._test_wall_placements("""
-            . .|1|. .
-              > >
-            . .|.|. .
-              > >
-            . . . . .
-            . . . . .
-            . . 2 . .
-        """)
-
-        self._test_wall_placements("""
-            . . 1v.v.
-            .v. .v. .
-            -+- -+-
-            . . . . .
-            . . . . .
-            . . 2 . .
-        """)
-
-        self._test_wall_placements("""
-            . . 1 . . .
-            . . . . . .
-            .v. .v. . .
-            -+- -+- >
-            . . . . . .
-            . . . . . .
-            . . 2 . . .
-        """)
-
-        self._test_wall_placements(
-            """
-            . 1 . . .
-              -+-
-            . . . . .
-              -+- -+-
-            .|. 2 . .
-                -+-
-            .|. .|.|.
-            >
-            . . .|.|.
-        """,
-            True,
-        )
diff --git a/deep_quoridor/test/quoridor_test.py b/deep_quoridor/test/quoridor_test.py
index d92f057d..a3c0846a 100644
--- a/deep_quoridor/test/quoridor_test.py
+++ b/deep_quoridor/test/quoridor_test.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
-from quoridor import Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
+from quoridor import ActionEncoder, Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
+from quoridor_env import env as quoridor_env
 
 
 def parse_board(board):
@@ -57,8 +58,8 @@ def parse_board(board):
                     # HACK: Since each wall has two "|" characters, this tries to place each wall twice, and
                     # the second time is off by one position. Checking whether we can place it stops us
                     # placing the second wall.
-                    if board.can_place_wall(np.array((row_n - 1, col_n)), WallOrientation.HORIZONTAL):
-                        board.add_wall(Player.ONE, np.array((row_n - 1, col_n)), WallOrientation.HORIZONTAL)
+                    if board.can_place_wall(Player.ONE, (row_n - 1, col_n), WallOrientation.HORIZONTAL):
+                        board.add_wall(Player.ONE, (row_n - 1, col_n), WallOrientation.HORIZONTAL)
 
                 if ch == ">" or ch == "-":
                     forbidden_walls.append((row_n - 1, col_n, WallOrientation.HORIZONTAL))
@@ -77,16 +78,16 @@ def parse_board(board):
                         # HACK: Since each wall has two "|" characters, this tries to place each wall twice, and
                         # the second time is off by one position. Checking whether we can place it stops us
                         # placing the second wall.
-                        if board.can_place_wall(np.array((row_n, col_n - 1)), WallOrientation.VERTICAL):
-                            board.add_wall(Player.ONE, np.array((row_n, col_n - 1)), WallOrientation.VERTICAL)
+                        if board.can_place_wall(Player.ONE, (row_n, col_n - 1), WallOrientation.VERTICAL):
+                            board.add_wall(Player.ONE, (row_n, col_n - 1), WallOrientation.VERTICAL)
                     continue
 
                 if ch == "*":
-                    potential_moves.append(np.array((row_n, col_n)))
+                    potential_moves.append((row_n, col_n))
                 elif ch == "1":
-                    board.move_player(Player.ONE, np.array((row_n, col_n)))
+                    board.move_player(Player.ONE, (row_n, col_n))
                 elif ch == "2":
-                    board.move_player(Player.TWO, np.array((row_n, col_n)))
+                    board.move_player(Player.TWO, (row_n, col_n))
 
                 col_positions[i] = col_n
                 col_n += 1
@@ -106,25 +107,116 @@ def _test_pawn_movements(self, s):
         game_moves = []
         for row in range(0, game.board.board_size):
             for col in range(0, game.board.board_size):
-                position = np.array((row, col))
+                position = (row, col)
                 if game.is_action_valid(MoveAction(position)):
                     game_moves.append(position)
 
-        np.testing.assert_equal(np.array(game_moves), np.array(potential_moves))
+        assert game_moves == potential_moves
 
-    def _test_wall_placements(self, s):
+        env = quoridor_env(board_size=game.board.board_size, max_walls=game.board.max_walls, game_start_state=game)
+        N = game.board.board_size
+        action_mask = env.observe("player_0")["action_mask"]
+
+        env_moves = []
+        for i, value in enumerate(action_mask[: N**2]):
+            if value == 1:
+                env_moves.append(divmod(i, N))
+
+        assert env_moves == potential_moves
+
+    def _test_wall_placements(self, s, just_highlighted=False):
         game, _, forbidden_walls = parse_board(s)
+        N = game.board.board_size
+        action_encoder = ActionEncoder(N)
         print(str(game))
 
         game_walls = []
         for orientation in [WallOrientation.HORIZONTAL, WallOrientation.VERTICAL]:
             for row in range(0, game.board.board_size - 1):
                 for col in range(0, game.board.board_size - 1):
-                    if not game.is_action_valid(WallAction(np.array((row, col)), orientation)):
+                    if not game.is_action_valid(WallAction((row, col), orientation)):
                         game_walls.append((row, col, orientation))
 
         assert set(game_walls) == set(forbidden_walls)
 
+        env = quoridor_env(board_size=game.board.board_size, max_walls=game.board.max_walls, game_start_state=game)
+        action_mask = env.observe("player_0")["action_mask"]
+
+        env_walls = []
+        for i in range(N**2, len(action_mask)):
+            if action_mask[i] == 0:
+                wall_action = action_encoder.index_to_action(i)
+
+                env_walls.append((wall_action.position[0], wall_action.position[1], wall_action.orientation))
+
+        if just_highlighted:
+            diff = set(forbidden_walls).difference(set(env_walls))
+            assert not diff
+        else:
+            assert set(env_walls) == set(forbidden_walls)
+
+    def _test_distance_to_target(self, s, moves_p1, moves_p2):
+        game, _, _ = parse_board(s)
+        N = game.board.board_size
+        player_one_position = game.board.get_player_position(Player.ONE)
+        assert moves_p1 == game.distance_to_target(player_one_position, N - 1)
+        player_two_position = game.board.get_player_position(Player.TWO)
+        assert moves_p2 == game.distance_to_target(player_two_position, 0)
+
+    def test_distance_to_goal(self):
+        self._test_distance_to_target(
+            """
+            1 . .
+            . . .
+            . . 2
+        """,
+            2,
+            2,
+        )
+
+        self._test_distance_to_target(
+            """
+            . * .
+            . . .
+            2 . 1
+        """,
+            0,
+            2,
+        )
+
+        self._test_distance_to_target(
+            """
+            1 . .
+            - -
+            . . .
+            . . 2
+        """,
+            4,
+            2,
+        )
+
+        self._test_distance_to_target(
+            """
+            1 .|.
+            - -+
+            . .|.
+            . . 2
+        """,
+            -1,
+            2,
+        )
+
+        self._test_distance_to_target(
+            """
+            1 . .
+            - -
+            . .|.
+            . .|2
+        """,
+            4,
+            2,
+        )
+
     def test_corner_movements(self):
         self._test_pawn_movements("""
             1 * .
@@ -265,7 +357,6 @@ def test_forbidden_walls(self):
             . . 2 . .
         """)
 
-    @pytest.mark.skip(reason="Disabled until Aaron implements blocking testing in Quoridor class")
     def test_forbidden_walls_due_to_blocking(self):
         self._test_wall_placements("""
             . .|1|. .


Reviews:
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-09T00:58:03Z
user: login: alejandromarcu
body: I still want to look a bit more, but I left some comments.  One thing that I don't like much is to use np.array for Position, I think a simple tuple would be nicer.
Also, I can help you to migrate the agents.
Does HumanAgent work? I think playing human vs human can be a good sanity test (you can do it by running `python src/play.py --players human human -r pygame`)
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T22:31:12Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T00:41:02Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T00:45:30Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T01:14:12Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T03:31:59Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T03:39:48Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-10T03:42:00Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T03:53:23Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T04:06:17Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T04:18:54Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T04:38:48Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-11T19:23:24Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-11T19:25:05Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-11T19:27:07Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-12T21:20:19Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-12T21:23:51Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-16T20:38:52Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-16T20:54:05Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T01:21:06Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T01:21:18Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T01:21:31Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T10:28:13Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-19T21:54:02Z
user: login: muralx
body: I think the approach you have taken will present some issues, in particular when env is transformed and agents need to inspect the environment or interact with it. For instance creating a new quoridor game won't work for an agent that is working on a transformed environment, it would need all the transformations for making the observation valid for constructing a new game. 

I left some initial round of comments, for arena and the interaction with the agents. I will look at it more later in the day.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T12:29:33Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-26T01:31:28Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-26T01:33:17Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:16:45Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:17:39Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:19:39Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:22:25Z
user: login: muralx
body: Thanks @jonbinney, I left some additional comments. Many of them are around the transformations, but you can add TODOs in there so we cover those when we add adapter or find a way of using per agent env wrappers.

I think there are a few things about naming that would be good to address, and some actual potential bug that I highlighted in the previous review (I might be missing something though).

I did not review greedy or simple (I will let @alejandromarcu to check those since he is more familiar with the code)
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:31:36Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T18:39:52Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T00:32:52Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T00:33:44Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T00:36:28Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T00:38:46Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:16:07Z
user: login: muralx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-27T19:21:22Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:05:20Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:09:19Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:11:22Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:12:08Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:31:14Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-27T23:41:02Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:35:42Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:52:54Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:53:20Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:55:38Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:56:24Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:56:57Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:57:37Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:58:49Z
user: login: jonbinney
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-28T00:59:27Z
user: login: muralx
body: Thanks @jonbinney, LGTM but make sure of getting @alejandromarcu to review the pieces he knows better, I have not looked at them.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-28T21:16:58Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T00:12:15Z
user: login: alejandromarcu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T00:35:32Z
user: login: alejandromarcu
body: Looks good! If you want just merge and I'll send a fix for the human player.

state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-30T00:51:39Z

Review comments:
diff_hunk: @@ -155,14 +191,17 @@ def is_wall_between(self, position1: Position, position2: Position) -> bool:
         assert position1_on_board or position2_on_board, "At least one position must be on the board"
 
         if position1_on_board and position2_on_board:
-            wall_position = position1 + position2
-            return self._grid[*wall_position] == Board.PLAYER1_WALL
+            wall_position = position1 + position2 + 2
+            return self._grid[*wall_position] == Board.WALL
         else:
             # By convention we treat the border as a "wall". This is makes checking jumps more convenient, since
             # players are allowed to jump diagonally if they are adjacent to another player and the border of the
             # board is on the other side of that player.
             return True
 
+    def get_old_style_walls(self):
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: For now the Board class keeps track of the walls in both the new combined grid as well as the `(board_size-1) x (board_size -1) x 2` grid so that agents can use whichever they want.
created_at: 2025-04-09T00:58:03Z
author_association: OWNER
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
         self._walls_remaining = [self.max_walls, self.max_walls]
 
+        self._potential_wall_neighbors = {
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: Can you add a comment explaining what this is?
created_at: 2025-04-09T21:44:55Z
author_association: COLLABORATOR
diff_hunk: @@ -109,7 +124,14 @@ def get_player_cell(self, position: Position) -> int:
         if not self.is_position_on_board(position):
             raise ValueError(f"Position {position} is out of bounds")
 
-        return self._grid[*position * 2]
+        return self._grid[*position * 2 + 2]
+
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        assert isinstance(player, Player)
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: I know that technically you could ignore the hint checker and pass whatever and the program would run and this would catch it, but I think that using both is redundant, I'd favor the type checker when possible (e.g. here we can assume that if we indicated that `player` is of type `Player`, this will be true)
created_at: 2025-04-09T21:47:08Z
author_association: COLLABORATOR
diff_hunk: @@ -181,22 +234,76 @@ def _get_wall_slice(self, position: Position, orientation: WallOrientation) -> t
         assert isinstance(orientation, WallOrientation)
 
         if orientation == WallOrientation.VERTICAL:
-            wall_slice = (slice(position[0] * 2, position[0] * 2 + 3), position[1] * 2 + 1)
+            wall_slice = (slice(position[0] * 2 + 2, position[0] * 2 + 5), position[1] * 2 + 3)
         elif orientation == WallOrientation.HORIZONTAL:
-            wall_slice = (position[0] * 2 + 1, slice(position[1] * 2, position[1] * 2 + 3))
+            wall_slice = (position[0] * 2 + 3, slice(position[1] * 2 + 2, position[1] * 2 + 5))
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: I think you could take advantage of the helper function and do something like this, so it's easier to read and make sense of it:
```
        r, c = self. _wall_position_to_grid_index(position, orientation)
        if orientation == WallOrientation.VERTICAL:
              wall_slice = (slice(r, r + 3), c)
        elif orientation == WallOrientation.HORIZONTAL:
               wall_slice = (r, slice(c, c + 3))
```



created_at: 2025-04-09T21:53:11Z
author_association: COLLABORATOR
diff_hunk: @@ -269,21 +377,44 @@ def is_action_valid(self, action: Action):
                 is_valid = False
 
         elif isinstance(action, WallAction):
-            is_valid = self.board.can_place_wall(action.position, action.orientation)
-
+            if self.board.can_place_wall(action.position, action.orientation):
+                if self.board._is_wall_potential_block(action.position, action.orientation):
+                    # Temporarily place the wall so that we can check player paths to their goals.
+                    self.board.add_wall(self.current_player, action.position, action.orientation)
+                    for p in [Player.ONE, Player.TWO]:
+                        if not self.can_reach(self.board.get_player_position(p), self.get_goal_row(p)):
+                            is_valid = False
+                            break
+                    # Restore the board to its previous state
+                    self.board.remove_wall(self.current_player, action.position, action.orientation)
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: Yey, this check looks very elegant now!
created_at: 2025-04-09T21:55:54Z
author_association: COLLABORATOR
diff_hunk: @@ -27,20 +27,37 @@
 import copy
 import functools
 import random
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
 from gymnasium import spaces
 from pettingzoo import AECEnv
 from pettingzoo.utils import wrappers
 
+from deep_quoridor.src.quoridor import Action, Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
path: deep_quoridor/src/quoridor_env.py
user: login: alejandromarcu
body: using `deep_quoridor.src.` breaks when I try to run it from the cmd line, I think it should work without it
```suggestion
from quoridor import Action, Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
```
created_at: 2025-04-09T21:57:16Z
author_association: COLLABORATOR
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: alejandromarcu
body: I think it would be cleaner to define an index method in `Action` (with NotImplementedError) and override it in each subclass, something like this:

```
class MoveAction(Action):
 # etc
   @override
   def index(self, board_size, wall_size)
      return self.destination[0] * board_size + self.destination[1] 
```


Then, action_params_to_index would just return `action.index(self.board_size, self.wall_size)` (or you can even not declare action_params_to_index and use that every time)
created_at: 2025-04-09T22:09:54Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: Should be removed, no?
created_at: 2025-04-09T22:21:26Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: I'm not a fan of position being an np array.  I know it lets you do some math, but I'd prefer to use a simple int,int tupe and use helper functions when needed, e.g. here it could be `self._set_player_cell(position, player)`.
created_at: 2025-04-09T22:23:01Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
         self._walls_remaining = [self.max_walls, self.max_walls]
 
+        self._potential_wall_neighbors = {
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: Also, why do you need np.array and not just lists? if you're just iterating a list would be enough
created_at: 2025-04-09T22:23:27Z
author_association: COLLABORATOR
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: jonbinney
body: The design now is that `quoridor.py` implements the game logic and then `quoridor_env.py` wraps it and knows about things like action masks and indexes into them. It seems icky to have `MoveAction` include a function which is specific to the way the Environment creates indices.
created_at: 2025-04-10T00:41:02Z
author_association: OWNER
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: I agree - since it would involve lots of little changes across `quoridor.py` I was avoiding including that change in this PR, but it's probably worth it. Do you like a tuple more than a `Position` class with row and col members?
created_at: 2025-04-10T00:45:30Z
author_association: OWNER
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: jonbinney
body: I've got an idea on how to do this nicely..... trying it out now.
created_at: 2025-04-10T01:14:12Z
author_association: OWNER
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: jonbinney
body: ok, here's what I'm thinking:

I don't think the Action class should be tied to a specific action->index encoding - that seems too tightly coupled. But agents shouldn't need to use a member function of the environment to do the translation between actions and indices either - that also seems like too tight of a coupling. How about having an ActionEncoder class, which takes the board size as an argument at initialization time, and has two functions: `action_to_index(self, action)` and `index_to_action(self, index)`. Each agent would instantiate one and use it. 

   The environment would also use this class out of convenience. But really, the interface between agents and the environment should be limited to a shared agreement on what the observation arrays mean and what each action index does - the fact that they both use the Quoridor class and the ActionEncoder is an implementation detail, but an agent could choose to use some other implementation of the game, so long as it had the same idea of what each action index does.

  The tight coupling, where agents are passed the entire environment and can call its functions, has made this PR more difficult, so I think it would be good to remove it.  I'll add a way to construct a Quoridor object from an observation, rather than passing the entire Quoridor object in the observation, or having the agent access it through the environment. Then get_action would only need to take the observation, instead of taking the environment and calling `.last()` on it. I'm not sure whether the agents need to know about termination and truncation.... right now they just return `None` in those cases and do nothing... but that might not always be the case?

What do you think?
created_at: 2025-04-10T03:31:59Z
author_association: OWNER
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
         self._walls_remaining = [self.max_walls, self.max_walls]
 
+        self._potential_wall_neighbors = {
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: Ah yes I'll add a comment. They're numpy arrays so that I can do slicing and math with them in the `_is_wall_potential_block` function. In this case it doesn't matter so much, but generally doing numpy operations on multiple values at once is much faster than iterating in python, since it all gets done in compiled C code in numpy.
created_at: 2025-04-10T03:39:48Z
author_association: OWNER
diff_hunk: @@ -109,7 +124,14 @@ def get_player_cell(self, position: Position) -> int:
         if not self.is_position_on_board(position):
             raise ValueError(f"Position {position} is out of bounds")
 
-        return self._grid[*position * 2]
+        return self._grid[*position * 2 + 2]
+
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        assert isinstance(player, Player)
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: Does the hint checker give some feedback in vscode? I might not have things configured correctly in my settings.
created_at: 2025-04-10T03:42:00Z
author_association: OWNER
diff_hunk: @@ -109,7 +124,14 @@ def get_player_cell(self, position: Position) -> int:
         if not self.is_position_on_board(position):
             raise ValueError(f"Position {position} is out of bounds")
 
-        return self._grid[*position * 2]
+        return self._grid[*position * 2 + 2]
+
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        assert isinstance(player, Player)
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: Yeah, it should give you, for example if I type

```
def meh(x : int):
    pass

meh("hey")
```
it underlines "hey" and says it's the wrong type. If it doesn't, try installing the pylance extension: https://marketplace.visualstudio.com/items/?itemName=ms-python.vscode-pylance 
created_at: 2025-04-10T03:53:22Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: I think the easiest would be just a simple tuple (r,c) (i.e `tuple[int, int]`).
created_at: 2025-04-10T04:06:17Z
author_association: COLLABORATOR
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: alejandromarcu
body: I like your idea of the ActionEncoder class, it will be elegant for decoupling things.  It's been frustrating that the agent receives an action_mask but it needs to understand what it means, or that it knows what move it wants to make but needs to encode it in an index. 
Also creating a Quoridor object from an observation would be great if it's not too complex to implement or too inefficient
created_at: 2025-04-10T04:18:54Z
author_association: COLLABORATOR
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: alejandromarcu
body: BTW, I'm not sure about returning None, I'd think that since in the Arena we check for that case already, we won't encounter it in the agents, I don't know why we are using it TBH, maybe we got it from an example that was doing things differently? or I can be missing something
created_at: 2025-04-10T04:38:48Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: A couple downsides with that:
- MoveAction, WallAction, Player are all classes. feels a bit odd for position to be treated differently
- can't do vector math, e.g. `position + offset`
- can't modify the `__str__(self)` function to pretty-print positions

But I'm not against it. I do think that change is outside of the scope of this PR though.
created_at: 2025-04-11T19:23:24Z
author_association: OWNER
diff_hunk: @@ -386,69 +268,42 @@ def observation_space(self, agent):
     def action_space(self, agent):
         return spaces.Discrete(self.board_size**2 + (self.wall_size**2) * 2, seed=random.randint(0, 2**32 - 1))
 
-    def action_params_to_index(self, row: int, col: int, action_type: int = 0) -> int:
+    def action_params_to_index(self, action) -> int:
         """
-        Takes action parameters (row, col, movement_type) to an action index
-        movement_type = 0 for moving, 1 for horizontal wall placement, 2 for vertical wall placement
+        Converts an action object to an action index
         """
-        if action_type == 0:  # Pawn movement
-            return row * self.board_size + col
-        elif action_type == 1:  # Placing a vertical wall
-            return self.board_size**2 + row * self.wall_size + col
-        elif action_type == 2:  # Placing a horizontal wall
-            return self.board_size**2 + self.wall_size**2 + row * self.wall_size + col
+        if isinstance(action, MoveAction):
+            return action.destination[0] * self.board_size + action.destination[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.VERTICAL:
+            return self.board_size**2 + action.position[0] * self.wall_size + action.position[1]
+        elif isinstance(action, WallAction) and action.orientation == WallOrientation.HORIZONTAL:
+            return self.board_size**2 + self.wall_size**2 + action.position[0] * self.wall_size + action.position[1]
path: deep_quoridor/src/quoridor_env.py
user: login: jonbinney
body: Yeah i think you're right that the agents don't need to worry about termination/truncation and that can be handled int the arena class. I've pushed changes that do that with the Greedy agent, and `play.py` works fine.
created_at: 2025-04-11T19:25:05Z
author_association: OWNER
diff_hunk: @@ -109,7 +124,14 @@ def get_player_cell(self, position: Position) -> int:
         if not self.is_position_on_board(position):
             raise ValueError(f"Position {position} is out of bounds")
 
-        return self._grid[*position * 2]
+        return self._grid[*position * 2 + 2]
+
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        assert isinstance(player, Player)
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: I do have pylance, but i've never seen that warning. Maybe i don't make that mistake? :-P
How would it catch more complicated problems though? Like:
```
x = None
if ...:
    x = right_type
else:
   x = wrong_type
foo(x)
```
created_at: 2025-04-11T19:27:07Z
author_association: OWNER
diff_hunk: @@ -175,35 +180,36 @@ def _log_action(
         path.insert(0, (int(opp_coords[0][0]), int(opp_coords[0][1])))
         self.action_log.path(path)
 
-    def get_action(self, game):
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
+    def get_action(self, observation, action_mask):
+        # Reconstruct the game from the observation.
+        game = Quoridor(Board(from_observation=observation))
 
         if random.random() < self.params.p_random:
             if self.action_log.is_enabled():
                 self.action_log.clear()
                 # TO DO: when we have the functionality in the log to output a message,
                 # use it to say that the move will be random.
-            return game.action_space(game.agent_selection).sample(observation["action_mask"])
+            return self.action_space.sample(action_mask)
 
-        goal_row = game.board_size - 1 if self.player_id == "player_0" else 0
-        opponent_row = game.board_size - 1 - goal_row
+        goal_row = self.board_size - 1 if self.player_id == "player_0" else 0
+        opponent_goal_row = self.board_size - 1 - goal_row
path: deep_quoridor/src/agents/greedy.py
user: login: alejandromarcu
body: we could use game.get_goal_row, no?
created_at: 2025-04-12T21:20:19Z
author_association: COLLABORATOR
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: alejandromarcu
body: The thing with using a class is that instead of just use for example `(1,2)` you need to use `Position(1,2)`. But it's nice to be able to have functions for offset or __str__.  We can discuss this more after this PR
created_at: 2025-04-12T21:23:51Z
author_association: COLLABORATOR
diff_hunk: @@ -109,7 +124,14 @@ def get_player_cell(self, position: Position) -> int:
         if not self.is_position_on_board(position):
             raise ValueError(f"Position {position} is out of bounds")
 
-        return self._grid[*position * 2]
+        return self._grid[*position * 2 + 2]
+
+    def get_walls_remaining(self, player: Player) -> int:
+        """
+        Get the number of walls remaining for the player.
+        """
+        assert isinstance(player, Player)
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: I went ahead and removed the type checking asserts. I'll have to figure out why pylance doesn't catch passing arguments with the wrong type.
created_at: 2025-04-16T20:38:52Z
author_association: OWNER
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
     FREE = -1
     # The player numbers start at 0 so that they can also be used as indices into the player positions array.
     PLAYER1 = 0
     PLAYER2 = 1
-    # Walls are stored as 10 + player number.We leave 2-9 free for implementing more than 2 players in the future.
-    PLAYER1_WALL = 10
-    PLAYER2_WALL = 11
+    WALL = 10
 
     def __init__(self, board_size: int = 9, max_walls: int = 10):
         self.board_size = board_size
         self.max_walls = max_walls
 
-        self._grid = np.full((board_size * 2 - 1, board_size * 2 - 1), Board.FREE, dtype=np.int8)
+        # We represent the board as a grid of cells alternating between wall cells and odd rows are player cells.
+        # To make some checks easier, we add a double border of walls around the grid.
+        self._grid = np.full((board_size * 2 + 3, board_size * 2 + 3), Board.FREE, dtype=np.int8)
+        self._grid[:2, :] = Board.WALL
+        self._grid[-2:, :] = Board.WALL
+        self._grid[:, :2] = Board.WALL
+        self._grid[:, -2:] = Board.WALL
+
+        # We also keep track of the walls in a (row x col x orientation) array so that we can give that
+        # representation to agents that expect it.
+        self._old_style_walls = np.zeros((self.board_size - 1, self.board_size - 1, 2), dtype=np.int8)
+
         self._players = [Player.ONE, Player.TWO]
         self._player_positions = [np.array([0, board_size // 2]), np.array([board_size - 1, board_size // 2])]
         for player, position in zip(self._players, self._player_positions):
-            self._grid[*position * 2] = player
+            self._grid[*(position * 2 + 2)] = player
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: ok i actually did go ahead and change from array to tuple. We can still use numpy functions, but have to be a bit careful because they change the type. For example `np.subtract(x, y)` returns an ndarray even if x and y are tuples.
created_at: 2025-04-16T20:54:05Z
author_association: OWNER
diff_hunk: @@ -181,22 +234,76 @@ def _get_wall_slice(self, position: Position, orientation: WallOrientation) -> t
         assert isinstance(orientation, WallOrientation)
 
         if orientation == WallOrientation.VERTICAL:
-            wall_slice = (slice(position[0] * 2, position[0] * 2 + 3), position[1] * 2 + 1)
+            wall_slice = (slice(position[0] * 2 + 2, position[0] * 2 + 5), position[1] * 2 + 3)
         elif orientation == WallOrientation.HORIZONTAL:
-            wall_slice = (position[0] * 2 + 1, slice(position[1] * 2, position[1] * 2 + 3))
+            wall_slice = (position[0] * 2 + 3, slice(position[1] * 2 + 2, position[1] * 2 + 5))
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: Great idea. I've updated the code.
created_at: 2025-04-17T01:21:06Z
author_association: OWNER
diff_hunk: @@ -27,20 +27,37 @@
 import copy
 import functools
 import random
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
 from gymnasium import spaces
 from pettingzoo import AECEnv
 from pettingzoo.utils import wrappers
 
+from deep_quoridor.src.quoridor import Action, Board, MoveAction, Player, Quoridor, WallAction, WallOrientation
path: deep_quoridor/src/quoridor_env.py
user: login: jonbinney
body: Fixed
created_at: 2025-04-17T01:21:18Z
author_association: OWNER
diff_hunk: @@ -49,35 +50,49 @@ def is_valid_position_type(position: Position) -> bool:
 
 
 class Board:
-    """
-    Grid which has both walls and possible pawn positions.
-
-    Pawn positions are on even rows and columns, walls are on odd rows and columns. We use negative
-    values to represent free cells and walls, 0 for player 1, and 1 for player 2.
-    integers.
-    """
-
     # Possible values for each cell in the grid.
     OUT_OF_BOUNDS = -2  # Returned if a move is off of the board.
path: deep_quoridor/src/quoridor.py
user: login: jonbinney
body: Yep, removed now.
created_at: 2025-04-17T01:21:31Z
author_association: OWNER
diff_hunk: @@ -94,6 +95,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: Just a comment in this. The encoder really depends on the observations, env wrapper being used. It should not really depend on the internal implementation of the game.
created_at: 2025-04-17T10:28:13Z
author_association: COLLABORATOR
diff_hunk: @@ -94,6 +95,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: True - it only depends on the rest of `quoridor.py` because it uses the MoveAction and WallAction classes. I'd like to keep it out of `quoridor_env.py` because it really shouldn't depend on anything in that file either, but maybe it should get its own file?
created_at: 2025-04-19T21:54:02Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
path: deep_quoridor/src/arena.py
user: login: muralx
body: I am not sure this would work with env wrappers. It is probably better to add a method in the wrapper and the env and force wrappers that modify the action mask to store the last action mask after transforming it.
created_at: 2025-04-25T12:01:07Z
author_association: COLLABORATOR
diff_hunk: @@ -104,7 +133,14 @@ def _play_games(self, players: list[str | Agent], times: int, mode: PlayMode):
                 agents.append(p)
             else:
                 agents.append(
-                    AgentRegistry.create_from_encoded_name(p, board_size=self.board_size, max_walls=self.max_walls)
+                    AgentRegistry.create_from_encoded_name(
+                        p,
+                        board_size=self.board_size,
+                        max_walls=self.max_walls,
+                        action_space=self.game.action_space(
+                            None
+                        ),  # Action space doesn't depend on the player so we just pass None
path: deep_quoridor/src/arena.py
user: login: muralx
body: It might too, for instance. If I want to represent movement using (1,0) (-1,0) ... to indicate the direction of the movement, action space could be different. I think for now we don't have agent wrappers, so it won't break anything.
created_at: 2025-04-25T12:04:08Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
path: deep_quoridor/src/arena.py
user: login: muralx
body: I am not sure this is correct, in general we want to have (s_n, action, s_n+1, rewards), but in here we have (s_n+1, action, rewards). 
created_at: 2025-04-25T12:13:02Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
path: deep_quoridor/src/arena.py
user: login: muralx
body: Since last item added to the observations is s_n+1, this is not the obsertation_before_action, but after action ?
created_at: 2025-04-25T12:13:51Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
path: deep_quoridor/src/arena.py
user: login: muralx
body: maybe make this name explicit as opponent_observation_after_action ?
created_at: 2025-04-25T12:14:35Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=observations[-1][3],
+                    action=action,
+                    agent_id=agent_id,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    observation_before_action=observations[-1][1],
path: deep_quoridor/src/arena.py
user: login: muralx
body: naming here will be confusing. Since we are calling the agent, we should name this from the POV of the agent being called.  So this should be opponent_observation_before_action, similar for the rest.
created_at: 2025-04-25T12:17:57Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=observations[-1][3],
+                    action=action,
+                    agent_id=agent_id,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
path: deep_quoridor/src/arena.py
user: login: muralx
body: This could be my_observation_after_opponents_actions?
created_at: 2025-04-25T12:18:27Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
path: deep_quoridor/src/arena.py
user: login: muralx
body: I don't think you need this. It is harder to see what's going on and in general you only need the current states and not past. See my other comment for game termination.
created_at: 2025-04-25T12:21:46Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
path: deep_quoridor/src/arena.py
user: login: muralx
body: For this case you can pass none for observation after_action, etc, since there is no Action in this  case. Maybe we can later refactor this to be a separate method that says, handle_opponent_victory which only receives the info it needs.
created_at: 2025-04-25T12:23:36Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=observations[-1][3],
+                    action=action,
+                    agent_id=agent_id,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    observation_before_action=observations[-1][1],
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Yes i was wondering about this! Ok that makes sense.
created_at: 2025-04-26T01:31:28Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: You mean the list of past observations? I added that to make it easier to figure out what went wrong when a game gets into a bad state and you catch it in the debugger. I'm ok with removing it though, and I think it repeats some information that is stored in other places.
created_at: 2025-04-26T01:33:16Z
author_association: OWNER
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: maybe also name it as opponent_observation_after_action
created_at: 2025-04-26T18:16:44Z
author_association: COLLABORATOR
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: This should not be passed ideally. Since this is the outcome self, self.player_id should be use internally.
created_at: 2025-04-26T18:17:39Z
author_association: COLLABORATOR
diff_hunk: @@ -185,24 +222,24 @@ def handle_step_outcome_all(self, observation_before_action, action, game, playe
                     return reward
             return 0
 
-        state_before_action = self.observation_to_tensor(observation_before_action, player_id)
-        state_after_action = self.observation_to_tensor(game.observe(player_id), player_id)
-        opponent_state = game.observe(self.get_opponent_player_id(player_id))
+        state_before_action = self.observation_to_tensor(observation_before_action["observation"], agent_id)
+        state_after_action = self.observation_to_tensor(observation_after_action["observation"], agent_id)
+        # opponent_state = game.observe(self.get_opponent_player_id(player_id))
+        opponent_state = opponent_observation
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: nit, just use opponent_observation.
created_at: 2025-04-26T18:19:39Z
author_association: COLLABORATOR
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: nit: note this is the player_id for the game and not the id of the agent. If we prefer using agent id, maybe it will be better to rename all other variables and methods to agent_id too. Though I think an id is kind of fixed, and should not change from match to match.
created_at: 2025-04-26T18:22:24Z
author_association: COLLABORATOR
diff_hunk: @@ -259,19 +296,13 @@ def update_target_network(self):
         """Copy parameters from online network to target network."""
         self.target_network.load_state_dict(self.online_network.state_dict())
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: is this action mask also in the observation? Note that the use of observation above includes the action mask. Not sure how to name the two version of observations so we can understand what we have in each one.. 
created_at: 2025-04-26T18:24:20Z
author_association: COLLABORATOR
diff_hunk: @@ -94,6 +95,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: I guess my point was similar to the other comments. If observation is changed, which currently should be handled OK by the agent (assuming wrapper has the action encoder function), it won't no longer work after this change since action encoder is different depending on the transformations being used.
created_at: 2025-04-26T18:26:50Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
path: deep_quoridor/src/arena.py
user: login: muralx
body: I thought you were using it for keeping state in between steps. 

It does add some readability challenges into what is being passed to the methods, so maybe if you want to keep it like this, you could extract the values to variables so it is easier to see the intention of the code and detect potential future bugs.


created_at: 2025-04-26T18:39:52Z
author_association: COLLABORATOR
diff_hunk: @@ -259,19 +296,13 @@ def update_target_network(self):
         """Copy parameters from online network to target network."""
         self.target_network.load_state_dict(self.online_network.state_dict())
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: This observations is `observation["observation"]`, which doesn't include the action mask. I don't like that pettingzoo allows environments to put the action mask in either the observation or the info dict. Agree that `observation["observation"]` is a confusing naming convention.
created_at: 2025-04-27T00:32:52Z
author_association: OWNER
diff_hunk: @@ -94,6 +95,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: Ah, yes that's a problem.
created_at: 2025-04-27T00:33:44Z
author_association: OWNER
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: Makes sense. `handle_step_outcome_all` does need to have the agent_id passed in though, right?

Side note: i've started trying to be more consistent about using `agent_id` for the string id of the agent, and `player` terminology within the Quoridor class.
created_at: 2025-04-27T00:36:28Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: In that case is`s_n+1` the state after this action, or the state after this action and the opponent's action?
created_at: 2025-04-27T00:38:46Z
author_association: OWNER
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: muralx
body: yes, since it is called from two different pov. Fine with any naming convention, lets just keep them consistent.
created_at: 2025-04-27T19:16:06Z
author_association: COLLABORATOR
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
path: deep_quoridor/src/arena.py
user: login: muralx
body: s_n+1 is the state after the action. For each step executed, there is always a s_i and s_i+1. In this case for the player that selected the action passed to the step method
created_at: 2025-04-27T19:21:22Z
author_association: COLLABORATOR
diff_hunk: @@ -175,35 +180,36 @@ def _log_action(
         path.insert(0, (int(opp_coords[0][0]), int(opp_coords[0][1])))
         self.action_log.path(path)
 
-    def get_action(self, game):
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
+    def get_action(self, observation, action_mask):
+        # Reconstruct the game from the observation.
+        game = Quoridor(Board(from_observation=observation))
 
         if random.random() < self.params.p_random:
             if self.action_log.is_enabled():
                 self.action_log.clear()
                 # TO DO: when we have the functionality in the log to output a message,
                 # use it to say that the move will be random.
-            return game.action_space(game.agent_selection).sample(observation["action_mask"])
+            return self.action_space.sample(action_mask)
 
-        goal_row = game.board_size - 1 if self.player_id == "player_0" else 0
-        opponent_row = game.board_size - 1 - goal_row
+        goal_row = self.board_size - 1 if self.player_id == "player_0" else 0
+        opponent_goal_row = self.board_size - 1 - goal_row
path: deep_quoridor/src/agents/greedy.py
user: login: jonbinney
body: Fixed now
created_at: 2025-04-27T23:05:20Z
author_association: OWNER
diff_hunk: @@ -104,7 +133,14 @@ def _play_games(self, players: list[str | Agent], times: int, mode: PlayMode):
                 agents.append(p)
             else:
                 agents.append(
-                    AgentRegistry.create_from_encoded_name(p, board_size=self.board_size, max_walls=self.max_walls)
+                    AgentRegistry.create_from_encoded_name(
+                        p,
+                        board_size=self.board_size,
+                        max_walls=self.max_walls,
+                        action_space=self.game.action_space(
+                            None
+                        ),  # Action space doesn't depend on the player so we just pass None
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: True. I've updated the comment to remind us to fix this if we add agent wrappers.
created_at: 2025-04-27T23:09:19Z
author_association: OWNER
diff_hunk: @@ -94,6 +95,7 @@ def __init__(
         self.training_mode = params.training_mode
         self.params = params
         self.action_size = self._calculate_action_size()
+        self.action_encoder = ActionEncoder(self.board_size)
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: As discussed below I'll leave this for now and we can change it when we start using wrappers.
created_at: 2025-04-27T23:11:22Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: As discussed below, I'll leave this in this PR and we can change it when we use wrappers.
created_at: 2025-04-27T23:12:08Z
author_association: OWNER
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: Done
created_at: 2025-04-27T23:31:14Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: ok, i understand. I've changed this to store `(s_n, action, s_n+1, rewards),`
created_at: 2025-04-27T23:41:01Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: I think you're right. I think I've fixed this during my switch back to individual variables instead of a big list.
created_at: 2025-04-28T00:35:41Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Good idea; done.
created_at: 2025-04-28T00:52:54Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=observations[-1][3],
+                    action=action,
+                    agent_id=agent_id,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    observation_before_action=observations[-1][1],
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Updated to use the POV of the agent being aclled.
created_at: 2025-04-28T00:53:20Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
 
             if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
-                agent.handle_step_outcome(observation, action, self.game)
+                agent.handle_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
+                    observation_after_action=self.game.observe(agent_id),
+                    reward=observations[-1][3],
+                    action=action,
+                    agent_id=agent_id,
+                    done=self.game.is_done(),
+                )
 
-            opponent_agent = agents[self.game.agent_selection]
-            if opponent_agent.is_trainable() and isinstance(opponent_agent, AbstractTrainableAgent):
-                opponent_agent.handle_opponent_step_outcome(observation, action, self.game)
+            if opponent_agent.is_trainable() and isinstance(opponent_agent_id, AbstractTrainableAgent):
+                opponent_agent.handle_opponent_step_outcome(
+                    observation_before_action=observations[-1][1],
+                    opponent_observation=self.game.observe(opponent_agent_id),
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Done
created_at: 2025-04-28T00:55:38Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Agreed about the readability problems; i've switched to individual variables instead of a list.
created_at: 2025-04-28T00:56:24Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: Done
created_at: 2025-04-28T00:56:56Z
author_association: OWNER
diff_hunk: @@ -50,32 +50,61 @@ def _play_game(self, agent1: Agent, agent2: Agent, game_id: str) -> GameResult:
         self.plugins.start_game(self.game, agent1, agent2)
         start_time = time.time()
         step = 0
-        for player_id in self.game.agent_iter():
+        observations = []  # list of tuples of (agent_id, observation, action)
+        for agent_id in self.game.agent_iter():
             observation, _, termination, truncation, _ = self.game.last()
-            agent = agents[player_id]
+            agent = agents[agent_id]
+            opponent_agent_id = self.game.get_opponent(agent_id)
+            opponent_agent = agents[agent_id]
+
             if termination or truncation:
                 if agent.is_trainable() and isinstance(agent, AbstractTrainableAgent):
                     # Handle end of game (in case winner was not this agent)
-                    agent.handle_step_outcome(observation, None, self.game)
+                    agent.handle_step_outcome(
+                        observation_before_action=observations[-2][1],
+                        opponent_observation=observations[-1][1],
+                        observation_after_action=observation,
+                        reward=self.game.rewards[agent_id],
+                        action=None,
+                        agent_id=agent_id,
+                        done=True,
+                    )
                 break
 
-            action = int(agent.get_action(self.game))
+            assert (observation["action_mask"] == self.game.last_action_mask[agent_id]).all()
+
+            action = int(agent.get_action(observation["observation"], observation["action_mask"]))
 
             self.plugins.before_action(self.game, agent)
             self.game.step(action)
+            observations.append((agent_id, self.game.observe(agent_id), action, self.game.rewards[agent_id]))
path: deep_quoridor/src/arena.py
user: login: jonbinney
body: I ended up using individual variables instead of a list.
created_at: 2025-04-28T00:57:37Z
author_association: OWNER
diff_hunk: @@ -185,24 +222,24 @@ def handle_step_outcome_all(self, observation_before_action, action, game, playe
                     return reward
             return 0
 
-        state_before_action = self.observation_to_tensor(observation_before_action, player_id)
-        state_after_action = self.observation_to_tensor(game.observe(player_id), player_id)
-        opponent_state = game.observe(self.get_opponent_player_id(player_id))
+        state_before_action = self.observation_to_tensor(observation_before_action["observation"], agent_id)
+        state_after_action = self.observation_to_tensor(observation_after_action["observation"], agent_id)
+        # opponent_state = game.observe(self.get_opponent_player_id(player_id))
+        opponent_state = opponent_observation
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body: Done
created_at: 2025-04-28T00:58:49Z
author_association: OWNER
diff_hunk: @@ -155,23 +157,58 @@ def get_opponent_player_id(self, player_id):
         """Get the opponent player ID."""
         return "player_1" if player_id == "player_0" else "player_0"
 
-    def handle_opponent_step_outcome(self, observation_before_action, action, game):
+    def handle_opponent_step_outcome(
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
+        done=False,
+    ):
         pass
 
-    def adjust_reward(self, r, game):
-        if game.is_done():
+    def adjust_reward(self, r, done):
+        if done:
             r *= self.final_reward_multiplier
         return r
 
-    def handle_step_outcome(self, observation_before_action, action, game):
+    def handle_step_outcome(
+        self,
+        observation_before_action,
+        opponent_observation,
+        observation_after_action,
+        reward,
+        action,
+        agent_id,
path: deep_quoridor/src/agents/core/trainable_agent.py
user: login: jonbinney
body:  Makes sense. I'm leaving it alone in this PR, but we will need to think a bit about what convention to use.
created_at: 2025-04-28T00:59:27Z
author_association: OWNER
diff_hunk: @@ -25,21 +27,22 @@ def name(self):
             return self.params.nick
         return "Human"
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
         # Import this here to avoid circular dependencies
         from renderers.pygame import PygameQuoridor
 
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
-
         # Transform the action mask into a set of valid moves
         valid_moves = set()
-        for action, value in enumerate(observation["action_mask"]):
+        for action, value in enumerate(action_mask):
             if value == 1:
-                r, c, type = game.action_index_to_params(action)
-                valid_moves.add((r, c, type))
+                valid_action = self.action_encoder.index_to_action(action)
+                if isinstance(valid_action, MoveAction):
+                    valid_moves.add((valid_action.destination[0], valid_action.destination[1], 0))
+                elif isinstance(valid_action, WallAction):
+                    valid_moves.add(
+                        (valid_action.position[0], valid_action.position[1], valid_action.orientation.value - 1)
+                    )
 
         result = PygameQuoridor.instance().get_human_input(valid_moves)
 
-        return None if result is None else game.action_params_to_index(*result)
+        return None if result is None else self.action_encoder.action_to_index(*result)
path: deep_quoridor/src/agents/human.py
user: login: alejandromarcu
body: When running the UI with human this fails:

```
(.venv) (base) amarcu@MacBookPro deep_rabbit_hole % /Users/amarcu/code/deep_rabbit_hole/.venv/bin/python /Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/play.py -yi human_vs_greedy                                                                                              
====< Running human_vs_greedy >====
Exception in thread Thread-1 (_play_games):
Traceback (most recent call last):
  File "/opt/homebrew/Cellar/python@3.12/3.12.9/Frameworks/Python.framework/Versions/3.12/lib/python3.12/threading.py", line 1075, in _bootstrap_inner
    self.run()
  File "/opt/homebrew/Cellar/python@3.12/3.12.9/Frameworks/Python.framework/Versions/3.12/lib/python3.12/threading.py", line 1012, in run
    self._target(*self._args, **self._kwargs)
  File "/Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/arena.py", line 160, in _play_games
    result = self._play_game(agent_1, agent_2, f"game_{match_id:04d}")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/arena.py", line 75, in _play_game
    agent.get_action(observation_before_action["observation"], observation_before_action["action_mask"])
  File "/Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/agents/human.py", line 48, in get_action
    return None if result is None else self.action_encoder.action_to_index(*result)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: ActionEncoder.action_to_index() takes 2 positional arguments but 4 were given
```

The problem is that result is a tuple of 3 elements (row, col, type) and we're expecting an Action.  Is there a way to convert those 3 elements into an Action or we don't use them anymore?  Otherwise, in pygame.py we should change the `handle_hover` method to create an Action (I hope this would be straightforward), and change the type of hover_action and click_action to be Action
created_at: 2025-04-30T00:12:15Z
author_association: COLLABORATOR
diff_hunk: @@ -25,21 +27,22 @@ def name(self):
             return self.params.nick
         return "Human"
 
-    def get_action(self, game):
+    def get_action(self, observation, action_mask):
         # Import this here to avoid circular dependencies
         from renderers.pygame import PygameQuoridor
 
-        observation, _, termination, truncation, _ = game.last()
-        if termination or truncation:
-            return None
-
         # Transform the action mask into a set of valid moves
         valid_moves = set()
-        for action, value in enumerate(observation["action_mask"]):
+        for action, value in enumerate(action_mask):
             if value == 1:
-                r, c, type = game.action_index_to_params(action)
-                valid_moves.add((r, c, type))
+                valid_action = self.action_encoder.index_to_action(action)
+                if isinstance(valid_action, MoveAction):
+                    valid_moves.add((valid_action.destination[0], valid_action.destination[1], 0))
+                elif isinstance(valid_action, WallAction):
+                    valid_moves.add(
+                        (valid_action.position[0], valid_action.position[1], valid_action.orientation.value - 1)
+                    )
 
         result = PygameQuoridor.instance().get_human_input(valid_moves)
 
-        return None if result is None else game.action_params_to_index(*result)
+        return None if result is None else self.action_encoder.action_to_index(*result)
path: deep_quoridor/src/agents/human.py
user: login: alejandromarcu
body: I have a fix for it, but not sure what's a good way to send it to you.  I'm afraid to push something and make a mess.  I tried making a patch but seems to have included other stuff.  Here's what I did.
Otherwise you could merge it and I'll send this in a separate PR.  
```
diff --git a/deep_quoridor/src/agents/human.py b/deep_quoridor/src/agents/human.py
index ceb0868..3e897b7 100644
--- a/deep_quoridor/src/agents/human.py
+++ b/deep_quoridor/src/agents/human.py
@@ -36,13 +36,8 @@ class HumanAgent(Agent):
         for action, value in enumerate(action_mask):
             if value == 1:
                 valid_action = self.action_encoder.index_to_action(action)
-                if isinstance(valid_action, MoveAction):
-                    valid_moves.add((valid_action.destination[0], valid_action.destination[1], 0))
-                elif isinstance(valid_action, WallAction):
-                    valid_moves.add(
-                        (valid_action.position[0], valid_action.position[1], valid_action.orientation.value - 1)
-                    )
+                valid_moves.add(valid_action)

         result = PygameQuoridor.instance().get_human_input(valid_moves)

-        return None if result is None else self.action_encoder.action_to_index(*result)
+        return None if result is None else self.action_encoder.action_to_index(result)
diff --git a/deep_quoridor/src/renderers/pygame.py b/deep_quoridor/src/renderers/pygame.py
index 9e73e8d..e916314 100644
--- a/deep_quoridor/src/renderers/pygame.py
+++ b/deep_quoridor/src/renderers/pygame.py
@@ -9,7 +9,7 @@ import pygame
 from agents import ActionLog, Agent
 from agents.human import HumanAgent
 from arena import GameResult
-from quoridor import Player
+from quoridor import Action, MoveAction, Player, WallAction, WallOrientation

 from renderers import Renderer

@@ -72,8 +72,8 @@ class BoardState:
 class HumanInput:
     enabled: bool = False
     valid_moves: set = field(default_factory=set)
-    hover_action: tuple[int, int, int] | None = None
-    click_action: tuple[int, int, int] | None = None
+    hover_action: Action | None = None
+    click_action: Action | None = None


 class PygameQuoridor:
@@ -275,11 +275,13 @@ class PygameQuoridor:
         if self.human_input.hover_action is None or not self.human_input.enabled:
             return

-        row, col, type = self.human_input.hover_action
-        if type == 0:
+        if isinstance(self.human_input.hover_action, MoveAction):
+            row, col = self.human_input.hover_action.destination
             self._draw_player(COLOR_PREVIEW, row, col, False)
-        else:
-            self._draw_wall(row, col, type == 2, is_preview=True)
+        elif isinstance(self.human_input.hover_action, WallAction):
+            row, col = self.human_input.hover_action.position
+            is_horizontal = self.human_input.hover_action.orientation == WallOrientation.HORIZONTAL
+            self._draw_wall(row, col, is_horizontal, is_preview=True)

     def draw_screen(self):
         if self.state is None:
@@ -408,16 +410,16 @@ class PygameQuoridor:

         type_action = 0
         if v_wall:
-            type_action = 1
+            action = WallAction((row, col), WallOrientation.VERTICAL)
         elif h_wall:
-            type_action = 2
-
-        action = (row, col, type_action)
+            action = WallAction((row, col), WallOrientation.HORIZONTAL)
+        else:
+            action = MoveAction((row, col))

         if action in self.human_input.valid_moves:
             self.human_input.hover_action = action

-    def get_human_input(self, valid_moves: set) -> Tuple[int, int, int] | None:
+    def get_human_input(self, valid_moves: set) -> Action | None:
         self.human_input.valid_moves = valid_moves
         self.human_input.click_action = None
```
created_at: 2025-04-30T00:35:32Z
author_association: COLLABORATOR

PR comments:
user: login: jonbinney
created_at: 2025-04-09T01:05:49Z
body: I'll debug the agents so that `play.py` works again in the next couple of days. It'll be good for me to catch up on that part of the code.
author_association: OWNER
user: login: jonbinney
created_at: 2025-04-11T19:29:48Z
body: I've pushed my implementation of `ActionEncoder` and a few other things:
- agents take observation and action_mask in get_action, instead of taking the entire environment
- Quoridor class can now construct itself from an observation, so that agents that want to try out actions and see how they affect the state can (Like Greedy)
- The (row, col, action_type) representation is gone - instead we use either "action index" (the index into the action_mask) or the MoveAction and WallAction classes from quoridor.py
- the mapping between action index and action classes is now its own "ActionEncoder" class in quoridor.py, so that agents can convert back and forth without needing access to the entire QuoridorEnv class

The easiest way to see the changes is to look at the greedy agent's `get_action()` function in this PR. I've got it running again with `play.py` but there may still be subtle bugs hiding.
author_association: OWNER
user: login: jonbinney
created_at: 2025-04-11T19:31:14Z
body: Let me know what you think about those changes. Now I'm going to update the other agents.
author_association: OWNER
user: login: alejandromarcu
created_at: 2025-04-12T21:24:28Z
body: I like the changes, it's cleaner using a Quoridor object than the mess of having the whole environment, thanks!
author_association: COLLABORATOR
user: login: alejandromarcu
created_at: 2025-04-15T18:48:36Z
body: It's looking good and tests are passing, yey! I guess dexp and ppo are missing, anything else? we probably should be running tests for them in CI, the hard part is making sure that the alias that we use is always working. e.g. we could use 'best', but if we bump the model version, we need to make sure that we create the tag immediately.  Or we could use `latest`, but there's the possibility that something fails and it would break the tests.  
author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-04-15T19:00:04Z
body: Yeah, I need to update dexp and ppo, and make sure that I didn't break incremental rewards. Working on that now.
author_association: OWNER
user: login: jonbinney
created_at: 2025-04-17T01:16:24Z
body: Almost there! Tests pass, agents run. Two problems:
- when running `python deep_quoridor/src/play.py -t 2 -yp` the human agent hangs. This might be because the agent its playing against hangs; i haven't dug into it yet.
- I get slightly different results when running the tournament, so something has changed.... a bit...

Here are the results of `python deep_quoridor/src/play.py -yi benchmark_B5W3 ` on this branch:
```
+------------+------------+--------+------------+--------+--------+
|  P1 \ P2   | dexp (111) | greedy | greedy-ish | sb3ppo | Total  |
+------------+------------+--------+------------+--------+--------+
| dexp (111) |     -      |  82%   |    78%     |  98%   |  86%   |
|   greedy   |     0%     |   -    |    98%     |  100%  |  66%   |
| greedy-ish |    22%     |  66%   |     -      |  100%  |  63%   |
|   sb3ppo   |     0%     |   0%   |     8%     |   -    |   3%   |
|   ======   |   ======   | ====== |   ======   | ====== | ====== |
|   Total    |     7%     |  49%   |    61%     |  99%   |  54%   |
+------------+------------+--------+------------+--------+--------+
```
And on `origin/main`:
```
+------------+------------+--------+------------+--------+--------+
|  P1 \ P2   | dexp (111) | greedy | greedy-ish | sb3ppo | Total  |
+------------+------------+--------+------------+--------+--------+
| dexp (111) |     -      |  94%   |    72%     |  98%   |  88%   |
|   greedy   |     4%     |   -    |    100%    |  100%  |  68%   |
| greedy-ish |    26%     |  74%   |     -      |  98%   |  66%   |
|   sb3ppo   |     2%     |   0%   |     4%     |   -    |   2%   |
|   ======   |   ======   | ====== |   ======   | ====== | ====== |
|   Total    |    11%     |  56%   |    59%     |  99%   |  56%   |
+------------+------------+--------+------------+--------+--------+
```

Anyone see something obvious I'm missing? If not I'll continue debugging tomorrow.
author_association: OWNER
user: login: jonbinney
created_at: 2025-04-24T21:28:32Z
body: @adamantivm @muralx @alejandromarcu It's past time that I stop tweaking this trying to get everything to work perfectly. Here are the highlights:


- We now pass the observations in to each agents `get_action` function, rather than passing in the entire environment and letting them do whatever they want with it
- Same for the `handle_step_*` functions in the trainable agent
- there's a function `construct_game_from_observation()` in quoridor.py that agents can use to construct a Quoridor game object from an observation. This allows us to keep the agent interfaces clean while still having a convenient way for agents like Greedy to try out actions
- Tests pass, profiling script works, play and train work (at least with the various argument combos I've been trying.

And the two things that have been holding me back from merging:
- Judging by the decreased win percentages, I've broken dexp in some way while changing the API. I've mucked about with it but I really need @muralx 's help fixing it - sorry for the trouble!
- Using classes for the actions makes things clearer, but also adds overhead for python to create the objects. This is about 3x slower, which matters most for agents like greedy that look at the actions instead of just the action mask. I think we should go back to tuples at some point, but I think we'd be better off getting all the other changes in first.

author_association: OWNER
user: login: adamantivm
created_at: 2025-04-26T00:34:55Z
body: Just want to say thanks for the hard work. I realize I'm too out of the loop on the evolution of all our abstractions, including the agent superclasses, the env and the arena, so I'm not sure how much I'd be able to contribute with a review. 
I'll support whatever you all decide and don't mind if a merge breaks things, I'll deal with it.
author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-04-26T01:37:56Z
body: @muralx Thank you for your review! I completely forgot to consider the case of wrapped environments, so I'll dig into that tomorrow. My first (probably wrong) thought is: shouldn't it just work? If the environment is wrapped, then when `Arena._play_game()` calls `.observe()` on the wrapped environment, it will get the observation with any transformations applied by the environment wrappers. Then it will pass that observation into the agent's `.get_action()` function, and the when the agent constructs a Quoridor object from it, it will be of the transformed game. Or do the transformations do things that lead to observations that don't actually represent a game anymore?
author_association: OWNER
user: login: muralx
created_at: 2025-04-26T18:05:22Z
body: @jonbinney I think the observation could be any format, or orientation, so it would be hard for Quoridor game to construct itself from it. Current agents that are using the Quoridor might anyway depends in a particular observation, so things might still work. But any other agent depending on a different observation style, would need some transformer to convert the observation back to the quoridor game accepted format. Maybe not important for now, but once I add the agent adapters it will become more important if game interaction is needed from the agent.
author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-04-27T00:42:46Z
body: @muralx I understand what you mean now about wrappers that modify the observation breaking my approach. I think the best way forward is for me to keep the part of this PR that uses the Quoridor class within the environment, but to revert the changes to the `Agent.get_action()` and `Agent.handle_step_outcome*` functions (go back to passing the environment in. Does that sound reasonable?
author_association: OWNER
user: login: muralx
created_at: 2025-04-27T19:18:55Z
body: @jonbinney I think nothing will break for now, and I will be adding per agent adapter or environment, so I will need to change that code anyway, so for me it is OK for you to submit and we can refactor things after when we add support for use of wrappers (only one using wrapper now is sb3 for training purposes).


author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-04-28T01:25:34Z
body: @muralx I made the naming changes you suggested and replaced the observations list with individual variables for clarity. It looks like I've still broken something in dexp, since its win rate has gone way down. If you're ok with merging and fixing in another PR, I think that might make things simpler.
author_association: OWNER
user: login: alejandromarcu
created_at: 2025-04-28T22:45:56Z
body: I'll take a look today or tomorrow, thanks guys!
author_association: COLLABORATOR
user: login: alejandromarcu
created_at: 2025-04-29T22:55:26Z
body: I just tried training dexp with the same parameters than I did in main (and I merged main here), and the metrics were significantly different. Are we ok with this?


<img width="1426" alt="image" src="https://github.com/user-attachments/assets/9f38f856-8bb8-486d-99f5-a16d87db9497" />


Training cmd:
```
/Users/amarcu/code/deep_rabbit_hole/.venv/bin/python /Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/train.py  -N 5 -W 3 -e 600 -i 43 -p dexp:rotate=true,split=true,turn=false,epsilon=1,epsilon_decay=0.9985,epsilon_min=0.1,final_reward_multiplier=20,training_mode=true,use_opponents_actions=true,update_target_every=100,gamma=0.95,target_as_source_for_opponent=true,use_negative_qvalue_function=true,batch_size=128,mask_targetq=true,nick=dexp_train greedy:p_random=0.3 greedy:p_random=0.1 simple -rp cucu -s
```

author_association: COLLABORATOR
user: login: muralx
created_at: 2025-04-30T00:08:59Z
body: If only dexp is the problem, I will take care of it once it is merged.
Thanks for double checking

On Tue, Apr 29, 2025, 11:55 PM Alejandro Marcu ***@***.***>
wrote:

> *alejandromarcu* left a comment (jonbinney/deep_rabbit_hole#158)
> <https://github.com/jonbinney/deep_rabbit_hole/pull/158#issuecomment-2840415678>
>
> I just tried training dexp with the same parameters than I did in main
> (and I merged main here), and the metrics were significantly different. Are
> we ok with this?
> image.png (view on web)
> <https://github.com/user-attachments/assets/9f38f856-8bb8-486d-99f5-a16d87db9497>
>
> Training cmd:
>
> /Users/amarcu/code/deep_rabbit_hole/.venv/bin/python /Users/amarcu/code/deep_rabbit_hole/deep_quoridor/src/train.py  -N 5 -W 3 -e 600 -i 43 -p dexp:rotate=true,split=true,turn=false,epsilon=1,epsilon_decay=0.9985,epsilon_min=0.1,final_reward_multiplier=20,training_mode=true,use_opponents_actions=true,update_target_every=100,gamma=0.95,target_as_source_for_opponent=true,use_negative_qvalue_function=true,batch_size=128,mask_targetq=true,nick=dexp_train greedy:p_random=0.3 greedy:p_random=0.1 simple -rp cucu -s
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/jonbinney/deep_rabbit_hole/pull/158#issuecomment-2840415678>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAHKU32SDTFNP6CT5BTTBGL2377PLAVCNFSM6AAAAAB2XUPLUWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQNBQGQYTKNRXHA>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

author_association: COLLABORATOR
user: login: jonbinney
created_at: 2025-04-30T16:05:56Z
body: Thanks guys! Merging!
author_association: OWNER