# Machine Learning Value Function for Shardok AI

## Overview

### The Problem

The current Shardok AI uses a hand-crafted scoring function to evaluate game states. This works, but has limitations:

1. **Variable score ranges**: Scores range from -500 to +50,000 depending on scenario and game phase
2. **Normalization challenges**: MCTS requires normalized scores, but our attempts over-compress differences
3. **Positional complexity**: Early-game evaluation requires sophisticated positional reasoning
4. **Maintenance burden**: Hand-tuned scoring is brittle and hard to improve

### The Solution

Replace the hand-crafted scoring function with a **learned value function** trained on self-play game outcomes:

```
Current: state → hand_crafted_score(state) → (needs complex normalization)
Proposed: state → neural_network(state) → win_probability ∈ [0,1]
```

**Key Benefits:**
- Naturally bounded output [0,1] perfect for MCTS
- Learns from aggregate patterns across thousands of games
- Can capture complex positional patterns automatically
- Iteratively improvable through self-play

### Approach

**Supervised Learning from Self-Play Outcomes**

1. Generate games using current AI (iterative deepening with current scoring)
2. Record every game state with label: "Did the current player win?"
3. Train neural network: `state → win_probability`
4. Use trained model as scoring function for MCTS
5. Generate new games with MCTS → retrain → improve (iterative)

This is similar to AlphaGo Zero but using **supervised learning only** (no policy head or reinforcement learning in v1).

---

## Phase 1: Data Collection

### Goal
Generate 50,000-100,000 game states with win/loss labels from self-play.

### Implementation Steps

#### 1.1 Extend Game State Serialization

**File:** Create `src/main/cpp/net/eagle0/shardok/ai/ml/StateRecorder.hpp`

```cpp
class StateRecorder {
public:
    struct StateRecord {
        std::vector<uint8_t> serializedStateView;  // FlatBuffer bytes (PLAYER VIEW)
        int currentPlayer;                          // 0 or 1
        int roundNumber;                            // Which round

        // Will be filled in after game ends
        bool currentPlayerWon;                      // Label: did current player win?
        int finalOutcome;                           // For future: outcome type
    };

    // Record the game state VIEW from current player's perspective
    // This should be called BEFORE the player makes their move
    void RecordStateView(const GameStateView& view, int currentPlayer);

    void SetGameOutcome(int winner);
    void WriteToFile(const std::string& filename);

private:
    std::vector<StateRecord> states_;
    int winner_ = -1;
};
```

**Critical Design Decisions:**

1. **Record PLAYER VIEW, not full game state**
   - Players have incomplete information (can see opponent units, but not exact stats)
   - Training on full state while inferring on partial view = distribution mismatch
   - Model must learn from what it can actually see

2. **Record AFTER each action**
   - Record the resulting position after applying each command
   - This teaches: "is reaching this position good for the current player?"
   - Important: Same player may make multiple moves in a row (only switches on END_TURN)
   - Don't record the state immediately after END_TURN (that's mid-transition to next player)

3. **Current player perspective**
   - Makes model generalizable to both sides
   - Label: "did current player win?" not "did P0 win?"

#### 1.2 Generate Self-Play Games

**File:** Create `src/main/cpp/net/eagle0/shardok/ai/ml/SelfPlayGenerator.cpp`

```cpp
void GenerateSelfPlayGames(int numGames, const std::string& outputDir) {
    for (int i = 0; i < numGames; ++i) {
        StateRecorder recorder;

        // Create random scenario (or use map pool)
        auto scenario = GenerateRandomScenario();

        // Play game: iterative deepening AI vs itself
        while (!game.IsTerminal()) {
            int currentPlayer = game.GetCurrentPlayer();

            // Current AI makes move
            auto move = IterativeDeepeningAI::Search(game, timeLimit);
            game.ApplyMove(move);

            // CRITICAL: Record AFTER applying the move
            // - This is the resulting position we want to evaluate
            // - Same player may make multiple moves in a row (only switches on END_TURN)
            // - Don't record immediately after END_TURN (mid-transition state)
            if (!move.IsEndTurn()) {
                int stillCurrentPlayer = game.GetCurrentPlayer();
                GameStateView playerView = game.GetPlayerView(stillCurrentPlayer);
                recorder.RecordStateView(playerView, stillCurrentPlayer);
            }
        }

        // Set outcome for all states
        recorder.SetGameOutcome(game.GetWinner());
        recorder.WriteToFile(outputDir + "/game_" + std::to_string(i) + ".bin");
    }
}
```

