README.md

---
license: mit
library_name: tensorflow
tags:
  - othello
  - reversi
  - game-ai
  - transformer
  - mixture-of-experts
  - rmsnorm
  - swiglu
---

## Transformer Reversi AI
This Reversi AI is a Transformer-based model that uses Mixture-of-Experts architecture for decision-making instead of CNNs.  
[Github - rsu-Suba/ReversiGPT](https://github.com/rsu-Suba/ReversiGPT)

## Architecture

### Normal model `M2.keras`
- Parameters: 1.98M
- Model size: 40.4MB
- Input: `(8, 8, 3)`
  - Own stones
  - Opponent stones
  - Legal moves
- Board representation: 64 tokens
- Embedding dimension: 128
- Transformer blocks: 4
- Attention heads: 4
- MHA experts: 3 (with optional 8-direction attention mask)
- FFN experts: 2 (SwiGLU)
- Routing: Independent Top-k sparse routing with load balancing loss
- Iterative routing steps: 2
- Normalization: Pre-LN RMSNorm
- FFN activation: SwiGLU

### Small model `M2_small.keras`
- Parameters: 756K
- Model size: 15.7MB
- Embedding dimension: 96
- Transformer blocks: 3
- Attention heads: 3
- MHA experts: 2
- FFN experts: 2
- Iterative routing steps: 1

### Diagram
```
Input (8×8×3)
    ↓
64 Tokens
    ↓
Dense 128 (or 96)
    ↓
Token + Row + Column + Move Embeddings
    ↓
DynamicAssembly × 4 (or 3)
  ├─ MHA Expert × 3 (or 2) [MHA Router (Top-1)]
  └─ FFN Expert × 2 (SwiGLU) [FFN Router (Top-1)]
    ↓
    ├─ Policy Head → RMSNorm → Dense → 64 moves
    └─ Value Head  → Attention Pooling → Shared Dense
                        ├─ Win Head   → Tanh (Win Rate [-1, 1])
                        └─ Score Head → Tanh (Score Diff [-1, 1])
```

### DynamicAssembly

Each DynamicAssembly block contains separate pools of MHA and FFN experts with independent routers. At each iteration step, the mean-pooled board representation is combined with a step embedding and passed to the respective router.

The router selects the top expert (Top-k sparse routing) and dispatches the computation, with auxiliary load-balancing loss during training to ensure uniform expert utilization.

### Policy Head

The policy head applies RMSNorm and dense layers to predict a probability distribution over the 64 board positions.

### Value Head

The value head uses **Attention Pooling** over the token sequence and predicts a 2-dimensional vector in `[-1, 1] × [-1, 1]`:
- **Win Rate**: Estimated game outcome `[-1, 1]`
- **Score Difference**: Estimated final stone difference `[-1, 1]`

## Training

This model was trained using self-play game records generated by [Reversi-Transformer-1](https://huggingface.co/rsu/Reversi-Transformer-1).

### Self-Play data generation
- MCTS simulations: 256
- MCTS batch size: 8
- PUCT constant: 2.5
- Total games: 30,000

### Model Training
- Optimizer: AdamW
- Clipnorm: 1.0
- Weight decay: 0.01
- Learning rate: 2.0e-4
- Learning rate scheduler: WarmupCosineDecay
- Learning rate warmup: 2 epochs
- Training batch size: 512
- Epochs: 30

## Evaluation

[Reversi-Transformer-1](https://huggingface.co/rsu/Reversi-Transformer-1) was used as the previous model.

### Win rate
- vs Random : 100% (100 / 100 games)
- vs Prev model : 54% (54 / 100 games)

### Average final stone count
- vs Random : 59
- vs Prev model : 38

## Usage / Quick start

```python
import numpy as np
import tensorflow as tf
from huggingface_hub import hf_hub_download
from model import RMSNorm, TokenAndPositionEmbedding, MHA, FFN, DynamicAssembly, AttentionPooling

model_path = hf_hub_download(repo_id="rsu/Reversi-Transformer-2", filename="M2.keras", repo_type="model")
# If you want to use M2_small
# model_path = hf_hub_download(repo_id="rsu/Reversi-Transformer-2", filename="M2_small.keras", repo_type="model")

custom_objects = {
    "RMSNorm": RMSNorm,
    "TokenAndPositionEmbedding": TokenAndPositionEmbedding,
    "MHA": MHA,
    "FFN": FFN,
    "DynamicAssembly": DynamicAssembly,
    "AttentionPooling": AttentionPooling,
}
model = tf.keras.models.load_model(model_path, custom_objects=custom_objects, compile=False)

board = np.zeros((1, 8, 8, 3), dtype=np.float32)
board[0, 3, 3, 1] = 1.0
board[0, 3, 4, 0] = 1.0
board[0, 4, 3, 0] = 1.0
board[0, 4, 4, 1] = 1.0

board[0, 2, 3, 2] = 1.0
board[0, 3, 2, 2] = 1.0
board[0, 4, 5, 2] = 1.0
board[0, 5, 4, 2] = 1.0

policy, value = model(board, training=False)
print("Policy (64 moves prob): ", policy.numpy())
print("Win rate est [-1, 1]: ", value[0, 0].numpy().item())
print("Score diff est [-1, 1]: ", value[0, 1].numpy().item())
```