예제 소스

rl_chart_trainer.py

"""
차트 봉 특징 → 매도/매수/유지 정책 (PyTorch MLP) + REINFORCE 학습.
매도 승률·누적 수익률로 에피소드 보상을 주고 재학습한다.
"""

from __future__ import annotations

import logging
import math
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

import main_llm_db_test_ollama as ollama_bt
from lib.rl_policy_teacher import (
    kl_student_from_teacher_logits,
    teacher_logits_from_chart,
    write_rl_training_meta,
)

logger = logging.getLogger(__name__)

# 가격·포지션 특징 + Market 극단 라벨(없으면 0)
STATE_BASE_DIM = 10
LABEL_FEAT_DIM = 7
STATE_DIM = STATE_BASE_DIM + LABEL_FEAT_DIM


def _safe_float(x, default: float = 0.0) -> float:
    try:
        v = float(x)
        if math.isnan(v) or math.isinf(v):
            return default
        return v
    except (TypeError, ValueError):
        return default


def _bool01(x) -> float:
    if x is None or (isinstance(x, float) and (math.isnan(x) or math.isinf(x))):
        return 0.0
    if isinstance(x, (bool, np.bool_)):
        return 1.0 if x else 0.0
    if isinstance(x, (int, np.integer)):
        return 1.0 if x else 0.0
    if isinstance(x, str):
        return 1.0 if x.lower() in ("true", "1", "yes") else 0.0
    return 1.0 if x else 0.0


def _label_feature_vector(row: pd.Series) -> list[float]:
    """merge_extreme_labels_from_csv 로 붙은 ext_* 가 없으면 0 벡터."""
    if "ext_anomaly_score" not in row.index:
        return [0.0] * LABEL_FEAT_DIM
    score = _safe_float(row.get("ext_anomaly_score"), 0.0)
    rid = _safe_float(row.get("ext_regime_id"), 0.0)
    return [
        min(max(score, 0.0), 1.0),
        _bool01(row.get("ext_crash_signal")),
        _bool01(row.get("ext_surge_signal")),
        _bool01(row.get("ext_price_spike")),
        _bool01(row.get("ext_volume_spike")),
        _bool01(row.get("ext_pattern_anomaly")),
        min(max(rid / 6.0, 0.0), 1.0),
    ]


def build_state_vector(
    df: pd.DataFrame,
    bar_idx: int,
    portfolio: ollama_bt.PaperLlmPortfolio,
    price: float,
) -> np.ndarray:
    """단일 봉 인덱스 기준 상태 벡터 (정규화된 스칼라들)."""
    row = df.iloc[bar_idx]
    close = _safe_float(row.get("close", price), price)
    feats: list[float] = []
    for w in (1, 3, 5, 10):
        j = bar_idx - w
        if j >= 0:
            p0 = _safe_float(df.iloc[j].get("close", close), close)
            feats.append(math.log(max(close / max(p0, 1e-9), 1e-9)))
        else:
            feats.append(0.0)
    ma20 = _safe_float(row.get("ma20", close), close)
    ma60 = _safe_float(row.get("ma60", close), close)
    feats.append((close - ma20) / max(close, 1e-9))
    feats.append((close - ma60) / max(close, 1e-9))
    feats.append(_safe_float(row.get("PB", 0.5), 0.5))
    feats.append(_safe_float(row.get("MFI10", 50.0), 50.0) / 100.0)
    feats.append(1.0 if portfolio.unit > 1e-12 else 0.0)
    bp = portfolio.buy_price
    if portfolio.unit > 1e-12 and bp > 0:
        feats.append((price - bp) / max(bp, 1e-9))
    else:
        feats.append(0.0)
    if len(feats) != STATE_BASE_DIM:
        raise RuntimeError(f"STATE_BASE_DIM mismatch: got {len(feats)}")
    feats.extend(_label_feature_vector(row))
    arr = np.asarray(feats, dtype=np.float32)
    if arr.shape[0] != STATE_DIM:
        raise RuntimeError(f"STATE_DIM mismatch: got {arr.shape[0]}, expected {STATE_DIM}")
    return arr


