migma 0.1.0__tar.gz

migma-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+vendor/
+__pycache__/
+build/
+*egg-info/
+dist/
+*.egg

migma-0.1.0/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: migma
+Version: 0.1.0
+Summary: Modular action policy components for Vision-Language-Action models
+Project-URL: Repository, https://github.com/lucasjinreal/migma
+License: Apache-2.0
+Keywords: action-policy,diffusion,flow-matching,robotics,vla
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: diffusers>=0.27.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: transformers>=4.40.0
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == 'dev'
+Requires-Dist: pytest>=7.0.0; extra == 'dev'
+Requires-Dist: ruff>=0.1.0; extra == 'dev'
+Requires-Dist: twine>=4.0.0; extra == 'dev'

migma-0.1.0/migma/__init__.py

@@ -0,0 +1,21 @@
+"""
+migma — Modular action policy components for Vision-Language-Action (VLA) models.
+
+Provides layerwise flow-matching action heads that attach to a frozen VLM backbone
+and decode robot actions from its per-layer KV cache or hidden states.
+"""
+
+from migma.models.action_heads import (
+    FlowmatchingActionHeadConfig,
+    LayerwiseFMDSCActionHead,
+    LayerwiseFMPi05ActionHead,
+    LayerwiseFMPiActionHead,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "FlowmatchingActionHeadConfig",
+    "LayerwiseFMPiActionHead",
+    "LayerwiseFMPi05ActionHead",
+    "LayerwiseFMDSCActionHead",
+]

migma-0.1.0/migma/models/__init__.py

@@ -0,0 +1,11 @@
+from migma.models.action_heads import (
+    FlowmatchingActionHeadConfig,
+    LayerwiseFMDSCActionHead,
+    LayerwiseFMPiActionHead,
+)
+
+__all__ = [
+    "FlowmatchingActionHeadConfig",
+    "LayerwiseFMPiActionHead",
+    "LayerwiseFMDSCActionHead",
+]

migma-0.1.0/migma/models/action_heads/__init__.py

@@ -0,0 +1,11 @@
+from migma.models.action_heads.configuration import FlowmatchingActionHeadConfig
+from migma.models.action_heads.layerwise_fm_dsc import LayerwiseFMDSCActionHead
+from migma.models.action_heads.layerwise_fm_pi import LayerwiseFMPiActionHead
+from migma.models.action_heads.layerwise_fm_pi05 import LayerwiseFMPi05ActionHead
+
+__all__ = [
+    "FlowmatchingActionHeadConfig",
+    "LayerwiseFMPiActionHead",
+    "LayerwiseFMDSCActionHead",
+    "LayerwiseFMPi05ActionHead",
+]

migma-0.1.0/migma/models/action_heads/configuration.py

