from __future__ import annotations

from typing import Any

import torch.nn as nn

from competitive_vla_model import CompetitiveVLAModel
from hybrid_vla_model import HybridCompetitiveVLAModel


class MultiBackboneRoutedVLAModel(nn.Module):
    """A compact Base router plus a DINOv2-Large specialist."""

    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__()
        self.policy_config = config
        self.base_model = HybridCompetitiveVLAModel(config["base_router_config"])
        self.large_model = CompetitiveVLAModel(config["large_config"])