scripts/Evo1_splitMeanflow.py

"""
EVO1 + SplitMeanFlow Action Head

compute_splitMeanflow_loss():
  JVP 없이 정적분 가법성 항등식으로 u_tgt 계산:

  r = t (flow matching, 50%):
      u_tgt = v_hat = w·v + (1-w)·u(z_r, t, t)   [CFG correction]

  r < t (SplitMeanFlow, 50%):
      s ~ Uniform(r, t)
      z_s = z_r + (s-r)·sg[u(z_r, s, r)]          [Euler step, stop-grad]
      u_tgt = [(s-r)·sg[u(z_r,s,r)] + (t-s)·sg[u(z_s,t,s)]] / (t-r)

  L = adaptive_l2_loss(u_θ(z_r,t,r) - sg(u_tgt))
"""

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from typing import List, Union, Tuple

import torch
import torch.nn as nn
from PIL import Image

from model.internvl3.internvl3_embedder import InternVL3Embedder
from model.action_head.flow_matching_splitMeanflow import (
    SplitMeanFlowActionHead,
    adaptive_l2_loss,
    stopgrad,
    sample_t_r_curriculum,
    get_max_gap,
)
import logging


class EVO1_SplitMeanFlow(nn.Module):
    """EVO1 모델 (VLM + SplitMeanFlow Action Head)"""

    def __init__(self, config: dict):
        super().__init__()
        self.config  = config
        self._device = config.get("device", "cuda")
        self.return_cls_only = config.get("return_cls_only", False)

        vlm_name = config.get("vlm_name", "OpenGVLab/InternVL3-1B")
        self.embedder = InternVL3Embedder(model_name=vlm_name, device=self._device)

        horizon        = config.get("action_horizon", config.get("horizon", 16))
        per_action_dim = config.get("per_action_dim", 7)
        action_dim     = horizon * per_action_dim

        config["horizon"]        = horizon
        config["per_action_dim"] = per_action_dim
        config["action_dim"]     = action_dim

        self.horizon        = horizon
        self.per_action_dim = per_action_dim
        self.action_dim     = action_dim
        self.w              = config.get("cfg_scale", 2.0)

        from types import SimpleNamespace
        self.action_head = SplitMeanFlowActionHead(
            config=SimpleNamespace(
                embed_dim        = config.get("embed_dim",        896),
                hidden_dim       = config.get("hidden_dim",       1024),
                action_dim       = action_dim,
                horizon          = horizon,
                per_action_dim   = per_action_dim,
                state_dim        = config.get("state_dim",        7),
                state_hidden_dim = config.get("state_hidden_dim", 1024),
                num_heads        = config.get("num_heads",        8),
                num_layers       = config.get("num_layers",       8),
                dropout          = config.get("dropout",          0.0),
                num_categories   = config.get("num_categories",   1),
                cfg_scale        = config.get("cfg_scale",        2.0),
            )
        ).to(self._device)

    # ── VLM 인코딩 ─────────────────────────────────────────────────────────────
    def get_vl_embeddings(
        self,
        images: List[Image.Image],
        image_mask: torch.Tensor,
        prompt: str = "",
        return_cls_only: Union[bool, None] = None,
    ) -> torch.Tensor:
        if return_cls_only is None:
            return_cls_only = self.return_cls_only
        if images is None or len(images) == 0:
            raise ValueError("Must provide at least one image.")
        return self.embedder.get_fused_image_text_embedding_from_tensor_images(
            image_tensors=images,
            image_mask=image_mask,
            text_prompt=prompt,
            return_cls_only=return_cls_only,
        )

    def prepare_state(self, state_input):
        if isinstance(state_input, list):
            state_tensor = torch.tensor(state_input)
        elif isinstance(state_input, torch.Tensor):
            state_tensor = state_input
        else:
            raise TypeError("Unsupported state input type")
        if state_tensor.ndim == 1:
            state_tensor = state_tensor.unsqueeze(0)
        return state_tensor.to(self._device)

    # ── 추론 ──────────────────────────────────────────────────────────────────
    def forward(
        self,
        fused_tokens: torch.Tensor,
        state: torch.Tensor = None,
        action_mask: torch.Tensor = None,
        embodiment_ids: torch.Tensor = None,
    ) -> torch.Tensor:
        return self.action_head.get_action(
            fused_tokens,
            state=state,
            action_mask=action_mask,
            embodiment_id=embodiment_ids,
        )

    @torch.no_grad()
    def run_inference(self, images, image_mask, prompt, state_input,
                      return_cls_only=None, action_mask=None) -> torch.Tensor:
        fused = self.get_vl_embeddings(images, image_mask, prompt, return_cls_only)
        state = self.prepare_state(state_input)
        return self.forward(fused, state=state, action_mask=action_mask)

    # ── SplitMeanFlow 학습 손실 ────────────────────────────────────────────────
    def compute_splitMeanflow_loss(
        self,
        fused_tokens: torch.Tensor,   # (B, seq_len, embed_dim)
        states: torch.Tensor,          # (B, state_dim)
        actions_gt: torch.Tensor,      # (B, data_H, per_action_dim) or (B, data_H*D)
        action_mask: torch.Tensor,     # (B, ...) or None
        step: int,
        total_steps: int,
    ) -> Tuple[torch.Tensor, dict]:
        """
        SplitMeanFlow 손실 — JVP 없음, 순수 대수 항등식 + Endpoint Supervision.

        수식:
          z_r = (1-r)·x + r·e            [보간: 시작점 at r]
          gap = t - r

          [r = t 샘플 (flow matching, ~50%)]:
            u_tgt = w·v + (1-w)·sg[u(z_r, t, t)]   [CFG correction]

          [r < t 샘플 (SplitMF, ~40%)]:
            s ~ U(r, t)
            z_s = z_r + (s-r)·sg[u(z_r, s, r)]
            u_tgt = [(s-r)·sg[u_rs] + (t-s)·sg[u_st]] / (t-r)

          [Endpoint 샘플 (t=1, r=0, ~10%)]:
            z_r = x  (r=0이므로 clean action에서 시작)
            s ~ U(0, 1)
            u_tgt = s·sg[u(x, s, 0)] + (1-s)·sg[u(z_s, 1, s)]
            → NFE=1 추론 조건을 직접 학습

          L = adaptive_l2_loss(u_θ(z_r, t, r) - sg(u_tgt))
        """
        head   = self.action_head
        B      = fused_tokens.shape[0]
        device = fused_tokens.device

        # endpoint_ratio: 배치에서 (t=1, r=0)으로 강제할 비율
        endpoint_ratio = self.config.get("endpoint_ratio", 0.1)

        embodiment_id = torch.zeros(B, dtype=torch.long, device=device)

        # ── State 임베딩 ──────────────────────────────────────────────────────
        state_emb = None
        if states is not None and head.state_encoder is not None:
            enc_dtype = next(head.state_encoder.parameters()).dtype
            state_emb = head.state_encoder(
                states.to(enc_dtype), embodiment_id
            ).to(fused_tokens.dtype)

        # ── 데이터 준비 ───────────────────────────────────────────────────────
        x_raw = actions_gt.to(fused_tokens.dtype)
        if x_raw.ndim == 2:
            x_raw = x_raw.view(B, -1, self.per_action_dim)
        x = x_raw[:, :self.horizon, :].contiguous()   # (B, H, D)
        e = torch.randn_like(x)

        if action_mask is not None:
            mask_raw = action_mask.view(B, -1, self.per_action_dim)
            mask = mask_raw[:, :self.horizon, :].to(x.dtype)
            e = e * mask

        v = e - x   # straight-path velocity (B, H, D)

        # ── 커리큘럼 (t, r) 샘플링 ───────────────────────────────────────────
        t, r = sample_t_r_curriculum(B, device, step, total_steps, flow_ratio=0.5)
        t = t.to(x.dtype)
        r = r.to(x.dtype)

        # ── Endpoint 강제: 일부 샘플을 NFE=1 추론 조건 (t=1, r=0)으로 고정 ──
        # SplitMF는 algebraic target이라 gap=1.0에서도 target이 안정적
        ep_mask = torch.rand(B, device=device) < endpoint_ratio
        t[ep_mask] = 1.0
        r[ep_mask] = 0.0

        t_ = t.view(B, 1, 1)
        r_ = r.view(B, 1, 1)

        # z_r: 시작점 at r
        z_r = (1.0 - r_) * x + r_ * e   # (B, H, D)

        gap     = (t - r)                 # (B,)
        fm_mask = gap < 1e-6              # r = t → flow matching
        mf_mask = ~fm_mask               # r < t → SplitMeanFlow

        # ── u_tgt 계산 (stop-grad, no JVP) ──────────────────────────────────
        u_tgt = torch.empty(B, self.action_dim, device=device, dtype=x.dtype)

        # [1] Flow matching 샘플 (r = t)
        if fm_mask.any():
            idx = fm_mask
            with torch.no_grad():
                # u(z_r, t, t): 순간 속도 예측
                u_t_fm = head.forward_meanflow(
                    fused_tokens[idx],
                    state_emb[idx] if state_emb is not None else None,
                    z_r[idx], t[idx], r[idx],
                    embodiment_id[idx],
                )
            u_t_fm_seq = u_t_fm.view(-1, self.horizon, self.per_action_dim)
            # CFG 보정: v_hat = w·v + (1-w)·u_t
            v_hat_fm = self.w * v[idx] + (1.0 - self.w) * u_t_fm_seq
            u_tgt[idx] = v_hat_fm.reshape(-1, self.action_dim)

        # [2] SplitMeanFlow 샘플 (r < t)
        if mf_mask.any():
            idx    = mf_mask
            B_mf   = idx.sum().item()
            z_r_mf = z_r[idx]
            t_mf   = t[idx]
            r_mf   = r[idx]
            gap_mf = gap[idx]
            ft_mf  = fused_tokens[idx]
            se_mf  = state_emb[idx] if state_emb is not None else None
            ei_mf  = embodiment_id[idx]

            # split point s ∈ (r, t)
            alpha = torch.rand(B_mf, device=device, dtype=x.dtype)
            s_mf  = r_mf + alpha * gap_mf   # (B_mf,)

            with torch.no_grad():
                # u(z_r, s, r) — 짧은 구간
                u_rs     = head.forward_meanflow(ft_mf, se_mf, z_r_mf, s_mf, r_mf, ei_mf)
                u_rs_seq = u_rs.view(B_mf, self.horizon, self.per_action_dim)

                # Euler step: z_s = z_r + (s-r)·u_rs
                dt_rs = (s_mf - r_mf).view(B_mf, 1, 1)
                z_s   = z_r_mf + dt_rs * u_rs_seq   # (B_mf, H, D)

                # u(z_s, t, s) — 나머지 구간
                u_st     = head.forward_meanflow(ft_mf, se_mf, z_s, t_mf, s_mf, ei_mf)
                u_st_seq = u_st.view(B_mf, self.horizon, self.per_action_dim)

                # 대수 target: [(s-r)·u_rs + (t-s)·u_st] / (t-r)
                dt_st   = (t_mf - s_mf).view(B_mf, 1, 1)
                dt_rt   = gap_mf.view(B_mf, 1, 1).clamp(min=1e-6)
                u_tgt_seq = (dt_rs * u_rs_seq + dt_st * u_st_seq) / dt_rt   # (B_mf, H, D)

            u_tgt[idx] = u_tgt_seq.reshape(B_mf, self.action_dim)

        # ── 현재 모델 예측 (gradient 흐름) ────────────────────────────────────
        u_pred = head.forward_meanflow(fused_tokens, state_emb, z_r, t, r, embodiment_id)

        # ── 손실 ─────────────────────────────────────────────────────────────
        error = u_pred - stopgrad(u_tgt)
        loss  = adaptive_l2_loss(error)

        loss_dict = {
            "smf_loss": loss.item(),
            "mse":      (stopgrad(error) ** 2).mean().item(),
            "max_gap":  get_max_gap(step, total_steps),
            "mean_gap": float(gap.mean().item()),
            "fm_ratio": float(fm_mask.float().mean().item()),
            "ep_ratio": float(ep_mask.float().mean().item()),
        }
        return loss, loss_dict

    # ── Finetune 플래그 ────────────────────────────────────────────────────────
    def _freeze_module(self, module: nn.Module, name: str):
        print(f"Freezing {name} parameters...")
        for p in module.parameters():
            p.requires_grad = False

    def set_finetune_flags(self):
        config = self.config
        if not config.get("finetune_vlm", False):
            self._freeze_module(self.embedder, "VLM (InternVL3)")
        else:
            print("Finetuning VLM (InternVL3)...")
        if not config.get("finetune_action_head", False):
            self._freeze_module(self.action_head, "Action Head (SplitMeanFlow)")
        else:
            print("Finetuning Action Head (SplitMeanFlow)...")