cida-plugin 1.0.0__py3-none-any.whl

cida_plugin/__init__.py

@@ -0,0 +1,62 @@
+"""
+CIDA-Plugin: Universal Evidence-Grounded Multi-Agent Deliberation Layer (v3).
+
+Использование (минимальное):
+    from cida_plugin import CIDAPlugin, CIDAPluginConfig
+
+    cfg = CIDAPluginConfig(d_input=768, num_classes=2)
+    plugin = CIDAPlugin(cfg)
+    out = plugin(pooled_output)   # pooled_output: (B, 768)
+    logits = out["p_final"]       # (B, 2)
+
+Использование (с seq_output для evidence pointers):
+    out = plugin(pooled_output, seq_output=hidden_states, mask=attention_mask)
+"""
+
+from .config import CIDAPluginConfig
+from .core import CIDAPlugin
+from .agent import (
+    AgentState,
+    apply_role_prior,
+    compute_role_orthogonality_loss,
+    RoleNames,
+)
+from .deliberation import (
+    AgentEvidenceExtractor,
+    MessageFormulator,
+    CounterargumentCommunication,
+    AgentUpdater,
+)
+from .consensus import ConsensusAggregator, HaltingPredictor
+from .losses import CIDALoss
+from .diagnostics import DebateDiagnostics
+from .ttt import TestTimeTrainer
+from .liquid_dynamics import LiquidDeliberationSolver
+from .hf import wrap_hf_model, HFModelWithCIDA
+
+__all__ = [
+    # Главный интерфейс
+    "CIDAPlugin",
+    "CIDAPluginConfig",
+    "wrap_hf_model",
+    "HFModelWithCIDA",
+    # Потери и диагностика
+    "CIDALoss",
+    "DebateDiagnostics",
+    # [v5] Test-Time Training
+    "TestTimeTrainer",
+    # [v5] Liquid Neural ODE Dynamics
+    "LiquidDeliberationSolver",
+    # Компоненты (для кастомных архитектур)
+    "AgentState",
+    "apply_role_prior",
+    "compute_role_orthogonality_loss",
+    "RoleNames",
+    "AgentEvidenceExtractor",
+    "MessageFormulator",
+    "CounterargumentCommunication",
+    "AgentUpdater",
+    "ConsensusAggregator",
+    "HaltingPredictor",
+]
+

cida_plugin/agent.py

