예제 소스

train_tetris_rl.py

import argparse
import os
from datetime import datetime
from pathlib import Path


def _require(pkg: str, import_name: str | None = None):
    try:
        return __import__(import_name or pkg)
    except Exception as e:  # pragma: no cover
        name = import_name or pkg
        raise RuntimeError(
            "\n".join(
                [
                    f"필수 패키지 임포트에 실패했습니다: {name}",
                    "",
                    "아래 패키지들을 설치한 뒤 다시 실행해 주세요.",
                    "  - gymnasium",
                    "  - stable-baselines3",
                    "  - numpy",
                    "",
                    "가상환경을 쓰는 경우:  .venv/bin/pip install gymnasium stable-baselines3 numpy",
                    "실행도 같은 파이썬으로:        .venv/bin/python train_tetris_rl.py ...",
                    "",
                    f"원래 에러: {type(e).__name__}: {e}",
                ]
            )
        ) from e


def make_env(
    env_id: str,
    seed: int,
    frame_stack: int,
    obs_size: int,
    use_macro: bool = False,
    stack_height_penalty_scale: float = 0.0,
    bumpiness_penalty_scale: float = 0.0,
    max_height_penalty_scale: float = 0.0,
    llm_reward_cfg: dict | None = None,
):
    # 요구사항: 이 레포의 tetris 엔진을 사용
    from local_tetris_env import LocalTetrisConfig, LocalTetrisEnv

    # env_id는 C API 라이브러리 경로로 사용 (기본값: tetris/build/... 추정)
    lib_path = env_id
    cfg = LocalTetrisConfig(
        stack_height_penalty_scale=float(stack_height_penalty_scale),
        bumpiness_penalty_scale=float(bumpiness_penalty_scale),
        max_height_penalty_scale=float(max_height_penalty_scale),
    )
    if use_macro:
        from macro_planning_env import MacroPlanningTetrisEnv

        env = MacroPlanningTetrisEnv(lib_path=lib_path, cfg=cfg)
    else:
        env = LocalTetrisEnv(lib_path=lib_path, cfg=cfg)

    if llm_reward_cfg and bool(llm_reward_cfg.get("enabled", False)):
        from llm_reward_wrapper import LLMRewardConfig, LLMRewardWrapper

        env = LLMRewardWrapper(env, LLMRewardConfig(**llm_reward_cfg))
    env.reset(seed=seed)
    return env