@@ -0,0 +1,121 @@
+"""
+Configuration dataclass for layerwise flow-matching action heads.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class FlowmatchingActionHeadConfig:
+    """
+    Unified configuration for all layerwise flow-matching action head variants.
+
+    The two concrete head classes (:class:`~migma.models.action_heads.LayerwiseFMPiActionHead`
+    and :class:`~migma.models.action_heads.LayerwiseFMDSCActionHead`) share this
+    configuration schema and ignore fields that are not relevant to them.
+
+    Core dimensions
+    ---------------
+    action_dim:
+        Dimension of a single action vector (e.g. 14 for bimanual euler, 20 for r6d).
+    action_horizon:
+        Total rollout length including the current step
+        (``future_action_window_size + 1`` in the original starVLA convention).
+    vl_hidden_dim:
+        Hidden dimension of the VLM backbone (e.g. 2048 for Qwen2.5-VL-2B).
+    num_inference_timesteps:
+        Number of Euler integration steps at inference time.
+
+    DiT architecture
+    ----------------
+    dit_cfg:
+        Dict of overrides merged on top of the default DiT hyperparameters.
+        Recognised keys: ``num_layers``, ``num_attention_heads``,
+        ``attention_head_dim``, ``dropout``, ``activation_fn``, …
+        Any key accepted by the underlying :class:`~migma.nn.DiT` /
+        :class:`~migma.nn.DSCDiT` constructor is valid here.
+    use_kv_cache_dit:
+        When ``True``, uses the KV-cache-capable DiT variant.
+    kv_hidden_size:
+        Hidden size of the KV tensors supplied by the VLM.
+    kv_compress_ratio:
+        If < 1.0, a :class:`~migma.nn.GatedMLPCompressor` is applied to the KV
+        cache before projection (DSC head only).
+    global_cond_dim:
+        Dimension of ``vlm_last_hidden_state``. Defaults to ``kv_hidden_size``
+        when ``None`` (DSC head only).
+
+    State encoder
+    -------------
+    state_dim:
+        When set, a two-layer MLP maps the proprioceptive state vector to
+        ``vl_hidden_dim`` before concatenating with action tokens.
+
+    Positional embedding
+    --------------------
+    add_pos_embed:
+        Whether to add a learnable positional embedding to action features.
+    max_seq_len:
+        Maximum sequence length for the positional embedding table.
+    num_target_vision_tokens:
+        Number of learnable "future" tokens prepended to the action sequence.
+
+    Flow-matching noise
+    -------------------
+    noise_beta_alpha / noise_beta_beta:
+        Parameters of the Beta distribution used to sample diffusion time.
+    noise_s:
+        Flow-matching schedule shift:  ``t = (s - Beta_sample) / s``.
+    num_timestep_buckets:
+        Number of discrete bins for the continuous time variable.
+    use_scaled_noise:
+        Whether to scale initial noise by per-dimension action statistics.
+    t_eps:
+        Minimum ``(1 - t)`` denominator to avoid division by zero in the AML
+        velocity formula (DSC head only).
+
+    Loss coefficients
+    -----------------
+    smoothness_loss_weight:
+        Weight for the action smoothness regulariser (0 = disabled).
+    enable_first_action_state_loss:
+        Anchor the first predicted action to the current proprioceptive state
+        (Pi head only).
+    first_action_state_loss_weight:
+        Weight for the first-action anchor loss.
+    """
+
+    # ── Core dimensions (required) ──────────────────────────────────────────
+    action_dim: int
+    action_horizon: int
+    vl_hidden_dim: int
+    num_inference_timesteps: int
+
+    # ── DiT architecture ────────────────────────────────────────────────────
+    dit_cfg: Dict[str, Any] = field(default_factory=dict)
+    use_kv_cache_dit: bool = False
+    kv_hidden_size: Optional[int] = None  # None means same as vl_hidden_dim (no projection)
+    kv_compress_ratio: float = 1.0
+    global_cond_dim: Optional[int] = None   # DSC-specific
+
+    # ── State encoder ───────────────────────────────────────────────────────
+    state_dim: Optional[int] = None
+
+    # ── Token sequence ──────────────────────────────────────────────────────
+    num_target_vision_tokens: int = 32
+    add_pos_embed: bool = True
+    max_seq_len: int = 1024
+
+    # ── Flow-matching noise ─────────────────────────────────────────────────
+    noise_beta_alpha: float = 1.5
+    noise_beta_beta: float = 1.0
+    noise_s: float = 0.999
+    num_timestep_buckets: int = 1000
+    use_scaled_noise: bool = True
+    t_eps: float = 5e-2  # DSC/AML-specific
+
+    # ── Loss coefficients ───────────────────────────────────────────────────
+    smoothness_loss_weight: float = 0.0
+    enable_first_action_state_loss: bool = False  # Pi-specific
+    first_action_state_loss_weight: float = 0.0

migma-0.1.0/migma/models/action_heads/layerwise_fm_dsc.py

@@ -0,0 +1,403 @@
+"""
+Layerwise Flow-Matching Action Head — AML / DSC variant.
+
+Velocity target:  ``v = (actions − x_t) / (1 − t)``  (AML / posterior-mean)
+KV conditioning:  list of ``(key, value)`` tuples or hidden-state tensors.
+Extra signal:     optional ``vlm_last_hidden_state`` drives dual-stream
+                  conditioning inside :class:`~migma.nn.DSCDiT`.
+
+References:
+  - Consistent Flow Matching (AML training objective)
+  - DSC-DiT dual-stream conditioning design
+"""
+
+import logging
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions import Beta
+
+from migma.models.action_heads.configuration import FlowmatchingActionHeadConfig
+from migma.nn.dit import DiT
+from migma.nn.dit_dsc import DSCDiT
+from migma.nn.encoders import ActionEncoder, MLP
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_DIT_CFG = {
+    "num_layers": 28,
+    "attention_head_dim": 64,
+    "dropout": 0.1,
+    "attention_bias": True,
+    "activation_fn": "gelu-approximate",
+    "norm_type": "ada_norm",
+    "norm_elementwise_affine": False,
+    "norm_eps": 1e-5,
+    "max_num_positional_embeddings": 512,
+    "final_dropout": True,
+    "positional_embeddings": "sinusoidal",
+    "interleave_self_attention": False,
+}
+
+
+def _build_dit_config(cfg: FlowmatchingActionHeadConfig) -> dict:
+    final = _DEFAULT_DIT_CFG.copy()
+    final.update(cfg.dit_cfg)
+
+    vl_dim = cfg.vl_hidden_dim
+    head_dim = final.get("attention_head_dim", 64)
+    final["num_attention_heads"] = final.get("num_attention_heads", vl_dim // head_dim)
+    final["output_dim"] = vl_dim
+    final["kv_hidden_size"] = final.get("kv_hidden_size", cfg.kv_hidden_size)
+    final["kv_compress_ratio"] = final.get("kv_compress_ratio", cfg.kv_compress_ratio)
+    final["global_cond_dim"] = final.get(
+        "global_cond_dim",
+        cfg.global_cond_dim if cfg.global_cond_dim is not None else cfg.kv_hidden_size,
+    )
+    # cross_attention_dim only used by the standard DiT fallback.
+    final["cross_attention_dim"] = final.get("cross_attention_dim", vl_dim)
+    return final
+
+
+class LayerwiseFMDSCActionHead(nn.Module):
+    """
+    AML flow-matching head with Dual-Stream Conditioned DiT (DSC-DiT).
+
+    **Training objective** (AML / consistent FM):
+
+    .. math::
+
+        x_t = t \\cdot x_0 + (1-t) \\cdot \\epsilon, \\quad
+        v_\\text{target} = \\frac{x_0 - x_t}{1-t}
+
+    where :math:`\\epsilon \\sim \\mathcal{N}(0, I)` (optionally scaled).
+
+    **Conditioning**:
+
+    * *Primary* — per-layer ``(key, value)`` tensors from the VLM or a list
+      of hidden-state tensors (standard cross-attention).
+    * *Secondary* — ``vlm_last_hidden_state`` fed into :class:`~migma.nn.DSCDiT`
+      for attention-pooled global adaLN and sequence-level cross-attention.
+
+    Args:
+        config: :class:`~migma.models.action_heads.FlowmatchingActionHeadConfig`
+
+    Example::
+
+        cfg = FlowmatchingActionHeadConfig(
+            action_dim=14, action_horizon=16,
+            vl_hidden_dim=2048, num_inference_timesteps=10,
+            use_kv_cache_dit=True,
+            kv_hidden_size=512,
+        )
+        head = LayerwiseFMDSCActionHead(cfg)
+
+        loss = head(context_list, actions, vlm_last_hidden_state=last_hs)
+        actions = head.predict_action(context_list, vlm_last_hidden_state=last_hs)
+    """
+
+    def __init__(self, config: FlowmatchingActionHeadConfig) -> None:
+        super().__init__()
+        self.config = config
+
+        dit_kwargs = _build_dit_config(config)
+        use_kv_cache = config.use_kv_cache_dit
+        if use_kv_cache:
+            self.model = DSCDiT(**dit_kwargs)
+            logger.info("LayerwiseFMDSCActionHead: using DSCDiT")
+        else:
+            self.model = DiT(**dit_kwargs)
+            logger.info("LayerwiseFMDSCActionHead: using standard DiT (no KV cache)")
+
+        self.use_kv_cache_dit = use_kv_cache
+        self.action_dim = config.action_dim
+        self.action_horizon = config.action_horizon
+        self.num_inference_timesteps = config.num_inference_timesteps
+        self.smoothness_loss_weight = config.smoothness_loss_weight
+        self.t_eps = config.t_eps
+
+        hidden_dim = config.vl_hidden_dim
+
+        self.state_encoder = (
+            MLP(input_dim=config.state_dim, output_dim=hidden_dim)
+            if config.state_dim
+            else None
+        )
+        self.action_encoder = ActionEncoder(
+            action_dim=config.action_dim, hidden_size=hidden_dim
+        )
+        self.action_decoder = MLP(
+            input_dim=hidden_dim, hidden_dim=1024, output_dim=config.action_dim
+        )
+
+        self.future_tokens = nn.Embedding(config.num_target_vision_tokens, hidden_dim)
+        nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02)
+
+        if config.add_pos_embed:
+            self.position_embedding = nn.Embedding(config.max_seq_len, hidden_dim)
+            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
+
+        self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta)
+        self.num_timestep_buckets = config.num_timestep_buckets
+        self.use_scaled_noise = config.use_scaled_noise
+        self.register_buffer("noise_scale", torch.ones(config.action_dim))
+
+        logger.info(
+            "LayerwiseFMDSCActionHead initialised | "
+            f"action_dim={config.action_dim}  horizon={config.action_horizon}  "
+            f"vl_hidden_dim={config.vl_hidden_dim}  "
+            f"params={sum(p.numel() for p in self.parameters()):,}"
+        )
+
+    # -----------------------------------------------------------------------
+    # Helpers
+    # -----------------------------------------------------------------------
+
+    def _sample_time(
+        self, batch_size: int, device: torch.device, dtype: torch.dtype
+    ) -> torch.Tensor:
+        sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype)
+        return (self.config.noise_s - sample) / self.config.noise_s
+
+    @staticmethod
+    def _is_kv_list(context: list) -> bool:
+        return (
+            isinstance(context, (list, tuple))
+            and len(context) > 0
+            and isinstance(context[0], (list, tuple))
+        )
+
+    def _build_sa_embs(
+        self,
+        action_features: torch.Tensor,
+        state_features: Optional[torch.Tensor],
+        B: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        future_tokens = self.future_tokens.weight.unsqueeze(0).expand(B, -1, -1)
+        if self.config.add_pos_embed:
+            pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+            action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0)
+        if state_features is not None:
+            return torch.cat([state_features, future_tokens, action_features], dim=1)
+        return torch.cat([future_tokens, action_features], dim=1)
+
+    def _run_model(
+        self,
+        sa_embs: torch.Tensor,
+        context_list: list,
+        t_disc: torch.Tensor,
+        is_kv: bool,
+        vlm_last_hidden_state: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        device, dtype = sa_embs.device, sa_embs.dtype
+
+        if self.use_kv_cache_dit:
+            return self.model(
+                hidden_states=sa_embs,
+                kv_cache_list=context_list,
+                timestep=t_disc,
+                encoder_attention_mask=None,
+                vlm_last_hidden_state=vlm_last_hidden_state,
+            )
+
+        # Standard DiT: build encoder_hidden_states from context.
+        if vlm_last_hidden_state is not None:
+            enc_hs = vlm_last_hidden_state.to(device=device, dtype=dtype)
+        else:
+            enc_hs = (context_list[-1][1] if is_kv else context_list[-1]).to(
+                device=device, dtype=dtype
+            )
+        return self.model(
+            hidden_states=sa_embs,
+            encoder_hidden_states=enc_hs,
+            timestep=t_disc,
+            encoder_attention_mask=None,
+            vlm_last_hidden_state=vlm_last_hidden_state,
+        )
+
+    # -----------------------------------------------------------------------
+    # Noise-scale calibration
+    # -----------------------------------------------------------------------
+
+    def set_noise_scale_from_stats(
+        self,
+        action_std: Union[list, "np.ndarray", torch.Tensor],
+    ) -> None:
+        """
+        Calibrate per-dimension noise scaling from dataset action statistics.
+
+        See :meth:`LayerwiseFMPiActionHead.set_noise_scale_from_stats` for the
+        full docstring — the logic is identical.
+        """
+        if isinstance(action_std, (list, np.ndarray)):
+            action_std = torch.tensor(action_std, dtype=torch.float32)
+
+        if len(action_std) == 14 and self.action_dim == 20:
+            action_std = _expand_14d_to_20d(action_std)
+
+        if len(action_std) != self.action_dim:
+            raise ValueError(
+                f"action_std length {len(action_std)} != action_dim {self.action_dim}."
+            )
+
+        TARGET = 0.3
+        scale = torch.clamp(TARGET / (action_std + 1e-6), max=1.0)
+        self.noise_scale.data.copy_((action_std * scale).to(self.noise_scale.device))
+        logger.info(
+            f"Noise scale updated | min={self.noise_scale.min():.4f}  "
+            f"max={self.noise_scale.max():.4f}"
+        )
+
+    # -----------------------------------------------------------------------
+    # Forward (training)
+    # -----------------------------------------------------------------------
+
+    def forward(
+        self,
+        context_list: List[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]],
+        actions: torch.Tensor,
+        state: Optional[torch.Tensor] = None,
+        vlm_last_hidden_state: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Compute the AML flow-matching velocity loss.
+
+        Args:
+            context_list:           Per-layer context — list of ``(key, value)``
+                                    tuples or list of hidden-state tensors.
+            actions:                Target trajectory ``(B, T, action_dim)``.
+            state:                  Optional proprioceptive state.
+            vlm_last_hidden_state:  Optional ``(B, S, D)`` — enables DSC dual-stream
+                                    conditioning when using :class:`~migma.nn.DSCDiT`.
+
+        Returns:
+            Scalar training loss.
+        """
+        device = actions.device
+        is_kv = self._is_kv_list(context_list)
+        B = context_list[0][0].shape[0] if is_kv else context_list[0].shape[0]
+
+        t = self._sample_time(B, device=device, dtype=actions.dtype)
+        t = t[:, None, None]  # (B, 1, 1)
+
+        noise = torch.randn(B, self.action_horizon, self.action_dim,
+                            dtype=actions.dtype, device=device)
+        if self.use_scaled_noise:
+            noise = noise * self.noise_scale[None, None, :]
+
+        # AML interpolation: x_t = t·x_0 + (1-t)·ε
+        noisy = t * actions + (1 - t) * noise
+        # AML velocity target: v = (x_0 - x_t) / (1 - t)
+        velocity = (actions - noisy) / (1 - t).clamp_min(self.t_eps)
+
+        t_disc = (t[:, 0, 0] * self.num_timestep_buckets).long()
+        action_features = self.action_encoder(noisy, t_disc)
+
+        state_features = (
+            self.state_encoder(state.to(device=device, dtype=actions.dtype))
+            if state is not None and self.state_encoder is not None
+            else None
+        )
+        sa_embs = self._build_sa_embs(action_features, state_features, B, device)
+
+        model_out = self._run_model(sa_embs, context_list, t_disc, is_kv, vlm_last_hidden_state)
+
+        # AML: model predicts action samples; velocity derived from predictions.
+        pred_actions = self.action_decoder(model_out)[:, -actions.shape[1]:]
+        pred_velocity = (pred_actions - noisy) / (1 - t).clamp_min(self.t_eps)
+
+        loss = ((pred_velocity - velocity) ** 2).mean()
+
+        if self.smoothness_loss_weight > 0 and pred_actions.shape[1] > 1:
+            smooth_loss = ((pred_actions[:, 1:] - pred_actions[:, :-1]) ** 2).mean()
+            loss = loss + self.smoothness_loss_weight * smooth_loss
+
+        return loss
+
+    # -----------------------------------------------------------------------
+    # Predict (inference)
+    # -----------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict_action(
+        self,
+        context_list: List[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]],
+        state: Optional[torch.Tensor] = None,
+        vlm_last_hidden_state: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Denoise from pure noise to a clean action rollout.
+
+        Args:
+            context_list:           Same format as :meth:`forward`.
+            state:                  Optional proprioceptive state.
+            vlm_last_hidden_state:  Optional last-layer VLM hidden states.
+
+        Returns:
+            ``(B, action_horizon, action_dim)``
+        """
+        is_kv = self._is_kv_list(context_list)
+        if is_kv:
+            B, device = context_list[0][0].shape[0], context_list[0][0].device
+        else:
+            B, device = context_list[0].shape[0], context_list[0].device
+        dtype = torch.float32
+
+        actions = torch.randn(B, self.action_horizon, self.action_dim, dtype=dtype, device=device)
+        if self.use_scaled_noise:
+            actions = actions * self.noise_scale[None, None, :].to(dtype=dtype)
+
+        state_features = None
+        if state is not None and self.state_encoder is not None:
+            state_features = self.state_encoder(state.to(device=device, dtype=dtype))
+
+        dt = 1.0 / self.num_inference_timesteps
+        for step in range(self.num_inference_timesteps):
+            t_cont = step / float(self.num_inference_timesteps)
+            t_disc_int = int(t_cont * self.num_timestep_buckets)
+            t_tensor = torch.full((B,), t_disc_int, device=device, dtype=torch.long)
+
+            action_features = self.action_encoder(actions, t_tensor)
+            sa_embs = self._build_sa_embs(action_features, state_features, B, device)
+
+            model_out = self._run_model(sa_embs, context_list, t_tensor, is_kv,
+                                        vlm_last_hidden_state)
+
+            pred_actions = self.action_decoder(model_out)[:, -self.action_horizon:]
+            # AML: velocity derived from predicted actions
+            pred_velocity = (pred_actions - actions) / (1.0 - t_cont)
+            actions = actions + dt * pred_velocity
+
+        return actions
+
+    # -----------------------------------------------------------------------
+    # Properties
+    # -----------------------------------------------------------------------
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(iter(self.parameters())).dtype
+
+
+# ---------------------------------------------------------------------------
+# Utilities
+# ---------------------------------------------------------------------------
+
+
+def _expand_14d_to_20d(action_std_14: torch.Tensor) -> torch.Tensor:
+    """Expand 14-D euler action statistics to 20-D r6d representation."""
+    out = torch.zeros(20, dtype=action_std_14.dtype)
+    out[0:3] = action_std_14[0:3]
+    out[3:9] = torch.repeat_interleave(action_std_14[3:6], 2)
+    out[9] = action_std_14[6]
+    out[10:13] = action_std_14[7:10]
+    out[13:19] = torch.repeat_interleave(action_std_14[10:13], 2)
+    out[19] = action_std_14[13]
+    return out