scripts/train_meanflow.py

"""
MeanFlow 학습 스크립트 (evo1_origin.sh 기준과 동일한 인프라)

train.py와의 차이점:
  1. EVO1_MeanFlow 모델 사용
  2. compute_meanflow_loss() 호출 (JVP + adaptive_l2_loss)
  3. step, max_steps 를 손실 함수에 전달 (커리큘럼용)
  4. autocast 없이 float32로 MeanFlow 손실 계산 (JVP 수치 안정성)
"""

import sys
import os
import math

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

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from torch.optim.lr_scheduler import LambdaLR
from torch.optim import AdamW
from accelerate import Accelerator, DistributedType
import logging
import argparse
import json
import shutil
import warnings

from Evo1_meanflow import EVO1_MeanFlow

accelerator = Accelerator()


# ── 공통 유틸 (train.py와 동일) ───────────────────────────────────────────────
def get_with_warning(config, key, default):
    if key in config:
        return config[key]
    warnings.warn(f"'{key}' not found in config, using default: {default!r}")
    return default


def custom_collate_fn(batch):
    prompts       = [item["prompt"]    for item in batch]
    images        = [item["images"]    for item in batch]
    states        = torch.stack([item["state"]         for item in batch])
    actions       = torch.stack([item["action"]        for item in batch])
    action_mask   = torch.stack([item["action_mask"]   for item in batch])
    image_masks   = torch.stack([item["image_mask"]    for item in batch])
    state_mask    = torch.stack([item["state_mask"]    for item in batch])
    embodiment_ids = torch.stack([item["embodiment_id"] for item in batch])
    return {
        "prompts": prompts, "images": images,
        "states": states, "actions": actions,
        "action_mask": action_mask, "state_mask": state_mask,
        "image_masks": image_masks, "embodiment_ids": embodiment_ids,
    }


def get_lr_lambda(warmup_steps, total_steps, resume_step=0):
    def lr_lambda(current_step):
        current_step += resume_step
        if current_step < warmup_steps:
            return current_step / max(1, warmup_steps)
        progress = (current_step - warmup_steps) / max(1, total_steps - warmup_steps)
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    return lr_lambda


def setup_logging(log_dir):
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, f"train_mf_{timestamp}.log")
    if accelerator.is_main_process:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s [%(levelname)s] %(message)s",
            handlers=[
                logging.FileHandler(log_path),
                logging.StreamHandler(),
            ],
        )
        logging.info(f"[MeanFlow] Logging to: {log_path}")
    return log_path


def prepare_dataset(config):
    dataset_type   = get_with_warning(config, "dataset_type", "lerobot")
    image_size     = get_with_warning(config, "image_size", 448)
    max_samples    = get_with_warning(config, "max_samples_per_file", None)
    horizon        = get_with_warning(config, "horizon", 50)
    binarize_gripper = get_with_warning(config, "binarize_gripper", False)
    use_augmentation = get_with_warning(config, "use_augmentation", False)

    if dataset_type == "lerobot":
        from dataset.lerobot_dataset_pretrain_mp import LeRobotDataset
        import yaml
        with open(config.get("dataset_config_path"), "r") as f:
            dataset_config = yaml.safe_load(f)
        dataset = LeRobotDataset(
            config=dataset_config,
            image_size=image_size,
            max_samples_per_file=max_samples,
            action_horizon=horizon,
            binarize_gripper=binarize_gripper,
            use_augmentation=use_augmentation,
        )
    else:
        raise ValueError(f"Unknown dataset_type: {dataset_type}")

    if accelerator.is_main_process:
        logging.info(f"Loaded {len(dataset)} samples ({dataset_type})")
    return dataset


def prepare_dataloader(dataset, config):
    batch_size  = get_with_warning(config, "batch_size", 8)
    num_workers = get_with_warning(config, "num_workers", 8)
    dataloader  = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=False,
        drop_last=True,
        collate_fn=custom_collate_fn,
    )
    if accelerator.is_main_process:
        logging.info(f"Dataloader: batch_size={batch_size}")
    return dataloader


def check_numerical_stability(step, **named_tensors):
    for name, tensor in named_tensors.items():
        if not torch.isfinite(tensor).all():
            logging.warning(f"[Step {step}] Non-finite in {name}")
            return False
    return True