**Key Flow:**
1. Get current player (P0 or P1)
2. AI selects and applies move
3. **Record resulting position** (unless it was END_TURN)
4. Get player's view of the new position
5. Same player may move again, or may have switched after END_TURN
6. At end, label all recorded positions with outcome

**Command to run:**
```bash
bazel run //src/main/cpp/net/eagle0/shardok/ai/ml:self_play_generator -- \
    --num_games 10000 \
    --output_dir /tmp/shardok_training_data \
    --time_per_move 500ms
```

#### 1.3 Data Format

**Output:** Binary files with repeated records
```
[StateRecord 1][StateRecord 2]...[StateRecord N]
```

Each StateRecord:
```
- Serialized FlatBuffer state (variable length)
- Current player: int32
- Round number: int32
- Label (current player won): bool
```

**Expected Data Volume:**
- 50,000 games × ~50 states/game = 2.5M training examples
- Each state ~10-50KB → ~50-250GB total
- Tip: Compress with gzip, decompress on-the-fly during training

---

## Phase 2: Model Training

### Goal
Train a neural network to predict win probability from game state.

### 2.1 Technology Stack

**Framework:** PyTorch (recommended for beginners, excellent documentation)

**Rationale:**
- Simpler API than TensorFlow
- Better debugging (eager execution by default)
- Excellent tutorials: https://pytorch.org/tutorials/
- Easy C++ integration via LibTorch

**Installation:**
```bash
pip install torch torchvision torchaudio
pip install numpy pandas tensorboard
```

### 2.2 State Feature Engineering

**Challenge:** FlatBuffer game states are complex (units, terrain, spells, etc.). Neural networks need fixed-size numerical inputs.

**Approach:** Extract features into fixed-size vector

**File:** Create `src/main/python/shardok/ml/feature_extractor.py`

```python
import numpy as np
from typing import List

class StateFeatureExtractor:
    """Convert FlatBuffer game state to fixed-size feature vector."""

    # Feature dimensions
    MAX_UNITS = 32  # Max units per side
    FEATURE_DIM_PER_UNIT = 20  # Stats per unit
    GLOBAL_FEATURES = 50  # Map, round, etc.

    TOTAL_DIM = (MAX_UNITS * 2 * FEATURE_DIM_PER_UNIT) + GLOBAL_FEATURES

    def extract_features(self, state_bytes: bytes) -> np.ndarray:
        """
        Extract features from serialized FlatBuffer state.

        Returns:
            np.ndarray of shape (TOTAL_DIM,)
        """
        # Parse FlatBuffer
        state = ParseGameState(state_bytes)

        features = []

        # Global features
        features.append(state.current_round / state.max_rounds)  # Normalized round
        features.append(len(state.p0_units) / self.MAX_UNITS)    # Army size ratio
        features.append(len(state.p1_units) / self.MAX_UNITS)
        # ... add ~50 global features ...

        # Per-unit features (current player)
        for i in range(self.MAX_UNITS):
            if i < len(state.current_player_units):
                unit = state.current_player_units[i]
                features.extend(self._extract_unit_features(unit))
            else:
                features.extend([0] * self.FEATURE_DIM_PER_UNIT)  # Padding

        # Per-unit features (opponent)
        for i in range(self.MAX_UNITS):
            if i < len(state.opponent_units):
                unit = state.opponent_units[i]
                features.extend(self._extract_unit_features(unit))
            else:
                features.extend([0] * self.FEATURE_DIM_PER_UNIT)

        return np.array(features, dtype=np.float32)

    def _extract_unit_features(self, unit) -> List[float]:
        """Extract normalized features for one unit."""
        return [
            unit.battalion_size / 1200.0,           # Normalized by max capacity
            unit.health / 100.0,
            unit.morale / 100.0,
            unit.x / MAP_WIDTH,                     # Normalized position
            unit.y / MAP_HEIGHT,
            unit.has_moved,                          # Boolean features
            unit.is_on_castle,
            # ... ~20 features total ...
        ]
```

**Alternative (Advanced):** Use raw state as tensor, let model learn features
- Requires more data and training time
- Better long-term, but harder to debug
- Start with hand-crafted features, upgrade later

### 2.3 Dataset and DataLoader

