File size: 5,635 Bytes
20c251e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | 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")
|