def build_param_groups(model, wd):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if n.endswith("bias") or "norm" in n.lower() or p.dim() == 1:
            no_decay.append(p)
        else:
            decay.append(p)
    return [{"params": decay, "weight_decay": wd},
            {"params": no_decay, "weight_decay": 0.0}]


def get_and_clip_grad_norm(accelerator, model, loss, max_norm=1.0):
    grad_norms = [p.grad.norm(2) for p in model.parameters() if p.grad is not None]
    if len(grad_norms) == 0:
        total_norm   = torch.tensor(0.0, device=loss.device)
        clipped_norm = torch.tensor(0.0, device=loss.device)
    else:
        total_norm = torch.norm(torch.stack(grad_norms), 2)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        clipped_norms = [p.grad.norm(2) for p in model.parameters() if p.grad is not None]
        clipped_norm  = torch.norm(torch.stack(clipped_norms), 2)
    return total_norm, clipped_norm


def save_checkpoint(save_dir, step, model_engine, loss, accelerator, config=None, norm_stats=None):
    tag           = f"step_{step}"
    checkpoint_dir = os.path.join(save_dir, tag)

    if accelerator.is_main_process and os.path.exists(checkpoint_dir):
        shutil.rmtree(checkpoint_dir)

    accelerator.wait_for_everyone()

    client_state = {
        "step": step,
        "best_loss": loss if isinstance(loss, float) else loss.item(),
        "config": config,
    } if accelerator.is_main_process else {}

    if hasattr(model_engine, "save_checkpoint"):
        model_engine.save_checkpoint(save_dir, tag=tag, client_state=client_state)
    else:
        unwrapped = accelerator.unwrap_model(model_engine)
        accelerator.save_model(unwrapped, os.path.join(save_dir, tag))
        if accelerator.is_main_process:
            torch.save(client_state, os.path.join(checkpoint_dir, "training_state.pt"))

    if accelerator.is_main_process:
        if config is not None:
            with open(os.path.join(checkpoint_dir, "config.json"), "w") as f:
                json.dump(config, f, indent=2)
        if norm_stats is not None:
            with open(os.path.join(checkpoint_dir, "norm_stats.json"), "w") as f:
                json.dump(norm_stats, f, indent=2)
        with open(os.path.join(checkpoint_dir, "checkpoint.json"), "w") as f:
            json.dump({"type": "ds_model", "version": 0.0,
                       "checkpoints": "mp_rank_00_model_states.pt"}, f, indent=2)
        logging.info(f"Saved checkpoint → {checkpoint_dir}")

        # 숫자 체크포인트 최대 4개 유지 (step_best, step_final, epoch_best 제외)
        # step_<N> 또는 epoch_<N> 형태 모두 처리
        if isinstance(step, int) or (isinstance(step, str) and step.startswith("epoch_") and step[6:].isdigit()):
            import pathlib
            root = pathlib.Path(save_dir)
            numeric_ckpts = []
            for d in root.iterdir():
                if not d.is_dir():
                    continue
                for prefix in ("step_", "epoch_"):
                    if d.name.startswith(prefix):
                        suffix = d.name[len(prefix):]
                        try:
                            numeric_ckpts.append((int(suffix), d))
                        except ValueError:
                            pass
            numeric_ckpts.sort(key=lambda x: x[0], reverse=True)
            for _, old_dir in numeric_ckpts[4:]:
                shutil.rmtree(old_dir)
                logging.info(f"Removed old checkpoint: {old_dir}")


def load_checkpoint_with_deepspeed(model_engine, load_dir, accelerator,
                                   tag="step_best", load_optimizer_states=True,
                                   resume_pretrain=False):
    try:
        load_path, client_state = model_engine.load_checkpoint(
            load_dir, tag=tag, load_module_strict=True,
            load_optimizer_states=load_optimizer_states and not resume_pretrain,
            load_lr_scheduler_states=load_optimizer_states and not resume_pretrain,
        )
        if accelerator.is_main_process:
            logging.info(f"Loaded checkpoint: {load_dir}/{tag}")
        return client_state.get("step", 0), client_state
    except Exception as e:
        if accelerator.is_main_process:
            logging.warning(f"Checkpoint load with optimizer failed: {e}. Retrying model-only...")
        load_path, client_state = model_engine.load_checkpoint(
            load_dir, tag=tag, load_module_strict=True,
            load_optimizer_states=False, load_lr_scheduler_states=False,
        )
        return client_state.get("step", 0), client_state


