[심층 강화학습] 16. Black-Box 최적화: 진화 전략과 유전 알고리즘

import torch
import torch.nn as nn
import numpy as np
import gymnasium as gym

class ESPolicy(nn.Module):
    """ES용 정책 네트워크"""
    def __init__(self, obs_size, act_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, 64),
            nn.Tanh(),
            nn.Linear(64, act_size),
        )

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

    def get_action(self, obs):
        with torch.no_grad():
            logits = self.forward(torch.FloatTensor(obs))
            return logits.argmax().item()

def evaluate_policy(env, policy, num_episodes=5):
    """정책의 평균 보상 계산"""
    total_reward = 0.0
    for _ in range(num_episodes):
        obs, _ = env.reset()
        done = False
        while not done:
            action = policy.get_action(obs)
            obs, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            done = terminated or truncated
    return total_reward / num_episodes

def train_es_cartpole(pop_size=50, num_generations=100,
                      sigma=0.1, lr=0.01):
    """CartPole에서 ES 학습"""
    env = gym.make('CartPole-v1')
    obs_size = env.observation_space.shape[0]
    act_size = env.action_space.n

    # 기본 정책
    policy = ESPolicy(obs_size, act_size)
    base_params = nn.utils.parameters_to_vector(
        policy.parameters()
    ).detach()

    for gen in range(num_generations):
        # 1. 무작위 변이 생성
        noise = torch.randn(pop_size, len(base_params))

        # 2. 각 변이에 대해 보상 평가
        rewards = []
        for i in range(pop_size):
            # 양방향 변이 (미러링)
            for sign in [1, -1]:
                perturbed = base_params + sign * sigma * noise[i]
                nn.utils.vector_to_parameters(perturbed, policy.parameters())
                reward = evaluate_policy(env, policy, num_episodes=3)
                rewards.append((reward, sign, i))

        # 3. 보상으로 가중된 업데이트 방향 계산
        rewards_array = np.array([r[0] for r in rewards])

        # 보상 정규화 (rank-based)
        ranked = np.argsort(np.argsort(-rewards_array)).astype(np.float32)
        ranked = (ranked - ranked.mean()) / (ranked.std() + 1e-8)

        # 그래디언트 추정
        grad = torch.zeros_like(base_params)
        for idx, (_, sign, noise_idx) in enumerate(rewards):
            grad += ranked[idx] * sign * noise[noise_idx]
        grad /= (2 * pop_size * sigma)

        # 4. 파라미터 업데이트
        base_params += lr * grad

        # 성능 확인
        nn.utils.vector_to_parameters(base_params, policy.parameters())
        avg_reward = evaluate_policy(env, policy, num_episodes=10)
        if gen % 10 == 0:
            print(f"세대 {gen}: 평균 보상 = {avg_reward:.1f}")

        if avg_reward >= 495:
            print(f"세대 {gen}에서 해결!")
            break

    return policy

class ESContinuousPolicy(nn.Module):
    def __init__(self, obs_size, act_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, 256),
            nn.Tanh(),
            nn.Linear(256, 256),
            nn.Tanh(),
            nn.Linear(256, act_size),
            nn.Tanh(),
        )

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

    def get_action(self, obs):
        with torch.no_grad():
            return self.forward(torch.FloatTensor(obs)).numpy()

