File size: 3,031 Bytes
bf6e2c9 | 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 | """Build the text and visual prefix used for waypoint prediction."""
from __future__ import annotations
from typing import List, Optional, Tuple
import torch
import torch.nn as nn
from .history_tokens import interleave_tvi_tokens
def build_visual_prefix(
llm: nn.Module,
tokenizer,
projector: nn.Module,
tvi: nn.Module,
coarse_tokens: torch.Tensor,
coarse_tidx: torch.Tensor,
fine_tokens: torch.Tensor,
fine_tidx: torch.Tensor,
instructions: List[str],
history_valid_mask: torch.Tensor,
*,
yaw_hist: Optional[torch.Tensor] = None,
yaw_curr: Optional[torch.Tensor] = None,
use_angle: bool = False,
include_history: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Build the text/visual prefix before the ACT token.
The current observation is always represented by ``fine_tokens``. Callers
must provide past observations separately as ``coarse_tokens``; an empty
history is represented by zero-filled coarse slots and a false validity
mask, rather than by duplicating the current image.
"""
if coarse_tokens.ndim != 3 or fine_tokens.ndim != 3:
raise ValueError("coarse_tokens and fine_tokens must be [B, N, C]")
batch_size = fine_tokens.size(0)
if coarse_tokens.size(0) != batch_size:
raise ValueError("coarse and fine batch sizes must match")
if (history_valid_mask is None or history_valid_mask.dtype != torch.bool
or history_valid_mask.ndim != 2):
raise ValueError("history_valid_mask must be a boolean [B, H] tensor")
if history_valid_mask.size(0) != batch_size:
raise ValueError("history_valid_mask batch size must match visual tokens")
device = next(projector.parameters()).device
vis_fine = projector(fine_tokens.to(device))
pieces = []
masks = []
if include_history:
vis_coarse = projector(coarse_tokens.to(device))
vis_coarse, coarse_attention = interleave_tvi_tokens(
tvi,
vis_coarse,
coarse_tidx.to(device),
kind_id=0,
yaw_per_frame=yaw_hist,
use_angle=use_angle,
frame_valid_mask=history_valid_mask.to(device),
)
pieces.append(vis_coarse)
masks.append(coarse_attention.to(dtype=torch.long))
vis_fine, fine_attention = interleave_tvi_tokens(
tvi,
vis_fine,
fine_tidx.to(device),
kind_id=1,
yaw_per_frame=yaw_curr,
use_angle=use_angle,
)
tok = tokenizer(
instructions,
return_tensors="pt",
padding=True,
truncation=True,
max_length=128,
)
tok = {key: value.to(device) for key, value in tok.items()}
text_emb = llm.get_input_embeddings()(tok["input_ids"])
pieces.append(vis_fine)
masks.extend([fine_attention.to(dtype=torch.long)])
sequence = torch.cat([text_emb, *pieces], dim=1).to(llm.dtype)
attention = torch.cat([tok["attention_mask"], *masks], dim=1)
return sequence, attention
|