gouhanke
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import abc
|
|
import torch
|
|
import torch.nn as nn
|
|
from typing import Dict, Any, Optional
|
|
|
|
class VLABackbone(nn.Module, abc.ABC):
|
|
"""
|
|
Contract for Vision/Language Backbones.
|
|
Must return a feature tensor of shape (B, Seq, Embed_Dim).
|
|
"""
|
|
@abc.abstractmethod
|
|
def forward(self, obs: Dict[str, torch.Tensor]) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
obs: Dictionary containing 'image' and optionally 'text'.
|
|
Returns:
|
|
features: (B, S, D) embedding.
|
|
"""
|
|
pass
|
|
|
|
|
|
class VLAProjector(nn.Module, abc.ABC):
|
|
"""
|
|
Contract for the adaptation layer (Projector).
|
|
Connects Backbone features to the Policy Head.
|
|
"""
|
|
@abc.abstractmethod
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
pass
|
|
|
|
|
|
class VLAHead(nn.Module, abc.ABC):
|
|
"""
|
|
Contract for Action Generation Heads (Policies).
|
|
Handles both training (loss calculation) and inference (action generation).
|
|
"""
|
|
@abc.abstractmethod
|
|
def forward(self, embeddings: torch.Tensor, actions: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
embeddings: (B, S, Hidden) from Projector.
|
|
actions: (B, Pred_Horizon, Action_Dim) - Ground truth for training.
|
|
Returns:
|
|
Dict containing 'loss' (if actions provided) or 'pred_actions'.
|
|
"""
|
|
pass |