import torch
import torch.nn as nn
from roboimi.vla.core.interfaces import VLAProjector

class MLPProjector(VLAProjector):
    """
    A simple Linear Projection layer.
    First-class citizen: Adapts Backbone dim -> Head dim.
    """
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.GELU(),
            nn.Linear(output_dim, output_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)