def train_es_continuous(env_name='HalfCheetah-v4', pop_size=200,
                        num_generations=500, sigma=0.02, lr=0.01):
    """연속 환경에서 ES 학습 (병렬화 가능)"""
    env = gym.make(env_name)
    obs_size = env.observation_space.shape[0]
    act_size = env.action_space.shape[0]

    policy = ESContinuousPolicy(obs_size, act_size)
    base_params = nn.utils.parameters_to_vector(
        policy.parameters()
    ).detach()

    # Adam 모멘텀 추적
    m = torch.zeros_like(base_params)  # 1차 모멘트
    v = torch.zeros_like(base_params)  # 2차 모멘트
    beta1, beta2 = 0.9, 0.999

    for gen in range(num_generations):
        noise = torch.randn(pop_size, len(base_params))
        rewards_pos = np.zeros(pop_size)
        rewards_neg = np.zeros(pop_size)

        # 병렬화 가능한 평가 루프
        for i in range(pop_size):
            # 양방향 변이 평가
            params_pos = base_params + sigma * noise[i]
            nn.utils.vector_to_parameters(params_pos, policy.parameters())
            rewards_pos[i] = evaluate_policy(env, policy, num_episodes=1)

            params_neg = base_params - sigma * noise[i]
            nn.utils.vector_to_parameters(params_neg, policy.parameters())
            rewards_neg[i] = evaluate_policy(env, policy, num_episodes=1)

        # 보상 차이로 그래디언트 추정
        all_rewards = np.concatenate([rewards_pos, rewards_neg])
        reward_mean = all_rewards.mean()
        reward_std = all_rewards.std() + 1e-8

        grad = torch.zeros_like(base_params)
        for i in range(pop_size):
            normalized = (rewards_pos[i] - rewards_neg[i]) / reward_std
            grad += normalized * noise[i]
        grad /= (2 * pop_size * sigma)

        # Adam 업데이트
        m = beta1 * m + (1 - beta1) * grad
        v = beta2 * v + (1 - beta2) * grad ** 2
        m_hat = m / (1 - beta1 ** (gen + 1))
        v_hat = v / (1 - beta2 ** (gen + 1))

        base_params += lr * m_hat / (torch.sqrt(v_hat) + 1e-8)

        nn.utils.vector_to_parameters(base_params, policy.parameters())
        if gen % 10 == 0:
            avg = evaluate_policy(env, policy, num_episodes=5)
            print(f"세대 {gen}: 평균 보상 = {avg:.1f}")

    return policy

import copy

def train_ga_cartpole(pop_size=100, num_generations=50,
                      elite_ratio=0.1, mutation_rate=0.1,
                      mutation_sigma=0.1):
    """유전 알고리즘으로 CartPole 학습"""
    env = gym.make('CartPole-v1')
    obs_size = env.observation_space.shape[0]
    act_size = env.action_space.n

    # 초기 인구 생성
    population = [ESPolicy(obs_size, act_size) for _ in range(pop_size)]
    num_elite = max(int(pop_size * elite_ratio), 1)

    for gen in range(num_generations):
        # 1. 적합도 평가
        fitness = []
        for individual in population:
            reward = evaluate_policy(env, individual, num_episodes=5)
            fitness.append(reward)

        # 정렬 (적합도 내림차순)
        sorted_indices = np.argsort(fitness)[::-1]
        best_fitness = fitness[sorted_indices[0]]

        if gen % 5 == 0:
            print(f"세대 {gen}: 최고 적합도 = {best_fitness:.1f}, "
                  f"평균 = {np.mean(fitness):.1f}")

        if best_fitness >= 495:
            print(f"세대 {gen}에서 해결!")
            return population[sorted_indices[0]]

        # 2. 엘리트 선택
        elites = [copy.deepcopy(population[i]) for i in sorted_indices[:num_elite]]

        # 3. 다음 세대 생성
        new_population = list(elites)  # 엘리트는 그대로 유지

        while len(new_population) < pop_size:
            # 토너먼트 선택
            parent1 = tournament_select(population, fitness)
            parent2 = tournament_select(population, fitness)

            # 교차
            child = crossover(parent1, parent2)

            # 변이
            mutate(child, mutation_rate, mutation_sigma)

            new_population.append(child)

        population = new_population

    return population[sorted_indices[0]]

