blokus/PLAN.md

# Blokus Gym Harness — Design & Implementation Plan

## 1. Overview

A Gymnasium-compatible RL environment for the board game **Blokus**, supporting both **single-agent** (agent vs. configurable bots) and **multi-agent** (PettingZoo-compatible) modes. The harness uses **numpy** for maximum training speed and broad framework compatibility, with **matplotlib** as an optional rendering backend.

## 2. File Structure

```
blokus/
├── pyproject.toml              # Package config, dependencies, entry points
├── README.md
├── LICENSE
├── PLAN.md                     # This file
├── src/blokus_gym/
│   ├── __init__.py             # Package exports
│   ├── envs/
│   │   ├── __init__.py         # Env exports
│   │   ├── blokus_env.py       # Main gymnasium.Env implementation
│   │   ├── multiagent.py       # PettingZoo-compatible multi-agent wrapper
│   │   └── registration.py     # gymnasium.register() calls for presets
│   ├── core/
│   │   ├── __init__.py
│   │   ├── pieces.py           # Piece definitions, orientation generation, custom piece sets
│   │   ├── board.py            # Board state (numpy array), placement validation, corner tracking
│   │   ├── game.py             # Game rules: turn order, move validation, scoring, game-over detection
│   │   └── bots.py             # Bot implementations: Random, Greedy, Minimax
│   ├── wrappers/
│   │   ├── __init__.py
│   │   └── action_mask.py      # Optional wrapper to expose action masks as part of observation
│   └── utils/
│       ├── __init__.py
│       └── render.py           # Text/ANSI and matplotlib rendering
├── tests/
│   ├── __init__.py
│   ├── conftest.py             # Shared fixtures
│   ├── test_pieces.py
│   ├── test_board.py
│   ├── test_game.py
│   ├── test_env.py
│   ├── test_multiagent.py
│   └── test_registration.py
└── examples/
    ├── basic_usage.py
    ├── train_sb3.py            # Stable-Baselines3 training example
    ├── train_rllib.py          # Ray RLlib training example
    └── multiagent_example.py
```

## 3. Core Design Decisions

### 3.1 Action Space: Discrete + Action Mask

**Design**: Pre-compute all unique `(piece, orientation, position)` combinations into a flat `Discrete(N)` action space. At each step, provide an action mask in `info["action_mask"]` indicating which of the N actions are currently valid.

**Why**:
- Works with any RL algorithm (masking-aware libs like SB3, RLlib use the mask; others can ignore it)
- Simple, standard API
- Action index is stable across episodes (important for policy networks)

**Action space size** (standard 20×20 board, 21 pieces): ~25,000 discrete actions. Smaller boards scale down proportionally.

**Action generation algorithm** (at `__init__` time):
1. For each piece in the piece set, generate all unique orientations (rotations × flips, deduplicated)
2. For each orientation, enumerate all positions where the piece fits on the board
3. Assign each `(piece, orientation, x, y)` a unique integer index
4. Store as a list of `Move` objects (piece_id, orientation_id, x, y) for O(1) lookup

**Action mask computation** (at each `step`):
- Start with all-False mask
- For each action index, check if the move is valid (piece available, in bounds, no overlap, corner rule satisfied)
- Set valid indices to True

### 3.2 Observation Space: Dict with Board, Pieces, and Corners

```python
observation_space = spaces.Dict({
    "board":       spaces.Box(0, num_players, (board_size, board_size), dtype=np.int8),
    "pieces":      spaces.MultiBinary(num_pieces),           # 1 = still available
    "my_turn":     spaces.Discrete(1),                        # Always 1 (agent's perspective)
    "corners":     spaces.Box(0, board_size, (board_size, board_size), dtype=bool),  # Valid corner cells
})
```

**Info dict**:
```python
info = {
    "action_mask": np.array([bool] * num_actions),  # Valid actions
    "current_player": int,
    "players_with_moves": [int, ...],  # Who still has valid moves
    "step_count": int,
}
```

### 3.3 Reward Design

**Single-agent mode** (agent is player 0):
- **Step reward**: `0` by default (sparse). Optional shaping: `+0.01 * piece_size` for each placed piece
- **Terminal reward**: Based on final score comparison
  - `+1.0` if agent wins (higher score than all opponents)
  - `0.0` if tie
  - `-1.0` if agent loses
  - Alternative: normalized score difference `agent_score / max_possible_score`

**Multi-agent mode** (PettingZoo):
- Each agent gets `reward = placed_squares / total_squares` at game end
- Penalty for unplaced squares: `-unplaced_squares / total_squares`
- `+15/total_squares` bonus for emptying all pieces
- `+5/total_squares` bonus for last piece being monomino