# ── 학습 메인 ─────────────────────────────────────────────────────────────────
def train(config):
    save_dir = get_with_warning(config, "save_dir", "checkpoints_mf")
    setup_logging(save_dir)

    if get_with_warning(config, "debug", False):
        torch.autograd.set_detect_anomaly(True)

    dataset    = prepare_dataset(config)
    dataloader = prepare_dataloader(dataset, config)

    model = EVO1_MeanFlow(config)
    model.train()
    model.set_finetune_flags()

    lr  = get_with_warning(config, "lr",            1e-5)
    wd  = get_with_warning(config, "weight_decay",  1e-5)
    optimizer = AdamW(build_param_groups(model, wd), lr=lr)
    if accelerator.is_main_process:
        logging.info(f"[MeanFlow] AdamW lr={lr}, wd={wd}")

    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
    model_engine = model

    # ── 스텝 수 계산 ──────────────────────────────────────────────────────────
    max_epochs = config.get("max_epochs", None)
    if max_epochs is not None:
        steps_per_epoch = len(dataloader)
        max_steps = int(math.ceil(max_epochs * steps_per_epoch))
        if accelerator.is_main_process:
            logging.info(f"{max_epochs} epochs × {steps_per_epoch} steps/epoch = {max_steps} total")
    else:
        max_steps = get_with_warning(config, "max_steps", 1000)

    warmup_steps       = get_with_warning(config, "warmup_steps",       300)
    log_interval       = get_with_warning(config, "log_interval",       100)
    ckpt_interval      = get_with_warning(config, "ckpt_interval",      1000)
    ckpt_epoch_interval = get_with_warning(config, "ckpt_epoch_interval", None)  # 에폭 기준 저장
    ckpt_epoch_start    = get_with_warning(config, "ckpt_epoch_start",    0.0)   # 이 에폭 이전엔 저장 안함
    max_norm           = get_with_warning(config, "grad_clip_norm",     1.0)

    # 에폭 기준이 설정된 경우 steps_per_epoch 기반으로 step 환산
    if ckpt_epoch_interval is not None:
        steps_per_epoch_for_ckpt = len(dataloader)
        ckpt_interval    = int(ckpt_epoch_interval * steps_per_epoch_for_ckpt)
        ckpt_start_step  = int(ckpt_epoch_start    * steps_per_epoch_for_ckpt)
        if accelerator.is_main_process:
            logging.info(
                f"[Checkpoint] 에폭 기준 저장: epoch {ckpt_epoch_start}부터 "
                f"매 {ckpt_epoch_interval} epoch = {ckpt_interval} steps"
            )
    else:
        ckpt_start_step = 0

    os.makedirs(save_dir, exist_ok=True)
    best_loss = float("inf")

    # ── Resume ────────────────────────────────────────────────────────────────
    resume        = get_with_warning(config, "resume",         False)
    resume_path   = get_with_warning(config, "resume_path",    None)
    resume_pretrain = get_with_warning(config, "resume_pretrain", False)

    if resume != bool(resume_path):
        raise ValueError("--resume and --resume_path must be set together.")

    if resume:
        resume_path = resume_path.rstrip("/")
        resume_dir, resume_tag = os.path.split(resume_path)
        step, client_state = load_checkpoint_with_deepspeed(
            model_engine, resume_dir, accelerator, resume_tag,
            load_optimizer_states=True, resume_pretrain=resume_pretrain,
        )
        best_loss = client_state.get("best_loss", float("inf"))
        if accelerator.is_main_process:
            logging.info(f"Resumed from {resume_path}, step={step}")
    else:
        step = 0
        if accelerator.is_main_process:
            logging.info("Starting fresh MeanFlow training")

    if resume_pretrain:
        step = 0

    scheduler = LambdaLR(optimizer, get_lr_lambda(warmup_steps, max_steps, resume_step=step))

    # ── 학습 루프 ─────────────────────────────────────────────────────────────
    run_name        = config.get("run_name", "MeanFlow")
    steps_per_epoch = len(dataloader)

    pbar = tqdm(
        total=max_steps, initial=step, desc=run_name,
        disable=not accelerator.is_main_process,
        dynamic_ncols=True, smoothing=0.1,
    )

    # unwrap once (single GPU per dataset, no DDP sync issue)
    unwrapped_model = accelerator.unwrap_model(model_engine)

    while step < max_steps:
        for batch in dataloader:
            if step >= max_steps:
                break

            prompts       = batch["prompts"]
            images_batch  = batch["images"]
            image_masks   = batch["image_masks"]
            states        = batch["states"]
            actions_gt    = batch["actions"]
            action_mask   = batch["action_mask"]

            # ── VLM 인코딩 (bfloat16, VLM frozen 여부에 따라 자동 처리) ──────
            fused_tokens_list = []
            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
                for prompt, images, image_mask in zip(prompts, images_batch, image_masks):
                    fused = unwrapped_model.get_vl_embeddings(
                        images=images, image_mask=image_mask, prompt=prompt,
                        return_cls_only=False,
                    )
                    fused_tokens_list.append(fused)

            # float32로 업캐스트 (JVP 수치 안정성)
            fused_tokens = torch.cat(fused_tokens_list, dim=0).float()
            states       = states.float()
            actions_gt   = actions_gt.float()

            # ── MeanFlow 손실 (JVP 포함, float32) ───────────────────────────
            loss, loss_dict = unwrapped_model.compute_meanflow_loss(
                fused_tokens  = fused_tokens,
                states        = states,
                actions_gt    = actions_gt,
                action_mask   = action_mask,
                step          = step,
                total_steps   = max_steps,
            )

            # ── 수치 안정성 체크 ─────────────────────────────────────────────
            if not check_numerical_stability(step, loss=loss):
                logging.warning(f"[Step {step}] Skipping due to non-finite loss")
                continue

            # ── Backward ────────────────────────────────────────────────────
            optimizer.zero_grad(set_to_none=True)
            accelerator.backward(loss)
            total_norm, _ = get_and_clip_grad_norm(accelerator, model, loss, max_norm)
            optimizer.step()
            scheduler.step()

            # ── 로깅 ────────────────────────────────────────────────────────
            loss_val = loss.item()
            if step % log_interval == 0 and accelerator.is_main_process:
                epoch = step / steps_per_epoch
                logging.info(
                    f"[Step {step}] Loss: {loss_val:.4f} | "
                    f"mf_loss={loss_dict['mf_loss']:.4f} | "
                    f"max_gap={loss_dict['max_gap']:.3f} | "
                    f"mean_gap={loss_dict['mean_gap']:.3f} | "
                    f"epoch={epoch:.2f} | "
                    f"lr={scheduler.get_last_lr()[0]:.2e} | "
                    f"grad_norm={total_norm:.3f}"
                )

            # ── Best 체크포인트 ──────────────────────────────────────────────
            if accelerator.is_main_process:
                is_best = loss_val < best_loss
                if is_best:
                    best_loss = loss_val
                is_best_tensor = torch.tensor(int(is_best), device=accelerator.device)
            else:
                is_best_tensor = torch.tensor(0, device=accelerator.device)

            if accelerator.distributed_type != DistributedType.NO:
                torch.distributed.broadcast(is_best_tensor, src=0)

            if is_best_tensor.item() == 1 and step > 1000:
                save_checkpoint(
                    save_dir, step="best", model_engine=model_engine,
                    loss=loss, accelerator=accelerator, config=config,
                    norm_stats=dataset.arm2stats_dict,
                )
                if accelerator.is_main_process:
                    logging.info(f"Saved best checkpoint at step {step}, loss={loss_val:.6f}")

            step += 1

            # ── 진행 바 ──────────────────────────────────────────────────────
            pbar.update(1)
            pbar.set_postfix(
                loss=f"{loss_val:.4f}",
                gap=f"{loss_dict['max_gap']:.2f}",
                lr=f"{scheduler.get_last_lr()[0]:.2e}",
                epoch=f"{step / steps_per_epoch:.2f}",
            )

            # ── 주기적 체크포인트 ─────────────────────────────────────────────
            if step % ckpt_interval == 0 and step > 0 and step >= ckpt_start_step:
                if ckpt_epoch_interval is not None:
                    # 에폭 번호로 태그 (epoch_5, epoch_10, ...)
                    current_epoch = int(step / steps_per_epoch)
                    ckpt_tag = f"epoch_{current_epoch}"
                else:
                    ckpt_tag = step  # step 번호로 태그
                save_checkpoint(
                    save_dir, step=ckpt_tag, model_engine=model_engine,
                    loss=loss, accelerator=accelerator, config=config,
                    norm_stats=dataset.arm2stats_dict,
                )
                if accelerator.is_main_process:
                    epoch_disp = step / steps_per_epoch
                    logging.info(f"Periodic checkpoint: tag={ckpt_tag} (epoch={epoch_disp:.1f})")

    pbar.close()

    # ── 최종 저장 ─────────────────────────────────────────────────────────────
    save_checkpoint(
        save_dir, step="final", model_engine=model_engine,
        loss=loss, accelerator=accelerator, config=config,
        norm_stats=dataset.arm2stats_dict,
    )
    if accelerator.is_main_process:
        logging.info(f"Training done. best_loss={best_loss:.6f}")