def tournament_select(population, fitness, k=3):
    """토너먼트 선택: k개 중 최고 적합도 선택"""
    indices = np.random.choice(len(population), k, replace=False)
    best_idx = indices[np.argmax([fitness[i] for i in indices])]
    return copy.deepcopy(population[best_idx])

def crossover(parent1, parent2):
    """균일 교차: 각 파라미터를 50% 확률로 선택"""
    child = copy.deepcopy(parent1)
    for c_param, p2_param in zip(child.parameters(), parent2.parameters()):
        mask = torch.rand_like(c_param) > 0.5
        c_param.data[mask] = p2_param.data[mask]
    return child

def mutate(individual, rate=0.1, sigma=0.1):
    """가우시안 변이"""
    for param in individual.parameters():
        mask = torch.rand_like(param) < rate
        noise = torch.randn_like(param) * sigma
        param.data += mask.float() * noise

class DeepGA:
    """Deep GA: 파라미터 인코딩으로 메모리 효율적 GA"""

    def __init__(self, policy_fn, pop_size=1000, seed_size=5):
        self.pop_size = pop_size
        self.policy_fn = policy_fn

        # 각 개체를 시드 시퀀스로 인코딩 (메모리 절약)
        self.population = []
        for _ in range(pop_size):
            seeds = [np.random.randint(0, 2**31) for _ in range(seed_size)]
            self.population.append(seeds)

    def decode(self, seed_sequence):
        """시드 시퀀스에서 정책 파라미터 복원"""
        policy = self.policy_fn()
        base_params = nn.utils.parameters_to_vector(
            policy.parameters()
        ).detach()

        params = torch.zeros_like(base_params)
        for seed in seed_sequence:
            rng = np.random.RandomState(seed)
            noise = torch.FloatTensor(
                rng.randn(len(base_params))
            ) * 0.01
            params += noise

        nn.utils.vector_to_parameters(params, policy.parameters())
        return policy

    def evolve(self, fitness_fn, num_generations=100):
        for gen in range(num_generations):
            # 적합도 평가
            fitness = []
            for seeds in self.population:
                policy = self.decode(seeds)
                fitness.append(fitness_fn(policy))

            # 변이: 시드 하나 추가
            sorted_idx = np.argsort(fitness)[::-1]
            elites = [self.population[i] for i in sorted_idx[:10]]

            new_pop = [list(e) for e in elites]
            while len(new_pop) < self.pop_size:
                parent = elites[np.random.randint(len(elites))]
                child = list(parent) + [np.random.randint(0, 2**31)]
                new_pop.append(child)

            self.population = new_pop

            if gen % 10 == 0:
                print(f"세대 {gen}: 최고 = {max(fitness):.1f}")

class NoveltySearch:
    """신규성 탐색: 행동 다양성으로 개체 선택"""

    def __init__(self, archive_size=1000, k_nearest=15):
        self.archive = []
        self.archive_size = archive_size
        self.k = k_nearest

    def behavior_characterization(self, trajectory):
        """에이전트의 행동을 특성 벡터로 요약"""
        # 예: 최종 위치, 방문한 상태의 통계 등
        if len(trajectory) == 0:
            return np.zeros(4)

        states = np.array(trajectory)
        return np.array([
            states[:, 0].mean(),  # 평균 x 위치
            states[:, 1].mean(),  # 평균 y 위치
            states[:, 0].std(),   # x 분산
            states[:, 1].std(),   # y 분산
        ])

    def novelty_score(self, behavior):
        """k-최근접 이웃과의 평균 거리 = 신규성 점수"""
        if len(self.archive) < self.k:
            return float('inf')

        distances = [np.linalg.norm(behavior - b)
                     for b in self.archive]
        distances.sort()
        return np.mean(distances[:self.k])

    def add_to_archive(self, behavior):
        self.archive.append(behavior)
        if len(self.archive) > self.archive_size:
            # 랜덤하게 하나 제거
            idx = np.random.randint(len(self.archive))
            self.archive.pop(idx)