예제 소스

llm_reward_wrapper.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import gymnasium as gym
import numpy as np

from llm_teacher import LLMTeacherConfig, LocalGGUFTeacher
from vlm_teacher import VLMTeacherConfig, LocalQwenVLViaTransformersTeacher


@dataclass(frozen=True)
class LLMRewardConfig:
    enabled: bool = False
    # "gguf" (text LLM) or "qwen2_5_vl" (image VLM via transformers)
    backend: str = "gguf"
    # env reward + alpha * llm_score
    alpha: float = 0.2
    # 매 스텝 호출은 비싸므로 N스텝마다 1회
    call_every: int = 4
    # gguf backend
    model_path: str = "llms/qwen2.5-3b-instruct-q4_k_m.gguf"
    # vlm backend (local HF directory)
    vlm_model_dir: str = "llms/Qwen2.5-VL-3B-Instruct"
    n_ctx: int = 2048
    n_threads: int = 8
    n_gpu_layers: int = 0
    temperature: float = 0.0
    top_p: float = 0.9
    max_tokens: int = 128
    # VLM render size (bigger = clearer, slower)
    vlm_cell_px: int = 20


class LLMRewardWrapper(gym.Wrapper):
    """
    DQN은 그대로 학습하되, step 후 관측(그리드)을 텍스트로 변환해 로컬 LLM이 점수를 매김.
    - next_action: LLM 추천 (info에만 기록)
    - llm_score: [-1,1] 점수
    - reward = env_reward + alpha * llm_score
    """

    def __init__(self, env: gym.Env, cfg: LLMRewardConfig):
        super().__init__(env)
        self.cfg = cfg
        if not cfg.enabled:
            raise ValueError("LLMRewardWrapper requires cfg.enabled=True")
        if cfg.backend == "qwen2_5_vl":
            self.teacher = LocalQwenVLViaTransformersTeacher(
                VLMTeacherConfig(
                    enabled=True,
                    model_dir=cfg.vlm_model_dir,
                    call_every=cfg.call_every,
                    cell_px=cfg.vlm_cell_px,
                    temperature=cfg.temperature,
                    top_p=cfg.top_p,
                    max_new_tokens=cfg.max_tokens,
                )
            )
        else:
            self.teacher = LocalGGUFTeacher(
                LLMTeacherConfig(
                    enabled=True,
                    model_path=cfg.model_path,
                    n_ctx=cfg.n_ctx,
                    n_threads=cfg.n_threads,
                    n_gpu_layers=cfg.n_gpu_layers,
                    temperature=cfg.temperature,
                    top_p=cfg.top_p,
                    max_tokens=cfg.max_tokens,
                    reward_alpha=cfg.alpha,
                    call_every=cfg.call_every,
                )
            )
        self._last_action = 0

    def step(self, action: int):
        self._last_action = int(action)
        obs, reward, terminated, truncated, info = self.env.step(action)

        # LLM judge uses obs after action
        next_action, score = self.teacher.judge_and_suggest(obs, self._last_action, dict(info))
        shaped = float(reward) + float(self.cfg.alpha) * float(score)

        info = dict(info)
        info["llm_score"] = float(score)
        info["llm_next_action"] = int(next_action)
        info["reward_env"] = float(reward)
        info["reward_shaped"] = float(shaped)

        return obs, shaped, terminated, truncated, info


__all__ = ["LLMRewardConfig", "LLMRewardWrapper"]