garvitsachdeva/SpindleFlow-RL

1"""2State encoder for the policy network.3MLP-based (replaces GNN from v3 design — too complex for hackathon timeline).4Document: GNN would be used in production for the delegation graph component.5"""6 7from __future__ import annotations8import torch9import torch.nn as nn10 11 12class StateEncoder(nn.Module):13    """14    Encodes the flat state vector into a compressed representation.15    The SB3 policy will use this as its feature extractor.16 17    Architecture:18      - Input: flat state vector (~1376 + N*768 dims)19      - Hidden: 512 → 256 → 12820      - Output: 128-dim feature vector21 22    Note: The MLP operates on the full flat vector including:23      - Task embedding (384)24      - Roster + called specialist embeddings (padded)25      - Graph adjacency vector (100)26      - Scratchpad summary (384)27      - Scalar features (8)28    This is the "MLP adjacency" approach that replaces the GNN.29    """30 31    def __init__(self, input_dim: int, output_dim: int = 128):32        super().__init__()33        self.network = nn.Sequential(34            nn.Linear(input_dim, 512),35            nn.LayerNorm(512),36            nn.ReLU(),37            nn.Dropout(0.1),38            nn.Linear(512, 256),39            nn.LayerNorm(256),40            nn.ReLU(),41            nn.Linear(256, output_dim),42            nn.ReLU(),43        )44 45    def forward(self, x: torch.Tensor) -> torch.Tensor:46        return self.network(x)47