**File:** `src/main/python/shardok/ml/dataset.py`

```python
import torch
from torch.utils.data import Dataset, DataLoader
import gzip

class ShardokDataset(Dataset):
    """Dataset for loading game states and win/loss labels."""

    def __init__(self, data_dir: str, feature_extractor: StateFeatureExtractor):
        self.data_dir = data_dir
        self.feature_extractor = feature_extractor

        # Load all game files
        self.game_files = list(Path(data_dir).glob("game_*.bin.gz"))

        # Build index: [(game_file, state_idx), ...]
        self.state_index = []
        for game_file in self.game_files:
            num_states = self._count_states(game_file)
            for i in range(num_states):
                self.state_index.append((game_file, i))

    def __len__(self):
        return len(self.state_index)

    def __getitem__(self, idx):
        game_file, state_idx = self.state_index[idx]

        # Load state (cached or streamed)
        state_bytes, label = self._load_state(game_file, state_idx)

        # Extract features
        features = self.feature_extractor.extract_features(state_bytes)

        return torch.tensor(features), torch.tensor(label, dtype=torch.float32)

# Usage
dataset = ShardokDataset('/tmp/shardok_training_data', StateFeatureExtractor())
train_loader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=4)
```

### 2.4 Model Architecture (v1: Simple Feedforward)

**File:** `src/main/python/shardok/ml/model.py`

```python
import torch
import torch.nn as nn

class ShardokValueNetwork(nn.Module):
    """
    Simple feedforward network for win probability prediction.

    Architecture:
        Input (1330 features)
        → FC(512) → ReLU → Dropout(0.3)
        → FC(256) → ReLU → Dropout(0.3)
        → FC(128) → ReLU
        → FC(1) → Sigmoid
        Output: win probability [0,1]
    """

    def __init__(self, input_dim=1330):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),

            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),

            nn.Linear(256, 128),
            nn.ReLU(),

            nn.Linear(128, 1),
            nn.Sigmoid()  # Output: [0,1] probability
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)  # Shape: (batch_size,)
```

**Why this architecture?**
- Simple to train and debug
- Proven effective for tabular game state data
- ~600K parameters (small by modern standards)

**Future improvements:**
- Add residual connections
- Use transformer for unit interactions
- See Phase 6 for details

### 2.5 Training Loop

**File:** `src/main/python/shardok/ml/train.py`

```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter

def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    num_epochs: int = 50,
    learning_rate: float = 0.001,
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
):
    model = model.to(device)

    # Binary cross-entropy loss (for win probability)
    criterion = nn.BCELoss()

    # Adam optimizer (good default choice)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Learning rate scheduler (reduce on plateau)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5
    )

    # TensorBoard for monitoring
    writer = SummaryWriter('runs/shardok_value_v1')

    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        # Training
        model.train()
        train_loss = 0.0
        train_acc = 0.0

        for features, labels in train_loader:
            features, labels = features.to(device), labels.to(device)

            # Forward pass
            predictions = model(features)
            loss = criterion(predictions, labels)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Metrics
            train_loss += loss.item()
            train_acc += ((predictions > 0.5) == labels).float().mean().item()

        train_loss /= len(train_loader)
        train_acc /= len(train_loader)

        # Validation
        model.eval()
        val_loss = 0.0
        val_acc = 0.0

        with torch.no_grad():
            for features, labels in val_loader:
                features, labels = features.to(device), labels.to(device)
                predictions = model(features)
                loss = criterion(predictions, labels)

                val_loss += loss.item()
                val_acc += ((predictions > 0.5) == labels).float().mean().item()

        val_loss /= len(val_loader)
        val_acc /= len(val_loader)

        # Logging
        print(f"Epoch {epoch+1}/{num_epochs}")
        print(f"  Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
        print(f"  Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")

        writer.add_scalar('Loss/train', train_loss, epoch)
        writer.add_scalar('Loss/val', val_loss, epoch)
        writer.add_scalar('Accuracy/train', train_acc, epoch)
        writer.add_scalar('Accuracy/val', val_acc, epoch)

        # Learning rate scheduling
        scheduler.step(val_loss)

        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'models/shardok_value_best.pth')
            print(f"  → Saved new best model (val_loss: {val_loss:.4f})")

    writer.close()

# Usage
if __name__ == '__main__':
    # Prepare data
    feature_extractor = StateFeatureExtractor()
    dataset = ShardokDataset('/tmp/shardok_training_data', feature_extractor)

    # 80/20 train/val split
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size]
    )

    train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=4)

    # Train model
    model = ShardokValueNetwork(input_dim=feature_extractor.TOTAL_DIM)
    train_model(model, train_loader, val_loader, num_epochs=50)
```

