Tetris AI Training

A Deep Reinforcement Learning agent for playing Tetris using Double DQN with Prioritized Experience Replay.

DEMO

airis.demo.mp4

Attribution

This implementation is based on:

• truonging/Tetris-A.I (Version2) - Double DQN Tetris implementation

What's the Same

• Core Algorithm: Double DQN with heap-based Prioritized Experience Replay
• State Representation: 6-feature state (heights, bumpiness, holes, lines, y_pos, pillars)
• Network Architecture: Simple MLP [32, 32, 32] hidden layers
• Training Schedule: Train every 200 steps, target update every 1000 steps
• State Enumeration: Numba JIT calc_all_states() for fast position computation

What's Different

Aspect	Reference	This Implementation
Loss Function	SmoothL1Loss (Huber)	MSELoss
Priority Storage	Positive priorities	Negative priorities (max-heap via min-heap)
PER Sampling	probabilities ** beta	Direct priority-weighted sampling
Priority Update	Per-index tuple replacement	Full heap rebuild (less efficient)
Importance Weights	(1/N * 1/P)^beta	(N * P)^(-beta)

Loss Function Note

The reference uses SmoothL1Loss (Huber loss) which is more robust to outliers. This implementation uses MSELoss. Both work well for this task.

Architecture

Double DQN with Heap-based PER

Following the reference implementation:

• Network: Simple feedforward MLP [32, 32, 32] hidden layers
• Training: Double DQN (action selection with online network, evaluation with target network)
• Memory: Heap-based Prioritized Experience Replay
• Target Updates: Hard copy every 1000 steps
• Training Schedule: Train every 200 steps

State Features (6 values)

1. Total Height - Sum of all column heights
2. Bumpiness - Sum of height differences between adjacent columns
3. Lines Cleared - Lines cleared by this move
4. Holes - Empty cells with blocks above
5. Y Position - Row where piece landed (normalized)
6. Pillars - Deep wells (≥3 depth) that trap pieces

Scoring (NES Tetris)

Lines Cleared	Base Points
1 line	40
2 lines	100
3 lines	300
4 lines (Tetris)	1200

Score = base_points × level, where level = total_lines // 10 + 1

Hyperparameters

Parameter	Value	Description
Network	[32, 32, 32]	Hidden layer sizes
Learning Rate	0.01	Adam optimizer
Discount (γ)	0.999	Future reward weight
Epsilon	0.3 → 0.01	Polynomial decay over games
Batch Size	128	Training batch
Memory	10,000	Replay buffer size
Train Every	200 steps	Training frequency
Target Update	1000 steps	Hard copy interval

Results

Training achieves excellent performance:

• Best Score: 5,650,660+ points
• Best Lines: 1,292 lines in a single game
• Training Time: ~10 minutes for 1500 games

Installation

git clone https://github.com/yourusername/TetrisBattleArena.git
cd TetrisBattleArena

# Create virtual environment
python3 -m venv venv
source venv/bin/activate

# Install dependencies
pip install torch numpy pygame numba

Usage

Training

# Train without visualization (faster)
python train.py --games 5000

# Train with visualization
python train.py --render --games 2000

# Load existing model and continue training
python train.py --render --load --games 1000

Visualization Controls

Key	Action
SPACE	Pause/unpause
UP/DOWN	Adjust speed (MAX → 60/s → 30/s → 15/s → 5/s → 1/s)
ESC	Save and exit

Genetic Algorithm (Optional)

Evolve reward weights:

python genetic_algo.py

Saves optimized weights to data/best_weights.json.

File Structure

TetrisBattleArena/
├── settings.py           # Tetrominos, Numba JIT functions, hyperparameters
├── tetris.py             # Game environment with calc_all_states()
├── agent.py              # Agent with heap-based PER
├── model.py              # Linear_QNet, QTrainer (Double DQN)
├── train.py              # Training loop with visualization
├── genetic_algo.py       # GA for reward weight tuning
├── data/
│   └── best_weights.json # Evolved reward weights
└── model/
    └── best_model.pth    # Saved checkpoint

How It Works

1. State Enumeration: For each piece, calc_all_states() (Numba JIT) computes all possible final positions and their resulting board states
2. Action Selection: Agent evaluates Q-values for all possible states and picks the best (or random with ε probability)
3. Reward Calculation: Uses evolved weights to compute reward based on state features
4. Experience Replay: Stores (state, next_state, reward, done) in prioritized heap
5. Training: Every 200 steps, samples batch weighted by TD-error priority
6. Target Network: Hard-copied every 1000 steps for stable Q-targets

References

• truonging/Tetris-A.I - Reference implementation