# ── argparse ──────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train EVO1 with MeanFlow")
    default_cfg = os.path.abspath(
        os.path.join(os.path.dirname(__file__), "..", "dataset", "config.yaml")
    )

    parser.add_argument("--device",              type=str,  default="cuda")
    parser.add_argument("--run_name",            type=str,  default="Evo1_meanflow")
    parser.add_argument("--vlm_name",            type=str,  default="OpenGVLab/InternVL3-1B")
    parser.add_argument("--dataset_config_path", type=str,  default=default_cfg)
    parser.add_argument("--finetune_vlm",        type=lambda x: x.lower() == "true", default=False)
    parser.add_argument("--finetune_action_head",type=lambda x: x.lower() == "true", default=True)
    parser.add_argument("--state_dim",           type=int,  default=8)
    parser.add_argument("--batch_size",          type=int,  default=4)
    parser.add_argument("--max_epochs",          type=int,  default=None)
    parser.add_argument("--max_steps",           type=int,  default=None)
    parser.add_argument("--warmup_steps",        type=int,  default=1000)
    parser.add_argument("--log_interval",        type=int,  default=10)
    parser.add_argument("--ckpt_interval",       type=int,  default=25000,
                        help="체크포인트 저장 간격 (step 기준). ckpt_epoch_interval 설정 시 무시됨")
    parser.add_argument("--ckpt_epoch_interval", type=float, default=None,
                        help="체크포인트 저장 간격 (에폭 기준). 예: 1.25 → 매 1.25 에폭마다 저장")
    parser.add_argument("--ckpt_epoch_start",   type=float, default=0.0,
                        help="이 에폭 이전에는 주기적 체크포인트 저장 안 함. 예: 20 → epoch 20부터 저장")
    parser.add_argument("--save_dir",            type=str,  default="checkpoints_mf")
    parser.add_argument("--lr",                  type=float,default=1e-5)
    parser.add_argument("--weight_decay",        type=float,default=1e-5)
    parser.add_argument("--grad_clip_norm",      type=float,default=1.0)
    parser.add_argument("--use_augmentation",    action="store_true")
    parser.add_argument("--resume",              action="store_true")
    parser.add_argument("--resume_pretrain",     action="store_true")
    parser.add_argument("--resume_path",         type=str,  default=None)
    parser.add_argument("--cfg_scale",           type=float,default=2.0,
                        help="CFG scale w for MeanFlow (MP1 default=2.0)")
    # MeanFlow 커리큘럼 파라미터 (현재는 하드코딩, 필요 시 확장 가능)
    # phase1: 0-40% → gap 0.05→0.40
    # phase2: 40-80% → gap 0.40→1.00
    # phase3: 80-100% → gap 1.00, large-gap bias 70%

    args = parser.parse_args()
    config = vars(args)
    config["dataset_type"] = "lerobot"

    train(config)