**Expected Training Time:**
- 2.5M examples, batch size 256, ~10K iterations/epoch
- 50 epochs ≈ 500K iterations
- On GPU: ~5-10 hours
- On CPU: ~2-3 days

**Monitoring Training:**
```bash
tensorboard --logdir runs/
# Open http://localhost:6006
```

### 2.6 Hyperparameter Tuning

**Key hyperparameters to tune:**

| Parameter | Initial Value | Search Range | Notes |
|-----------|--------------|--------------|-------|
| Learning rate | 0.001 | [0.0001, 0.01] | Most important |
| Batch size | 256 | [128, 512] | Larger = more stable |
| Dropout | 0.3 | [0.0, 0.5] | Prevent overfitting |
| Hidden dims | [512, 256, 128] | Vary | Bigger = more capacity |
| Weight decay | 0.0001 | [0, 0.001] | L2 regularization |

**Simple grid search:**
```python
for lr in [0.0001, 0.001, 0.01]:
    for batch_size in [128, 256, 512]:
        model = ShardokValueNetwork()
        train_model(model, lr=lr, batch_size=batch_size)
```

**Or use Optuna for automated hyperparameter search:**
https://optuna.readthedocs.io/

---

## Phase 3: Model Integration (C++ Inference)

### Goal
Use the trained PyTorch model from C++ game code.

### 3.1 Export Model to TorchScript

**File:** `src/main/python/shardok/ml/export_model.py`

```python
import torch

# Load trained model
model = ShardokValueNetwork(input_dim=1330)
model.load_state_dict(torch.load('models/shardok_value_best.pth'))
model.eval()

# Convert to TorchScript (C++ compatible format)
example_input = torch.randn(1, 1330)
traced_model = torch.jit.trace(model, example_input)

# Save
traced_model.save('models/shardok_value_v1.pt')
print("Exported to shardok_value_v1.pt")
```

### 3.2 C++ Inference with LibTorch

**Dependencies:** Add to WORKSPACE:
```python
# Download LibTorch
http_archive(
    name = "libtorch",
    urls = ["https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.0%2Bcpu.zip"],
    strip_prefix = "libtorch",
)
```

**File:** Create `src/main/cpp/net/eagle0/shardok/ai/ml/MLValueFunction.hpp`

```cpp
#pragma once

#include <torch/script.h>
#include <torch/torch.h>
#include <memory>
#include <vector>

namespace shardok::ml {

class MLValueFunction {
public:
    MLValueFunction(const std::string& modelPath);

    // Evaluate win probability for current player
    double EvaluateState(const GameStateW& state);

private:
    torch::jit::script::Module model_;

    // Convert game state to feature vector
    std::vector<float> ExtractFeatures(const GameStateW& state);
};

} // namespace shardok::ml
```

**File:** `src/main/cpp/net/eagle0/shardok/ai/ml/MLValueFunction.cpp`

```cpp
#include "MLValueFunction.hpp"

namespace shardok::ml {

MLValueFunction::MLValueFunction(const std::string& modelPath) {
    try {
        model_ = torch::jit::load(modelPath);
        model_.eval();
    } catch (const c10::Error& e) {
        throw std::runtime_error("Failed to load model: " + std::string(e.what()));
    }
}

double MLValueFunction::EvaluateState(const GameStateW& state) {
    // Extract features (matches Python feature extractor)
    std::vector<float> features = ExtractFeatures(state);

    // Convert to tensor
    auto options = torch::TensorOptions().dtype(torch::kFloat32);
    torch::Tensor input = torch::from_blob(
        features.data(),
        {1, static_cast<int64_t>(features.size())},
        options
    ).clone();  // Clone to own the data

    // Forward pass
    torch::NoGradGuard no_grad;  // Disable gradient computation
    torch::Tensor output = model_.forward({input}).toTensor();

    // Extract scalar value
    return output.item<double>();
}

std::vector<float> MLValueFunction::ExtractFeatures(const GameStateW& state) {
    std::vector<float> features;
    features.reserve(1330);  // Must match Python TOTAL_DIM

    // IMPORTANT: Must exactly match Python feature extraction!

    // Global features
    features.push_back(state->current_round() / (float)state.max_rounds());
    // ... (match Python's extract_features exactly) ...

    // Per-unit features (current player)
    for (int i = 0; i < 32; ++i) {
        if (i < current_player_units.size()) {
            const auto* unit = current_player_units[i];
            features.push_back(unit->battalion().size() / 1200.0f);
            features.push_back(unit->health() / 100.0f);
            // ... (match Python's _extract_unit_features exactly) ...
        } else {
            // Padding
            for (int j = 0; j < 20; ++j) features.push_back(0.0f);
        }
    }

    // Per-unit features (opponent)
    // ... (same pattern) ...

    return features;
}

} // namespace shardok::ml
```

