transfer_agent.py

"""
transfer_agent.py — Architecture-aware Transfer Learning for LoopUnrollEnv

- arch별로 x86 / arm64 등을 선택적으로 지원
- Backbone: 기존 {arch}_base 모델의 일부 레이어를 백본으로 사용
- Adapter: 새 환경(또는 새 CPU)에 맞게 소형 레이어만 재학습
"""

import os
import glob
import sys
import argparse

import numpy as np
import torch
import torch.nn as nn
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

from compiler_env import LoopUnrollEnv


# ─────────────────────────────────────────────────────────────
# 유틸: 경로 및 기본 설정
# ─────────────────────────────────────────────────────────────

PROJECT_ROOT = os.path.expanduser("~/projects/machineai")
MODELS_DIR = os.path.join(PROJECT_ROOT, "models")
BENCH_DIR = os.path.join(PROJECT_ROOT, "benchmarks")


def get_model_paths(arch: str):
    """
    아키텍처별 기본 모델/전이 모델 경로 생성
    - base:     models/model_{arch}_base.zip
    - transfer: models/model_{arch}_transfer.zip
    """
    base = os.path.join(MODELS_DIR, f"model_{arch}_base.zip")
    transfer = os.path.join(MODELS_DIR, f"model_{arch}_transfer.zip")
    return base, transfer


# ─────────────────────────────────────────────────────────────
# Backbone 가중치 추출
# ─────────────────────────────────────────────────────────────

def extract_backbone_weights(model_path: str) -> dict:
    """
    기존 PPO 모델에서 mlp_extractor의 일부 레이어를 백본으로 추출
    - 현재는 policy_net의 첫 두 레이어를 백본으로 사용
    """
    print(f"[Backbone] 로드: {model_path}")
    model = PPO.load(model_path)
    state_dict = model.policy.state_dict()
    backbone = {}

    for k, v in state_dict.items():
        if "mlp_extractor.policy_net.0" in k or "mlp_extractor.policy_net.2" in k:
            backbone[k] = v.clone()

    print(f"[Backbone] 추출 레이어:")
    for k in backbone.keys():
        print(f"  - {k}")
    return backbone


# ─────────────────────────────────────────────────────────────
# Transfer PPO 빌더
# ─────────────────────────────────────────────────────────────

def build_transfer_model(env, backbone_weights: dict | None, freeze_backbone: bool = True):
    """
    Backbone 동결 + Adapter 레이어 추가한 PPO 모델 구성
    - backbone_weights가 None이면 순수 새 모델로 시작
    """
    print("[Model] Transfer PPO 생성 중...")
    model = PPO(
        policy="MlpPolicy",
        env=env,
        learning_rate=1e-4,      # 전이학습은 낮은 lr
        n_steps=256,
        batch_size=64,
        n_epochs=10,
        gamma=0.99,
        verbose=1,
        policy_kwargs=dict(net_arch=[64, 64, 32]),  # +32 adapter layer
    )

    # 백본 가중치 주입
    if backbone_weights is not None:
        print("[Model] Backbone 가중치 주입...")
        state_dict = model.policy.state_dict()
        injected, skipped = 0, 0
        for k, v in backbone_weights.items():
            if k in state_dict and state_dict[k].shape == v.shape:
                state_dict[k] = v
                injected += 1
                print(f"  ✔ 주입: {k}")
            else:
                skipped += 1
                print(f"  ✗ 스킵: {k} (shape mismatch or not found)")
        model.policy.load_state_dict(state_dict)
        print(f"[Model] 주입 완료: {injected}개, 스킵: {skipped}개")
    else:
        print("[Model] Backbone 없이 새 모델로 시작")

    # 백본 동결
    if freeze_backbone and backbone_weights is not None:
        print("[Model] Backbone 파라미터 동결...")
        for name, param in model.policy.named_parameters():
            if "mlp_extractor.policy_net.0" in name or "mlp_extractor.policy_net.2" in name:
                param.requires_grad = False
                print(f"  🔒 동결: {name}")

    trainable = sum(p.numel() for p in model.policy.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.policy.parameters())
    print(f"\n[Model] 파라미터: {trainable}/{total} 학습가능 ({trainable/total*100:.1f}%)")
    return model


# ─────────────────────────────────────────────────────────────
# 메인 전이학습 실행
# ─────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Architecture-aware transfer learning for LoopUnrollEnv")
    parser.add_argument("--arch", type=str, default="x86", help="타겟 아키텍처 (예: x86, arm64)")
    parser.add_argument("--timesteps", type=int, default=2000, help="전이학습 스텝 수")
    parser.add_argument("--load-base", action="store_true", help="기존 base 모델에서 backbone을 로드할지 여부")
    parser.add_argument("--base-path", type=str, default="", help="직접 base 모델 경로 지정 (옵션)")
    parser.add_argument("--out-path", type=str, default="", help="전이 결과 저장 경로 직접 지정 (옵션)")
    parser.add_argument("--repeat-runs", type=int, default=3, help="실행 시간 측정 반복 횟수")
    parser.add_argument("--freeze-backbone", action="store_true", help="Backbone 레이어를 동결할지 여부")
    parser.add_argument("--clang-bin", type=str, default="", help="사용할 clang 바이너리 (비우면 기본값)")
    parser.add_argument("--opt-bin", type=str, default="", help="사용할 opt 바이너리 (비우면 기본값)")
    parser.add_argument("--source-files", type=str, nargs="+", default=[], help="학습에 사용할 소스 파일 목록")
    args = parser.parse_args()

    arch = args.arch
    print(f"[Config] arch={arch}")

    # 경로 설정
    os.makedirs(MODELS_DIR, exist_ok=True)
    default_base, default_transfer = get_model_paths(arch)

    base_model_path = args.base_path or default_base
    transfer_model_path = args.out_path or default_transfer

    print(f"[Config] base_model_path    = {base_model_path}")
    print(f"[Config] transfer_model_path= {transfer_model_path}")

    # 학습 대상 소스 파일
    if args.source_files:
        source_files = [os.path.abspath(f) for f in args.source_files]
    else:
        source_files = sorted(glob.glob(os.path.join(BENCH_DIR, "*.c")))
    print(f"[Data] 학습 대상: {source_files}")

    # Backbone 로드 (옵션)
    backbone = None
    if args.load_base:
        if not os.path.exists(base_model_path):
            raise FileNotFoundError(f"Base 모델을 찾을 수 없습니다: {base_model_path}")
        backbone = extract_backbone_weights(base_model_path)
    else:
        print("[Backbone] base 모델 로드 생략 (순수 새 모델로 시작)")

    # Env 생성 함수
    def make_env():
        return LoopUnrollEnv(
            source_files=source_files,
            repeat_runs=args.repeat_runs,
            arch=arch,
            clang_bin=args.clang_bin or None,
            opt_bin=args.opt_bin or None,
        )

    vec_env = make_vec_env(make_env, n_envs=1)

    # Transfer 모델 빌드
    print("\n=== Transfer 모델 빌드 ===")
    model = build_transfer_model(vec_env, backbone, freeze_backbone=args.freeze_backbone)

    # 학습
    print(f"\n=== Adapter 학습 ({args.timesteps} 스텝) ===")
    model.learn(total_timesteps=args.timesteps, progress_bar=True)

    # 저장
    model.save(transfer_model_path.replace(".zip", ""))
    print(f"\n저장 완료: {transfer_model_path}")


if __name__ == "__main__":
    main()