強化学習完全ガイド: DQN・PPOからRLHF・DPOによるLLMアライメントまで

import torch
import torch.nn as nn
import numpy as np
from collections import deque
import random

class DQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim)
        )

    def forward(self, x):
        return self.net(x)

class ReplayBuffer:
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return (
            torch.FloatTensor(np.array(states)),
            torch.LongTensor(actions),
            torch.FloatTensor(rewards),
            torch.FloatTensor(np.array(next_states)),
            torch.FloatTensor(dones)
        )

    def __len__(self):
        return len(self.buffer)

def train_dqn_step(online_net, target_net, optimizer, buffer,
                   batch_size=64, gamma=0.99):
    if len(buffer) < batch_size:
        return
    states, actions, rewards, next_states, dones = buffer.sample(batch_size)

    # 現在のQ値
    current_q = online_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)

    # ターゲットQ値 (ターゲットネットワーク使用)
    with torch.no_grad():
        max_next_q = target_net(next_states).max(1)[0]
        target_q = rewards + gamma * max_next_q * (1 - dones)

    loss = nn.MSELoss()(current_q, target_q)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

class DuelingDQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.feature = nn.Sequential(
            nn.Linear(state_dim, 128), nn.ReLU()
        )
        self.value_stream = nn.Sequential(
            nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 1)
        )
        self.advantage_stream = nn.Sequential(
            nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, action_dim)
        )

    def forward(self, x):
        feat = self.feature(x)
        value = self.value_stream(feat)
        advantage = self.advantage_stream(feat)
        return value + advantage - advantage.mean(dim=1, keepdim=True)

from stable_baselines3 import PPO
import gymnasium as gym

env = gym.make("CartPole-v1")
model = PPO(
    "MlpPolicy",
    env,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    gae_lambda=0.95,
    clip_range=0.2,       # epsilon: クリップ比率
    ent_coef=0.01,        # エントロピーボーナス係数
    verbose=1
)
model.learn(total_timesteps=100_000)

obs, _ = env.reset()
for _ in range(1000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        obs, _ = env.reset()

from stable_baselines3 import SAC
import gymnasium as gym

env = gym.make("HalfCheetah-v4")
model = SAC(
    "MlpPolicy", env,
    learning_rate=3e-4,
    buffer_size=1_000_000,
    learning_starts=10_000,
    batch_size=256,
    tau=0.005,
    gamma=0.99,
    train_freq=1,
    gradient_steps=1,
    verbose=1
)
model.learn(total_timesteps=1_000_000)

from stable_baselines3 import HerReplayBuffer, SAC

env = gym.make("FetchReach-v2")
model = SAC(
    "MultiInputPolicy",
    env,
    replay_buffer_class=HerReplayBuffer,
    replay_buffer_kwargs=dict(
        n_sampled_goal=4,
        goal_selection_strategy="future",
    ),
    verbose=1,
)
model.learn(total_timesteps=100_000)

import torch
import torch.nn as nn

class RewardModel(nn.Module):
    def __init__(self, base_model_name="gpt2"):
        super().__init__()
        from transformers import AutoModel
        self.backbone = AutoModel.from_pretrained(base_model_name)
        hidden_size = self.backbone.config.hidden_size
        self.reward_head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden = outputs.last_hidden_state[:, -1, :]
        return self.reward_head(last_hidden).squeeze(-1)

def compute_reward_model_loss(reward_model, chosen_ids, chosen_mask,
                               rejected_ids, rejected_mask):
    """Bradley-Terryモデルベースの選好損失"""
    chosen_reward = reward_model(chosen_ids, chosen_mask)
    rejected_reward = reward_model(rejected_ids, rejected_mask)
    loss = -torch.log(torch.sigmoid(chosen_reward - rejected_reward)).mean()
    return loss

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer

ppo_config = PPOConfig(
    model_name="gpt2",
    learning_rate=1.41e-5,
    batch_size=128,
    mini_batch_size=16,
    gradient_accumulation_steps=4,
    target_kl=0.1,
    kl_penalty="kl",
)

model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_config.model_name)
tokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name)
ppo_trainer = PPOTrainer(ppo_config, model, ref_model=None, tokenizer=tokenizer)

for batch in dataloader:
    query_tensors = batch["input_ids"]
    response_tensors = ppo_trainer.generate(query_tensors, max_new_tokens=200)
    rewards = [reward_model(q, r) for q, r in zip(query_tensors, response_tensors)]
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)

import torch.nn.functional as F

def dpo_loss(pi_logps_chosen, pi_logps_rejected,
             ref_logps_chosen, ref_logps_rejected, beta=0.1):
    """
    DPO損失関数
    Args:
        pi_logps_chosen:    学習モデルの選好応答log-prob
        pi_logps_rejected:  学習モデルの非選好応答log-prob
        ref_logps_chosen:   参照モデルの選好応答log-prob
        ref_logps_rejected: 参照モデルの非選好応答log-prob
        beta: KLペナルティ強度
    """
    pi_log_ratio = pi_logps_chosen - pi_logps_rejected
    ref_log_ratio = ref_logps_chosen - ref_logps_rejected
    # 暗黙的な報酬の差
    logits = beta * (pi_log_ratio - ref_log_ratio)
    loss = -F.logsigmoid(logits).mean()
    return loss

class MADDPGCritic(nn.Module):
    """全エージェントの情報を活用する中央集権クリティック"""
    def __init__(self, n_agents, obs_dim, action_dim):
        super().__init__()
        input_dim = n_agents * (obs_dim + action_dim)
        self.net = nn.Sequential(
            nn.Linear(input_dim, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 1)
        )

    def forward(self, all_obs, all_actions):
        x = torch.cat([all_obs, all_actions], dim=-1)
        return self.net(x)

import pyspiel

game = pyspiel.load_game("tic_tac_toe")
state = game.new_initial_state()

while not state.is_terminal():
    legal_actions = state.legal_actions()
    action = legal_actions[0]  # 実際には方策で選択
    state.apply_action(action)

returns = state.returns()
print(f"プレイヤー0の報酬: {returns[0]}")
print(f"プレイヤー1の報酬: {returns[1]}")

import gymnasium as gym
import numpy as np

class NormalizedObsWrapper(gym.ObservationWrapper):
    """観測値を [-1, 1] に正規化するラッパー"""
    def __init__(self, env):
        super().__init__(env)
        self.obs_low = env.observation_space.low
        self.obs_high = env.observation_space.high

    def observation(self, obs):
        normalized = (
            2.0 * (obs - self.obs_low)
            / (self.obs_high - self.obs_low + 1e-8)
            - 1.0
        )
        return normalized.astype(np.float32)

class RewardShapingWrapper(gym.RewardWrapper):
    """報酬スケーリングラッパー"""
    def __init__(self, env, scale=0.01):
        super().__init__(env)
        self.scale = scale

    def reward(self, reward):
        return reward * self.scale

base_env = gym.make("LunarLander-v2")
env = RewardShapingWrapper(NormalizedObsWrapper(base_env))
obs, info = env.reset(seed=42)

# Isaac Lab の例（簡略化）
from isaaclab.envs import DirectRLEnvCfg

class CartpoleEnvCfg(DirectRLEnvCfg):
    num_envs = 4096          # 4096個の並列環境
    episode_length_s = 5.0
    decimation = 2
    action_scale = 100.0

# GPU上で4096個のCartPoleを並列実行
# obs.shape: [4096, obs_dim]