**Critical:** Feature extraction in C++ must **exactly match** Python. Mismatch = garbage predictions!

**Testing feature parity:**
```python
# Generate test data
state_bytes = ... # From real game
python_features = feature_extractor.extract_features(state_bytes)

# Save to file
np.save('test_features.npy', python_features)
```

```cpp
// Load and compare
std::vector<float> cpp_features = ExtractFeatures(state);
std::vector<float> python_features = LoadNpy("test_features.npy");

for (size_t i = 0; i < cpp_features.size(); ++i) {
    if (std::abs(cpp_features[i] - python_features[i]) > 1e-5) {
        std::cerr << "Mismatch at index " << i << std::endl;
    }
}
```

### 3.3 Integrate with MCTS

**File:** Modify `src/main/cpp/net/eagle0/shardok/ai/mcts/adapters/ShardokGameState.cpp`

```cpp
#include "src/main/cpp/net/eagle0/shardok/ai/ml/MLValueFunction.hpp"

double ShardokGameState::score(MCTSPlayerId playerId) const {
    // Option 1: Use ML value function
    static auto mlValueFunc = std::make_shared<ml::MLValueFunction>(
        "models/shardok_value_v1.pt"
    );

    return mlValueFunc->EvaluateState(state_);

    // Option 2: Hybrid approach
    // Use ML for win probability, scale by hand-crafted value for granularity
    // double win_prob = mlValueFunc->EvaluateState(state_);
    // double heuristic_value = AIScoreCalculator::GuessedStateScore(...);
    // return win_prob * 1000.0 + heuristic_value * 0.1;
}
```

**For MCTS:** Win probability [0,1] is perfect! No normalization needed!

### 3.4 Performance Considerations

**Inference speed:**
- Forward pass: ~1-5ms on CPU
- MCTS does thousands of evaluations → optimize!

**Optimizations:**
1. **Batch inference**: Accumulate states, evaluate in batches
2. **Model quantization**: Convert to int8 (4x faster, minimal accuracy loss)
3. **ONNX Runtime**: Faster than LibTorch for inference
4. **GPU**: If available, 10-100x faster

**Example batch inference:**
```cpp
std::vector<double> EvaluateBatch(const std::vector<GameStateW>& states) {
    // Extract all features
    std::vector<std::vector<float>> all_features;
    for (const auto& state : states) {
        all_features.push_back(ExtractFeatures(state));
    }

    // Stack into single tensor
    torch::Tensor batch = torch::zeros({states.size(), 1330});
    for (size_t i = 0; i < states.size(); ++i) {
        batch[i] = torch::from_blob(all_features[i].data(), {1330});
    }

    // Single forward pass
    torch::Tensor output = model_.forward({batch}).toTensor();

    // Extract results
    std::vector<double> results;
    for (int i = 0; i < states.size(); ++i) {
        results.push_back(output[i].item<double>());
    }
    return results;
}
```

---

## Phase 4: Iterative Improvement (Self-Play Loop)

### Goal
Use the ML model to generate better training data, creating a virtuous improvement cycle.

### 4.1 Self-Play with MCTS

Once the ML model is integrated:

```cpp
// Configure MCTS to use ML value function
MCTSConfig config;
config.simulationPolicy = MCTSSimulationPolicy::BEST_IMMEDIATE;
config.useMLValueFunction = true;  // New flag

// Generate self-play games
GenerateSelfPlayGames(
    num_games = 10000,
    ai_type = "MCTS",  // Use MCTS instead of iterative deepening
    output_dir = "/tmp/shardok_training_data_v2"
);
```