### 3.4 Dependencies

| Dependency | Purpose | Required? |
|---|---|---|
| `numpy` | Board state, coordinate math, action masks | Yes |
| `gymnasium` | RL environment API | Yes |
| `pettingzoo` | Multi-agent wrapper | Yes (for multi-agent mode) |
| `matplotlib` | Visual rendering | No (optional) |

**Rationale**: numpy is the universal standard for array operations in Python RL. Every framework (SB3, RLlib, Tianshou, etc.) consumes numpy arrays natively. No torch/tensorflow dependency means faster installation and broader compatibility.

## 4. Component Designs

### 4.1 Pieces (`core/pieces.py`)

**Piece definition**: A piece is a set of `(x, y)` integer coordinates relative to a reference point.

```python
@dataclass
class Piece:
    name: str
    squares: frozenset[tuple[int, int]]  # Relative coordinates
    size: int                             # Number of squares

@dataclass
class OrientedPiece:
    piece_id: int
    orientation_id: int
    squares: list[tuple[int, int]]        # Absolute coordinates on board
    corners: list[tuple[int, int]]        # Corner cells for placement rule
```

**Standard piece set** (`STANDARD_PIECES`): All 21 free polyominoes (1 monomino, 1 domino, 2 trominoes, 5 tetrominoes, 12 pentominoes).

**Orientation generation**:
1. Start with original coordinates
2. Apply 4 rotations (0°, 90°, 180°, 270°) using rotation matrix `(x, y) → (y, -x)`
3. For each rotation, apply 2 flips: original and horizontal-flip `(x, y) → (-x, y)`
4. Normalize each result (shift so min x = 0, min y = 0)
5. Deduplicate by comparing sorted coordinate sets

**Custom piece sets**:
```python
# Option 1: Define custom pieces
custom_pieces = [
    Piece("A", frozenset([(0,0), (1,0), (0,1)])),  # L-tromino
    Piece("B", frozenset([(0,0), (1,0), (2,0)])),  # I-tromino
]

# Option 2: Use a preset
from blokus_gym.core.pieces import STANDARD_PIECES, DUO_PIECES, JUNIOR_PIECES

# Option 3: Load from JSON
pieces = PieceSet.from_json("my_pieces.json")

# Option 4: Filter standard set
pentominoes_only = [p for p in STANDARD_PIECES if p.size == 5]
```

### 4.2 Board (`core/board.py`)

```python
class Board:
    def __init__(self, size: int):
        self.size = size
        self.grid = np.zeros((size, size), dtype=np.int8)  # 0 = empty, 1-4 = player
    
    def place(self, player_idx: int, squares: list[tuple[int, int]]) -> None
    def is_empty(self, x: int, y: int) -> bool
    def in_bounds(self, x: int, y: int) -> bool
    def has_overlap(self, squares: list[tuple[int, int]]) -> bool
    def get_player_corners(self, player_idx: int) -> set[tuple[int, int]]
```

**Key methods**:
- `place()`: Sets grid cells to player index
- `has_overlap()`: Checks if any square is already occupied
- `get_player_corners()`: Returns all corner cells adjacent to player's pieces (for placement rule)

### 4.3 Game Logic (`core/game.py`)

```python
class BlokusGame:
    def __init__(self, board: Board, pieces: list[Piece], num_players: int = 4):
        self.board = board
        self.pieces = pieces
        self.num_players = num_players
        self.current_player = 0
        self.player_pieces: list[set[str]]  # Available piece names per player
        self.player_corners: list[set[tuple[int, int]]]
        self.rounds = 0
        self.game_over = False
    
    def valid_move(self, player_idx: int, move: Move) -> bool
    def apply_move(self, player_idx: int, move: Move) -> None
    def get_valid_moves(self, player_idx: int) -> list[int]  # Returns action indices
    def get_action_mask(self, player_idx: int) -> np.ndarray
    def next_player(self) -> int  # Cycles to next player with moves
    def is_game_over(self) -> bool
    def get_scores(self) -> list[int]  # Negative = unplaced squares
    def get_winners(self) -> list[int]
```

**Move validation rules**:
1. Piece must be in player's available pieces
2. All squares must be in bounds
3. No overlap with existing pieces
4. **Corner rule**: At least one square of the placed piece must touch a same-color piece at a corner (diagonal adjacency)
5. **No edge adjacency**: No square of the placed piece may share an edge with a same-color piece
6. **First move**: Must be placed in a corner (if `corner_rule=True`)

### 4.4 Bots (`core/bots.py`)

