ACTmodel.py

# 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