### 4.2 Iterative Training Loop

**Pseudo-code:**
```
Model₀ = train_on_data(iterative_deepening_games)

for iteration in [1, 2, 3, ...]:
    # Generate new self-play data
    games_i = generate_games(MCTS with Model_{i-1})

    # Combine with previous data (optional)
    data_i = games_i + sample(previous_data, 50%)

    # Retrain
    Model_i = train_on_data(data_i)

    # Evaluate improvement
    elo_rating = tournament(Model_i vs Model_{i-1} vs iterative_deepening)

    if no_improvement_for_3_iterations:
        break
```

### 4.3 Evaluation

**How to measure improvement:**

1. **Head-to-head tournaments**
   ```cpp
   RunTournament({
       {"Model_v1", mcts_with_v1},
       {"Model_v2", mcts_with_v2},
       {"Baseline", iterative_deepening}
   }, num_games=1000);
   ```

2. **Elo ratings**: Track relative strength
   ```
   Model_v1: 1500 Elo (baseline)
   Model_v2: 1580 Elo (+80)
   Model_v3: 1650 Elo (+70)
   Iterative Deepening: 1450 Elo (reference)
   ```

3. **Win rate on test scenarios**
   - Create 100 fixed test scenarios
   - Measure win % for each model
   - Track improvement over iterations

---

## Phase 5: Production Deployment

### 5.1 Model Versioning

**Directory structure:**
```
models/
├── v1/
│   ├── shardok_value.pt          # TorchScript model
│   ├── feature_extractor.py      # For reproducibility
│   └── metadata.json              # Training info
├── v2/
│   └── ...
└── production -> v2/              # Symlink to current production model
```

### 5.2 Model Loading

```cpp
class MLModelManager {
public:
    static MLModelManager& Instance() {
        static MLModelManager instance;
        return instance;
    }

    MLValueFunction* GetModel(const std::string& version = "production") {
        std::lock_guard lock(mutex_);

        if (models_.find(version) == models_.end()) {
            std::string path = "models/" + version + "/shardok_value.pt";
            models_[version] = std::make_unique<MLValueFunction>(path);
        }

        return models_[version].get();
    }

private:
    std::mutex mutex_;
    std::map<std::string, std::unique_ptr<MLValueFunction>> models_;
};
```

### 5.3 Fallback Strategy

```cpp
double EvaluateState(const GameStateW& state) {
    try {
        return MLModelManager::Instance().GetModel()->EvaluateState(state);
    } catch (const std::exception& e) {
        // Fallback to hand-crafted scoring
        LOG(WARNING) << "ML evaluation failed: " << e.what();
        return AIScoreCalculator::GuessedStateScore(...);
    }
}
```

---

## Phase 6: Richer Outcomes (Beyond Win/Loss)

### Motivation

In the larger strategic game, not all outcomes are equal:
- **Crushing Victory** (all enemies dead) > **Standard Victory** (enemies fled)
- **Tactical Retreat** (preserved army) > **Pyrrhic Victory** (won but lost army)
- **Total Defeat** (all units captured) is catastrophic

### 6.1 Outcome Taxonomy

```cpp
enum class BattleOutcome {
    // Victories
    CRUSHING_VICTORY,      // All enemies eliminated, minimal losses
    STANDARD_VICTORY,      // Enemies fled/routed, moderate losses
    PYRRHIC_VICTORY,       // Won but heavy casualties

    // Defeats
    TACTICAL_RETREAT,      // Lost but preserved army (successful flee)
    HEAVY_DEFEAT,          // Lost with heavy casualties
    TOTAL_DEFEAT,          // Complete annihilation/capture
};

struct BattleResult {
    BattleOutcome outcome;
    float units_preserved_ratio;   // [0,1] % of starting units preserved
    float casualties_inflicted;     // Damage dealt to opponent
    int rounds_survived;           // How long we lasted
};
```

### 6.2 Multi-Task Learning Architecture

Instead of single output (win probability), train **multiple outputs**:

```python
class ShardokValueNetworkV2(nn.Module):
    def __init__(self, input_dim=1330):
        super().__init__()

        # Shared trunk (feature extractor)
        self.shared = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
        )

        # Task-specific heads
        self.win_prob_head = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # Win probability [0,1]
        )

        self.outcome_head = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 6),  # 6 outcome classes
            nn.Softmax(dim=-1)
        )

        self.preservation_head = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # Units preserved [0,1]
        )

    def forward(self, x):
        shared_features = self.shared(x)

        return {
            'win_prob': self.win_prob_head(shared_features),
            'outcome': self.outcome_head(shared_features),
            'preservation': self.preservation_head(shared_features),
        }
```

