InstructGPTとは？

import os
from google.colab import drive
drive.mount('/content/drive')

SAVE_ROOT = "/content/drive/MyDrive/instructgpt"
os.makedirs(SAVE_ROOT, exist_ok=True)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from torch.utils.data import DataLoader
from peft import LoraConfig, get_peft_model
from tqdm import tqdm

# ==========================================
# 1. 初期設定とデータの準備
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "llm-jp/llm-jp-3-1.8b"
MAX_LEN = 512                          # llm-jp-3 の上限は 4096 だが、メモリ都合でパイプライン全体を 512 に統一

print("トークナイザーを準備中...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
# llm-jp-3 は <PAD|LLM-jp> を持つため通常ここは通らないが、保険として残す
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("データセットを準備中...")
raw_dataset = load_dataset("llm-jp/hh-rlhf-12k-ja", split="train")

def build_prompt(conversations):
    """
    対話履歴からプロンプトを構築します。
    末尾にスペースを付けない点が重要です（トークン境界のずれを防ぐため）。
    回答と連結する際は join_prompt_response() を使ってください。
    """
    prompt = ""
    for turn in conversations:
        if turn["from"] == "human":
            prompt += f"Human: {turn['value']}\n"
        elif turn["from"] == "gpt":
            prompt += f"Assistant: {turn['value']}\n"
    return prompt + "Assistant:"        # 末尾スペースを削除


def join_prompt_response(prompt, response):
    """プロンプトと回答を連結します。スペースは回答側の先頭に付けます。"""
    return prompt + " " + response.strip()


# 長すぎるデータを切り捨て（truncation）ではなく除外（filter）する理由：
# 切り捨てると末尾のEOSトークンが失われ、モデルが「どこで発話を終えるか」を学べず、
# 後段のPPOで生成が止まらない原因になるため、完結したデータのみを使います。
def filter_length_function(example):
    prompt = build_prompt(example["conversations"])
    full_text = join_prompt_response(prompt, example["chosen"]) + tokenizer.eos_token
    return len(tokenizer(full_text)["input_ids"]) <= MAX_LEN


print(f"{MAX_LEN}トークン以内の高品質データのみにフィルタリング中...")
filtered_dataset = raw_dataset.shuffle(seed=42).filter(filter_length_function)
dataset = filtered_dataset.select(range(min(1000, len(filtered_dataset))))

def preprocess_function(example):
    """
    プロンプト部分は -100 でマスクし、回答（chosen）とEOSだけを学習対象にします。
    CrossEntropyLoss はデフォルトで -100 の要素を無視して計算します。
    """
    prompt = build_prompt(example["conversations"])
    full_text = join_prompt_response(prompt, example["chosen"]) + tokenizer.eos_token

    prompt_ids = tokenizer(prompt)["input_ids"]          # BOS込み・末尾に孤立▁は入らない
    tokenized = tokenizer(full_text, truncation=True, max_length=MAX_LEN, padding="max_length")
    input_ids = tokenized["input_ids"]

    # 境界が正しいことを保証する（ずれていれば即座に気付ける）
    n_prompt = len(prompt_ids)
    assert input_ids[:n_prompt] == prompt_ids, "プロンプトのトークン境界がずれています"

    labels = [
        -100 if (i < n_prompt or tok == tokenizer.pad_token_id) else tok
        for i, tok in enumerate(input_ids)
    ]

    return {
        "input_ids": input_ids,
        "attention_mask": tokenized["attention_mask"],
        "labels": labels,
    }


print("データセットをトークナイズ中...")
tokenized_dataset = dataset.map(preprocess_function, remove_columns=dataset.column_names)
tokenized_dataset.set_format(type="torch")

# 1.8B モデルではバッチを小さくし、勾配累積で実効バッチサイズを稼ぎます
BATCH_SIZE = 2
ACCUM_STEPS = 4   # 実効バッチサイズ = 2 × 4 = 8
dataloader = DataLoader(tokenized_dataset, batch_size=BATCH_SIZE, shuffle=True)

# ==========================================
# 2. モデルの準備（bf16 + LoRA）
# ==========================================
print("SFT用のベースモデルを初期化中...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,   # 1.8B を fp32 で載せると重み・勾配・Adam状態で約29GB必要になる
).to(device)
model.config.pad_token_id = tokenizer.pad_token_id
model.config.use_cache = False  # メモリ節約のため学習中はKVキャッシュを無効化
model.gradient_checkpointing_enable()
model.enable_input_require_grads()   # 勾配チェックポイント + LoRA の併用に必要

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    # llm-jp-3 は Llama 系アーキテクチャなのでモジュール名が GPT-2 と異なります
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()   # 全体の 1% 未満だけを学習します

# LoRA は全体の学習率より高め（1e-4 前後）にするのが定石です
optimizer = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad], lr=1e-4
)

