예제 소스

rl_vision_viewer.py

"""
viewer.VisionBacktestViewer 확장: 체결 신호는 RL 정책, LLM·MA는 참고 표시.
"""

from __future__ import annotations

import logging
import math
import time

import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
from pathlib import Path

from test_hf_img import _build_analysis_prompt, _parse_opinion, _png_bytes_to_jpg_bytes

import main_llm_db_test_ollama as ollama_bt
from viewer import ACCENT, VisionBacktestViewer, extract_llm_reason, opinion_to_signal_label
from lib.rl_chart_trainer import (
    RlChartPolicy,
    action_to_opinion,
    build_state_vector,
    select_action_deterministic,
)
from lib.rl_policy_teacher import (
    kl_student_from_teacher_logits,
    opinion_to_action_index,
    teacher_logits_from_chart,
    write_rl_training_meta,
)

logger = logging.getLogger(__name__)


class RlVisionBacktestViewer(VisionBacktestViewer):
    """차트·포트폴리오 UI는 동일, 실행 신호만 RL 정책(argmax)으로 결정."""

    def __init__(
        self,
        *,
        rl_policy: RlChartPolicy,
        online_train: bool = False,
        online_sample_action: bool = True,
        online_max_episodes: int | None = None,
        online_save_every_ep: int = 1,
        online_clear_text_each_ep: bool = False,
        checkpoint_path: str | None = None,
        train_lr: float = 1e-3,
        train_weight_decay: float = 1e-5,
        reward_win_weight: float = 2.0,
        reward_profit_scale: float = 5.0,
        initial_krw: float | None = None,
        min_trade_bars: int | None = None,
        teacher_kl_coef: float = 0.0,
        teacher_temperature: float = 1.0,
        llm_align_weight: float = 0.0,
        vision_policy_doc: str | None = None,
        policy_doc_path_for_meta: str | None = None,
        save_train_meta_with_checkpoint: bool = False,
        **kwargs,
    ) -> None:
        self._rl_policy = rl_policy
        self._online_train = bool(online_train)
        self._online_sample_action = bool(online_sample_action)
        self._online_max_episodes = online_max_episodes if online_max_episodes is None else int(online_max_episodes)
        self._online_save_every_ep = max(1, int(online_save_every_ep))
        self._online_clear_text_each_ep = bool(online_clear_text_each_ep)
        self._checkpoint_path = checkpoint_path
        self._reward_win_weight = float(reward_win_weight)
        self._reward_profit_scale = float(reward_profit_scale)
        self._teacher_kl_coef = float(teacher_kl_coef)
        self._teacher_temperature = float(teacher_temperature)
        self._llm_align_weight = float(llm_align_weight)
        self._vision_policy_doc = (vision_policy_doc or "").strip()
        self._policy_doc_path_for_meta = policy_doc_path_for_meta
        self._save_train_meta_with_checkpoint = bool(save_train_meta_with_checkpoint)

        self._train_opt: optim.Optimizer | None = None
        if self._online_train:
            self._train_opt = optim.Adam(
                self._rl_policy.parameters(),
                lr=float(train_lr),
                weight_decay=float(train_weight_decay),
            )

        self._episode = 1
        self._traj_log_probs: list[torch.Tensor] = []
        self._traj_teacher_kl: list[torch.Tensor] = []
        self._traj_llm_ce: list[torch.Tensor] = []
        self._initial_krw = float(initial_krw) if initial_krw is not None else None
        self._min_trade_bars = int(min_trade_bars) if min_trade_bars is not None else None
        self._reset_to_idx = None  # 생성자에서 확정
        super().__init__(**kwargs)
        # VisionBacktestViewer 가 self.idx를 min_bars 기준으로 세팅한 뒤 기록
        self._reset_to_idx = int(self.idx)

    def _worker_step(self) -> None:
        err: str | None = None
        png_bytes: bytes | None = None
        reply = ""
        llm_opinion: str | None = None
        opinion: str | None = None
        ma_opinion: str | None = None
        ma_cross_opinion: str | None = None
        ma60_down = False
        rl_opinion: str | None = None
        reason = ""
        price = 0.0
        t_llm = 0.0
        bar = self.idx
        trade_executed = False

        try:
            df_hist = self.df.iloc[: bar + 1].copy()
            chart_idx = len(df_hist)
            price = float(df_hist.iloc[-1]["close"])

            t0 = time.perf_counter()
            base_png, _parts = self.chart_gen.render_and_slice_by_x(
                df_hist,
                idx=chart_idx,
                max_data=min(ollama_bt.CHART_MAX_DATA, len(df_hist)),
                ticker=ollama_bt.TICKER,
                n_slices=ollama_bt.CHART_N_SLICES,
                overlap_px=ollama_bt.CHART_OVERLAP_PX,
                scale=ollama_bt.CHART_SCALE,
            )
            png_bytes = base_png
            img_bytes = _png_bytes_to_jpg_bytes(base_png)
            prompt = _build_analysis_prompt(
                ollama_bt.TICKER,
                policy_reference_doc=self._vision_policy_doc or None,
            )

            if self.dry_run:
                reply = "참고의견: 유지\n근거: dry-run 모드입니다."
                llm_opinion = "유지"
                reason = extract_llm_reason(reply)
            else:
                t2 = time.perf_counter()
                reply = self._call_vision_with_retries(img_bytes, prompt)
                t_llm = time.perf_counter() - t2

                llm_opinion, raw = _parse_opinion(reply)
                reason = extract_llm_reason(raw or reply)
                if llm_opinion is None:
                    llm_opinion = "유지"

            if "ma20" in df_hist.columns and len(df_hist) >= 2:
                ma20 = float(df_hist.iloc[-1]["ma20"])
                ma_signal = self._ma20_cross_signal(close=price, ma20=ma20)
                if ma_signal == 0:
                    ma_cross_opinion = "매도"
                elif ma_signal == 1:
                    ma_cross_opinion = "매수"
                else:
                    ma_cross_opinion = "유지"
            else:
                ma_cross_opinion = None

            ma60_down = self._ma60_is_downward_curve(df_hist)
            ma_opinion = ma_cross_opinion
            if ma_opinion == "매수" and ma60_down:
                ma_opinion = "유지"

            st_vec = build_state_vector(self.df, bar, self.portfolio, price)
            if self._online_train and self._online_sample_action:
                x = torch.from_numpy(st_vec).unsqueeze(0)
                logits = self._rl_policy(x)
                dist = Categorical(logits=logits)
                a = dist.sample()
                self._traj_log_probs.append(dist.log_prob(a))
                if self._teacher_kl_coef > 0:
                    t_np = teacher_logits_from_chart(self.df, bar, self.portfolio, price)
                    self._traj_teacher_kl.append(
                        kl_student_from_teacher_logits(
                            logits.squeeze(0),
                            t_np,
                            temperature=self._teacher_temperature,
                        )
                    )
                if self._llm_align_weight > 0 and llm_opinion:
                    log_p = F.log_softmax(logits.squeeze(0), dim=-1)
                    ai = opinion_to_action_index(llm_opinion)
                    self._traj_llm_ce.append(-log_p[ai])
                act = int(a.item())
            else:
                act = select_action_deterministic(self._rl_policy, st_vec)
            rl_opinion = action_to_opinion(act)
            opinion = rl_opinion

            signal, probabilities = ollama_bt._opinion_to_signal(opinion)

            trade_executed = self.portfolio.execute_llm_trade(
                signal,
                price,
                current_bar=bar,
                probabilities=probabilities,
            )

            elapsed = time.perf_counter() - t0
        except Exception as e:
            err = str(e)
            logger.exception("step 실패 bar=%s", bar)
            elapsed = 0.0

        self.after(
            0,
            lambda: self._apply_step_ui_rl(
                bar=bar,
                png_bytes=png_bytes,
                llm_opinion=llm_opinion,
                opinion=opinion,
                rl_opinion=rl_opinion,
                ma_opinion=ma_opinion,
                ma_cross_opinion=ma_cross_opinion,
                ma60_down=ma60_down,
                reason=reason,
                reply=reply,
                price=price,
                err=err,
                t_llm=t_llm,
                elapsed=elapsed,
                trade_executed=trade_executed,
            ),
        )

    def _apply_step_ui_rl(
        self,
        *,
        bar: int,
        png_bytes: bytes | None,
        llm_opinion: str | None,
        opinion: str | None,
        rl_opinion: str | None,
        ma_opinion: str | None,
        ma_cross_opinion: str | None = None,
        ma60_down: bool = False,
        reason: str,
        reply: str,
        price: float,
        err: str | None,
        t_llm: float,
        elapsed: float,
        trade_executed: bool = False,
    ) -> None:
        try:
            if png_bytes:
                self._photo = self._png_to_tk(png_bytes)
                self.lbl_chart.configure(image=self._photo)

            if err:
                self.lbl_judgment.configure(text="오류", fg="#b2182b")
                self.txt_reason.delete("1.0", "end")
                self.txt_reason.insert("1.0", err)
                hist_line = f"--- bar={bar} 가격={price:,.0f} ERROR ---\n{err}\n\n"
            else:
                label = opinion_to_signal_label(opinion)
                llm_label = opinion_to_signal_label(llm_opinion) if llm_opinion else "N/A"
                ma_label = opinion_to_signal_label(ma_opinion) if ma_opinion else "N/A"
                cross_label = (
                    opinion_to_signal_label(ma_cross_opinion) if ma_cross_opinion else "N/A"
                )
                ma60_txt = "예" if ma60_down else "아니오"
                rl_label = opinion_to_signal_label(rl_opinion) if rl_opinion else label
                self.lbl_judgment.configure(
                    text=f"{label} (체결=RL:{rl_label}, MA20돌파:{cross_label}, MA60↓:{ma60_txt}, LLM참고:{llm_label})",
                    fg=ACCENT,
                )
                self.txt_reason.delete("1.0", "end")
                self.txt_reason.insert("1.0", reason or "(근거 없음)")
                hist_line = (
                    f"--- bar={bar}/{self.n - 1} 가격={price:,.0f} "
                    f"판단={label} (체결=RL, MA참고:{ma_label}, MA20돌파:{cross_label}, MA60↓:{ma60_txt}, "
                    f"LLM참고:{llm_label}) llm={t_llm:.2f}s total={elapsed:.2f}s ---\n"
                    f"{reply.strip()}\n\n"
                )
                self._refresh_portfolio_summary(price)
                self._append_profit_history(
                    bar=bar,
                    price=price,
                    opinion_label=label,
                    trade_executed=trade_executed,
                )

            self.txt_history.insert("end", hist_line)
            self.txt_history.see("end")

            if err is None:
                nxt = bar + self.bar_step
                if nxt >= self.n:
                    if self._online_train:
                        self._finish_episode_and_reset(last_price=price)
                    else:
                        self.idx = self.n
                        self.btn_next.configure(state="disabled")
                        self.auto_mode.set(False)
                        self._cancel_auto_timer()
                        self.btn_stop.configure(state="disabled")
                else:
                    self.idx = nxt

            tag = f" · {self.status_backend_label}" if self.status_backend_label else ""
            if err:
                self._set_status(f"오류 bar={bar} — {err[:100]}")
            else:
                st = self.portfolio.get_status(price)
                self._set_status(
                    f"bar={self.idx}/{self.n - 1} 가격={price:,.0f} "
                    f"판단={opinion_to_signal_label(opinion)} · 평가 {st['portfolio_value']:,.0f}원 "
                    f"({st['profit_rate'] * 100:+.2f}%) · 매도승률 {st['win_rate']*100:.1f}% · llm={t_llm:.2f}s{tag}"
                )
        finally:
            self._set_busy(False)

        if self.auto_mode.get() and err is None:
            self._schedule_auto_next()

    def _finish_episode_and_reset(self, *, last_price: float) -> None:
        """
        한 에피소드 종료: 포트폴리오 결과로 보상 계산 → policy 업데이트 → (선택) 저장 → 리셋 후 재도전.
        UI 스레드에서만 호출됨.
        """
        st = self.portfolio.get_status(float(last_price))
        wr = float(st["win_rate"])
        pr = float(st["profit_rate"])
        reward = self._reward_win_weight * wr + math.tanh(pr * self._reward_profit_scale)
        # footer 에 최근 학습 리워드 표시
        try:
            self.set_training_reward(float(reward))
            self._refresh_portfolio_summary(float(last_price))
        except Exception:
            pass

        if self._train_opt is not None and self._traj_log_probs:
            self._train_opt.zero_grad()
            lp = torch.stack(self._traj_log_probs).sum()
            loss = -lp * reward
            if self._traj_teacher_kl:
                loss = loss + self._teacher_kl_coef * torch.stack(self._traj_teacher_kl).mean()
            if self._traj_llm_ce:
                loss = loss + self._llm_align_weight * torch.stack(self._traj_llm_ce).mean()
            loss.backward()
            self._train_opt.step()

        # 로그
        self.txt_history.insert(
            "end",
            (
                f"=== EPISODE {self._episode} END | reward={reward:.4f} "
                f"win_rate={wr*100:.1f}% pnl={pr*100:+.2f}% trades={st['total_trades']} ===\n\n"
            ),
        )
        self.txt_history.see("end")

        # 저장
        if self._checkpoint_path and (self._episode % self._online_save_every_ep == 0):
            try:
                Path(self._checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
                torch.save(self._rl_policy.state_dict(), self._checkpoint_path)
                if self._save_train_meta_with_checkpoint and self._policy_doc_path_for_meta:
                    write_rl_training_meta(
                        self._checkpoint_path,
                        policy_doc_path=str(self._policy_doc_path_for_meta),
                        extra={
                            "online_train": True,
                            "episode": int(self._episode),
                            "teacher_kl_coef": self._teacher_kl_coef,
                            "teacher_temperature": self._teacher_temperature,
                            "llm_align_weight": self._llm_align_weight,
                        },
                    )
            except Exception as e:
                logger.warning("온라인 RL 체크포인트 저장 실패: %s", e)

        # 종료 조건
        if self._online_max_episodes is not None and self._episode >= self._online_max_episodes:
            self.auto_mode.set(False)
            self._cancel_auto_timer()
            self.btn_stop.configure(state="disabled")
            self.btn_next.configure(state="disabled")
            self._set_status(f"온라인 학습 종료: episodes={self._episode}")
            self.idx = self.n
            self._traj_log_probs.clear()
            self._traj_teacher_kl.clear()
            self._traj_llm_ce.clear()
            return

        # 리셋(재도전)
        self._episode += 1
        self._traj_log_probs.clear()
        self._traj_teacher_kl.clear()
        self._traj_llm_ce.clear()
        self._prev_close = None
        self._prev_ma20 = None

        if self._online_clear_text_each_ep:
            self.txt_profit_hist.delete("1.0", "end")
            self.txt_history.delete("1.0", "end")

        init_krw = self._initial_krw if self._initial_krw is not None else float(getattr(self.portfolio, "orgin_money", 5_000_000.0))
        new_port = ollama_bt.PaperLlmPortfolio(init_krw)
        if self._min_trade_bars is not None:
            new_port.min_trade_bars = max(0, int(self._min_trade_bars))
        self.portfolio = new_port

        self.idx = max(50, int(self._reset_to_idx or 100))
        self._set_status(f"온라인 학습: EP {self._episode} 시작 (자동 진행 유지)")