def action_to_opinion(action: int) -> str:
    if action == 0:
        return "매도"
    if action == 1:
        return "매수"
    return "유지"


class RlChartPolicy(nn.Module):
    def __init__(self, state_dim: int = STATE_DIM, hidden: int = 64) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.Tanh(),
            nn.Linear(hidden, hidden),
            nn.Tanh(),
            nn.Linear(hidden, 3),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def _run_teacher_pretrain(
    policy: RlChartPolicy,
    opt: optim.Optimizer,
    df: pd.DataFrame,
    *,
    n: int,
    min_bars: int,
    bar_step: int,
    initial_krw: float,
    min_trade_bars: int,
    pretrain_episodes: int,
    teacher_temperature: float,
) -> None:
    """규칙 teacher 행동으로 환경을 굴리며 KL(teacher||policy)만 최소화 (모방·워밍업)."""
    if pretrain_episodes <= 0:
        return
    desc = "RL teacher 사전학습"
    for ep in range(pretrain_episodes):
        port = ollama_bt.PaperLlmPortfolio(initial_krw)
        port.min_trade_bars = min_trade_bars
        kls: list[torch.Tensor] = []
        idx = min_bars
        while idx < n:
            price = _safe_float(df.iloc[idx]["close"])
            st = build_state_vector(df, idx, port, price)
            x = torch.from_numpy(st).unsqueeze(0)
            logits = policy(x).squeeze(0)
            t_np = teacher_logits_from_chart(df, idx, port, price)
            kls.append(
                kl_student_from_teacher_logits(
                    logits,
                    t_np,
                    temperature=teacher_temperature,
                )
            )
            act = int(np.argmax(t_np))
            sig, probs = ollama_bt._opinion_to_signal(action_to_opinion(act))
            port.execute_llm_trade(sig, price, current_bar=idx, probabilities=probs)
            nxt = idx + bar_step
            if nxt >= n:
                break
            idx = nxt
        if not kls:
            continue
        opt.zero_grad()
        loss = torch.stack(kls).mean()
        loss.backward()
        opt.step()
        logger.info("%s ep %s/%s kl_mean=%.6f", desc, ep + 1, pretrain_episodes, float(loss.item()))


def _steps_per_episode(n: int, min_bars: int, bar_step: int) -> int:
    """train_reinforce_episodes 의 내부 while 과 동일한 스텝 수."""
    cnt = 0
    idx = min_bars
    while idx < n:
        cnt += 1
        nxt = idx + bar_step
        if nxt >= n:
            break
        idx = nxt
    return cnt


