# ACT (Action Chunking Transformer) Policy для LeRobot. from dataclasses import dataclass, field from typing import Optional, Any import torch @dataclass class ACTConfig: # Архитектура трансформера dim_model: int = 512 n_heads: int = 8 dim_feedforward: int = 3200 n_encoder_layers: int = 4 n_decoder_layers: int = 1 # ACT-специфичные chunk_size: int = 100 latent_dim: int = 32 use_vae: bool = True temporal_ensemble_coeff: Optional[float] = 0.1 n_action_steps: int = 1 # Параметры VAE (именно их не хватает) n_vae_encoder_layers = 4 # <-- обязательно # use_vae = False # <-- важно, если используешь VAE-ветку # Активации и нормализация feedforward_activation: str = "relu" pre_norm: bool = True dropout: float = 0.1 # Признаки (должны быть тензорами, чтобы инициализировать Linear) robot_state_feature: Optional[torch.Tensor] = None env_state_feature: Optional[torch.Tensor] = None action_feature: Optional[torch.Tensor] = None # Изображения (можно оставить None для теста без картинок) image_features: bool = True vision_backbone: str = "resnet18" pretrained_backbone_weights: str | None = "ResNet18_Weights.IMAGENET1K_V1" replace_final_stride_with_dilation: bool = False def validate_features(self): # заглушка, чтобы не ломался вызов config.validate_features() в ACTPolicy pass import math from collections import deque from collections.abc import Callable from itertools import chain import einops import numpy as np import torch import torch.nn.functional as F import torchvision from torch import Tensor, nn from torchvision.models._utils import IntermediateLayerGetter from torchvision.ops.misc import FrozenBatchNorm2d OBS_STATE = "observation.state" ACTION = "action" OBS_ENV_STATE = "env_state" OBS_IMAGES = "observation.images.realsense" def create_sinusoidal_pos_embedding(num_positions: int, dimension: int) -> Tensor: def get_position_angle_vec(position: int): # Вычисляем углы для каждой размерности return [ position / np.power(10000, 2 * (hid_j // 2) / dimension) for hid_j in range(dimension) ] # Строим таблицу: каждая строка — позиция, каждый столбец — размерность sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(num_positions)]) # Применяем sin к чётным индексам, cos — к нечётным sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) return torch.from_numpy(sinusoid_table).float() class ACTSinusoidalPositionEmbedding2d(nn.Module): def __init__(self, dimension: int): super().__init__() self.dimension = dimension self._two_pi = 2 * math.pi self._eps = 1e-6 self._temperature = 10000 # стандартная константа для частоты синусоид def forward(self, x: Tensor) -> Tensor: # Создаём маску-заглушку, чтобы получить размеры H и W not_mask = torch.ones_like(x[0, :1]) # (1, H, W) # Накопительная сумма даёт индексы строк/столбцов: [1..H], [1..W] y_range = not_mask.cumsum(1, dtype=torch.float32) x_range = not_mask.cumsum(2, dtype=torch.float32) # Нормализуем в [0, 2π] y_range = y_range / (y_range[:, -1:, :] + self._eps) * self._two_pi x_range = x_range / (x_range[:, :, -1:] + self._eps) * self._two_pi # Частоты синусоид (как в оригинальной статье про трансформеры) inverse_frequency = self._temperature ** ( 2 * (torch.arange(self.dimension, dtype=torch.float32, device=x.device) // 2) / self.dimension ) # Добавляем измерение для broadcasting x_range = x_range.unsqueeze(-1) / inverse_frequency # (1, H, W, C//2) y_range = y_range.unsqueeze(-1) / inverse_frequency # Чередуем sin/cos для чётных/нечётных размерностей pos_embed_x = torch.stack( (x_range[..., 0::2].sin(), x_range[..., 1::2].cos()), dim=-1 ).flatten(3) pos_embed_y = torch.stack( (y_range[..., 0::2].sin(), y_range[..., 1::2].cos()), dim=-1 ).flatten(3) # Объединяем и переставляем оси: (1, C, H, W) pos_embed = torch.cat((pos_embed_y, pos_embed_x), dim=3).permute(0, 3, 1, 2) return pos_embed def get_activation_fn(activation: str) -> Callable: """Возвращает функцию активации по имени (для конфигурации модели).""" if activation == "relu": return F.relu if activation == "gelu": return F.gelu if activation == "glu": return F.glu raise RuntimeError(f"activation should be relu/gelu/glu, not {activation}.") class ACTEncoderLayer(nn.Module): """Один слой энкодера трансформера с self-attention и feed-forward.""" def __init__(self, config: ACTConfig): super().__init__() # Multihead self-attention self.self_attn = nn.MultiheadAttention( config.dim_model, config.n_heads, dropout=config.dropout, batch_first=True ) # Feed-forward сеть self.linear1 = nn.Linear(config.dim_model, config.dim_feedforward) self.dropout = nn.Dropout(config.dropout) self.linear2 = nn.Linear(config.dim_feedforward, config.dim_model) # LayerNorm'ы и dropout'ы self.norm1 = nn.LayerNorm(config.dim_model) self.norm2 = nn.LayerNorm(config.dim_model) self.dropout1 = nn.Dropout(config.dropout) self.dropout2 = nn.Dropout(config.dropout) self.activation = get_activation_fn(config.feedforward_activation) self.pre_norm = config.pre_norm # если True, то LayerNorm перед attention def forward( self, x: Tensor, pos_embed: Tensor | None = None, key_padding_mask: Tensor | None = None, return_internals: bool = False ) -> Tensor| tuple[Tensor, dict]: skip = x internals = {} # Pre-Norm: нормализация перед attention if self.pre_norm: x = self.norm1(x) q = k = x if pos_embed is None else x + pos_embed attn_out, attn_weights = self.self_attn(q, k, value=x, key_padding_mask=key_padding_mask) x = skip + self.dropout1(attn_out) if return_internals: internals["self_attn_weights"] = attn_weights # (B, H, T_enc, T_enc) internals["self_attn_out"] = attn_out # (B, T_enc, D) # Дальше — feed-forward if self.pre_norm: skip = x x = self.norm2(x) else: x = self.norm1(x) skip = x ff_pre = x x = self.linear2(self.dropout(self.activation(self.linear1(x)))) x = skip + self.dropout2(x) if not self.pre_norm: x = self.norm2(x) if return_internals: internals["ff_pre"] = ff_pre internals["ff_out"] = x return (x, internals) if return_internals else x class ACTEncoder(nn.Module): """Стек слоёв энкодера. Может быть использован как VAE-энкодер или обычный энкодер.""" def __init__(self, config: ACTConfig, is_vae_encoder: bool = False): super().__init__() num_layers = ( config.n_vae_encoder_layers if is_vae_encoder else config.n_encoder_layers ) self.layers = nn.ModuleList([ACTEncoderLayer(config) for _ in range(num_layers)]) # Последний LayerNorm self.norm = nn.LayerNorm(config.dim_model) if config.pre_norm else nn.Identity() self.is_vae_encoder = is_vae_encoder def forward( self, x: Tensor, pos_embed: Tensor | None = None, key_padding_mask: Tensor | None = None ) -> Tensor: for layer in self.layers: x = layer(x, pos_embed=pos_embed, key_padding_mask=key_padding_mask) x = self.norm(x) return x class ACTDecoderLayer(nn.Module): """Один слой декодера трансформера: self-attention + cross-attention + feed-forward.""" def __init__(self, config: ACTConfig): super().__init__() self.self_attn = nn.MultiheadAttention(config.dim_model, config.n_heads, dropout=config.dropout, batch_first=True ) self.multihead_attn = nn.MultiheadAttention(config.dim_model, config.n_heads, dropout=config.dropout, batch_first=True ) self.linear1 = nn.Linear(config.dim_model, config.dim_feedforward) self.dropout = nn.Dropout(config.dropout) self.linear2 = nn.Linear(config.dim_feedforward, config.dim_model) self.norm1 = nn.LayerNorm(config.dim_model) self.norm2 = nn.LayerNorm(config.dim_model) self.norm3 = nn.LayerNorm(config.dim_model) self.dropout1 = nn.Dropout(config.dropout) self.dropout2 = nn.Dropout(config.dropout) self.dropout3 = nn.Dropout(config.dropout) self.activation = get_activation_fn(config.feedforward_activation) self.pre_norm = config.pre_norm def maybe_add_pos_embed(self, tensor: Tensor, pos_embed: Tensor | None) -> Tensor: return tensor if pos_embed is None else tensor + pos_embed def forward( self, x: Tensor, encoder_out: Tensor, decoder_pos_embed: Tensor | None = None, encoder_pos_embed: Tensor | None = None ) -> Tensor: skip = x if self.pre_norm: x = self.norm1(x) q = k = self.maybe_add_pos_embed(x, decoder_pos_embed) attn_self, _ = self.self_attn(q, k, value=x) x = skip + self.dropout1(attn_self) if self.pre_norm: skip = x x = self.norm2(x) else: x = self.norm1(x) skip = x #print ("x.dec: ", x.shape) #print("decoder_pos_embed: ", decoder_pos_embed) #print("encoder_pos_embed: ", encoder_pos_embed) # Cross-attention к энкодеру query = self.maybe_add_pos_embed(x, decoder_pos_embed) key = self.maybe_add_pos_embed(encoder_out, encoder_pos_embed) value = encoder_out attn_cross, attn_weights = self.multihead_attn( query=query, key=key, value=value, ) #print("attn_cross.shape:", attn_cross.shape) #print("attn_weights.shape:", attn_weights.shape) x = skip + self.dropout2(attn_cross) if self.pre_norm: skip = x x = self.norm3(x) else: x = self.norm2(x) skip = x x = self.linear2(self.dropout(self.activation(self.linear1(x)))) x = skip + self.dropout3(x) if not self.pre_norm: x = self.norm3(x) return x class ACTDecoder(nn.Module): def __init__(self, config: ACTConfig): super().__init__() self.layers = nn.ModuleList([ACTDecoderLayer(config) for _ in range(config.n_decoder_layers)]) self.norm = nn.LayerNorm(config.dim_model) def forward( self, x: Tensor, encoder_out: Tensor, decoder_pos_embed: Tensor | None = None, encoder_pos_embed: Tensor | None = None ) -> Tensor: for layer in self.layers: x = layer( x, encoder_out, decoder_pos_embed=decoder_pos_embed, encoder_pos_embed=encoder_pos_embed ) if self.norm is not None: x = self.norm(x) return x class ACTTemporalEnsembler: def __init__(self, temporal_ensemble_coeff: float, chunk_size: int) -> None: self.chunk_size = chunk_size # Предвычисляем веса и их кумулятивную сумму self.ensemble_weights = torch.exp(-temporal_ensemble_coeff * torch.arange(chunk_size)) self.ensemble_weights_cumsum = torch.cumsum(self.ensemble_weights, dim=0) self.reset() def reset(self): """Сброс состояния при ресете среды (новый эпизод).""" self.ensembled_actions = None self.ensembled_actions_count = None def update(self, actions: Tensor) -> Tensor: # Переносим веса на устройство тензора действий self.ensemble_weights = self.ensemble_weights.to(device=actions.device) self.ensemble_weights_cumsum = self.ensemble_weights_cumsum.to(device=actions.device) if self.ensembled_actions is None: # Инициализация: первое предсказание становится начальным ансамблем self.ensembled_actions = actions.clone() self.ensembled_actions_count = torch.ones( (self.chunk_size, 1), dtype=torch.long, device=self.ensembled_actions.device ) else: # Онлайн-обновление: взвешенное усреднение с учётом новых предсказаний # Это эквивалентно пересчёту среднего с новыми весами без хранения всей истории self.ensembled_actions *= self.ensemble_weights_cumsum[self.ensembled_actions_count - 1] self.ensembled_actions += actions[:, :-1] * self.ensemble_weights[self.ensembled_actions_count] self.ensembled_actions /= self.ensemble_weights_cumsum[self.ensembled_actions_count] # Увеличиваем счётчик количества шагов в ансамбле (не более chunk_size) self.ensembled_actions_count = torch.clamp( self.ensembled_actions_count + 1, max=self.chunk_size ) # Добавляем последнее действие из нового чанка как новый элемент в конец последовательности self.ensembled_actions = torch.cat([self.ensembled_actions, actions[:, -1:]], dim=1) self.ensembled_actions_count = torch.cat( [ self.ensembled_actions_count, torch.ones_like(self.ensembled_actions_count[-1:]) ] ) # «Потребляем» первое действие из ансамбля и сдвигаем окно вперёд action, self.ensembled_actions, self.ensembled_actions_count = ( self.ensembled_actions[:, 0], self.ensembled_actions[:, 1:], self.ensembled_actions_count[1:] ) return action class ACT(nn.Module): def __init__(self, config): super().__init__() self.config = config if self.config.use_vae: # Трансформер-энкодер, который кодирует [cls, robot_state, action_sequence] self.vae_encoder = ACTEncoder(config, is_vae_encoder=True) # Эмбеддинг для токена CLS — из него извлекаются параметры латента self.vae_encoder_cls_embed = nn.Embedding(1, config.dim_model) # Проекция состояния робота в пространство модели if self.config.robot_state_feature is not None: self.vae_encoder_robot_state_input_proj = nn.Linear( self.config.robot_state_feature.shape[0], config.dim_model ) # Проекция целевых действий в пространство модели (для VAE-энкодера) self.vae_encoder_action_input_proj = nn.Linear( self.config.action_feature.shape[0], config.dim_model, ) # Проекция выхода VAE-энкодера в параметры латентного распределения # (mean и log_var для каждой размерности латента) self.vae_encoder_latent_output_proj = nn.Linear( config.dim_model, config.latent_dim * 2 ) # Позиционные эмбеддинги для VAE-энкодера num_input_token_encoder = 1 + config.chunk_size # 1 cls + chunk_size действий if self.config.robot_state_feature is not None: num_input_token_encoder += 1 # +1 для состояния робота self.register_buffer( "vae_encoder_pos_enc", create_sinusoidal_pos_embedding(num_input_token_encoder, config.dim_model).unsqueeze(0), persistent=False, ) if self.config.image_features: # Загружаем предобученную модель (например, ResNet) backbone_model = getattr(torchvision.models, config.vision_backbone)( replace_stride_with_dilation=[False, False, config.replace_final_stride_with_dilation], weights=config.pretrained_backbone_weights, norm_layer=FrozenBatchNorm2d, ) # Получаем feature map из последнего блока (layer4) self.backbone = IntermediateLayerGetter(backbone_model, return_layers={"layer4": "feature_map"}) # --------------------------------------------------------------------- # Трансформер (основной энкодер и декодер) # --------------------------------------------------------------------- self.encoder = ACTEncoder(config) self.decoder = ACTDecoder(config) # Проекции входов в пространство трансформера if self.config.robot_state_feature is not None: self.encoder_robot_state_input_proj = nn.Linear( self.config.robot_state_feature.shape[0], config.dim_model ) if self.config.env_state_feature is not None: self.encoder_env_state_input_proj = nn.Linear( self.config.env_state_feature.shape[0], config.dim_model ) # Проекция латентной переменной (z) в пространство модели self.encoder_latent_input_proj = nn.Linear(config.latent_dim, config.dim_model) if self.config.image_features: # Проецируем каналы feature map в dim_model # backbone_model.fc.in_features — это количество каналов layer4 у выбранной модели self.encoder_img_feat_input_proj = nn.Conv2d( backbone_model.fc.in_features, config.dim_model, kernel_size=1 ) # Позиционные эмбеддинги n_1d_tokens = 1 # для латентной переменной if self.config.robot_state_feature is not None: n_1d_tokens += 1 if self.config.env_state_feature is not None: n_1d_tokens += 1 self.encoder_1d_feature_pos_embed = nn.Embedding(n_1d_tokens, config.dim_model) if self.config.image_features: self.encoder_cam_feat_pos_embed = ACTSinusoidalPositionEmbedding2d(config.dim_model // 2) # Learnable positional embeddings для декодера (как в DETR) self.decoder_pos_embed = nn.Embedding(config.chunk_size, config.dim_model) # Head для предсказания действий self.action_head = nn.Linear(config.dim_model, self.config.action_feature.shape[0]) self._reset_parameters() def _reset_parameters(self): """Инициализация весов трансформера (Xavier uniform).""" for p in chain(self.encoder.parameters(), self.decoder.parameters()): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward( self, batch: dict[str, Tensor] ) -> tuple[Tensor, tuple[Tensor, Tensor] | tuple[None, None]]: # --------------------------------------------------------------------- # Шаг 1: VAE-энкодер — получаем латентную переменную z # --------------------------------------------------------------------- mu_hat = log_sigma_x2_hat = None if self.config.use_vae and ACTION in batch: # Собираем входные токены для VAE-энкодера: [cls, robot_state (если есть), action_sequence] tokens = [] # Токен CLS cls_token = self.vae_encoder_cls_embed(torch.zeros(batch[OBS_STATE].size(0), dtype=torch.long, device=batch[OBS_STATE].device)) tokens.append(cls_token) # Состояние робота (если есть) if self.config.robot_state_feature is not None: robot_state_proj = self.vae_encoder_robot_state_input_proj(batch[OBS_STATE]) tokens.append(robot_state_proj) # Целевые действия (для VAE) action_proj = self.vae_encoder_action_input_proj(batch[ACTION]) # action_proj имеет форму (B, chunk_size, dim_model), разворачиваем в последовательность tokens.append(action_proj.transpose(1, 0)) # (chunk_size, B, dim_model) -> потом сложим # Объединяем токены в одну последовательность (B, seq_len, dim_model) vae_input = torch.cat(tokens, dim=1) # (B, 1+state+chunk_size, dim_model) # Добавляем позиционные эмбеддинги vae_input += self.vae_encoder_pos_enc[:, :vae_input.size(1), :] # Пропускаем через VAE-энкодер vae_output = self.vae_encoder(vae_input) # (B, seq_len, dim_model) # Берём только выход токена CLS (первый токен) cls_output = vae_output[:, 0, :] # (B, dim_model) # Преобразуем в параметры распределения латента latent_params = self.vae_encoder_latent_output_proj(cls_output) # (B, latent_dim*2) mu_hat, log_sigma_x2_hat = latent_params.chunk(2, dim=-1) # разделяем mean и log_var # Сэмплируем латент z через трюк репараметризации std = torch.exp(0.5 * log_sigma_x2_hat) eps = torch.randn_like(std) z = mu_hat + eps * std # (B, latent_dim) else: # Если VAE не используется (инференс или обучение без VAE), z = 0 batch_size = batch[OBS_STATE].dim() z = torch.zeros(batch_size, self.config.latent_dim, device=batch[OBS_STATE].device) # --------------------------------------------------------------------- # Шаг 2: Извлечение признаков изображений (если есть) # --------------------------------------------------------------------- img_features = None if self.config.image_features: images = batch[OBS_IMAGES] # images может быть списком камер или одним тензором (B, C, H, W) или (B, N, C, H, W) # В LeRobot обычно это список тензоров (B, C, H, W) для каждой камеры if isinstance(images, list): # Для каждой камеры получаем feature map img_feats = [] for img in images: feat_dict = self.backbone(img) img_feats.append(feat_dict["feature_map"]) # (B, C_feat, H_feat, W_feat) # Склеиваем по каналам: (B, N*C_feat, H, W) img_features = torch.cat(img_feats, dim=1) else: # Если уже один тензор feat_dict = self.backbone(images) img_features = feat_dict["feature_map"] # Проецируем в dim_model и добавляем 2D позиционные эмбеддинги img_features = self.encoder_img_feat_input_proj(img_features) # (B, dim_model, H, W) pos_embed_2d = self.encoder_cam_feat_pos_embed(img_features) img_features += pos_embed_2d # Разворачиваем feature map в последовательность токенов (B, H*W, dim_model) B, C, H, W = img_features.shape img_tokens = einops.rearrange(img_features, "b c h w -> b (h w) c") else: img_tokens = None # --------------------------------------------------------------------- # Шаг 3: Сбор токенов для основного трансформера-энкодера # --------------------------------------------------------------------- encoder_tokens = [] # Латент z z_proj = self.encoder_latent_input_proj(z) # (B, dim_model) z_proj_token = z_proj.unsqueeze(1) # (B, 1, D) encoder_tokens.append(z_proj_token) # (B, 1, dim_model) # Состояние робота if self.config.robot_state_feature is not None: robot_state = batch[OBS_STATE] # если это один пример, делаем батч размера 1 robot_state = robot_state.unsqueeze(0) robot_state_proj = self.encoder_robot_state_input_proj(robot_state) # (B, 512) robot_state_proj_token = robot_state_proj.unsqueeze(1) # (B, 1, D) encoder_tokens.append(robot_state_proj_token) #print ("robot_state_proj_t: ", robot_state_proj_token.shape) # Состояние среды if self.config.env_state_feature: env_state_proj = self.encoder_env_state_input_proj(batch[OBS_ENV_STATE]) # (B, dim_model) env_state_proj_token = env_state_proj.unsqueeze(1) # (B, 1, D) encoder_tokens.append(env_state_proj_token) #print ("env_state_proj: ", env_state_proj_token.shape) # Признаки изображений (если есть) if img_tokens is not None: encoder_tokens.append(img_tokens) # (B, H*W, dim_model) for i, t in enumerate(encoder_tokens): print(i, t.shape, t.dtype) # Объединяем все токены encoder_input = torch.cat(encoder_tokens, dim=1) # (B, num_tokens, dim_model) # Количество 1D токенов = len(encoder_tokens) - (1 если есть img_tokens else 0) num_1d_tokens = len(encoder_tokens) if img_tokens is not None: num_1d_tokens -= 1 # image tokens используют 2D pos embed, не 1D encoder_input[:, :num_1d_tokens, :] += self.encoder_1d_feature_pos_embed.weight.unsqueeze(0) # Пропускаем через энкодер трансформера encoder_output = self.encoder(encoder_input) # (B, num_tokens, dim_model) # Query для декодера — learnable pos embeddings длины chunk_size decoder_queries = self.decoder_pos_embed.weight.unsqueeze(0).expand(batch_size, -1, -1) # Декодер делает cross-attention к выходу энкодера decoder_output = self.decoder( decoder_queries, encoder_output, decoder_pos_embed=None, # уже в decoder_queries encoder_pos_embed=None ) # (B, chunk_size, dim_model) actions_hat = self.action_head(decoder_output) # (B, chunk_size, action_dim) return actions_hat, (mu_hat, log_sigma_x2_hat) class ACTPolicy(nn.Module): name = "act" def __init__( self, config, **kwargs, ): super().__init__() config.validate_features() self.config = config # Инициализируем ядро модели ACT self.model = ACT(config) # Если включён temporal ensembling — создаём сглаживатель if config.temporal_ensemble_coeff is not None: self.temporal_ensembler = ACTTemporalEnsembler( config.temporal_ensemble_coeff, config.chunk_size ) else: # Иначе используем простую очередь действий (deque) self._action_queue = deque([], maxlen=self.config.n_action_steps) self.reset() def get_optim_params(self) -> dict: return [ { "params": [ p for n, p in self.named_parameters() if not n.startswith("model.backbone") and p.requires_grad ] }, { "params": [ p for n, p in self.named_parameters() if n.startswith("model.backbone") and p.requires_grad ], "lr": self.config.optimizer_lr_backbone, }, ] def reset(self): if self.config.temporal_ensemble_coeff is not None: self.temporal_ensembler.reset() else: # Очищаем очередь действий, если не используем temporal ensembling self._action_queue.clear() @torch.no_grad() def select_action(self, batch: dict[str, Tensor]) -> Tensor: batch_size = batch[OBS_STATE].size(0) if self.config.temporal_ensemble_coeff is not None: # Temporal ensembling: обновляем сглаженный прогноз и берём первое действие actions_chunk, _ = self.model(batch) # (B, chunk_size, action_dim) action = self.temporal_ensembler.update(actions_chunk) return action # Обычная логика с очередью (без ensembling) if len(self._action_queue) == 0: # Очередь пуста: делаем предсказание и кладём чанк целиком actions_chunk, _ = self.model(batch) # (B, chunk_size, action_dim) # Разворачиваем чанк по батчу и добавляем все действия в очередь for t in range(actions_chunk.size(1)): self._action_queue.append(actions_chunk[:, t, :]) # Берём первое действие из очереди action = self._action_queue.popleft() # (B, action_dim) return action