```python
class Bot(ABC):
    @abstractmethod
    def select_action(self, game: BlokusGame, player_idx: int, valid_actions: np.ndarray) -> int

class RandomBot(Bot):
    """Randomly selects from valid actions."""

class GreedyBot(Bot):
    """Selects the move that places the most squares (largest piece)."""

class GreedyCornersBot(Bot):
    """Greedy + prefers moves that open up more corners for future play."""

class MinimaxBot(Bot):
    """Minimax search with configurable depth. Primarily for 2-player games."""
```

### 4.5 Gymnasium Environment (`envs/blokus_env.py`)

```python
class BlokusEnv(gymnasium.Env):
    metadata = {"render_modes": ["human", "ansi", "rgb_array"]}
    
    def __init__(
        self,
        num_players: int = 4,
        board_size: int = 20,
        pieces: list[Piece] | None = None,  # Default: STANDARD_PIECES
        render_mode: str | None = None,
        bot_type: type[Bot] = RandomBot,
        bot_strength: int = 1,
        reward_shaping: bool = False,
        corner_rule: bool = True,
        max_steps: int | None = None,
        seed: int | None = None,
    ):
        super().__init__()
        # Initialize game, spaces, action lookup table
    
    def reset(self, seed=None, options=None) -> tuple[dict, dict]
    def step(self, action: int) -> tuple[dict, float, bool, bool, dict]
    def render(self) -> str | np.ndarray | None
    def close(self)
```

**Key implementation details**:
- `__init__` generates the full action lookup table (piece, orientation, position → index)
- `reset()` initializes a new game, places bots, optionally lets bots take first turns
- `step()` applies the agent's action, then lets bots play until it's the agent's turn again
- Returns `(observation, reward, terminated, truncated, info)` where `info` contains the action mask

### 4.6 Multi-Agent Wrapper (`envs/multiagent.py`)

PettingZoo-compatible wrapper that exposes all players as separate agents:

```python
class BlokusMultiAgentEnv(AECEnv):
    def __init__(self, **kwargs):
        # Same config as BlokusEnv, but wraps it
    
    def reset(self, seed=None, options=None) -> tuple[dict, dict]
    def step(self, action: int) -> None
    # PettingZoo API: observe(), close(), render(), etc.
    
    @property
    def agents(self) -> list[str]  # ["player_0", "player_1", ...]
    @property
    def agent_selection(self) -> str
```

**Behavior**:
- Cycles through players in order (0, 1, 2, 3)
- Each agent gets its own observation (board + own pieces + own corners + action mask)
- Rewards are per-agent: negative for unplaced squares, bonus for emptying pieces
- Agents with no valid moves are marked as done (truncated) but still receive final rewards

### 4.7 Registration (`envs/registration.py`)

Predefined environment configurations registered with `gymnasium.register()`:

```python
# Standard 4-player Blokus
"Blokus-v0":        num_players=4, board_size=20, pieces=STANDARD_PIECES

# Blokus Duo (2 players, 14x14)
"BlokusDuo-v0":     num_players=2, board_size=14, pieces=STANDARD_PIECES

# Blokus Junior (2 players, 14x14, simplified pieces)
"BlokusJunior-v0":  num_players=2, board_size=14, pieces=JUNIOR_PIECES

# Simple test env (2 players, 7x7, only pieces < 5 squares)
"BlokusSimple-v0":  num_players=2, board_size=7, pieces=[p for p in STANDARD_PIECES if p.size < 5]

# Greedy bot variants
"BlokusGreedy-v0":  bot_type=GreedyBot
"BlokusDuoGreedy-v0": bot_type=GreedyBot, num_players=2, board_size=14
```

### 4.8 Rendering (`utils/render.py`)

**Text/ANSI mode** (`render_mode="human"`):
```
  0 1 2 3 4 5 6 7 8 9 ...
0 . . . . . . . . . .
1 . . . . . . . . . .
2 . . . . . . . . . .
3 . . . . . . . . . .
4 . . . . . . . . . .
...
Player 0 (Red): I1 I2 I3 L4 O4 I4 S4 T4 I5 L5 N P T U V W X Y Z
Player 1 (Blue): I1 I2 I3 L4 O4 I4 S4 T4 I5 L5 N P T U V W X Y Z
Current player: 0
Valid moves: 152
```

**RGB array mode** (`render_mode="rgb_array"`):
- Matplotlib figure with colored squares (red, blue, yellow, green, grey)
- Grid lines
- Legend showing remaining pieces per player
- Valid move indicators (semi-transparent overlays)

### 4.9 Action Mask Wrapper (`wrappers/action_mask.py`)

Optional wrapper that moves the action mask from `info` into the observation space:

