Logic
231 lines
8.8 KiB
Python
231 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Mapping, Sequence
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from roboimi.vla.core.interfaces import VLABackbone
|
|
|
|
|
|
class _LEWMProjector(nn.Module):
|
|
"""LEWM projector MLP: 192 -> 2048 -> 192 with BatchNorm1d + GELU."""
|
|
|
|
def __init__(self, input_dim: int = 192, hidden_dim: int = 2048, output_dim: int = 192) -> None:
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(input_dim, hidden_dim),
|
|
nn.BatchNorm1d(hidden_dim),
|
|
nn.GELU(),
|
|
nn.Linear(hidden_dim, output_dim),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.net(x)
|
|
|
|
|
|
class LEWMViTBackbone(VLABackbone):
|
|
"""Frozen LEWM joint-multiview ViT backbone.
|
|
|
|
The backbone fuses the three camera views into a single LEWM-style image,
|
|
runs a ViT-tiny encoder plus the LEWM projector, and returns one joint
|
|
192-d embedding per timestep.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
checkpoint_path: str | Path | None = None,
|
|
*,
|
|
checkpoint: Mapping[str, Any] | None = None,
|
|
camera_names: Sequence[str] = ("r_vis", "top", "front"),
|
|
fused_camera_names: Sequence[str] = ("front", "top", "r_vis"),
|
|
num_cameras: int | None = None,
|
|
dataset_image_resize_shape: Sequence[int] | None = None,
|
|
eval_image_resize_shape: Sequence[int] | None = (256, 256),
|
|
freeze_backbone: bool = True,
|
|
joint_output_dim: int = 192,
|
|
image_size: int = 224,
|
|
output_dim: int = 192,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.camera_names = tuple(camera_names)
|
|
self.fused_camera_names = tuple(fused_camera_names)
|
|
self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names)
|
|
self.freeze_backbone = bool(freeze_backbone)
|
|
self.joint_output_dim = int(joint_output_dim)
|
|
self.image_size = int(image_size)
|
|
self._output_dim = int(output_dim)
|
|
self.dataset_image_resize_shape = (
|
|
tuple(int(v) for v in dataset_image_resize_shape)
|
|
if dataset_image_resize_shape is not None else None
|
|
)
|
|
self.eval_image_resize_shape = (
|
|
tuple(int(v) for v in eval_image_resize_shape)
|
|
if eval_image_resize_shape is not None else None
|
|
)
|
|
if self.num_cameras != len(self.camera_names):
|
|
raise ValueError(
|
|
f"num_cameras({self.num_cameras}) must match len(camera_names)({len(self.camera_names)})"
|
|
)
|
|
if set(self.fused_camera_names) != set(self.camera_names):
|
|
raise ValueError(
|
|
"fused_camera_names must contain the same cameras as camera_names. "
|
|
f"got camera_names={list(self.camera_names)}, fused_camera_names={list(self.fused_camera_names)}"
|
|
)
|
|
|
|
self.encoder = self._build_encoder(self.image_size)
|
|
self.projector = _LEWMProjector(
|
|
input_dim=self.encoder.config.hidden_size,
|
|
hidden_dim=2048,
|
|
output_dim=self.joint_output_dim,
|
|
)
|
|
|
|
self.register_buffer(
|
|
"mean",
|
|
torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(1, 3, 1, 1),
|
|
)
|
|
self.register_buffer(
|
|
"std",
|
|
torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(1, 3, 1, 1),
|
|
)
|
|
|
|
if checkpoint_path is not None and checkpoint is not None:
|
|
raise ValueError("checkpoint_path and checkpoint cannot both be provided")
|
|
if checkpoint_path is not None:
|
|
self.load_lewm_checkpoint(checkpoint_path)
|
|
elif checkpoint is not None:
|
|
self.load_lewm_checkpoint(checkpoint)
|
|
|
|
if self.freeze_backbone:
|
|
self._freeze_encoder_and_projector()
|
|
|
|
@staticmethod
|
|
def _build_encoder_config(image_size: int):
|
|
from transformers import ViTConfig
|
|
|
|
return ViTConfig(
|
|
image_size=image_size,
|
|
patch_size=14,
|
|
num_channels=3,
|
|
hidden_size=192,
|
|
intermediate_size=768,
|
|
num_hidden_layers=12,
|
|
num_attention_heads=3,
|
|
qkv_bias=True,
|
|
hidden_dropout_prob=0.0,
|
|
attention_probs_dropout_prob=0.0,
|
|
)
|
|
|
|
@classmethod
|
|
def _build_encoder(cls, image_size: int) -> nn.Module:
|
|
from transformers import ViTModel
|
|
|
|
return ViTModel(cls._build_encoder_config(image_size), add_pooling_layer=False)
|
|
|
|
@staticmethod
|
|
def _unwrap_state_dict(payload: Mapping[str, Any]) -> Mapping[str, torch.Tensor]:
|
|
state_dict = payload.get("state_dict", payload)
|
|
if not isinstance(state_dict, Mapping):
|
|
raise TypeError("checkpoint payload must contain a mapping state_dict")
|
|
return state_dict
|
|
|
|
@staticmethod
|
|
def _extract_prefixed_state_dict(
|
|
state_dict: Mapping[str, torch.Tensor],
|
|
prefix: str,
|
|
) -> Dict[str, torch.Tensor]:
|
|
extracted = {
|
|
key[len(prefix) :]: value
|
|
for key, value in state_dict.items()
|
|
if key.startswith(prefix)
|
|
}
|
|
if not extracted:
|
|
raise KeyError(f"checkpoint missing parameters with prefix {prefix!r}")
|
|
return extracted
|
|
|
|
def load_lewm_checkpoint(self, checkpoint_or_path: str | Path | Mapping[str, Any]) -> None:
|
|
if isinstance(checkpoint_or_path, (str, Path)):
|
|
payload = torch.load(Path(checkpoint_or_path), map_location="cpu", weights_only=False)
|
|
else:
|
|
payload = checkpoint_or_path
|
|
|
|
state_dict = self._unwrap_state_dict(payload)
|
|
encoder_state_dict = self._extract_prefixed_state_dict(state_dict, "model.encoder.")
|
|
projector_state_dict = self._extract_prefixed_state_dict(state_dict, "model.projector.")
|
|
|
|
self.encoder.load_state_dict(encoder_state_dict, strict=True)
|
|
self.projector.load_state_dict(projector_state_dict, strict=True)
|
|
|
|
def _freeze_encoder_and_projector(self) -> None:
|
|
for module in (self.encoder, self.projector):
|
|
module.eval()
|
|
for parameter in module.parameters():
|
|
parameter.requires_grad = False
|
|
|
|
def train(self, mode: bool = True) -> "LEWMViTBackbone":
|
|
super().train(mode)
|
|
if self.freeze_backbone:
|
|
self._freeze_encoder_and_projector()
|
|
return self
|
|
|
|
def _ordered_images(self, images: Dict[str, torch.Tensor]) -> list[torch.Tensor]:
|
|
missing = [camera_name for camera_name in self.camera_names if camera_name not in images]
|
|
if missing:
|
|
raise ValueError(
|
|
f"image input missing required cameras. missing={missing}, expected={list(self.camera_names)}"
|
|
)
|
|
|
|
ordered = [images[camera_name] for camera_name in self.camera_names]
|
|
reference_shape = ordered[0].shape
|
|
if len(reference_shape) != 5:
|
|
raise ValueError(f"expected image tensors shaped (B, T, C, H, W), got {reference_shape}")
|
|
|
|
for camera_name, image in zip(self.camera_names[1:], ordered[1:]):
|
|
if image.shape != reference_shape:
|
|
raise ValueError(
|
|
f"camera {camera_name!r} shape {tuple(image.shape)} does not match {tuple(reference_shape)}"
|
|
)
|
|
|
|
return ordered
|
|
|
|
def _prepare_pixels(self, images: Dict[str, torch.Tensor]) -> tuple[torch.Tensor, int, int]:
|
|
self._ordered_images(images)
|
|
fused = torch.cat([images[camera_name] for camera_name in self.fused_camera_names], dim=-2)
|
|
bsz, steps = fused.shape[:2]
|
|
fused = fused.reshape(bsz * steps, *fused.shape[2:]).contiguous().float()
|
|
|
|
fused = fused.clamp(0.0, 1.0)
|
|
fused = (fused - self.mean) / self.std
|
|
|
|
height, width = fused.shape[-2:]
|
|
short_side = min(height, width)
|
|
if short_side <= 0:
|
|
raise ValueError(f"invalid fused image shape: {tuple(fused.shape)}")
|
|
scale = self.image_size / float(short_side)
|
|
resized_height = int(round(height * scale))
|
|
resized_width = int(round(width * scale))
|
|
if (resized_height, resized_width) != (height, width):
|
|
fused = F.interpolate(
|
|
fused,
|
|
size=(resized_height, resized_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
antialias=True,
|
|
)
|
|
return fused, bsz, steps
|
|
|
|
def forward(self, images: Dict[str, torch.Tensor]) -> torch.Tensor:
|
|
pixels, bsz, steps = self._prepare_pixels(images)
|
|
with torch.set_grad_enabled(torch.is_grad_enabled() and not self.freeze_backbone):
|
|
output = self.encoder(pixel_values=pixels, interpolate_pos_encoding=True)
|
|
cls = output.last_hidden_state[:, 0]
|
|
embedding = self.projector(cls)
|
|
return embedding.view(bsz, steps, self.joint_output_dim)
|
|
|
|
@property
|
|
def output_dim(self) -> int:
|
|
return self._output_dim
|