Upload 2 files

ACTmodel.py

@@ -0,0 +1,725 @@
+
+# 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

so101_scene.xml

@@ -0,0 +1,198 @@
+<mujoco model="so101_pick_place">
+  <compiler angle="radian" autolimits="true"/>
+
+  <option timestep="0.002" gravity="0 0 -9.81" integrator="Euler">
+    <flag contact="enable"/>
+  </option>
+
+  <default>
+    <joint armature="0.02" damping="1.5"/>
+    <geom condim="4" friction="1.0 0.005 0.001" solimp="0.95 0.99 0.001 0.5 2" solref="0.02 1.0"/>
+    <position kp="100" ctrlrange="-3.14159 3.14159"/>
+  </default>
+
+  <visual>
+    <global offwidth="960" offheight="720"/>
+    <quality shadowsize="2048"/>
+    <map znear="0.01" zfar="10"/>
+  </visual>
+
+  <!-- ==================== ASSETS ==================== -->
+  <asset>
+    <texture name="grid" type="2d" builtin="checker" rgb1="0.85 0.85 0.85" rgb2="0.7 0.7 0.7" width="256" height="256"/>
+    <material name="floor_mat" texture="grid" texrepeat="8 8" reflectance="0.1"/>
+    <texture name="wood" type="2d" builtin="flat" rgb1="0.65 0.45 0.25" rgb2="0.55 0.35 0.18" width="256" height="256" mark="random" markrgb="0.6 0.4 0.2"/>
+    <material name="wood_mat" texture="wood" reflectance="0.05"/>
+    <material name="arm_mat" rgba="0.25 0.25 0.28 1" specular="0.5" shininess="0.5"/>
+    <material name="finger_mat" rgba="0.15 0.15 0.15 1" specular="0.3"/>
+    <material name="box_mat" rgba="0.85 0.15 0.15 1" specular="0.3" shininess="0.3"/>
+    <material name="bowl_mat" rgba="0.2 0.4 0.85 1" specular="0.4" shininess="0.5"/>
+    <material name="base_mat" rgba="0.3 0.3 0.35 1" specular="0.5" shininess="0.5"/>
+  </asset>
+
+  <!-- ==================== WORLD ==================== -->
+  <worldbody>
+    <!-- Floor -->
+    <geom name="floor" type="plane" size="1 1 0.01" pos="0 0 0" material="floor_mat" contype="1" conaffinity="1"/>
+
+    <!-- Ambient and directional lights -->
+    <light name="overhead" directional="true" diffuse="0.8 0.8 0.8" specular="0.3 0.3 0.3" pos="0 0 2" dir="0 0 -1" castshadow="true"/>
+    <light name="side_fill" directional="false" diffuse="0.4 0.4 0.4" specular="0.1 0.1 0.1" pos="0.5 -0.5 1.2"/>
+    <light name="back_fill" directional="false" diffuse="0.2 0.2 0.2" specular="0.05 0.05 0.05" pos="-0.3 0.4 1.0"/>
+
+    <!-- Overhead camera -->
+    <!-- Overhead camera — full view of table, arm, and objects -->
+    <camera name="webcam" pos="0.0 -0.5 1.1" quat="0.92 0.39 0 0" fovy="60"/>
+
+    <!-- ==================== TABLE ==================== -->
+    <body name="table" pos="0 0 0.2">
+      <geom name="table_top" type="box" size="0.30 0.25 0.02" pos="0 0 0.20" material="wood_mat" contype="1" conaffinity="1" mass="10"/>
+      <!-- Table legs -->
+      <geom type="cylinder" size="0.02 0.10" pos="0.25 0.20 0.10" material="wood_mat"/>
+      <geom type="cylinder" size="0.02 0.10" pos="-0.25 0.20 0.10" material="wood_mat"/>
+      <geom type="cylinder" size="0.02 0.10" pos="0.25 -0.20 0.10" material="wood_mat"/>
+      <geom type="cylinder" size="0.02 0.10" pos="-0.25 -0.20 0.10" material="wood_mat"/>
+    </body>
+
+    <!-- ==================== ROBOT ARM ==================== -->
+    <body name="base_link" pos="0 -0.15 0.42">
+      <geom name="base" type="cylinder" size="0.04 0.02" material="base_mat" mass="1.0"/>
+
+      <!-- Joint 1: Shoulder Pan (yaw) -->
+      <body name="shoulder_link" pos="0 0 0.03">
+        <joint name="shoulder_pan" type="hinge" axis="0 0 1" range="-3.14159 3.14159"/>
+        <geom name="shoulder_geom" type="capsule" fromto="0 0 0 0 0 0.06" size="0.025" material="arm_mat" mass="0.3"/>
+
+        <!-- Joint 2: Shoulder Lift (pitch) -->
+        <body name="upper_arm_link" pos="0 0 0.06">
+          <joint name="shoulder_lift" type="hinge" axis="0 1 0" range="-1.5708 1.5708"/>
+          <geom name="upper_arm_geom" type="capsule" fromto="0 0 0 0.10 0 0.04" size="0.02" material="arm_mat" mass="0.2"/>
+
+          <!-- Joint 3: Elbow Flex (pitch) -->
+          <body name="forearm_link" pos="0.10 0 0.04">
+            <joint name="elbow_flex" type="hinge" axis="0 1 0" range="-3.14159 0"/>
+            <geom name="forearm_geom" type="capsule" fromto="0 0 0 0.10 0 -0.02" size="0.018" material="arm_mat" mass="0.15"/>
+
+            <!-- Joint 4: Wrist Flex (pitch) -->
+            <body name="wrist_flex_link" pos="0.10 0 -0.02">
+              <joint name="wrist_flex" type="hinge" axis="0 1 0" range="-1.5708 1.5708"/>
+              <geom name="wrist_flex_geom" type="capsule" fromto="0 0 0 0.04 0 0" size="0.015" material="arm_mat" mass="0.1"/>
+
+              <!-- Joint 5: Wrist Roll (roll) -->
+              <body name="wrist_roll_link" pos="0.04 0 0">
+                <joint name="wrist_roll" type="hinge" axis="1 0 0" range="-3.14159 3.14159"/>
+                <geom name="wrist_roll_geom" type="cylinder" size="0.016 0.01" material="arm_mat" mass="0.05"/>
+
+                <!-- Wrist-mounted camera -->
+                <!-- Wrist camera: mounted above gripper, looking down at the fingers/workspace -->
+                <!-- Wrist camera: looks along -local_z toward the gripper/workspace -->
+                <!-- Wrist camera: looks down toward gripper/workspace -->
+                <!-- Wrist camera: mounted on gripper mount, looking forward/down through the fingers -->
+                <camera name="arm_cam" pos="0 0 -0.025" xyaxes="0 1 0 0 0 1" fovy="90"/>
+
+                <!-- End-effector site (at grasp center between finger pads) -->
+                <!-- EE site at center of finger pads (invisible — for IK only) -->
+                <!-- EE site at grasp center (invisible) -->
+                <site name="ee_site" pos="0.0 0 -0.046" size="0.001" rgba="0 0 0 0"/>
+                <!-- finger pads center at body(-0.024) + geom(0, +-0.006, -0.022) = wrist + (0, +-0.006, -0.046) -->
+
+                <!-- ==================== GRIPPER ==================== -->
+                <!-- Parallel jaw gripper. Fixed mount plate at wrist, two sliding
+                     finger prongs. qpos=0 -> closed, qpos=0.04 -> open. -->
+
+                <!-- Fixed gripper mount plate (always attached to wrist) -->
+                <geom name="gripper_mount" type="box" size="0.012 0.028 0.006" pos="0 0 -0.018"
+                      material="finger_mat" mass="0.02"/>
+
+                <!-- Left finger: slides along +Y -->
+                <body name="left_finger" pos="0.0 0.0 -0.024">
+                  <joint name="finger_left" type="slide" axis="0 1 0" range="0 0.04" damping="0.3"/>
+                  <!-- Horizontal slider bar (stays connected to mount visually) -->
+                  <geom name="left_finger_slider" type="box" size="0.006 0.018 0.004" pos="0 0.008 0"
+                        material="finger_mat" mass="0.003" contype="0" conaffinity="0"/>
+                  <!-- Vertical finger pad (gripping surface) -->
+                  <geom name="left_finger_pad" type="box" size="0.008 0.005 0.022" pos="0 0.006 -0.022"
+                        material="finger_mat" mass="0.01"
+                        condim="4" friction="5.0 0.01 0.001"
+                        solimp="0.95 0.99 0.001 0.5 2" solref="0.02 1.0"/>
+                </body>
+
+                <!-- Right finger: slides along -Y -->
+                <body name="right_finger" pos="0.0 0.0 -0.024">
+                  <joint name="finger_right" type="slide" axis="0 -1 0" range="0 0.04" damping="0.3"/>
+                  <!-- Horizontal slider bar -->
+                  <geom name="right_finger_slider" type="box" size="0.006 0.018 0.004" pos="0 -0.008 0"
+                        material="finger_mat" mass="0.003" contype="0" conaffinity="0"/>
+                  <!-- Vertical finger pad -->
+                  <geom name="right_finger_pad" type="box" size="0.008 0.005 0.022" pos="0 -0.006 -0.022"
+                        material="finger_mat" mass="0.01"
+                        condim="4" friction="5.0 0.01 0.001"
+                        solimp="0.95 0.99 0.001 0.5 2" solref="0.02 1.0"/>
+                </body>
+
+              </body> <!-- wrist_roll_link -->
+            </body> <!-- wrist_flex_link -->
+          </body> <!-- forearm_link -->
+        </body> <!-- upper_arm_link -->
+      </body> <!-- shoulder_link -->
+    </body> <!-- base_link -->
+
+    <!-- ==================== OBJECTS ==================== -->
+
+    <!-- Graspable box -->
+    <body name="box" pos="0.15 0.0 0.44">
+      <freejoint name="box_joint"/>
+      <geom name="box_geom" type="box" size="0.015 0.015 0.015"
+            material="box_mat" mass="0.03"
+            condim="4" friction="1.5 0.005 0.001"
+            solimp="0.95 0.99 0.001 0.5 2" solref="0.02 1.0"
+            contype="1" conaffinity="1"/>
+    </body>
+
+    <!-- Bowl (static) -->
+    <body name="bowl" pos="-0.15 0.0 0.42">
+      <!-- Bowl base (thin disc) -->
+      <geom name="bowl_base" type="cylinder" size="0.045 0.003" pos="0 0 0" material="bowl_mat" contype="1" conaffinity="1"/>
+      <!-- Bowl walls (ring of thin boxes approximating a cylinder shell) -->
+      <geom name="bowl_wall_front" type="box" size="0.003 0.045 0.02" pos="0.045 0 0.02" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_back" type="box" size="0.003 0.045 0.02" pos="-0.045 0 0.02" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_left" type="box" size="0.045 0.003 0.02" pos="0 0.045 0.02" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_right" type="box" size="0.045 0.003 0.02" pos="0 -0.045 0.02" material="bowl_mat" contype="1" conaffinity="1"/>
+      <!-- Diagonal walls for rounder shape -->
+      <geom name="bowl_wall_fl" type="box" size="0.003 0.045 0.02" pos="0.032 0.032 0.02" euler="0 0 0.785" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_fr" type="box" size="0.003 0.045 0.02" pos="0.032 -0.032 0.02" euler="0 0 -0.785" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_bl" type="box" size="0.003 0.045 0.02" pos="-0.032 0.032 0.02" euler="0 0 -0.785" material="bowl_mat" contype="1" conaffinity="1"/>
+      <geom name="bowl_wall_br" type="box" size="0.003 0.045 0.02" pos="-0.032 -0.032 0.02" euler="0 0 0.785" material="bowl_mat" contype="1" conaffinity="1"/>
+    </body>
+
+  </worldbody>
+
+  <!-- ==================== ACTUATORS ==================== -->
+  <actuator>
+    <position name="act_shoulder_pan"  joint="shoulder_pan"  kp="80" ctrlrange="-3.14159 3.14159"/>
+    <position name="act_shoulder_lift" joint="shoulder_lift" kp="80" ctrlrange="-1.5708 1.5708"/>
+    <position name="act_elbow_flex"    joint="elbow_flex"    kp="80" ctrlrange="-3.14159 0"/>
+    <position name="act_wrist_flex"    joint="wrist_flex"    kp="60"  ctrlrange="-1.5708 1.5708"/>
+    <position name="act_wrist_roll"    joint="wrist_roll"    kp="40"  ctrlrange="-3.14159 3.14159"/>
+    <position name="act_finger_left"   joint="finger_left"   kp="15"  ctrlrange="0 0.04"/>
+    <position name="act_finger_right"  joint="finger_right"  kp="15"  ctrlrange="0 0.04"/>
+  </actuator>
+
+  <equality>
+    <!-- Couple gripper fingers to move symmetrically -->
+    <joint joint1="finger_left" joint2="finger_right" polycoef="0 1 0 0 0"/>
+  </equality>
+
+  <!-- ==================== SENSORS ==================== -->
+  <sensor>
+    <jointpos name="sen_shoulder_pan"  joint="shoulder_pan"/>
+    <jointpos name="sen_shoulder_lift" joint="shoulder_lift"/>
+    <jointpos name="sen_elbow_flex"    joint="elbow_flex"/>
+    <jointpos name="sen_wrist_flex"    joint="wrist_flex"/>
+    <jointpos name="sen_wrist_roll"    joint="wrist_roll"/>
+    <jointpos name="sen_finger_left"   joint="finger_left"/>
+    <framepos  name="sen_ee_pos"  objtype="site" objname="ee_site"/>
+    <framequat name="sen_ee_quat" objtype="site" objname="ee_site"/>
+  </sensor>
+
+</mujoco>