```python
class ActionMaskWrapper(Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.observation_space = spaces.Dict({
            "observation": env.observation_space,
            "action_mask": spaces.Box(0, 1, (env.action_space.n,), dtype=bool),
        })
    
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        return {
            "observation": obs,
            "action_mask": info["action_mask"],
        }, reward, terminated, truncated, info
```

## 5. Customization Options

All options are passed as keyword arguments to `BlokusEnv()` or via `gymnasium.make()`:

| Option | Type | Default | Description |
|---|---|---|---|
| `num_players` | int | 4 | Number of players (2-4) |
| `board_size` | int | 20 | Board dimension (standard: 20, Duo: 14) |
| `pieces` | list[Piece] | STANDARD_PIECES | Custom piece set |
| `bot_type` | type[Bot] | RandomBot | Opponent bot type |
| `bot_strength` | int | 1 | Bot search depth (for Minimax) |
| `reward_shaping` | bool | False | Enable intermediate rewards |
| `corner_rule` | bool | True | Require first move in corner |
| `max_steps` | int | None | Max steps before truncation |
| `render_mode` | str | None | "human", "ansi", "rgb_array" |
| `seed` | int | None | Random seed |
| `starting_corners` | list | [(0,0), (0,19), (19,0), (19,19)] | Custom starting corners |
| `scoring_mode` | str | "standard" | "standard", "coverage", "winloss" |
| `first_player` | int | 0 | Which player goes first (agent is always 0) |

## 6. Implementation Phases

### Phase 1: Core Game Logic (Days 1-3)
- [ ] `core/pieces.py` — Piece definitions, orientation generation, custom piece support
- [ ] `core/board.py` — Board state, placement validation, corner tracking
- [ ] `core/game.py` — Game rules, move validation, scoring, game-over detection
- [ ] Unit tests for all three modules

### Phase 2: Bots (Days 4-5)
- [ ] `core/bots.py` — Random, Greedy, GreedyCorners, Minimax bots
- [ ] Tests for bot behavior

### Phase 3: Gymnasium Environment (Days 6-8)
- [ ] `envs/blokus_env.py` — Main environment class
- [ ] Action space generation and masking
- [ ] Observation and reward design
- [ ] `envs/registration.py` — Preset registrations
- [ ] `wrappers/action_mask.py` — Optional action mask wrapper
- [ ] `gymnasium.utils.env_checker.check_env()` passes

### Phase 4: Multi-Agent (Days 9-10)
- [ ] `envs/multiagent.py` — PettingZoo AECEnv wrapper
- [ ] Per-agent observations and rewards
- [ ] Tests for multi-agent mode

### Phase 5: Rendering (Days 11-12)
- [ ] `utils/render.py` — Text/ANSI and matplotlib rendering
- [ ] Integration with environment

### Phase 6: Testing & Examples (Days 13-14)
- [ ] Comprehensive test suite (unit + integration)
- [ ] `examples/basic_usage.py` — Basic usage
- [ ] `examples/train_sb3.py` — SB3 training
- [ ] `examples/train_rllib.py` — RLlib training
- [ ] `examples/multiagent_example.py` — Multi-agent
- [ ] `pyproject.toml` with all dependencies
- [ ] `README.md` with full documentation

## 7. Testing Strategy

### Unit Tests
- **Pieces**: Orientation generation correctness, deduplication, custom piece loading
- **Board**: Placement, overlap detection, bounds checking, corner computation
- **Game**: Move validation (all rules), scoring, game-over detection, turn cycling
- **Bots**: Random bot returns valid moves, greedy bot prefers larger pieces

### Integration Tests
- **Environment**: `check_env()` passes, reset/step return correct types
- **Action mask**: All masked actions are valid, all valid actions are masked
- **Game flow**: Full game from start to end, correct winner detection
- **Multi-agent**: All agents receive observations, rewards, and termination signals

### Edge Cases
- Player with no valid moves (should be skipped)
- Game ends with all players blocked
- Last piece is monomino (bonus scoring)
- Custom piece sets with unusual shapes
- Small boards where pieces don't fit

## 8. Key Design Tradeoffs

| Decision | Choice | Rationale |
|---|---|---|
| Gym library | Gymnasium | Maintained fork, latest API |
| Agent mode | Single + Multi-agent | Covers both use cases |
| Action space | Discrete + mask | Broadest RL framework compatibility |
| Dependencies | numpy only | Fastest, broadest compatibility |
| Rendering | Text + matplotlib | No Pygame dependency, optional |
| Piece definition | Coordinate sets | Simple, flexible for custom pieces |
| Action pre-computation | At init time | Stable action indices, fast step() |
| Multi-agent | PettingZoo AECEnv | Standard API, compatible with RLlib/PettingZoo algorithms |