AI Alignment

"""
RLHF training loop với PPO — phiên bản rút gọn cho giáo dục.
Dùng thư viện TRL của HuggingFace.
"""
import torch
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ───────────────────────────────────────────────────────
# 1. Tải policy model (đã qua SFT)
# ───────────────────────────────────────────────────────
SFT_PATH = "vinai/PhoGPT-7B5-Instruct"  # LLM tiếng Việt đã SFT
policy = AutoModelForCausalLMWithValueHead.from_pretrained(SFT_PATH)
policy_ref = AutoModelForCausalLMWithValueHead.from_pretrained(SFT_PATH)
tok = AutoTokenizer.from_pretrained(SFT_PATH)

# ───────────────────────────────────────────────────────
# 2. Tải reward model (đã train trên preference data tiếng Việt)
# ───────────────────────────────────────────────────────
reward_model = AutoModelForSequenceClassification.from_pretrained(
    "custom/vietnamese-rm-hhh",
    num_labels=1,
)
reward_model.eval()

def score_reward(texts: list[str]) -> torch.Tensor:
    """Trả về reward scalar cho mỗi phản hồi."""
    with torch.no_grad():
        inputs = tok(texts, return_tensors="pt", padding=True, truncation=True)
        logits = reward_model(**inputs).logits.squeeze(-1)
    return logits  # (batch,)

# ───────────────────────────────────────────────────────
# 3. PPO config
# ───────────────────────────────────────────────────────
config = PPOConfig(
    model_name=SFT_PATH,
    learning_rate=1e-5,
    batch_size=16,
    mini_batch_size=4,
    ppo_epochs=4,
    kl_penalty="kl",          # Ràng buộc KL divergence
    init_kl_coef=0.2,         # beta trong công thức
    target_kl=6.0,            # KL target; adaptive beta sẽ điều chỉnh
    cliprange=0.2,            # PPO clipping
    cliprange_value=0.2,
    gamma=1.0,                # Discount factor
    lam=0.95,                 # GAE lambda
)

trainer = PPOTrainer(
    config=config,
    model=policy,
    ref_model=policy_ref,
    tokenizer=tok,
)

# ───────────────────────────────────────────────────────
# 4. Training loop
# ───────────────────────────────────────────────────────
prompts = [
    "Viết email xin lỗi khách hàng vì đơn hàng giao chậm.",
    "Giải thích cơ chế hoạt động của đèn LED cho học sinh lớp 9.",
    # ... hàng chục nghìn prompt đa dạng
]

for epoch in range(3):
    for batch_prompts in chunk(prompts, config.batch_size):
        # (a) Encode prompts
        query_tensors = [
            tok.encode(p, return_tensors="pt").squeeze() for p in batch_prompts
        ]

        # (b) Policy sinh phản hồi
        response_tensors = []
        for q in query_tensors:
            resp = policy.generate(
                q.unsqueeze(0),
                max_new_tokens=256,
                do_sample=True,
                top_p=0.9,
                temperature=0.7,
            ).squeeze()
            response_tensors.append(resp[q.shape[0]:])  # chỉ phần sinh mới

        # (c) Decode và chấm reward
        response_texts = [tok.decode(r, skip_special_tokens=True) for r in response_tensors]
        rewards = score_reward(response_texts)

        # (d) PPO step: update policy để tăng reward với KL constraint
        stats = trainer.step(
            queries=query_tensors,
            responses=response_tensors,
            scores=[r for r in rewards],
        )

        print(f"Epoch {epoch} | reward mean: {rewards.mean():.3f} | kl: {stats['objective/kl']:.3f}")

# ───────────────────────────────────────────────────────
# 5. Lưu mô hình aligned
# ───────────────────────────────────────────────────────
policy.save_pretrained("./phogpt-aligned-vi")
tok.save_pretrained("./phogpt-aligned-vi")

"""
Phát hiện reward hacking đơn giản:
- Detect repetition (agent farm 1 chỗ)
- Detect stagnation (không tiến bộ về intended goal)
- Detect reward spike bất thường (jump >3 SD so với baseline)

Dùng trong training pipeline để dừng sớm hoặc rollback.
"""
from dataclasses import dataclass
import numpy as np
from collections import Counter


@dataclass
class HackingReport:
    is_hacking: bool
    reason: str
    confidence: float
    suggested_action: str


def detect_reward_hacking(
    episode_trajectory: list[tuple[int, int]],  # (x, y) positions
    rewards_history: list[float],
    intended_goal_reached: bool,
    baseline_reward_stats: tuple[float, float],  # (mean, std)
) -> HackingReport:
    """Phát hiện dấu hiệu reward hacking trong một tập episode."""
    # 1. Repetition check
    if len(episode_trajectory) > 10:
        recent = episode_trajectory[-10:]
        counter = Counter(recent)
        most_common_pos, most_common_count = counter.most_common(1)[0]
        if most_common_count >= 6:  # ≥60% thời gian ở 1 ô
            return HackingReport(
                is_hacking=True,
                reason=f"Agent loop ở ô {most_common_pos} ({most_common_count}/10 bước)",
                confidence=most_common_count / 10,
                suggested_action="Thêm KL penalty cao hơn, hoặc reward shaping negative cho repetition",
            )

    # 2. Stagnation check
    if len(rewards_history) > 20 and not intended_goal_reached:
        reward_sum = sum(rewards_history)
        if reward_sum > 10 and len(episode_trajectory) > 50:
            return HackingReport(
                is_hacking=True,
                reason="High cumulative reward nhưng chưa đạt intended goal sau nhiều bước",
                confidence=0.8,
                suggested_action="Thêm terminal reward lớn khi đạt goal, giảm per-step reward",
            )

    # 3. Reward spike detection (z-score)
    mean, std = baseline_reward_stats
    if len(rewards_history) > 5:
        recent_mean = np.mean(rewards_history[-5:])
        if std > 0:
            z = (recent_mean - mean) / std
            if z > 3.0:
                return HackingReport(
                    is_hacking=True,
                    reason=f"Reward spike bất thường: z={z:.2f} so với baseline",
                    confidence=min(z / 5, 0.95),
                    suggested_action="Audit reward function, kiểm tra xem có exploit cell không",
                )

    return HackingReport(
        is_hacking=False,
        reason="Hành vi trong ngưỡng bình thường",
        confidence=0.0,
        suggested_action="Tiếp tục training",
    )


# ─────────────────────────────────────────────────────
# Ví dụ gọi trong training loop
# ─────────────────────────────────────────────────────
baseline = (mean_reward, std_reward)  # lấy từ early epochs
for episode in range(N_EPISODES):
    traj, rewards, reached = run_episode(policy)
    report = detect_reward_hacking(traj, rewards, reached, baseline)
    if report.is_hacking:
        print(f"⚠ Detected: {report.reason}")
        print(f"   Action: {report.suggested_action}")
        # Rollback policy hoặc điều chỉnh training
        rollback_or_adjust(report)