""" 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)...")