def train(
    env_id: str,
    total_timesteps: int,
    seed: int,
    model_dir: Path,
    learning_rate: float,
    buffer_size: int,
    batch_size: int,
    train_freq: int,
    target_update_interval: int,
    exploration_fraction: float,
    exploration_final_eps: float,
    frame_stack: int,
    obs_size: int,
    save_name: str | None,
    load_model: Path | None = None,
    fixed_model_name: str = "dqn",
    use_macro: bool = False,
    stack_height_penalty_scale: float = 0.0,
    bumpiness_penalty_scale: float = 0.0,
    max_height_penalty_scale: float = 0.0,
    llm_reward_cfg: dict | None = None,
):
    _require("stable_baselines3", "stable_baselines3")

    from stable_baselines3 import DQN
    from stable_baselines3.common.callbacks import BaseCallback
    from stable_baselines3.common.monitor import Monitor
    from stable_baselines3.common.vec_env import DummyVecEnv

    model_dir.mkdir(parents=True, exist_ok=True)

    def _thunk():
        return Monitor(
            make_env(
                env_id=env_id,
                seed=seed,
                frame_stack=frame_stack,
                obs_size=obs_size,
                use_macro=use_macro,
                stack_height_penalty_scale=stack_height_penalty_scale,
                bumpiness_penalty_scale=bumpiness_penalty_scale,
                max_height_penalty_scale=max_height_penalty_scale,
                llm_reward_cfg=llm_reward_cfg,
            )
        )

    vec_env = DummyVecEnv([_thunk])
    # policy 자동 선택: SB3의 is_image_space()로 판단 (단순 3D 체크는 불충분)
    from stable_baselines3.common.preprocessing import is_image_space
    policy = "CnnPolicy" if is_image_space(vec_env.observation_space) else "MlpPolicy"
    # tensorboard가 없으면 SB3가 ImportError를 내므로, 설치 여부에 따라 로그를 켬/끔
    tb_log_dir = model_dir / "tb"
    try:
        __import__("tensorboard")
        tensorboard_log = str(tb_log_dir)
    except Exception:
        tensorboard_log = None

    if load_model is not None:
        # 기존 모델 이어받기: 파일이 .zip 없이 저장된 경우도 처리
        load_path = load_model if load_model.exists() else Path(str(load_model) + ".zip")
        print(f"기존 모델 로드: {load_path}")
        # 먼저 env 없이 로드해서 action_space로 매크로 여부를 추정 (40=매크로, 7=저수준)
        with open(load_path, "rb") as _f:
            tmp = DQN.load(_f)
        try:
            n = int(getattr(tmp.action_space, "n", 0))
        except Exception:
            n = 0
        if n == 40:
            use_macro = True
        elif n == 7:
            use_macro = False

        # 매크로 여부가 바뀌었을 수 있으니 env를 다시 만든다
        def _thunk2():
            return Monitor(
                make_env(
                    env_id=env_id,
                    seed=seed,
                    frame_stack=frame_stack,
                    obs_size=obs_size,
                    use_macro=use_macro,
                    stack_height_penalty_scale=stack_height_penalty_scale,
                    bumpiness_penalty_scale=bumpiness_penalty_scale,
                    max_height_penalty_scale=max_height_penalty_scale,
                    llm_reward_cfg=llm_reward_cfg,
                )
            )

        vec_env = DummyVecEnv([_thunk2])

        with open(load_path, "rb") as _f:
            model = DQN.load(_f, env=vec_env, tensorboard_log=tensorboard_log)
        # 하이퍼파라미터 덮어쓰기 (커맨드라인 값 반영)
        model.learning_rate = learning_rate
        model.exploration_fraction = exploration_fraction
        model.exploration_final_eps = exploration_final_eps
    else:
        model = DQN(
            policy=policy,
            env=vec_env,
            learning_rate=learning_rate,
            buffer_size=buffer_size,
            batch_size=batch_size,
            train_freq=train_freq,
            target_update_interval=target_update_interval,
            exploration_fraction=exploration_fraction,
            exploration_final_eps=exploration_final_eps,
            verbose=1,
            seed=seed,
            tensorboard_log=tensorboard_log,
        )

    class LLMStatsCallback(BaseCallback):
        """
        llm_reward_wrapper가 info에 넣는 llm_score를 집계해 로깅.
        - rollout/llm_score_mean
        - rollout/llm_score_std
        """

        def __init__(self, print_every: int = 1000):
            super().__init__()
            self.print_every = int(print_every)
            self.n = 0
            self.mean = 0.0
            self.m2 = 0.0

        def _on_step(self) -> bool:
            infos = self.locals.get("infos", None)
            if isinstance(infos, list):
                for info in infos:
                    if not isinstance(info, dict):
                        continue
                    if "llm_score" not in info:
                        continue
                    try:
                        x = float(info["llm_score"])
                    except Exception:
                        continue
                    self.n += 1
                    # Welford online variance
                    delta = x - self.mean
                    self.mean += delta / self.n
                    delta2 = x - self.mean
                    self.m2 += delta * delta2

            if self.n > 1:
                var = self.m2 / (self.n - 1)
                std = float(var ** 0.5)
                self.logger.record("rollout/llm_score_mean", float(self.mean))
                self.logger.record("rollout/llm_score_std", std)

                # 콘솔에서도 주기적으로 확인
                if self.print_every > 0 and (self.num_timesteps % self.print_every == 0):
                    print(
                        f"[llm] steps={self.num_timesteps:,} n={self.n:,} "
                        f"mean={self.mean:+.4f} std={std:.4f}"
                    )
            return True

    # tqdm/rich가 없으면 progress bar 사용 불가
    try:
        __import__("tqdm")
        __import__("rich")
        use_progress_bar = True
    except Exception:
        use_progress_bar = False

    model.learn(
        total_timesteps=total_timesteps,
        progress_bar=use_progress_bar,
        callback=LLMStatsCallback(print_every=2000),
    )

    # 기본: model/<fixed_model_name>.zip 에만 저장 (타임스탬프 파일 누적 방지)
    # 기존 파일이 있으면 덮어쓰기 전에 model/<fixed_model_name>.zip.bak 으로 1회 백업(항상 1개만 유지)
    import re as _re

    base = save_name or fixed_model_name
    base = _re.sub(r"[^\w\-]+", "_", str(base)).strip("_") or "dqn"
    out_path = model_dir / base
    zip_path = out_path.with_suffix(".zip")
    if zip_path.exists():
        backup_path = zip_path.with_name(f"{zip_path.name}.bak")
        if backup_path.exists():
            backup_path.unlink()
        zip_path.replace(backup_path)
        print(f"\n이전 모델 백업: {backup_path}")
    model.save(str(out_path))
    print(f"\n저장 완료: {zip_path}")


def main():
    # 옵션은 기본적으로 config/rl_config.py에서 관리
    p = argparse.ArgumentParser(description="Tetris 환경으로 DQN 학습 (config 기반)")
    p.add_argument(
        "--config",
        default="config/rl_config.py",
        help="설정 파일 경로 (기본: config/rl_config.py)",
    )
    p.add_argument("--save-name", default=None, help="(옵션) 이번 실행에서만 저장 이름 override")
    args = p.parse_args()

    # config 로드
    import importlib.util
    import sys

    cfg_path = Path(args.config).resolve()
    spec = importlib.util.spec_from_file_location("rl_config", str(cfg_path))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"config 로드 실패: {cfg_path}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = mod
    spec.loader.exec_module(mod)  # type: ignore[union-attr]
    cfg = getattr(mod, "CONFIG")

    model_dir = Path(cfg.io.model_dir).resolve()
    os.environ.setdefault("PYTHONHASHSEED", str(cfg.env.seed))

    load_model = Path(cfg.io.load_model).resolve() if cfg.io.load_model else None

    train(
        env_id=cfg.env.env_id,
        total_timesteps=cfg.dqn.timesteps,
        seed=cfg.env.seed,
        model_dir=model_dir,
        learning_rate=cfg.dqn.learning_rate,
        buffer_size=cfg.dqn.buffer_size,
        batch_size=cfg.dqn.batch_size,
        train_freq=cfg.dqn.train_freq,
        target_update_interval=cfg.dqn.target_update_interval,
        exploration_fraction=cfg.dqn.exploration_fraction,
        exploration_final_eps=cfg.dqn.exploration_final_eps,
        frame_stack=4,
        obs_size=84,
        save_name=args.save_name,
        load_model=load_model,
        fixed_model_name=cfg.io.fixed_model_name,
        use_macro=bool(cfg.env.macro),
        stack_height_penalty_scale=float(getattr(cfg.env, "stack_height_penalty_scale", 0.0)),
        bumpiness_penalty_scale=float(getattr(cfg.env, "bumpiness_penalty_scale", 0.0)),
        max_height_penalty_scale=float(getattr(cfg.env, "max_height_penalty_scale", 0.0)),
        llm_reward_cfg=getattr(cfg, "llm_reward").__dict__ if hasattr(cfg, "llm_reward") else None,
    )


if __name__ == "__main__":
    main()