tests/test_openvla_adapter.py

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace

import pytest

torch = pytest.importorskip("torch")

from dovla_cil.data.schema import ActionChunk
from dovla_cil.models.dovla import DoVLAConfig, DoVLAModel
from dovla_cil.models.openvla_adapter import (
    ExternalOpenVLAAdapter,
    PretrainedCLIPBackbone,
    ToyVLABackbone,
    VLABackbone,
)


def _config() -> DoVLAConfig:
    return DoVLAConfig(
        obs_dim=10,
        lang_dim=16,
        action_dim=8,
        hidden_dim=32,
        action_horizon=3,
        effect_dim=7,
        intervention_dim=24,
    )


def test_toy_vla_backbone_works() -> None:
    config = _config()
    backbone = ToyVLABackbone(config)
    observation = torch.randn(2, config.obs_dim)
    instructions = ["pick the mug", "open the drawer"]
    action = torch.randn(2, config.action_horizon, config.action_dim)

    context = backbone.encode_observation_language(observation, instructions)
    action_z = backbone.encode_action(action)
    policy = backbone.forward_policy(observation, instructions)
    intervention = backbone.forward_intervention(observation, instructions, action)
    decoded = backbone.decode_action(policy)

    assert isinstance(backbone, VLABackbone)
    assert context.shape == (2, config.hidden_dim)
    assert action_z.shape == (2, config.hidden_dim)
    assert policy.shape == (2, config.action_horizon, config.action_dim)
    assert intervention.shape == (2, config.intervention_dim)
    assert isinstance(decoded, ActionChunk)


def test_external_openvla_adapter_requires_configuration() -> None:
    with pytest.raises(NotImplementedError) as exc_info:
        ExternalOpenVLAAdapter()

    assert "checkpoint_path" in str(exc_info.value)


def test_external_openvla_adapter_methods_are_placeholders() -> None:
    adapter = ExternalOpenVLAAdapter(checkpoint_path=Path("missing-openvla-checkpoint"))

    with pytest.raises(NotImplementedError) as exc_info:
        adapter.forward_policy(None, "pick the mug")

    assert "extension point only" in str(exc_info.value)


def test_dovla_model_can_use_injected_backbone() -> None:
    config = _config()
    backbone = ToyVLABackbone(config)
    model = DoVLAModel(config, backbone=backbone)
    observation = torch.randn(2, config.obs_dim)
    instructions = ["pick the mug", "open the drawer"]
    action = torch.randn(2, config.action_horizon, config.action_dim)

    policy = model.forward_policy(observation, instructions)
    effect = model.forward_effect(observation, instructions, action)
    reward = model.forward_reward(observation, instructions, action)
    z = model.encode_intervention(observation, instructions, action)

    assert model.backbone is backbone
    assert policy.shape == (2, config.action_horizon, config.action_dim)
    assert effect["effect_vector"].shape == (2, config.effect_dim)
    assert reward.shape == (2,)
    assert z.shape == (2, config.intervention_dim)


class _FakeCLIP(torch.nn.Module):
    def __init__(self, projection_dim: int = 12) -> None:
        super().__init__()
        self.anchor = torch.nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(
            projection_dim=projection_dim,
            vision_config=SimpleNamespace(image_size=8),
        )

    def get_image_features(self, *, pixel_values):
        values = pixel_values.mean(dim=(1, 2, 3), keepdim=False).unsqueeze(1)
        return values.repeat(1, self.config.projection_dim) + self.anchor

    def get_text_features(self, *, input_ids, **_kwargs):
        values = input_ids.float().mean(dim=1, keepdim=True)
        return values.repeat(1, self.config.projection_dim) + self.anchor


class _FakeProcessor:
    tokenizer = None

    def __init__(self) -> None:
        self.tokenizer = self

    def __call__(self, texts, **_kwargs):
        return {
            "input_ids": torch.tensor(
                [[len(word) for word in text.split()][:4] for text in texts],
                dtype=torch.long,
            )
        }


def test_pretrained_clip_backbone_supports_raw_and_cached_features() -> None:
    config = DoVLAConfig(
        obs_dim=10,
        lang_dim=16,
        action_dim=8,
        hidden_dim=32,
        action_horizon=3,
        effect_dim=7,
        intervention_dim=24,
        observation_mode="rgb",
        backbone_type="native",
    )
    backbone = PretrainedCLIPBackbone(
        config,
        clip_model=_FakeCLIP(),
        processor=_FakeProcessor(),
    )
    images = torch.randint(0, 255, (2, 12, 10, 3), dtype=torch.uint8)
    instructions = ["pick red cube", "push blue cube"]
    action = torch.randn(2, config.action_horizon, config.action_dim)

    cached = backbone.encode_pretrained_features(images, instructions)
    raw_context = backbone.encode_observation_language(images, instructions)
    cached_context = backbone.encode_observation_language(cached, instructions)
    intervention = backbone.forward_intervention(cached, instructions, action)

    assert cached.shape == (2, 24)
    assert torch.allclose(raw_context, cached_context)
    assert raw_context.shape == (2, config.hidden_dim)
    assert intervention.shape == (2, config.intervention_dim)
    assert not any(parameter.requires_grad for parameter in backbone.clip_model.parameters())


def test_clip_model_config_requires_rgb_and_model_path() -> None:
    with pytest.raises(ValueError, match="observation_mode"):
        DoVLAConfig(backbone_type="clip", backbone_model="local-clip")
    with pytest.raises(ValueError, match="backbone_model"):
        DoVLAConfig(observation_mode="rgb", backbone_type="clip")