# ==========================================
# 3. SFT（教師あり微調整）の学習ループ
# ==========================================
epochs = 1
print("\n" + "=" * 40)
print("SFT (ステップ1) の学習開始")
print("=" * 40)

model.train()

for epoch in range(epochs):
    total_loss = 0.0
    optimizer.zero_grad()
    progress_bar = tqdm(dataloader, desc=f"SFT Epoch {epoch + 1}")

    for i, batch in enumerate(progress_bar):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss

        # 勾配累積のため損失をスケールしてから逆伝播
        (loss / ACCUM_STEPS).backward()

        if (i + 1) % ACCUM_STEPS == 0 or (i + 1) == len(dataloader):
            torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], 1.0
            )
            optimizer.step()
            optimizer.zero_grad()

        total_loss += loss.item()
        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})

    print(f"\nSFT Epoch {epoch + 1} 完了 | 平均Loss: {total_loss / len(dataloader):.4f}")

print("\nSFTモデルの学習が完了しました！")

# ==========================================
# 4. LoRA をマージして「普通のモデル」として保存
# ==========================================
# 後続のステップ2・3が AutoModel...from_pretrained() でそのまま読めるようにします
# アダプターをベースモデルに統合（マージ）します。
# これにより、ステップ2・3で AutoModelForCausalLM 等を使って
# 「ごく普通の1つのモデル」として読み込み直すことができます。
model = model.merge_and_unload()
model.config.use_cache = True  # PPOの大量生成（ロールアウト）に備えて有効に戻して保存
model.save_pretrained(f"{SAVE_ROOT}/trained_sft_model")
tokenizer.save_pretrained(f"{SAVE_ROOT}/trained_sft_model")

import random
import numpy as np
import torch
import torch.nn.functional as F
from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          get_cosine_schedule_with_warmup)
from datasets import load_dataset, Dataset
from torch.utils.data import DataLoader
from peft import LoraConfig, get_peft_model
from tqdm import tqdm

# ==========================================
# 0. 共通定義（ステップ1・3と完全に同一）
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sft_model_path = f"{SAVE_ROOT}/trained_sft_model"
MAX_LEN = 512

def build_prompt(conversations):
    prompt = ""
    for turn in conversations:
        if turn["from"] == "human":
            prompt += f"Human: {turn['value']}\n"
        elif turn["from"] == "gpt":
            prompt += f"Assistant: {turn['value']}\n"
    return prompt + "Assistant:"

def join_prompt_response(prompt, response):
    return prompt + " " + response.strip()

print("トークナイザーを準備中...")
tokenizer = AutoTokenizer.from_pretrained(sft_model_path)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# ==========================================
# 1. データの準備（元ペア + 合成ネガティブ）
# ==========================================
random.seed(42)

N_BASE  = 4000    # 元のペア（丁寧さ・無害性の軸）
N_SYNTH = 2000    # 合成ネガティブ（完結性の軸）
#    インデックス範囲を明確に分割し、検証データの混入を防ぐ
#    [0, 4000)      : 元ペア
#    [4000, 6000)   : 合成ネガティブの生成元
#    [11000, 11200) : 検証用（学習には一切使わない）

print("データセットを準備中...")
raw_dataset = load_dataset("llm-jp/hh-rlhf-12k-ja", split="train")
shuffled = raw_dataset.shuffle(seed=42)

# --- 元のペア ---
pairs = []
for ex in shuffled.select(range(N_BASE)):
    pairs.append({
        "prompt": build_prompt(ex["conversations"]),
        "chosen": ex["chosen"].strip(),
        "rejected": ex["rejected"].strip(),
    })

# --- 合成ネガティブ ---
FILLER = ["はい。", "そうですね。", "うん。", "なるほど。", "わかりました。", "そうかもしれません。"]