**Training with multiple losses:**
```python
def compute_loss(predictions, targets):
    # Binary cross-entropy for win probability
    win_loss = F.binary_cross_entropy(predictions['win_prob'], targets['win'])

    # Cross-entropy for outcome classification
    outcome_loss = F.cross_entropy(predictions['outcome'], targets['outcome'])

    # MSE for unit preservation
    preservation_loss = F.mse_loss(predictions['preservation'], targets['preservation'])

    # Weighted combination
    total_loss = (
        1.0 * win_loss +
        0.5 * outcome_loss +
        0.3 * preservation_loss
    )

    return total_loss
```

### 6.3 Using Richer Outcomes

**For strategic AI (Eagle layer):**
```cpp
// When deciding which battle to fight
auto outcome_prediction = ml_model->EvaluateRichOutcome(state);

if (outcome_prediction.win_prob > 0.6 &&
    outcome_prediction.preservation_ratio > 0.7) {
    // Good odds, low losses → fight
} else if (outcome_prediction.win_prob < 0.3 &&
           outcome_prediction.outcome == BattleOutcome::TACTICAL_RETREAT) {
    // Can retreat successfully → flee to fight another day
} else {
    // Avoid battle
}
```

**For MCTS (tactical layer):**
```cpp
// Use win probability for tree search
double ucb_score = predictions.win_prob + exploration_bonus;

// But use richer outcomes for final move selection
// Prefer moves that preserve units even if win prob is similar
if (move_a.win_prob ≈ move_b.win_prob) {
    choose move with higher preservation_ratio
}
```

---

## Phase 7: Advanced Topics

### 7.1 Curriculum Learning

Start with simpler scenarios, gradually increase difficulty:

```python
# Iteration 1: Small armies (5-10 units)
GenerateSelfPlayGames(scenario_size='small', num_games=10000)

# Iteration 2: Medium armies (10-20 units)
GenerateSelfPlayGames(scenario_size='medium', num_games=10000)

# Iteration 3: Full complexity
GenerateSelfPlayGames(scenario_size='any', num_games=20000)
```

### 7.2 Data Augmentation

Generate more training data by:
- **Symmetry**: Flip map horizontally/vertically
- **Player swap**: Train on both P0 and P1 perspectives
- **Noise injection**: Small perturbations to unit stats

### 7.3 Ensemble Models

Train multiple models, combine predictions:
```cpp
double EvaluateEnsemble(const GameStateW& state) {
    double sum = 0.0;
    for (auto& model : ensemble) {
        sum += model->Evaluate(state);
    }
    return sum / ensemble.size();
}
```

Reduces variance, improves robustness.

### 7.4 Explainability

**Attention mechanisms** to understand what model is looking at:
```python
class AttentionValueNetwork(nn.Module):
    def forward(self, x):
        # ...
        attention_weights = self.attention(x)
        # attention_weights shows which features are important
```

**Visualization:** Heatmap of board showing which positions matter most.

---

## Learning Resources

### PyTorch Tutorials
- **Official Tutorial**: https://pytorch.org/tutorials/beginner/basics/intro.html
- **60-minute blitz**: https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html
- **Training a classifier**: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html

### Game AI Specific
- **AlphaGo Zero paper**: https://www.nature.com/articles/nature24270
  - Skip the RL parts, focus on value network architecture
- **Leela Chess Zero**: https://lczero.org/dev/wiki/
  - Open-source chess engine using similar approach
- **MuZero**: https://arxiv.org/abs/1911.08265
  - More advanced, but good for ideas

### LibTorch (C++ Inference)
- **LibTorch Tutorial**: https://pytorch.org/tutorials/advanced/cpp_frontend.html
- **Loading models in C++**: https://pytorch.org/tutorials/advanced/cpp_export.html

### General ML
- **Fast.ai course** (practical, top-down): https://course.fast.ai/
- **Deep Learning book** (theoretical): https://www.deeplearningbook.org/

---

## Implementation Timeline

