Logic
396 lines
18 KiB
Python
396 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from dataclasses import fields
|
|
from typing import Dict, Optional, Sequence
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from roboimi.vla.models.normalization import NormalizationModule
|
|
|
|
|
|
class SmolVLANativeAgent(nn.Module):
|
|
"""RoboIMI wrapper for the native SmolVLA flow-matching model.
|
|
|
|
The wrapper owns RoboIMI-facing concerns only: dataset normalization,
|
|
camera ordering, language fallback/tokenization, rollout queues, and action
|
|
denormalization. The native SmolVLA model itself is imported lazily so unit
|
|
tests can inject fakes without downloading or loading a real VLM.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Optional[nn.Module] = None,
|
|
tokenizer=None,
|
|
action_dim: int = 16,
|
|
obs_dim: int = 16,
|
|
chunk_size: int = 16,
|
|
n_action_steps: int = 8,
|
|
obs_horizon: int = 1,
|
|
action_chunk_start: int = 0,
|
|
num_cams: int = 3,
|
|
camera_names: Optional[Sequence[str]] = None,
|
|
dataset_stats=None,
|
|
normalization_type: str = 'min_max',
|
|
task_description: Optional[str] = None,
|
|
model_config: Optional[dict] = None,
|
|
tokenizer_name: Optional[str] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__()
|
|
del kwargs
|
|
self.action_dim = int(action_dim)
|
|
self.obs_dim = int(obs_dim)
|
|
self.chunk_size = int(chunk_size)
|
|
self.pred_horizon = self.chunk_size
|
|
self.n_action_steps = int(n_action_steps)
|
|
self.num_action_steps = self.n_action_steps
|
|
self.obs_horizon = int(obs_horizon)
|
|
self.action_chunk_start = int(action_chunk_start)
|
|
self.num_cams = int(num_cams)
|
|
self.camera_names = tuple(camera_names) if camera_names is not None else None
|
|
if self.camera_names is not None and len(self.camera_names) != self.num_cams:
|
|
raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})')
|
|
if self.n_action_steps < 1:
|
|
raise ValueError('n_action_steps must be >= 1')
|
|
if self.action_chunk_start < 0:
|
|
raise ValueError('action_chunk_start must be >= 0')
|
|
if self.action_chunk_start + self.n_action_steps > self.chunk_size:
|
|
raise ValueError('action_chunk_start + n_action_steps must be <= chunk_size')
|
|
|
|
self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
|
|
self.task_description = task_description
|
|
self.model_config = dict(model_config or {})
|
|
self.max_state_dim = int(self.model_config.get('max_state_dim', self.obs_dim))
|
|
self.max_action_dim = int(self.model_config.get('max_action_dim', self.action_dim))
|
|
self.resize_imgs_with_padding = self._normalize_resize_shape(
|
|
self.model_config.get('resize_imgs_with_padding', self.model_config.get('image_resize_shape', (512, 512)))
|
|
)
|
|
self.tokenizer_max_length = int(self.model_config.get('tokenizer_max_length', 48))
|
|
self.pad_language_to = str(self.model_config.get('pad_language_to', 'longest'))
|
|
self.tokenizer_name = tokenizer_name
|
|
self.tokenizer = tokenizer if tokenizer is not None else self._build_tokenizer(tokenizer_name)
|
|
self.model = model if model is not None else self._build_model()
|
|
self.reset()
|
|
|
|
@staticmethod
|
|
def _normalize_resize_shape(shape):
|
|
if shape is None:
|
|
return None
|
|
normalized = tuple(int(v) for v in shape)
|
|
if len(normalized) != 2:
|
|
raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}')
|
|
return normalized
|
|
|
|
def _build_tokenizer(self, tokenizer_name):
|
|
if tokenizer_name is None:
|
|
tokenizer_name = self.model_config.get('tokenizer_name') or self.model_config.get('vlm_model_name')
|
|
if tokenizer_name is None:
|
|
raise ImportError('A tokenizer or tokenizer_name/model_config.vlm_model_name is required for SmolVLANativeAgent')
|
|
try:
|
|
from transformers import AutoTokenizer
|
|
except ImportError as exc:
|
|
raise ImportError('SmolVLANativeAgent requires transformers to construct the real tokenizer') from exc
|
|
return AutoTokenizer.from_pretrained(tokenizer_name)
|
|
|
|
def _native_config_kwargs(self) -> dict:
|
|
from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
|
|
|
|
cfg_kwargs = dict(self.model_config)
|
|
|
|
# Backwards-compatible aliases from earlier RoboIMI-facing drafts. The
|
|
# native model core intentionally keeps only SmolVLA model fields.
|
|
if 'image_resize_shape' in cfg_kwargs and 'resize_imgs_with_padding' not in cfg_kwargs:
|
|
cfg_kwargs['resize_imgs_with_padding'] = cfg_kwargs['image_resize_shape']
|
|
if 'freeze_vlm' in cfg_kwargs and 'train_expert_only' not in cfg_kwargs:
|
|
cfg_kwargs['train_expert_only'] = bool(cfg_kwargs['freeze_vlm'])
|
|
|
|
for wrapper_only_key in (
|
|
'state_dim',
|
|
'action_dim',
|
|
'tokenizer_name',
|
|
'freeze_vlm',
|
|
'image_resize_shape',
|
|
'num_cameras',
|
|
):
|
|
cfg_kwargs.pop(wrapper_only_key, None)
|
|
|
|
cfg_kwargs.setdefault('max_state_dim', self.max_state_dim)
|
|
cfg_kwargs.setdefault('max_action_dim', self.max_action_dim)
|
|
cfg_kwargs.setdefault('chunk_size', self.chunk_size)
|
|
cfg_kwargs.setdefault('n_action_steps', self.n_action_steps)
|
|
cfg_kwargs['resize_imgs_with_padding'] = self._normalize_resize_shape(
|
|
cfg_kwargs.get('resize_imgs_with_padding', self.resize_imgs_with_padding)
|
|
)
|
|
|
|
allowed_keys = {field.name for field in fields(NativeSmolVLAConfig)}
|
|
return {key: value for key, value in cfg_kwargs.items() if key in allowed_keys}
|
|
|
|
def _build_model(self):
|
|
try:
|
|
from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
|
|
from roboimi.vla.models.smolvla.modeling import VLAFlowMatching
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
'Native SmolVLA model modules are unavailable. Pass injected model/tokenizer fakes in tests, '
|
|
'or add roboimi.vla.models.smolvla.configuration/modeling for real construction.'
|
|
) from exc
|
|
config = NativeSmolVLAConfig(**self._native_config_kwargs())
|
|
return VLAFlowMatching(config=config)
|
|
|
|
def _get_model_device(self) -> torch.device:
|
|
try:
|
|
return next(self.model.parameters()).device
|
|
except (StopIteration, AttributeError):
|
|
return torch.device('cpu')
|
|
|
|
def _move_to_device(self, data, device: torch.device):
|
|
if torch.is_tensor(data):
|
|
return data.to(device)
|
|
if isinstance(data, dict):
|
|
return {k: self._move_to_device(v, device) for k, v in data.items()}
|
|
if isinstance(data, list):
|
|
return [self._move_to_device(v, device) for v in data]
|
|
if isinstance(data, tuple):
|
|
return tuple(self._move_to_device(v, device) for v in data)
|
|
return data
|
|
|
|
def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
if self.camera_names is None:
|
|
names = tuple(sorted(images.keys()))
|
|
if len(names) != self.num_cams:
|
|
raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})')
|
|
return {name: images[name] for name in names}
|
|
missing = [name for name in self.camera_names if name not in images]
|
|
if missing:
|
|
raise ValueError(f'image batch missing required cameras. missing={missing}, expected={list(self.camera_names)}')
|
|
return {name: images[name] for name in self.camera_names}
|
|
|
|
@staticmethod
|
|
def _is_missing_task(item) -> bool:
|
|
return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
|
|
|
|
def _resolve_task(self, task, batch_size: int):
|
|
def fallback_or(item):
|
|
if self._is_missing_task(item):
|
|
if self.task_description is not None:
|
|
return self.task_description
|
|
if item is None:
|
|
return ''
|
|
return item
|
|
|
|
if task is None:
|
|
task = self.task_description
|
|
if task is None:
|
|
return [''] * batch_size
|
|
if isinstance(task, str):
|
|
return [fallback_or(task)] * batch_size
|
|
task = list(task)
|
|
if len(task) != batch_size:
|
|
raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})')
|
|
return [fallback_or(item) for item in task]
|
|
|
|
def _tokenize_tasks(self, task, batch_size: int, device: torch.device) -> dict:
|
|
tasks = [str(text) + '\n' for text in self._resolve_task(task, batch_size)]
|
|
old_padding_side = getattr(self.tokenizer, 'padding_side', None)
|
|
if old_padding_side is not None:
|
|
self.tokenizer.padding_side = 'right'
|
|
try:
|
|
tokenized = self.tokenizer(
|
|
tasks,
|
|
padding=self.pad_language_to,
|
|
max_length=self.tokenizer_max_length,
|
|
truncation=True,
|
|
return_tensors='pt',
|
|
)
|
|
finally:
|
|
if old_padding_side is not None:
|
|
self.tokenizer.padding_side = old_padding_side
|
|
return {
|
|
'lang_tokens': tokenized['input_ids'].to(device=device),
|
|
'lang_masks': tokenized['attention_mask'].to(device=device, dtype=torch.bool),
|
|
}
|
|
|
|
@staticmethod
|
|
def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
|
|
current_dim = vector.shape[-1]
|
|
if current_dim == new_dim:
|
|
return vector
|
|
if current_dim > new_dim:
|
|
raise ValueError(f'cannot pad vector with dim {current_dim} to smaller dim {new_dim}')
|
|
padded_shape = list(vector.shape)
|
|
padded_shape[-1] = int(new_dim)
|
|
padded = torch.zeros(*padded_shape, dtype=vector.dtype, device=vector.device)
|
|
padded[..., :current_dim] = vector
|
|
return padded
|
|
|
|
@staticmethod
|
|
def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor:
|
|
if img.ndim != 4:
|
|
raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}')
|
|
import torch.nn.functional as F
|
|
|
|
cur_height, cur_width = img.shape[2:]
|
|
ratio = max(cur_width / width, cur_height / height)
|
|
resized_height = int(cur_height / ratio)
|
|
resized_width = int(cur_width / ratio)
|
|
resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
|
|
pad_height = max(0, int(height - resized_height))
|
|
pad_width = max(0, int(width - resized_width))
|
|
return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value)
|
|
|
|
def _prepare_native_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
|
|
ordered_images = self._order_images(images)
|
|
prepared_images: list[torch.Tensor] = []
|
|
img_masks: list[torch.Tensor] = []
|
|
reference_batch_size: int | None = None
|
|
for camera_name, image in ordered_images.items():
|
|
if image.ndim == 5:
|
|
image = image[:, -1]
|
|
elif image.ndim != 4:
|
|
raise ValueError(
|
|
f'image for camera {camera_name!r} must be shaped (B,T,C,H,W) or (B,C,H,W), got {tuple(image.shape)}'
|
|
)
|
|
if reference_batch_size is None:
|
|
reference_batch_size = int(image.shape[0])
|
|
elif int(image.shape[0]) != reference_batch_size:
|
|
raise ValueError(f'image batch size mismatch for camera {camera_name!r}')
|
|
|
|
image = image.contiguous().float().clamp(0.0, 1.0)
|
|
if self.resize_imgs_with_padding is not None:
|
|
image = self._resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
|
|
image = image * 2.0 - 1.0
|
|
prepared_images.append(image)
|
|
img_masks.append(torch.ones(image.shape[0], dtype=torch.bool, device=image.device))
|
|
return prepared_images, img_masks
|
|
|
|
def _prepare_native_state(self, states: torch.Tensor) -> torch.Tensor:
|
|
states = self.normalization.normalize_qpos(states)
|
|
if states.ndim > 2:
|
|
states = states[:, -1, :]
|
|
return self._pad_vector(states.float(), self.max_state_dim)
|
|
|
|
def _prepare_native_actions(self, actions: torch.Tensor) -> torch.Tensor:
|
|
actions = self.normalization.normalize_action(actions)
|
|
return self._pad_vector(actions.float(), self.max_action_dim)
|
|
|
|
def _prepare_model_inputs(self, batch: Dict[str, torch.Tensor], include_actions: bool = False) -> dict:
|
|
states = batch['qpos']
|
|
batch_size = states.shape[0]
|
|
device = states.device
|
|
images, img_masks = self._prepare_native_images(batch['images'])
|
|
inputs = {
|
|
'images': images,
|
|
'img_masks': img_masks,
|
|
'state': self._prepare_native_state(states),
|
|
}
|
|
inputs.update(self._tokenize_tasks(batch.get('task', None), batch_size, device))
|
|
if include_actions:
|
|
inputs['actions'] = self._prepare_native_actions(batch['action'])
|
|
return inputs
|
|
|
|
def _reduce_native_losses(self, losses: torch.Tensor, action_is_pad: torch.Tensor | None = None) -> torch.Tensor:
|
|
if losses.ndim == 0:
|
|
return losses
|
|
if losses.shape[-1] < self.action_dim:
|
|
raise RuntimeError(f'loss action dim mismatch: got {losses.shape[-1]}, expected at least {self.action_dim}')
|
|
losses = losses[..., : self.action_dim]
|
|
if action_is_pad is None:
|
|
return losses.mean()
|
|
if losses.ndim != 3:
|
|
raise RuntimeError(f'action padding mask requires per-element losses shaped (B,H,A), got {tuple(losses.shape)}')
|
|
if tuple(action_is_pad.shape) != tuple(losses.shape[:2]):
|
|
raise RuntimeError(
|
|
f'action_is_pad shape {tuple(action_is_pad.shape)} does not match loss prefix {tuple(losses.shape[:2])}'
|
|
)
|
|
mask = (~action_is_pad).to(device=losses.device, dtype=losses.dtype).unsqueeze(-1)
|
|
denom = (mask.sum() * losses.shape[-1]).clamp_min(1.0)
|
|
return (losses * mask).sum() / denom
|
|
|
|
def compute_loss(self, batch):
|
|
inputs = self._prepare_model_inputs(batch, include_actions=True)
|
|
action_is_pad = batch.get('action_is_pad', None)
|
|
output = self.model(**inputs)
|
|
if isinstance(output, dict):
|
|
if 'loss' not in output:
|
|
raise RuntimeError('SmolVLA model forward returned a dict without loss')
|
|
return output['loss']
|
|
if torch.is_tensor(output):
|
|
return self._reduce_native_losses(output, action_is_pad=action_is_pad)
|
|
if hasattr(output, 'loss'):
|
|
return output.loss
|
|
raise RuntimeError(f'Unsupported SmolVLA forward output type: {type(output)!r}')
|
|
|
|
@torch.no_grad()
|
|
def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
|
|
inputs = self._prepare_model_inputs(batch, include_actions=False)
|
|
if not hasattr(self.model, 'sample_actions'):
|
|
raise RuntimeError('SmolVLA model must implement sample_actions for inference')
|
|
actions = self.model.sample_actions(**inputs)
|
|
if not torch.is_tensor(actions) or actions.ndim != 3:
|
|
raise RuntimeError(f'sample_actions must return (B,H,A), got {type(actions)!r} {getattr(actions, "shape", None)}')
|
|
if actions.shape[-1] < self.action_dim:
|
|
raise RuntimeError(f'action dim mismatch: got {actions.shape[-1]}, expected at least {self.action_dim}')
|
|
actions = actions[..., : self.action_dim]
|
|
return self.normalization.denormalize_action(actions)
|
|
|
|
def reset(self):
|
|
self._queues = {
|
|
'qpos': deque(maxlen=self.obs_horizon),
|
|
'images': deque(maxlen=self.obs_horizon),
|
|
'task': deque(maxlen=self.obs_horizon),
|
|
'action': deque(maxlen=self.n_action_steps),
|
|
}
|
|
|
|
def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
|
|
if 'qpos' in observation:
|
|
self._queues['qpos'].append(observation['qpos'].clone())
|
|
if 'images' in observation:
|
|
ordered_images = self._order_images(observation['images'])
|
|
self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()})
|
|
if 'task' in observation:
|
|
self._queues['task'].append(observation['task'])
|
|
|
|
def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
|
|
qpos_list = list(self._queues['qpos'])
|
|
if not qpos_list:
|
|
raise ValueError('observation queue is empty.')
|
|
while len(qpos_list) < self.obs_horizon:
|
|
qpos_list.append(qpos_list[-1])
|
|
batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
|
|
|
|
images_list = list(self._queues['images'])
|
|
if not images_list:
|
|
raise ValueError('image queue is empty.')
|
|
while len(images_list) < self.obs_horizon:
|
|
images_list.append(images_list[-1])
|
|
names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
|
|
batch_images = {name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0) for name in names}
|
|
batch = {'qpos': batch_qpos, 'images': batch_images}
|
|
if self._queues['task']:
|
|
batch['task'] = [list(self._queues['task'])[-1]]
|
|
elif self.task_description is not None:
|
|
batch['task'] = [self.task_description]
|
|
return batch
|
|
|
|
@torch.no_grad()
|
|
def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
|
|
device = self._get_model_device()
|
|
observation = self._move_to_device(observation, device)
|
|
self._populate_queues(observation)
|
|
if len(self._queues['action']) == 0:
|
|
batch = self._prepare_observation_batch()
|
|
actions = self.predict_action_chunk(batch)
|
|
start = self.action_chunk_start
|
|
end = start + self.n_action_steps
|
|
if actions.shape[1] < end:
|
|
raise RuntimeError(f'action chunk too short: got {actions.shape[1]}, need at least {end}')
|
|
executable_actions = actions[:, start:end]
|
|
for i in range(executable_actions.shape[1]):
|
|
self._queues['action'].append(executable_actions[:, i].squeeze(0))
|
|
return self._queues['action'].popleft()
|
|
|
|
def get_normalization_stats(self):
|
|
return self.normalization.get_stats()
|