Add multimodal observation schema and DQN encoders

Author: sbischoff-ai

README.md

@@ -70,3 +70,30 @@ uv run black --check adeptly tests
 uv run mypy adeptly
 uv run pytest
 ```
+
+
+### Multimodal DQN inference (batched)
+```python
+import torch
+from adeptly.observations import MultimodalQNetwork, ObservationBatch
+
+batch_size = 8
+num_actions = 4
+model = MultimodalQNetwork(
+    image_channels=3,
+    telemetry_dim=6,
+    action_size=num_actions,
+    sequence_vocab_size=256,
+)
+
+# Minimal batched environment loop for inference-only usage.
+for _ in range(5):
+    observations = ObservationBatch(
+        image_frames=torch.randint(0, 256, (batch_size, 3, 84, 84), dtype=torch.uint8),
+        scalar_telemetry=torch.randn(batch_size, 6),
+        events_or_text=torch.randint(0, 256, (batch_size, 12), dtype=torch.int64),
+    )
+    q_values = model(observations)
+    actions = torch.argmax(q_values, dim=1)
+    # send `actions` back to your vectorized environment
+```

adeptly/observations/__init__.py

@@ -0,0 +1,22 @@
+"""Observation schemas and multimodal DQN building blocks."""
+
+from adeptly.observations.dqn_multimodal import MultimodalQNetwork
+from adeptly.observations.encoders import EventSequenceEncoder, ModalityFusion, TelemetryEncoder, VisionEncoder
+from adeptly.observations.schema import (
+    ObservationBatch,
+    normalize_image_frames,
+    normalize_scalar_telemetry,
+    validate_observation_batch,
+)
+
+__all__ = [
+    "EventSequenceEncoder",
+    "ModalityFusion",
+    "MultimodalQNetwork",
+    "ObservationBatch",
+    "TelemetryEncoder",
+    "VisionEncoder",
+    "normalize_image_frames",
+    "normalize_scalar_telemetry",
+    "validate_observation_batch",
+]

adeptly/observations/dqn_multimodal.py

@@ -0,0 +1,71 @@
+"""Multimodal DQN network composed from observation encoders and fusion."""
+
+from __future__ import annotations
+
+from torch import Tensor, nn
+
+from adeptly.observations.encoders import EventSequenceEncoder, ModalityFusion, TelemetryEncoder, VisionEncoder
+from adeptly.observations.schema import (
+    ObservationBatch,
+    normalize_image_frames,
+    normalize_scalar_telemetry,
+    validate_observation_batch,
+)
+
+
+class MultimodalQNetwork(nn.Module):
+    """DQN action-value network for multimodal observations."""
+
+    def __init__(
+        self,
+        image_channels: int,
+        telemetry_dim: int,
+        action_size: int,
+        *,
+        vision_embedding_dim: int = 128,
+        telemetry_embedding_dim: int = 64,
+        sequence_vocab_size: int | None = None,
+        sequence_embedding_dim: int = 32,
+        sequence_output_dim: int = 64,
+        fused_dim: int = 256,
+        use_attention_fusion: bool = False,
+    ) -> None:
+        super().__init__()
+
+        if action_size <= 0:
+            raise ValueError("action_size must be > 0")
+
+        self.vision_encoder = VisionEncoder(image_channels, vision_embedding_dim)
+        self.telemetry_encoder = TelemetryEncoder(telemetry_dim, telemetry_embedding_dim)
+
+        self.sequence_encoder: EventSequenceEncoder | None = None
+        encoder_dims = [vision_embedding_dim, telemetry_embedding_dim]
+        if sequence_vocab_size is not None:
+            self.sequence_encoder = EventSequenceEncoder(
+                vocab_size=sequence_vocab_size,
+                embedding_dim=sequence_embedding_dim,
+                output_dim=sequence_output_dim,
+            )
+            encoder_dims.append(sequence_output_dim)
+
+        self.fusion = ModalityFusion(encoder_dims, fused_dim=fused_dim, use_attention=use_attention_fusion)
+        self.action_head = nn.Linear(fused_dim, action_size)
+
+    def forward(self, observations: ObservationBatch) -> Tensor:
+        validate_observation_batch(observations)
+
+        image_frames = normalize_image_frames(observations.image_frames)
+        scalar_telemetry = normalize_scalar_telemetry(observations.scalar_telemetry)
+
+        embeddings = [
+            self.vision_encoder(image_frames),
+            self.telemetry_encoder(scalar_telemetry),
+        ]
+
+        if observations.events_or_text is not None:
+            if self.sequence_encoder is None:
+                raise ValueError("events_or_text provided but sequence encoder was not configured")
+            embeddings.append(self.sequence_encoder(observations.events_or_text))
+
+        fused = self.fusion(embeddings)
+        return self.action_head(fused)

adeptly/observations/encoders.py

@@ -0,0 +1,88 @@
+"""Composable multimodal encoders for DQN observations."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+
+class VisionEncoder(nn.Module):
+    """Small CNN encoder for image frame observations."""
+
+    def __init__(self, in_channels: int, embedding_dim: int) -> None:
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Conv2d(in_channels, 32, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, stride=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(),
+            nn.Linear(64, embedding_dim),
+            nn.ReLU(),
+        )
+
+    def forward(self, image_frames: Tensor) -> Tensor:
+        return self.backbone(image_frames)
+
+
+class TelemetryEncoder(nn.Module):
+    """MLP encoder for scalar telemetry channels."""
+
+    def __init__(self, input_dim: int, embedding_dim: int, hidden_dim: int = 128) -> None:
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, embedding_dim),
+            nn.ReLU(),
+        )
+
+    def forward(self, scalar_telemetry: Tensor) -> Tensor:
+        return self.model(scalar_telemetry)
+
+
+class EventSequenceEncoder(nn.Module):
+    """Embedding + GRU encoder for event/text token sequences."""
+
+    def __init__(self, vocab_size: int, embedding_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.gru = nn.GRU(embedding_dim, output_dim, batch_first=True)
+
+    def forward(self, events_or_text: Tensor) -> Tensor:
+        embedded = self.embedding(events_or_text.long())
+        _, hidden = self.gru(embedded)
+        return hidden.squeeze(0)
+
+
+class ModalityFusion(nn.Module):
+    """Fuse modality embeddings via concatenation projection or attention pooling."""
+
+    def __init__(self, input_dims: list[int], fused_dim: int, use_attention: bool = False) -> None:
+        super().__init__()
+        self.use_attention = use_attention
+        self.fused_dim = fused_dim
+
+        if use_attention:
+            if len(set(input_dims)) != 1:
+                raise ValueError("All modality dims must match when use_attention=True")
+            self.attention = nn.MultiheadAttention(embed_dim=input_dims[0], num_heads=1, batch_first=True)
+            self.output_projection = nn.Linear(input_dims[0], fused_dim)
+        else:
+            self.output_projection = nn.Linear(sum(input_dims), fused_dim)
+
+    def forward(self, modality_embeddings: list[Tensor]) -> Tensor:
+        if not modality_embeddings:
+            raise ValueError("modality_embeddings cannot be empty")
+
+        if self.use_attention:
+            stacked = torch.stack(modality_embeddings, dim=1)
+            attended, _ = self.attention(stacked, stacked, stacked)
+            pooled = attended.mean(dim=1)
+            return self.output_projection(pooled)
+
+        concatenated = torch.cat(modality_embeddings, dim=1)
+        return self.output_projection(concatenated)

adeptly/observations/schema.py

@@ -0,0 +1,57 @@
+"""Observation schema and validation utilities for multimodal DQN inputs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+
+@dataclass(slots=True)
+class ObservationBatch:
+    """Batched multimodal observations consumed by multimodal DQN networks."""
+
+    image_frames: Tensor
+    scalar_telemetry: Tensor
+    events_or_text: Tensor | None = None
+
+
+def validate_observation_batch(observations: ObservationBatch) -> None:
+    """Validate batch shapes early to prevent downstream runtime shape errors."""
+
+    image = observations.image_frames
+    telemetry = observations.scalar_telemetry
+    sequence = observations.events_or_text
+
+    if image.ndim != 4:
+        raise ValueError("image_frames must have shape [batch, channels, height, width]")
+    if telemetry.ndim != 2:
+        raise ValueError("scalar_telemetry must have shape [batch, features]")
+
+    batch_size = image.shape[0]
+    if telemetry.shape[0] != batch_size:
+        raise ValueError("image_frames and scalar_telemetry must share the same batch dimension")
+
+    if sequence is not None:
+        if sequence.ndim != 2:
+            raise ValueError("events_or_text must have shape [batch, sequence_length]")
+        if sequence.shape[0] != batch_size:
+            raise ValueError("events_or_text must share the same batch dimension as image_frames")
+
+
+def normalize_image_frames(image_frames: Tensor) -> Tensor:
+    """Convert image frames to float tensors in [0, 1]."""
+
+    if not torch.is_floating_point(image_frames):
+        image_frames = image_frames.to(torch.float32)
+    return image_frames / 255.0 if image_frames.max().item() > 1.0 else image_frames
+
+
+def normalize_scalar_telemetry(scalar_telemetry: Tensor, eps: float = 1e-6) -> Tensor:
+    """Per-batch z-score normalization for scalar telemetry channels."""
+
+    telemetry = scalar_telemetry.to(torch.float32)
+    mean = telemetry.mean(dim=0, keepdim=True)
+    std = telemetry.std(dim=0, keepdim=True).clamp_min(eps)
+    return (telemetry - mean) / std

tests/test_multimodal_observations.py

@@ -0,0 +1,62 @@
+import pytest
+import torch
+
+from adeptly.observations import (
+    MultimodalQNetwork,
+    ObservationBatch,
+    normalize_image_frames,
+    normalize_scalar_telemetry,
+    validate_observation_batch,
+)
+
+
+def test_validate_observation_batch_rejects_mismatched_batch_dims():
+    observations = ObservationBatch(
+        image_frames=torch.zeros(2, 3, 84, 84),
+        scalar_telemetry=torch.zeros(3, 4),
+    )
+
+    with pytest.raises(ValueError, match="share the same batch dimension"):
+        validate_observation_batch(observations)
+
+
+def test_normalization_helpers_return_expected_ranges_and_stats():
+    image = torch.randint(0, 256, (4, 3, 16, 16), dtype=torch.uint8)
+    normalized_image = normalize_image_frames(image)
+    assert normalized_image.dtype == torch.float32
+    assert float(normalized_image.max()) <= 1.0
+
+    telemetry = torch.tensor([[1.0, 2.0], [3.0, 6.0], [5.0, 10.0]])
+    normalized_telemetry = normalize_scalar_telemetry(telemetry)
+    assert torch.allclose(normalized_telemetry.mean(dim=0), torch.zeros(2), atol=1e-5)
+
+
+def test_multimodal_q_network_outputs_action_values_for_batched_input():
+    model = MultimodalQNetwork(
+        image_channels=3,
+        telemetry_dim=5,
+        action_size=4,
+        sequence_vocab_size=128,
+        use_attention_fusion=False,
+    )
+
+    observations = ObservationBatch(
+        image_frames=torch.randint(0, 256, (6, 3, 84, 84), dtype=torch.uint8),
+        scalar_telemetry=torch.randn(6, 5),
+        events_or_text=torch.randint(0, 128, (6, 10), dtype=torch.int64),
+    )
+
+    output = model(observations)
+    assert output.shape == (6, 4)
+
+
+def test_multimodal_q_network_raises_if_sequence_not_configured():
+    model = MultimodalQNetwork(image_channels=3, telemetry_dim=3, action_size=2)
+    observations = ObservationBatch(
+        image_frames=torch.randint(0, 256, (2, 3, 84, 84), dtype=torch.uint8),
+        scalar_telemetry=torch.randn(2, 3),
+        events_or_text=torch.randint(0, 100, (2, 8), dtype=torch.int64),
+    )
+
+    with pytest.raises(ValueError, match="sequence encoder"):
+        model(observations)