@@ -0,0 +1,198 @@
+"""
+agent.py — CIDA-Plugin Agent (упрощённая версия)
+
+Ключевое изменение:
+    БЫЛО:  RoleEmbeddings (аддитивный bias) + debate_loss + role_spec_loss + orth_loss
+    СТАЛО: ROLE_PRIORS — структурные априорные убеждения агентов
+
+Почему это лучше:
+    RoleEmbeddings создаёт СТАТИСТИЧЕСКОЕ различие (агенты начинают чуть
+    по-разному и потом схожи). Role priors создают СТРУКТУРНОЕ различие:
+
+    Прокурор  → P(y=1) = 0.85 до просмотра данных
+    Защитник  → P(y=1) = 0.15 до просмотра данных
+    Скептик   → P(y=1) = 0.50 (максимальная неопределённость)
+    Интегратор → P(y=1) = данные (без prior, объективен)
+
+    Это байесовски обоснованная альтернатива: агенты представляют разные
+    prior beliefs, а консенсус — это posterior после объединения.
+
+    Математически гарантировано:
+        |b_0 - b_1| ≥ BLEND * |P_0 - P_1| = 0.3 * 0.7 = 0.21 > 0
+    То есть несогласие является СВОЙСТВОМ архитектуры, а не результатом loss.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import Optional
+
+
+# ── Роли агентов: структурные prior beliefs ───────────────────────────────────
+#
+# Для binary/multi-label: prior = вероятность positive класса
+# Для multi-class: prior = вектор вероятностей классов
+#
+# BLEND = насколько сильно prior смешивается с данными
+#   0.0 → prior игнорируется (все агенты одинаковы)
+#   0.3 → 30% prior, 70% данные (рекомендуется)
+#   1.0 → агент игнорирует данные (только prior)
+
+ROLE_PRIOR_POSITIVE = [
+    0.85,  # Прокурор: склонен считать что патология есть
+    0.15,  # Защитник: склонен считать что патологии нет
+    0.50,  # Скептик: максимальная неопределённость
+    None,  # Интегратор: без prior, смотрит объективно
+]
+
+ROLE_PRIOR_BLEND = 0.30  # Сила prior относительно данных
+
+
+@dataclass
+class AgentState:
+    """Явное представление состояния агента."""
+    s: torch.Tensor       # (B, d)   - Hidden state
+    b: torch.Tensor       # (B, K)   - Belief (posterior после prior blending)
+    u: torch.Tensor       # (B, 1)   - Total uncertainty (backward compat)
+    p: torch.Tensor       # (B, n)   - Evidence pointer
+    e: torch.Tensor       # (B, d)   - Evidence vector
+    alpha: torch.Tensor   # (B, K)   - Dirichlet parameters
+    # v4: Разделение неопределённости (Josang 2002 + Meinert 2024)
+    u_epi: torch.Tensor = None    # (B, 1)   - Epistemic: K²/alpha_sum² (vacuity)
+    u_alea: torch.Tensor = None   # (B, 1)   - Aleatoric: 1 - max(b) (data noise)
+
+
+def apply_role_prior(
+    b_data: torch.Tensor,
+    role_idx: int,
+    multi_label: bool = False,
+    positive_class_idx: int = 1,
+) -> torch.Tensor:
+    """
+    Смешивает data-driven belief с role-specific prior (обобщённая версия для любого K).
+    """
+    role = role_idx % 4
+    if role == 3: # Интегратор: без prior
+        return b_data
+
+    K = b_data.size(-1)
+    device = b_data.device
+    dtype = b_data.dtype
+
+    if multi_label or K == 1:
+        prior_val = ROLE_PRIOR_POSITIVE[role]
+        prior = torch.full_like(b_data, prior_val)
+    elif K == 2:
+        prior_val = ROLE_PRIOR_POSITIVE[role]
+        prior = torch.tensor(
+            [1.0 - prior_val, prior_val],
+            device=device,
+            dtype=dtype,
+        ).expand_as(b_data)
+    else:
+        # Multi-class
+        if role == 0:   # Prosecutor: biased to positive class
+            prior_val = 0.85
+            uniform_val = (1.0 - prior_val) / max(K - 1, 1)
+            prior = torch.full_like(b_data, uniform_val)
+            prior = prior.clone()
+            prior[:, positive_class_idx] = prior_val
+        elif role == 1: # Defender: biased against positive class (pro-alternative classes)
+            prior_val = 0.15
+            uniform_val = (1.0 - prior_val) / max(K - 1, 1)
+            prior = torch.full_like(b_data, uniform_val)
+            prior = prior.clone()
+            prior[:, positive_class_idx] = prior_val
+        elif role == 2: # Skeptic: neutral / high uncertainty (uniform)
+            prior = torch.full_like(b_data, 1.0 / K)
+            
+        prior = prior / prior.sum(dim=-1, keepdim=True)
+
+    return (1.0 - ROLE_PRIOR_BLEND) * b_data + ROLE_PRIOR_BLEND * prior
+
+
+def apply_role_priors_batched(
+    b_data: torch.Tensor,
+    M: int,
+    multi_label: bool = False,
+    positive_class_idx: int = 1,
+) -> torch.Tensor:
+    """
+    Векторизованная версия apply_role_prior для батча агентов сразу (B, M, K).
+    """
+    B, _, K = b_data.shape
+    device = b_data.device
+    dtype = b_data.dtype
+    
+    priors = torch.zeros((M, K), device=device, dtype=dtype)
+    blend_mask = torch.full((M, 1), ROLE_PRIOR_BLEND, device=device, dtype=dtype)
+    
+    for i in range(M):
+        role = i % 4
+        if role == 3: # Интегратор: без prior
+            blend_mask[i, 0] = 0.0
+            continue
+            
+        if multi_label or K == 1:
+            priors[i, :] = ROLE_PRIOR_POSITIVE[role]
+        elif K == 2:
+            priors[i, 0] = 1.0 - ROLE_PRIOR_POSITIVE[role]
+            priors[i, 1] = ROLE_PRIOR_POSITIVE[role]
+        else:
+            # Multi-class
+            if role == 0:   # Prosecutor: biased to positive class
+                prior_val = 0.85
+                uniform_val = (1.0 - prior_val) / max(K - 1, 1)
+                priors[i, :] = uniform_val
+                priors[i, positive_class_idx] = prior_val
+            elif role == 1: # Defender: biased against positive class
+                prior_val = 0.15
+                uniform_val = (1.0 - prior_val) / max(K - 1, 1)
+                priors[i, :] = uniform_val
+                priors[i, positive_class_idx] = prior_val
+            elif role == 2: # Skeptic: neutral / uniform
+                priors[i, :] = 1.0 / K
+                
+            priors[i] = priors[i] / priors[i].sum()
+            
+    priors = priors.unsqueeze(0)        # (1, M, K)
+    blend_mask = blend_mask.unsqueeze(0) # (1, M, 1)
+    
+    return (1.0 - blend_mask) * b_data + blend_mask * priors
+
+
+
+def compute_role_orthogonality_loss(agent_states: torch.Tensor) -> torch.Tensor:
+    """
+    Мягкая ортогональность через ПРЕДСТАВЛЕНИЯ, а не через веса.
+
+    БЫЛО: orth_loss через Gram(weights) — веса могут быть ортогональны,
+          но representations при этом коллапсируют (если вход низкоранговый).
+
+    СТАЛО: orth_loss через Gram(representations) — напрямую измеряем
+           то, что нас реально интересует.
+
+    agent_states: (B, M, d) — hidden states агентов
+    Returns: scalar — потери (0 при полной ортогональности)
+
+    Математически:
+        G_ij = <s_i, s_j> / (||s_i|| · ||s_j||) — косинусное сходство
+        Loss = mean((G - I)²) → 0 при G = I
+    """
+    B, M, d = agent_states.shape
+    # Усредняем по батчу, нормализуем
+    s_mean = agent_states.mean(0)           # (M, d)
+    s_norm = F.normalize(s_mean, p=2, dim=-1)  # (M, d)
+    gram = s_norm @ s_norm.T                # (M, M)
+    eye = torch.eye(M, device=gram.device)
+    return ((gram - eye) ** 2).mean()
+
+
+class RoleNames:
+    """Имена агентов для диагностики."""
+    NAMES = ["Prosecutor", "Defender", "Skeptic", "Integrator"]
+
+    @classmethod
+    def get(cls, idx: int) -> str:
+        return cls.NAMES[idx % 4]

cida_plugin/config.py

@@ -0,0 +1,228 @@
+"""
+CIDAPluginConfig — единая точка конфигурации для CIDA-Plugin v3.
+"""
+from dataclasses import dataclass, asdict
+import json
+import os
+
+
+@dataclass
+class CIDAPluginConfig:
+    """
+    Конфигурация универсального CIDA-Plugin слоя (v3 — упрощённая архитектура).
+
+    Параметры ядра
+    --------------
+    d_input : int
+        Размерность pooled_output от внешнего энкодера.
+        Примеры: 128 (bert-tiny), 768 (bert-base / distilbert), 1024 (DenseNet121).
+    d_hidden : int
+        Внутренняя размерность агентов плагина.
+    d_message : int
+        Размерность сообщений при коммуникации агентов.
+    num_agents : int
+        Количество агентов (рекомендуется 4: Prosecutor, Defender, Skeptic, Integrator).
+    num_classes : int
+        Число целевых классов.
+    max_rounds : int
+        Максимальное число раундов deliberation.
+    multi_label : bool
+        True для multi-label задач (BCE), False для single-label (CE).
+
+    Архитектурные параметры
+    -----------------------
+    num_attn_heads : int
+        Число голов в TransformerAgentUpdater (cross-attention).
+        Должен делить d_hidden.
+
+    Role Priors
+    -----------
+    role_prior_blend : float
+        Сила смешивания роли с данными (0.0 = без priors, 1.0 = только priors).
+        Рекомендуется 0.30 (30% prior, 70% данные).
+    positive_class_idx : int
+        Индекс "положительного" класса для multi-class задач (K>2).
+        По умолчанию 1, но может быть 0, 2 и т.д. в зависимости от датасета.
+
+    Loss гиперпараметры
+    -------------------
+    lambda_cal : float
+        Вес calibration loss (Brier score). Рекомендуется 0.3–0.5.
+    lambda_ac : float
+        Вес anti-collapse loss. Рекомендуется 0.1–0.3.
+    min_disagreement : float
+        Минимальный порог несогласия для anti-collapse.
+
+    Регуляризация
+    -------------
+    comm_dropout : float
+        Dropout на канале коммуникации между агентами.
+    early_stop_threshold : float or None
+        Порог уверенности для досрочной остановки deliberation (inference only).
+        Если None — всегда идёт max_rounds раундов.
+    freeze_input_proj : bool
+        Если True — входной проекционный слой не обучается.
+
+    Ablation Flags
+    --------------
+    abl_no_pointers : bool
+        Отключить evidence pointers.
+    abl_no_messages : bool
+        Отключить формулировку сообщений.
+    abl_no_communication : bool
+        Полностью отключить коммуникацию.
+
+    Perspective Projector
+    ---------------------
+    use_perspective_projector : bool
+        Если True — каждый агент имеет свою проекцию входа (per-agent view).
+        Если False — все агенты начинают с одного представления (expand).
+        По умолчанию False для упрощённой архитектуры v3.
+
+    Test-Time Training (TTT)
+    -------------------------
+    use_ttt : bool
+        Агенты адаптируют свои веса к каждому входу перед deliberation.
+    ttt_steps : int
+        Число шагов внутренней оптимизации (K).
+    ttt_lr : float
+        Learning rate для inner Adam.
+    ttt_mask_ratio : float
+        Доля маскируемых компонент скрытого состояния.
+
+    Liquid Neural ODE Dynamics
+    --------------------------
+    use_liquid_dynamics : bool
+        Заменяет дискретные раунды на непрерывное ODE: ds/dt = -s/τ + F(s,r,e).
+    liquid_solver : str
+        Метод интегрирования: 'euler', 'dopri5' и др.
+    liquid_atol : float
+        Абсолютная толерантность для адаптивных решателей.
+    liquid_rtol : float
+        Относительная толерантность для адаптивных решателей.
+    trajectory_save_every : int
+        Сохранять каждое N-е состояние в trajectory для ODE решателя.
+        Для адаптивных решателей (dopri5) рекомендуется 10-50 для экономии памяти.
+        Для euler можно использовать 1.
+    """
+    # ─── Размерности ────────────────────────────────────────────────────────────
+    d_input: int = 128
+    d_hidden: int = 128
+    d_message: int = 128
+
+    # ─── Агенты и deliberation ───────────────────────────────────────────────────
+    num_agents: int = 4
+    num_classes: int = 2
+    max_rounds: int = 3
+    multi_label: bool = False
+
+    # ─── Архитектурные параметры ─────────────────────────────────────────────────
+    num_attn_heads: int = 4
+
+    # ─── Role Priors ─────────────────────────────────────────────────────────────
+    role_prior_blend: float = 0.30
+    positive_class_idx: int = 1
+
+    # ─── Loss гиперпараметры ─────────────────────────────────────────────────────
+    lambda_cal: float = 0.4
+    lambda_ac: float = 0.2
+    min_disagreement: float = 0.08
+
+    # ─── Регуляризация ──────────────────────────────────────────────────────────
+    comm_dropout: float = 0.2
+    early_stop_threshold: float = 0.90
+    freeze_input_proj: bool = False
+
+    # ─── Ablation flags ─────────────────────────────────────────────────────────
+    abl_no_pointers: bool = False
+    abl_no_messages: bool = False
+    abl_no_communication: bool = False
+
+    # ─── Perspective Projector ─────────────────────────────────────────────────
+    use_perspective_projector: bool = True
+
+    # ─── Test-Time Training ──────────────────────────────────────────────────────
+    use_ttt: bool = False
+    ttt_steps: int = 3
+    ttt_lr: float = 1e-3
+    ttt_mask_ratio: float = 0.15
+
+    # ─── Liquid Neural ODE Dynamics ──────────────────────────────────────────────
+    use_liquid_dynamics: bool = False
+    liquid_solver: str = "euler"
+    liquid_atol: float = 1e-3
+    liquid_rtol: float = 1e-3
+    trajectory_save_every: int = 10
+
+    # ─── CIDA v4: Architectural Improvements ─────────────────────────────────────
+    # Disagreement-Routed Sparse Communication (Li et al. EMNLP 2024)
+    # Каждый агент слушает только top-K по disagreement, не всех.
+    sparse_comm_k: int = 2
+
+    # Anonymous Message Passing (Choi et al. ACL 2025)
+    # Убирает s (identity) из сообщений → агент слышит аргумент, а не авторитет.
+    anonymous_messages: bool = True
+
+    # Dynamic Trust Weights (CortexDebate MDM, ACL 2025)
+    # EMA trust вместо статических ROLE_WEIGHTS.
+    use_dynamic_trust: bool = False
+    trust_ema_gamma: float = 0.9
+
+    # ─── Kairos Dynamic Deliberation (KDD) ───────────────────────────────────────
+    # Time-Travel (Rollback) mechanics
+    enable_rollback: bool = False
+    rollback_noise_std: float = 0.05
+    
+    # Latent Superposition
+    # If final disagreement is > threshold, output Superposition State
+    superposition_threshold: float = 0.7
+
+    def __post_init__(self):
+        assert self.d_input > 0, "d_input must be positive"
+        assert self.num_agents >= 2, "Need at least 2 agents"
+        assert self.num_classes >= 2, "Need at least 2 classes"
+        assert self.max_rounds >= 1, "Need at least 1 deliberation round"
+        assert self.d_hidden % self.num_attn_heads == 0, (
+            f"d_hidden ({self.d_hidden}) must be divisible by "
+            f"num_attn_heads ({self.num_attn_heads})"
+        )
+        assert 0.0 <= self.role_prior_blend <= 1.0, (
+            "role_prior_blend must be in [0.0, 1.0]"
+        )
+        if self.early_stop_threshold is not None:
+            assert 0.5 < self.early_stop_threshold <= 1.0, (
+                "early_stop_threshold must be in (0.5, 1.0]"
+            )
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, config_dict):
+        # Filter out unknown keys for forward compatibility
+        valid_keys = {f.name for f in cls.__dataclass_fields__.values()}
+        filtered = {k: v for k, v in config_dict.items() if k in valid_keys}
+        return cls(**filtered)
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        config_file = os.path.join(save_directory, "config.json")
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(self.to_dict(), f, indent=2, sort_keys=True)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str):
+        if os.path.isdir(pretrained_model_name_or_path):
+            config_file = os.path.join(
+                pretrained_model_name_or_path, "config.json"
+            )
+        else:
+            from huggingface_hub import hf_hub_download
+            config_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="config.json",
+            )
+
+        with open(config_file, "r", encoding="utf-8") as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)

cida_plugin/consensus.py

@@ -0,0 +1,178 @@
+"""
+consensus.py — Упрощённый ConsensusAggregator
+
+Ключевые исправления:
+
+1. УБРАН ReliabilityTracker
+   Причина: круговая зависимость rho ← accuracy ← predictions ← weights(rho).
+   На малых датасетах вырождается. Без теоретической гарантии сходимости.
+   Замена: статические веса ролей (обоснованные a priori).
+
+2. ИСПРАВЛЕН ConsensusAggregator
+   БЫЛО: PoE = log p ∝ Σ w_i · log(b_i)
+   ПРОБЛЕМА: PoE предполагает НЕЗАВИСИМЫХ экспертов.
+   Агенты CIDA КОРРЕЛИРОВАНЫ — общий энкодер + communication между агентами.
+   PoE на коррелированных агентах → усиление shared bias → overconfidence.
+
+   СТАЛО: Weighted Mean + Disagreement-as-Uncertainty
+   p_final = Σ w_i · b_i  (взвешенное среднее)
+   uncertainty = std(b) по агентам
+
+   Математическое обоснование:
+   Для коррелированных экспертов с ковариационной матрицей Σ:
+       Var(Σ w_i · b_i) = w^T · Σ · w  (правильно учитывает корреляции)
+       Var(PoE(b_i)) ≠ w^T · Σ · w     (PoE неверно предполагает Σ = σ²I)
+   При положительной корреляции: PoE под-оценивает variance → overconfidence.
+
+3. УПРОЩЁН интерфейс
+   Убраны: epoch_fraction, temperature scheduling, normalize_components флаг.
+   Всё это компенсировало нестабильность PoE на коррелированных агентах.
+   При правильной агрегации это не нужно.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Статические веса ролей (a priori обоснованы ролями)
+# Интегратор = наибольший вес (его роль — синтезировать)
+# Скептик = наименьший вес (его роль — поднимать вопросы, не решать)
+ROLE_WEIGHTS = torch.tensor([0.28, 0.28, 0.12, 0.32])  # P, D, S, I (сумма = 1.0)
+
+
+class ConsensusAggregator(nn.Module):
+    """
+    Weighted Mean консенсус для КОРРЕЛИРОВАННЫХ агентов.
+
+    v4: Dynamic Trust Weights (CortexDebate MDM, ACL 2025).
+    Статические веса заменяются на EMA trust — вес агента растёт
+    когда его предсказания совпадают с финальным консенсусом.
+
+    trust_i += γ·(1[b_i_correct] - trust_i) per batch, detached from main graph
+    """
+
+    def __init__(
+        self,
+        num_agents: int = 4,
+        multi_label: bool = False,
+        use_dynamic_trust: bool = False,
+        trust_ema_gamma: float = 0.9,
+    ):
+        super().__init__()
+        self.multi_label = multi_label
+        self.use_dynamic_trust = use_dynamic_trust
+        self.trust_ema_gamma = trust_ema_gamma
+
+        # Статические веса ролей — обучаемые для адаптации к задаче
+        init_weights = ROLE_WEIGHTS[:num_agents].clone()
+        init_weights = init_weights / init_weights.sum()
+        self.role_weights = nn.Parameter(init_weights)
+
+        # v4: Dynamic trust (EMA buffer, не участвует в градиентах)
+        if use_dynamic_trust:
+            self.register_buffer(
+                'trust_scores',
+                torch.ones(num_agents) / num_agents
+            )
+
+    @property
+    def weights(self) -> torch.Tensor:
+        """Нормализованные веса ролей (в simplex)."""
+        if self.use_dynamic_trust and hasattr(self, 'trust_scores'):
+            return self.trust_scores
+        return F.softmax(self.role_weights, dim=0)  # (M,)
+
+    def update_trust(self, b: torch.Tensor, p_final: torch.Tensor):
+        """
+        Обновляет trust scores на основе согласованности агента с консенсусом.
+        Вызывается ПОСЛЕ forward(), detached от основного графа.
+
+        b: (B, M, K) — убеждения агентов
+        p_final: (B, K) — финальный консенсус
+        """
+        if not self.use_dynamic_trust:
+            return
+        with torch.no_grad():
+            # Per-agent correctness = 1 - L1(b_i, p_final)
+            correctness = 1.0 - (b - p_final.unsqueeze(1)).abs().mean(dim=(0, 2))  # (M,)
+            gamma = self.trust_ema_gamma
+            self.trust_scores = gamma * self.trust_scores + (1 - gamma) * correctness
+            self.trust_scores = self.trust_scores / self.trust_scores.sum()
+
+    def forward(
+        self,
+        b: torch.Tensor,          # (B, M, K)
+        u: torch.Tensor = None,   # (B, M, 1) — опционально
+        u_epi: torch.Tensor = None,   # (B, M, 1) — v4: epistemic
+        u_alea: torch.Tensor = None,  # (B, M, 1) — v4: aleatoric
+    ) -> tuple:
+        """
+        Returns:
+            p_final:     (B, K)  — консенсус
+            uncertainty: (B,)   — мера неопределённости
+            disagreement:(B,)   — мера несогласия агентов
+        """
+        B, M, K = b.shape
+        w = self.weights  # (M,) — нормализованные role weights
+
+        # ── Взвешенное среднее (корректно для коррелированных агентов) ────────
+        p_final = (w.view(1, M, 1) * b).sum(dim=1)  # (B, K)
+
+        if not self.multi_label:
+            # FIX v4: Нормализация через log-softmax — численно стабильна
+            p_final = F.softmax(p_final.log().clamp(-10, 10), dim=-1)
+
+        # ── Несогласие = std убеждений по агентам (со стабильным градиентом) ──
+        var = b.var(dim=1, unbiased=True)
+        disagreement = torch.sqrt(var + 1e-8).mean(dim=-1)  # (B,)
+
+        # ── Неопределённость = несогласие + собственная uncertainty агентов ──
+        if u is not None:
+            agent_u = (w.view(1, M, 1) * u).sum(dim=1).squeeze(-1)  # (B,)
+            uncertainty = 0.5 * disagreement + 0.5 * agent_u
+        else:
+            uncertainty = disagreement
+
+        # v4: Обновляем dynamic trust (если включён)
+        if self.training and self.use_dynamic_trust:
+            self.update_trust(b, p_final)
+
+        return p_final, uncertainty, disagreement
+
+
+class HaltingPredictor(nn.Module):
+    """
+    Предсказатель остановки (упрощённый).
+    Используется для ACT: когда агенты достаточно согласны — останавливаемся.
+
+    В отличие от оригинала, решение основано на НЕСОГЛАСИИ (наблюдаемо),
+    а не на выученном MLP поверх всех состояний (трудно интерпретировать).
+    """
+
+    def __init__(self, agreement_threshold: float = 0.15):
+        """
+        agreement_threshold: при несогласии ниже этого порога — останавливаемся.
+        """
+        super().__init__()
+        self.threshold = agreement_threshold
+        # Небольшой выученный bias для подстройки порога под задачу
+        self.bias = nn.Parameter(torch.tensor(0.0))
+
+    def should_halt(self, disagreement: torch.Tensor) -> torch.Tensor:
+        """
+        disagreement: (B,)
+        Returns: (B,) bool — True если стоит остановиться
+        """
+        effective_threshold = self.threshold + self.bias.sigmoid() * 0.1
+        return disagreement < effective_threshold
+
+    def halting_probability(self, disagreement: torch.Tensor) -> torch.Tensor:
+        """
+        Мягкая версия: вероятность остановки ∈ [0, 1].
+        disagreement: (B,)
+        Returns: (B,)
+        """
+        effective_threshold = self.threshold + self.bias.sigmoid() * 0.1
+        # Sigmoid(-scale * (d - threshold)): высокая вероятность при d < threshold
+        return torch.sigmoid(-10.0 * (disagreement - effective_threshold))