def train_reinforce_episodes(
    df: pd.DataFrame,
    *,
    episodes: int,
    lr: float,
    hidden_dim: int,
    weight_decay: float,
    min_bars: int,
    bar_step: int,
    initial_krw: float,
    min_trade_bars: int,
    checkpoint_path: str | Path,
    reward_win_weight: float,
    reward_profit_scale: float,
    load_if_exists: bool = True,
    policy_doc_path: str | Path | None = None,
    teacher_kl_coef: float = 0.0,
    teacher_temperature: float = 1.0,
    teacher_pretrain_episodes: int = 0,
    save_train_meta: bool = True,
) -> RlChartPolicy:
    df = ollama_bt._ensure_pb_mfi10(df)
    n = len(df)
    if n < min_bars + bar_step:
        raise RuntimeError("학습용 데이터가 부족합니다.")

    ckpt = Path(checkpoint_path)
    policy = RlChartPolicy(STATE_DIM, hidden_dim)
    opt = optim.Adam(policy.parameters(), lr=lr, weight_decay=weight_decay)

    if load_if_exists and ckpt.exists():
        try:
            state = torch.load(ckpt, map_location="cpu")
            policy.load_state_dict(state)
            logger.info("RL 체크포인트 로드: %s", ckpt)
        except Exception as e:
            logger.warning("체크포인트 로드 실패(차원 변경 등), 처음부터 학습: %s", e)

    if teacher_pretrain_episodes > 0:
        _run_teacher_pretrain(
            policy,
            opt,
            df,
            n=n,
            min_bars=min_bars,
            bar_step=bar_step,
            initial_krw=initial_krw,
            min_trade_bars=min_trade_bars,
            pretrain_episodes=teacher_pretrain_episodes,
            teacher_temperature=teacher_temperature,
        )

    steps_ep = _steps_per_episode(n, min_bars, bar_step)
    total_steps = max(1, steps_ep * episodes)
    pbar = tqdm(
        total=total_steps,
        desc="RL 강화학습",
        unit="step",
        mininterval=0.3,
    )

    for ep in range(episodes):
        port = ollama_bt.PaperLlmPortfolio(initial_krw)
        port.min_trade_bars = min_trade_bars

        log_probs: list[torch.Tensor] = []
        aux_kl: list[torch.Tensor] = []
        idx = min_bars
        while idx < n:
            price = _safe_float(df.iloc[idx]["close"])
            st = build_state_vector(df, idx, port, price)
            x = torch.from_numpy(st).unsqueeze(0)
            logits = policy(x)
            dist = Categorical(logits=logits)
            action = dist.sample()
            log_probs.append(dist.log_prob(action))
            if teacher_kl_coef > 0:
                t_np = teacher_logits_from_chart(df, idx, port, price)
                aux_kl.append(
                    kl_student_from_teacher_logits(
                        logits.squeeze(0),
                        t_np,
                        temperature=teacher_temperature,
                    )
                )

            sig, probs = ollama_bt._opinion_to_signal(action_to_opinion(int(action.item())))
            port.execute_llm_trade(sig, price, current_bar=idx, probabilities=probs)

            pbar.update(1)
            pbar.set_postfix(ep=f"{ep + 1}/{episodes}", refresh=False)

            nxt = idx + bar_step
            if nxt >= n:
                break
            idx = nxt

        last_price = _safe_float(df.iloc[idx]["close"])
        status = port.get_status(last_price)
        wr = float(status["win_rate"])
        pr = float(status["profit_rate"])
        reward = reward_win_weight * wr + math.tanh(pr * reward_profit_scale)

        opt.zero_grad()
        if log_probs:
            lp = torch.stack(log_probs).sum()
            loss = -lp * reward
            if aux_kl and teacher_kl_coef > 0:
                loss = loss + teacher_kl_coef * torch.stack(aux_kl).mean()
            loss.backward()
            opt.step()
        else:
            reward = 0.0

        logger.info(
            "RL ep %s/%s reward=%.4f win_rate=%.3f pnl=%.4f%% trades=%s",
            ep + 1,
            episodes,
            reward,
            wr,
            pr * 100.0,
            status["total_trades"],
        )

    pbar.close()

    ckpt.parent.mkdir(parents=True, exist_ok=True)
    torch.save(policy.state_dict(), ckpt)
    logger.info("RL 정책 저장: %s", ckpt)
    if save_train_meta and policy_doc_path:
        write_rl_training_meta(
            ckpt,
            policy_doc_path=str(policy_doc_path),
            extra={
                "teacher_kl_coef": teacher_kl_coef,
                "teacher_temperature": teacher_temperature,
                "teacher_pretrain_episodes": teacher_pretrain_episodes,
                "episodes": episodes,
                "reward_win_weight": reward_win_weight,
                "reward_profit_scale": reward_profit_scale,
            },
        )
    policy.eval()
    return policy


def load_policy(
    checkpoint_path: str | Path,
    *,
    hidden_dim: int,
) -> RlChartPolicy:
    ckpt = Path(checkpoint_path)
    policy = RlChartPolicy(STATE_DIM, hidden_dim)
    if ckpt.exists():
        policy.load_state_dict(torch.load(ckpt, map_location="cpu"))
    policy.eval()
    return policy


def select_action_deterministic(policy: RlChartPolicy, state: np.ndarray) -> int:
    with torch.no_grad():
        logits = policy(torch.from_numpy(state).unsqueeze(0))
        return int(torch.argmax(logits, dim=1).item())