synth = []
for ex in shuffled.select(range(N_BASE, N_BASE + N_SYNTH)):
    prompt = build_prompt(ex["conversations"])
    full = ex["chosen"].strip()
    ids = tokenizer(full, add_special_tokens=False)["input_ids"]
    if len(ids) < 16:
        continue

    mode = random.choice(["truncate", "repeat", "filler"])

    if mode == "truncate":
        # 途中で切れた断片（「シベリア」型の失敗）
        k = random.randint(2, max(4, len(ids) // 5))
        neg = tokenizer.decode(ids[:k]).strip()
    elif mode == "repeat":
        # 同一文の反復（長いが中身がない失敗）
        head = tokenizer.decode(ids[:max(8, len(ids) // 6)]).strip()
        neg = "  ".join([head] * random.randint(3, 6)).strip()
    else:
        # 中身のない相槌
        neg = random.choice(FILLER)

    if neg and neg != full:
        synth.append({"prompt": prompt, "chosen": full, "rejected": neg})

all_pairs = pairs + synth
random.shuffle(all_pairs)
print(f"元ペア: {len(pairs)} 件 / 合成ネガティブ: {len(synth)} 件 / 合計: {len(all_pairs)} 件")

dataset = Dataset.from_list(all_pairs)

def preprocess_function(example):
    # padding しない。バッチ単位で動的にパディングする
    c = tokenizer(join_prompt_response(example["prompt"], example["chosen"]),
                  truncation=True, max_length=MAX_LEN)
    r = tokenizer(join_prompt_response(example["prompt"], example["rejected"]),
                  truncation=True, max_length=MAX_LEN)
    return {"input_ids_chosen": c["input_ids"], "input_ids_rejected": r["input_ids"]}

print("データセットをトークナイズ中...")
tokenized_dataset = dataset.map(preprocess_function,
                                remove_columns=dataset.column_names)

def collate_fn(batch):
    """バッチ内の最大長に合わせて右パディングする"""
    def pad(seqs):
        m = max(len(s) for s in seqs)
        ids = torch.full((len(seqs), m), tokenizer.pad_token_id, dtype=torch.long)
        mask = torch.zeros((len(seqs), m), dtype=torch.long)
        for i, s in enumerate(seqs):
            ids[i, :len(s)] = torch.tensor(s, dtype=torch.long)
            mask[i, :len(s)] = 1
        return ids, mask

    ic, mc = pad([b["input_ids_chosen"] for b in batch])
    ir, mr = pad([b["input_ids_rejected"] for b in batch])
    return {"input_ids_chosen": ic, "attention_mask_chosen": mc,
            "input_ids_rejected": ir, "attention_mask_rejected": mr}

BATCH_SIZE = 4     # 動的パディングにより 2 → 4 に増量可能。OOM が出たら 2 に戻す
ACCUM_STEPS = 2    # 実効バッチサイズ = 4 × 2 = 8（従来と同じ）
dataloader = DataLoader(tokenized_dataset, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=collate_fn)

# ==========================================
# 2. モデルの準備
# ==========================================
print("報酬モデル（スカラー値出力の分類器）を初期化中...")
reward_# SFTモデルをベースに、出力層を1次元（num_labels=1）の分類ヘッドに置き換えます。
# これにより、文全体の「報酬」を示す連続値（スカラー）を出力させます。
model = AutoModelForSequenceClassification.from_pretrained(
    sft_model_path, num_labels=1, torch_dtype=torch.bfloat16,
).to(device)

reward_model.config.pad_token_id = tokenizer.pad_token_id
reward_model.config.use_cache = False  # メモリ節約のため学習中はKVキャッシュを無効化
reward_model.gradient_checkpointing_enable()
reward_model.enable_input_require_grads()

lora_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="SEQ_CLS",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    modules_to_save=["score"],  # 【必須】新しく追加された分類ヘッド（score）をランダムなままにせず、学習対象に含める
)
reward_model = get_peft_model(reward_model, lora_config)
reward_model.print_trainable_parameters()

optimizer = torch.optim.AdamW(
    [p for p in reward_model.parameters() if p.requires_grad], lr=1e-4
)

epochs = 2
total_steps = (len(dataloader) // ACCUM_STEPS + 1) * epochs
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=int(total_steps * 0.03), num_training_steps=total_steps
)

# ==========================================
# 3. 学習ループ
# ==========================================
print("\n" + "=" * 40)
print("報酬モデル(RM)の学習開始")
print("=" * 40)

reward_model.train()

for epoch in range(epochs):
    total_loss, total_acc, n = 0.0, 0.0, 0
    optimizer.zero_grad()
    progress_bar = tqdm(dataloader, desc=f"Epoch {epoch + 1}")

    for i, batch in enumerate(progress_bar):
        input_ids_c = batch["input_ids_chosen"].to(device)
        attn_mask_c = batch["attention_mask_chosen"].to(device)
        input_ids_r = batch["input_ids_rejected"].to(device)
        attn_mask_r = batch["attention_mask_rejected"].to(device)

        reward_chosen = reward_model(input_ids=input_ids_c,
                                     attention_mask=attn_mask_c).logits.squeeze(-1)
        reward_rejected = reward_model(input_ids=input_ids_r,
                                       attention_mask=attn_mask_r).logits.squeeze(-1)

        diff = (reward_chosen - reward_rejected).float()
        loss = -F.logsigmoid(diff).mean()

        (loss / ACCUM_STEPS).backward()

        if (i + 1) % ACCUM_STEPS == 0 or (i + 1) == len(dataloader):
            torch.nn.utils.clip_grad_norm_(
                [p for p in reward_model.parameters() if p.requires_grad], 1.0
            )
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        total_loss += loss.item()
        total_acc += (diff > 0).float().mean().item()
        n += 1
        progress_bar.set_postfix({
            "loss": f"{loss.item():.4f}",
            "acc": f"{total_acc / n:.3f}",     # 学習中の判別正解率
        })

    print(f"\nEpoch {epoch + 1} 完了 | 平均Loss: {total_loss / n:.4f} | 学習時正解率: {total_acc / n:.3f}")

    # エポックごとに保存（Colab の切断対策）
    reward_model.save_pretrained(f"{SAVE_ROOT}/rm_adapter_epoch{epoch + 1}")

print("\n報酬モデルの学習が完了しました！")

# LoRA と score ヘッドをマージして保存
reward_model = reward_model.merge_and_unload()
reward_model.config.pad_token_id = tokenizer.pad_token_id
reward_model.save_pretrained(f"{SAVE_ROOT}/trained_reward_model")
tokenizer.save_pretrained(f"{SAVE_ROOT}/trained_reward_model")

import os
import warnings
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (AutoTokenizer, AutoModelForCausalLM,
                          AutoModelForSequenceClassification,
                          StoppingCriteria, StoppingCriteriaList)
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
from tqdm import tqdm

warnings.filterwarnings("ignore", message=".*use_return_dict.*")

# ==========================================
# 0. 共通定義（ステップ1・2と同一）
# ==========================================
def build_prompt(conversations):
    prompt = ""
    for turn in conversations:
        if turn["from"] == "human":
            prompt += f"Human: {turn['value']}\n"
        elif turn["from"] == "gpt":
            prompt += f"Assistant: {turn['value']}\n"
    return prompt + "Assistant:"


def join_prompt_response(prompt, response):
    return prompt + " " + response.strip()


class StopOnStrings(StoppingCriteria):
    """生成テキストに指定文字列が現れたら停止する。
    系列を打ち切るだけでトークンごとの分布は変えないため、オンポリシー性は保たれる。"""
    def __init__(self, tokenizer, stops, prompt_len):
        self.tokenizer, self.stops, self.prompt_len = tokenizer, stops, prompt_len

    def __call__(self, input_ids, scores, **kwargs):
        text = self.tokenizer.decode(input_ids[0, self.prompt_len:], skip_special_tokens=True)
        return any(s in text for s in self.stops)

# ==========================================
# 1. Policy + Value モデル
# ==========================================
class PolicyValueLM(nn.Module):
    def __init__(self, base_model_name, lora_config, dtype=torch.bfloat16):
        super().__init__()
        self.lm = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=dtype)
        hidden_size = self.lm.config.hidden_size
        self.lm = get_peft_model(self.lm, lora_config)
        self.value_head = nn.Linear(hidden_size, 1, dtype=dtype)

    def forward(self, input_ids, attention_mask=None):
        outputs = self.lm(input_ids, attention_mask=attention_mask, output_hidden_states=True)
        values = self.value_head(outputs.hidden_states[-1]).squeeze(-1)
        return outputs.logits, values

    def ref_logits(self, input_ids, attention_mask=None):
        """LoRA を無効化＝SFT初期モデル（リファレンス）"""
        with self.lm.disable_adapter():
            return self.lm(input_ids, attention_mask=attention_mask).logits

# ==========================================
# 2. GAE と PPO 損失
# ==========================================
def compute_advantages(rewards, values, gamma=0.99, lam=0.95):
    seq_len = rewards.size(1)
    advantages = torch.zeros_like(rewards)
    lastgaelam = 0
    next_values = torch.zeros_like(values[:, -1])

    for t in reversed(range(seq_len)):
        if t == seq_len - 1:
            next_non_terminal, next_val = 0.0, next_values
        else:
            next_non_terminal, next_val = 1.0, values[:, t + 1]

        delta = rewards[:, t] + gamma * next_val * next_non_terminal - values[:, t]
        advantages[:, t] = lastgaelam = delta + gamma * lam * next_non_terminal * lastgaelam

    return advantages, advantages + values


def ppo_loss(model, input_ids, old_log_probs, advantages, returns,
             clip_ratio=0.2, c_vf=0.1):
    """損失のみ返す。正規化と optimizer.step() は呼び出し側の責任。"""
    logits, values = model(input_ids)
    target_ids = input_ids[:, 1:]
    log_probs = F.log_softmax(logits[:, :-1, :].float(), dim=-1)
    new_log_probs = log_probs.gather(dim=-1, index=target_ids.unsqueeze(-1)).squeeze(-1)

    gen_len = old_log_probs.size(1)
    new_log_probs = new_log_probs[:, -gen_len:]
    values = values[:, :-1][:, -gen_len:]

    ratio = torch.exp(new_log_probs - old_log_probs.float())
    surr1 = ratio * advantages.float()
    surr2 = torch.clamp(ratio, 1.0 - clip_ratio, 1.0 + clip_ratio) * advantages.float()
    policy_loss = -torch.min(surr1, surr2).mean()
    value_loss = F.mse_loss(values.float(), returns.float())

    return policy_loss + c_vf * value_loss, policy_loss.item(), value_loss.item()

# ==========================================
# 3. 初期化
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sft_model_path = f"{SAVE_ROOT}/trained_sft_model"
MAX_LEN = 512

print("モデルとトークナイザーを準備中...")
tokenizer = AutoTokenizer.from_pretrained(sft_model_path)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

lora_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.0, bias="none", task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)

model = PolicyValueLM(sft_model_path, lora_config).to(device)
model.lm.config.pad_token_id = tokenizer.pad_token_id

print("学習済みの報酬モデルを読み込んでいます...")
reward_model = AutoModelForSequenceClassification.from_pretrained(
    f"{SAVE_ROOT}/trained_reward_model", num_labels=1, torch_dtype=torch.bfloat16
).to(device)
reward_model.config.pad_token_id = tokenizer.pad_token_id
reward_model.eval()

# ==========================================
# 4. ハイパーパラメータ
# ==========================================
beta = 0.05             # KLペナルティ係数
CKPT_EVERY = 20
ppo_epochs = 4          # 1バッチあたりの更新回数（InstructGPT準拠）
ROLLOUT_BATCH = 8       # このプロンプト数を集めてから1回更新
num_steps = 160         # ロールアウトの総数
lr = 3e-5

min_new_tokens = 8
max_new_tokens = 96

optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)

print("データセットを準備中...")
raw_dataset = load_dataset("llm-jp/hh-rlhf-12k-ja", split="train")
dataset = raw_dataset.shuffle(seed=42).select(range(num_steps))

history = {"reward": [], "kl": [], "len": [], "ploss": [], "vloss": []}
WATCH = [build_prompt(ex["conversations"])
         for ex in raw_dataset.shuffle(seed=42).select(range(11000, 11002))]

buffer = []
os.makedirs(f"{SAVE_ROOT}/ppo_ckpt", exist_ok=True)

# ==========================================
# 5. 学習ループ
# ==========================================
print("\n" + "=" * 40)
print("RLHF (PPO) 学習ループ開始")
print("=" * 40)

for step, data in enumerate(tqdm(dataset, desc="PPO Steps")):

    prompt = build_prompt(data["conversations"])
    inputs = tokenizer(prompt, return_tensors="pt",
                       truncation=True, max_length=MAX_LEN).to(device)
    prompt_length = inputs["input_ids"].size(1)

    # --- ロールアウト（オンポリシー）---
    model.eval()
    with torch.no_grad():
        output_ids = model.lm.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            min_new_tokens=min_new_tokens,
            do_sample=True,
            top_p=1.0,          # 分布を歪めない
            top_k=0,
            temperature=1.0,
            pad_token_id=tokenizer.pad_token_id,  # 【重要】文の末尾位置を特定してスコアを取り出すために必須
            eos_token_id=tokenizer.eos_token_id,
            stopping_criteria=StoppingCriteriaList(
                [StopOnStrings(tokenizer, ["\nHuman:"], prompt_length)]
            ),
        )

        gen_len = output_ids.size(1) - prompt_length
        if gen_len <= 1:
            continue

        logits, values = model(output_ids)
        target_ids = output_ids[:, 1:]
        log_probs = F.log_softmax(logits[:, :-1, :].float(), dim=-1).gather(
            dim=-1, index=target_ids.unsqueeze(-1)).squeeze(-1)

        ref_logits = model.ref_logits(output_ids)
        ref_log_probs = F.log_softmax(ref_logits[:, :-1, :].float(), dim=-1).gather(
            dim=-1, index=target_ids.unsqueeze(-1)).squeeze(-1)

        gen_log_probs = log_probs[:, -gen_len:]
        ref_gen_log_probs = ref_log_probs[:, -gen_len:]
        gen_values = values[:, :-1][:, -gen_len:].float()

        # 報酬：RM スコアを終端に、KL ペナルティを各トークンに
        response_text = tokenizer.decode(output_ids[0, prompt_length:], skip_special_tokens=True)
        reward_inputs = tokenizer(
            join_prompt_response(prompt, response_text),
            return_tensors="pt", truncation=True, max_length=MAX_LEN
        ).to(device)
        final_score = reward_model(**reward_inputs).logits[0, 0].float().item()

        step_rewards = -beta * (gen_log_probs - ref_gen_log_probs)
        step_rewards = step_rewards.clone()
        step_rewards[:, -1] += final_score

        advantages, returns = compute_advantages(step_rewards, gen_values, gamma=1.0)

        buffer.append({
            "input_ids": output_ids,
            "old_log_probs": gen_log_probs,
            "advantages": advantages,
            "returns": returns,
        })

        history["reward"].append(final_score)
        history["kl"].append((gen_log_probs - ref_gen_log_probs).mean().item())
        history["len"].append(gen_len)

    # --- バッファが溜まったら更新 ---
    if len(buffer) < ROLLOUT_BATCH:
        continue

    all_adv = torch.cat([b["advantages"].flatten() for b in buffer])
    adv_mean, adv_std = all_adv.mean(), all_adv.std()
    if torch.isnan(adv_std) or adv_std == 0:
        adv_std = torch.tensor(1.0, device=all_adv.device)

    model.train()
    for _ in range(ppo_epochs):
        optimizer.zero_grad()
        pl_sum, vl_sum = 0.0, 0.0
        for b in buffer:
            adv = (b["advantages"] - adv_mean) / (adv_std + 1e-8)
            loss, pl, vl = ppo_loss(model, b["input_ids"], b["old_log_probs"], adv, b["returns"])
            (loss / len(buffer)).backward()
            pl_sum += pl / len(buffer)
            vl_sum += vl / len(buffer)
        torch.nn.utils.clip_grad_norm_(
            [p for p in model.parameters() if p.requires_grad], 1.0
        )
        optimizer.step()

    history["ploss"].append(pl_sum)
    history["vloss"].append(vl_sum)
    buffer = []

    # --- チェックポイント保存と経過表示 ---
    if (step + 1) % CKPT_EVERY == 0:
        model.lm.save_pretrained(f"{SAVE_ROOT}/ppo_ckpt/step{step+1}")

        w = CKPT_EVERY
        print(
            f"\n[step {step+1}] "
            f"reward={np.mean(history['reward'][-w:]):+.3f} | "
            f"KL={np.mean(history['kl'][-w:]):+.3f} | "
            f"生成長={np.mean(history['len'][-w:]):.1f}"
        )

        model.eval()
        for wp in WATCH:
            inp = tokenizer(wp, return_tensors="pt", truncation=True, max_length=MAX_LEN).to(device)
            pl = inp["input_ids"].size(1)
            with torch.no_grad():
                o = model.lm.generate(**inp, max_new_tokens=64, do_sample=False,
                                      repetition_penalty=1.15,
                                      pad_token_id=tokenizer.pad_token_id,  # 【重要】文の末尾位置を特定してスコアを取り出すために必須
                                      eos_token_id=tokenizer.eos_token_id)
            print(f"  [watch] {tokenizer.decode(o[0, pl:], skip_special_tokens=True).strip()!r}")

print("\n学習ループが完了しました！")