### Week 1-2: Data Collection
- [ ] Implement `StateRecorder`
- [ ] Implement `SelfPlayGenerator`
- [ ] Generate 50K games
- [ ] Verify data quality

### Week 3-4: Feature Engineering & Training
- [ ] Implement `StateFeatureExtractor` (Python)
- [ ] Implement `ShardokDataset`
- [ ] Train v1 model
- [ ] Achieve >70% validation accuracy

### Week 5-6: C++ Integration
- [ ] Implement `MLValueFunction` (C++)
- [ ] Verify feature parity (Python vs C++)
- [ ] Integrate with MCTS
- [ ] Benchmark inference speed

### Week 7-8: Evaluation & Iteration
- [ ] Run tournaments (ML vs baseline)
- [ ] Generate self-play data with MCTS
- [ ] Retrain v2 model
- [ ] Measure improvement

### Week 9+: Production & Advanced Features
- [ ] Deploy to production
- [ ] Implement richer outcomes (multi-task learning)
- [ ] Curriculum learning
- [ ] Continuous improvement loop

---

## Success Metrics

### Phase 1 (MVP)
- ✅ Train model with >70% win prediction accuracy
- ✅ MCTS with ML model matches or beats baseline AI
- ✅ Inference time <5ms per evaluation

### Phase 2 (Iteration)
- ✅ Self-play iteration improves Elo by +50 points
- ✅ Model generalizes across different map types

### Phase 3 (Production)
- ✅ Stable inference in production
- ✅ <1% fallback rate to hand-crafted scoring
- ✅ Reduced maintenance burden (fewer hand-tuned parameters)

---

## Open Questions & Future Work

1. **Feature engineering vs raw state:**
   - Start with hand-crafted features
   - Long-term: Learn features end-to-end (transformer on unit embeddings)

2. **Model architecture:**
   - Start with feedforward network
   - Upgrade to attention/transformer when data is abundant

3. **Transfer learning:**
   - Can model trained on small maps transfer to large maps?
   - Fine-tuning vs training from scratch

4. **Opponent modeling:**
   - Current approach assumes symmetric gameplay
   - Could train on games with varied AI opponents

5. **Multi-agent learning:**
   - P0 and P1 have different objectives (attacker vs defender)
   - Separate models or shared with position encoding?

---

## Appendix: Quick Reference

### File Structure
```
eagle0/
├── src/main/cpp/net/eagle0/shardok/ai/ml/
│   ├── StateRecorder.hpp/cpp          # Record game states
│   ├── SelfPlayGenerator.cpp          # Generate training data
│   ├── MLValueFunction.hpp/cpp        # C++ inference
│   └── BUILD.bazel
├── src/main/python/shardok/ml/
│   ├── feature_extractor.py           # State → features
│   ├── dataset.py                     # PyTorch dataset
│   ├── model.py                       # Neural network
│   ├── train.py                       # Training loop
│   └── export_model.py                # PyTorch → TorchScript
└── models/
    ├── v1/shardok_value.pt
    └── production -> v1/
```

### Commands Cheat Sheet
```bash
# Generate training data (C++)
bazel run //src/main/cpp/net/eagle0/shardok/ai/ml:self_play_generator -- \
    --num_games 50000 --output_dir /tmp/training_data

# Train model (Python)
python src/main/python/shardok/ml/train.py \
    --data_dir /tmp/training_data \
    --epochs 50 \
    --batch_size 256

# Export model (Python)
python src/main/python/shardok/ml/export_model.py \
    --input models/shardok_value_best.pth \
    --output models/v1/shardok_value.pt

# Run with ML model (C++)
bazel run //src/main/cpp/net/eagle0/shardok:shardok-server -- \
    --use_ml_value_function=true \
    --ml_model_path=models/v1/shardok_value.pt
```

### Common Pitfalls
- ❌ **Feature mismatch**: Python and C++ extract different features
- ❌ **Label leakage**: Training on future information
- ❌ **Overfitting**: Model memorizes positions instead of learning strategy
- ❌ **Distribution shift**: Training on one map type, testing on another
- ❌ **Slow inference**: Not batching evaluations, using unoptimized model

### Debugging Checklist
- [ ] Features in C++ exactly match Python?
- [ ] Model loaded correctly (no errors)?
- [ ] Validation accuracy >60%?
- [ ] Predictions in [0,1] range?
- [ ] Inference time <10ms?
- [ ] Model improves over baseline in tournaments?
