Entrenando un agente de Deep Reinforcement Learning

# Crear un entorno virtual (recomendado)
python -m venv dqn_env
source dqn_env/bin/activate   # Linux/Mac
# dqn_env\Scripts\activate    # Windows

# Instalar dependencias
pip install torch torchvision
pip install "gymnasium[atari,accept-rom-license]"
pip install ale-py
pip install matplotlib numpy

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import random
import collections
import matplotlib.pyplot as plt
from collections import deque

# ── Hiperparámetros ──────────────────────────────────────
BATCH_SIZE       = 32        # Tamaño del mini-batch
GAMMA            = 0.99      # Factor de descuento
LR               = 1e-4      # Learning rate del optimizador
BUFFER_SIZE      = 100_000   # Capacidad del replay buffer
MIN_BUFFER       = 10_000    # Transiciones antes de empezar a entrenar
EPSILON_START    = 1.0       # Epsilon inicial (100% exploración)
EPSILON_END      = 0.02      # Epsilon mínimo
EPSILON_DECAY    = 100_000   # Steps para decaer de start a end
TARGET_UPDATE    = 1_000     # Cada cuántos steps sincronizar target net
TAU              = 0.005     # Coeficiente para soft update (alternativa)
NUM_EPISODES     = 800       # Episodios de entrenamiento
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Usando dispositivo: {DEVICE}")

SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed(SEED)

# Crear entorno sin frameskip (lo manejaremos nosotros)
env = gym.make("PongNoFrameskip-v4", render_mode="rgb_array")

print(f"Espacio de observación: {env.observation_space.shape}")  # (210, 160, 3)
print(f"Espacio de acciones:    {env.action_space.n}")           # 6 acciones

import cv2

class PreprocessFrame(gym.ObservationWrapper):
    """Convierte a escala de grises, recorta y redimensiona a 84×84."""

    def __init__(self, env):
        super().__init__(env)
        self.observation_space = gym.spaces.Box(
            low=0, high=255, shape=(84, 84, 1), dtype=np.uint8
        )

    def observation(self, obs):
        # Escala de grises
        gray = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
        # Recortar zona de juego (eliminar marcador superior)
        cropped = gray[34:194, :]   # 160×160
        # Redimensionar a 84×84
        resized = cv2.resize(cropped, (84, 84), interpolation=cv2.INTER_AREA)
        return resized.reshape(84, 84, 1)


class FrameStack(gym.Wrapper):
    """Apila los últimos `k` frames como canales (captura movimiento)."""

    def __init__(self, env, k=4):
        super().__init__(env)
        self.k = k
        self.frames = deque(maxlen=k)
        shp = env.observation_space.shape
        self.observation_space = gym.spaces.Box(
            low=0, high=255,
            shape=(k, shp[0], shp[1]),
            dtype=np.uint8
        )

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        for _ in range(self.k):
            self.frames.append(obs.squeeze(-1))
        return np.array(self.frames), info

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        self.frames.append(obs.squeeze(-1))
        return np.array(self.frames), reward, terminated, truncated, info


class MaxAndSkip(gym.Wrapper):
    """Repite cada acción `skip` frames y devuelve el máximo de los
    últimos 2 frames (elimina parpadeo de sprites Atari)."""

    def __init__(self, env, skip=4):
        super().__init__(env)
        self.skip = skip
        self._obs_buffer = deque(maxlen=2)

    def step(self, action):
        total_reward = 0.0
        terminated = truncated = False
        for _ in range(self.skip):
            obs, reward, terminated, truncated, info = self.env.step(action)
            self._obs_buffer.append(obs)
            total_reward += reward
            if terminated or truncated:
                break
        max_frame = np.max(np.stack(self._obs_buffer), axis=0)
        return max_frame, total_reward, terminated, truncated, info


class ScaleObservation(gym.ObservationWrapper):
    """Normaliza píxeles de [0, 255] a [0.0, 1.0]."""

    def observation(self, obs):
        return np.array(obs, dtype=np.float32) / 255.0

def make_env(env_name="PongNoFrameskip-v4", render_mode=None):
    """Crea el entorno con toda la cadena de preprocesamiento."""
    env = gym.make(env_name, render_mode=render_mode)
    env = MaxAndSkip(env, skip=4)
    env = PreprocessFrame(env)
    env = FrameStack(env, k=4)
    env = ScaleObservation(env)
    return env

# Crear el entorno de entrenamiento
env = make_env()
obs, info = env.reset(seed=SEED)
print(f"Observación preprocesada: shape={obs.shape}, dtype={obs.dtype}")
print(f"Rango de valores: [{obs.min():.2f}, {obs.max():.2f}]")

class DQN(nn.Module):
    """
    Red Deep Q-Network.
    Entrada: tensor (batch, 4, 84, 84) — 4 frames en escala de grises.
    Salida:  tensor (batch, n_actions)  — Q-value para cada acción.
    """

    def __init__(self, n_actions):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(4, 32, kernel_size=8, stride=4),  # → (32, 20, 20)
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2), # → (64, 9, 9)
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1), # → (64, 7, 7)
            nn.ReLU(),
        )
        # Tras las convoluciones: 64 × 7 × 7 = 3136
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions),
        )

    def forward(self, x):
        features = self.conv(x)
        return self.fc(features)

n_actions = env.action_space.n

# Red principal (la que se entrena)
policy_net = DQN(n_actions).to(DEVICE)

# Red target (copia para estabilidad — la veremos en el Paso 8)
target_net = DQN(n_actions).to(DEVICE)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()  # No se entrena directamente

# Optimizador
optimizer = optim.Adam(policy_net.parameters(), lr=LR)

# Función de pérdida: Huber Loss (más robusta que MSE ante outliers)
loss_fn = nn.SmoothL1Loss()

# Contar parámetros
total_params = sum(p.numel() for p in policy_net.parameters())
print(f"Parámetros totales: {total_params:,}")

class ReplayBuffer:
    """
    Buffer circular de experiencias para Experience Replay.
    Almacena transiciones (state, action, reward, next_state, done)
    y permite muestreo aleatorio de mini-batches.
    """

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

    def push(self, state, action, reward, next_state, done):
        """Almacena una transición."""
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        """Muestrea un mini-batch aleatorio y lo convierte a tensores."""
        batch = random.sample(self.buffer, batch_size)

        states, actions, rewards, next_states, dones = zip(*batch)

        states      = torch.FloatTensor(np.array(states)).to(DEVICE)
        actions     = torch.LongTensor(actions).to(DEVICE)
        rewards     = torch.FloatTensor(rewards).to(DEVICE)
        next_states = torch.FloatTensor(np.array(next_states)).to(DEVICE)
        dones       = torch.FloatTensor(dones).to(DEVICE)

        return states, actions, rewards, next_states, dones

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

buffer = ReplayBuffer(BUFFER_SIZE)
print(f"Buffer creado con capacidad: {BUFFER_SIZE:,}")

def select_action(state, step_count):
    """
    Política epsilon-greedy con decaimiento lineal.

    Args:
        state:      observación actual (numpy array 4×84×84)
        step_count: número total de steps de entrenamiento

    Returns:
        acción (int)
    """
    # Calcular epsilon actual (decaimiento lineal)
    epsilon = max(
        EPSILON_END,
        EPSILON_START - step_count / EPSILON_DECAY
    )

    if random.random() < epsilon:
        # Exploración: acción aleatoria
        return env.action_space.sample(), epsilon
    else:
        # Explotación: acción con mayor Q-value
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            q_values = policy_net(state_tensor)
        return q_values.argmax(dim=1).item(), epsilon

def optimize_model():
    """
    Un paso de optimización DQN:
    1. Muestrear mini-batch del buffer
    2. Calcular Q(s, a) con la policy net
    3. Calcular el target: r + γ · max Q_target(s', a') · (1 - done)
    4. Calcular la pérdida (Huber) y hacer backprop
    """
    if len(buffer) < MIN_BUFFER:
        return None  # No entrenar hasta tener suficientes experiencias

    # 1. Muestrear mini-batch
    states, actions, rewards, next_states, dones = buffer.sample(BATCH_SIZE)

    # 2. Q(s, a) — seleccionamos solo el Q-value de la acción tomada
    #    policy_net(states) → shape (batch, n_actions)
    #    .gather(1, ...) selecciona el Q de cada acción
    q_values = policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)

    # 3. Target: r + γ · max_a' Q_target(s', a') · (1 - done)
    with torch.no_grad():
        # max Q-value del siguiente estado, según la target net
        next_q_values = target_net(next_states).max(dim=1).values
        targets = rewards + GAMMA * next_q_values * (1 - dones)

    # 4. Calcular pérdida y actualizar pesos
    loss = loss_fn(q_values, targets)

    optimizer.zero_grad()
    loss.backward()
    # Gradient clipping para estabilidad
    torch.nn.utils.clip_grad_norm_(policy_net.parameters(), max_norm=10.0)
    optimizer.step()

    return loss.item()

def train():
    """Bucle principal de entrenamiento DQN."""
    episode_rewards = []
    step_count = 0
    best_reward = -float('inf')

    for episode in range(NUM_EPISODES):
        state, info = env.reset()
        episode_reward = 0
        done = False

        while not done:
            # 1. Seleccionar acción con ε-greedy
            action, epsilon = select_action(state, step_count)

            # 2. Ejecutar acción en el entorno
            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated

            # 3. Almacenar transición en el buffer
            buffer.push(state, action, reward, next_state, float(done))

            # 4. Optimizar la red (si hay suficientes datos)
            loss = optimize_model()

            # 5. Actualizar la target network (hard update)
            if step_count % TARGET_UPDATE == 0:
                target_net.load_state_dict(policy_net.state_dict())

            state = next_state
            episode_reward += reward
            step_count += 1

        episode_rewards.append(episode_reward)

        # ── Logging cada 10 episodios ─────────────────────────
        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            print(
                f"Ep {episode+1:4d} | "
                f"Reward: {episode_reward:6.1f} | "
                f"Avg(10): {avg_reward:6.1f} | "
                f"ε: {epsilon:.3f} | "
                f"Buffer: {len(buffer):,} | "
                f"Steps: {step_count:,}"
            )

            # Guardar mejor modelo
            if avg_reward > best_reward:
                best_reward = avg_reward
                torch.save(policy_net.state_dict(), "dqn_pong_best.pth")

    print(f"\n✅ Entrenamiento completado. Mejor reward promedio: {best_reward:.1f}")
    return episode_rewards

# Lanzar el entrenamiento
rewards_history = train()

def soft_update(policy_net, target_net, tau=TAU):
    """
    Soft update (Polyak averaging) de la target network.
    θ_target = τ · θ_policy + (1 - τ) · θ_target
    """
    for target_param, policy_param in zip(
        target_net.parameters(), policy_net.parameters()
    ):
        target_param.data.copy_(
            tau * policy_param.data + (1.0 - tau) * target_param.data
        )

# ── En el training loop, reemplazar: ─────────────────────
# ANTES (hard update):
# if step_count % TARGET_UPDATE == 0:
#     target_net.load_state_dict(policy_net.state_dict())

# DESPUÉS (soft update — en cada step):
# soft_update(policy_net, target_net, tau=TAU)

# ── Versión híbrida (recomendada para empezar): ─────────
# Usar hard update cada N steps + soft update continuo

def train_with_soft_update():
    """Training loop con soft update."""
    episode_rewards = []
    step_count = 0
    best_reward = -float('inf')

    for episode in range(NUM_EPISODES):
        state, info = env.reset()
        episode_reward = 0
        done = False

        while not done:
            action, epsilon = select_action(state, step_count)
            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated

            buffer.push(state, action, reward, next_state, float(done))
            loss = optimize_model()

            # ★ Soft update cada step en vez de hard update periódico
            soft_update(policy_net, target_net, tau=TAU)

            state = next_state
            episode_reward += reward
            step_count += 1

        episode_rewards.append(episode_reward)

        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            print(
                f"Ep {episode+1:4d} | "
                f"Reward: {episode_reward:6.1f} | "
                f"Avg(10): {avg_reward:6.1f} | "
                f"ε: {epsilon:.3f} | "
                f"Steps: {step_count:,}"
            )
            if avg_reward > best_reward:
                best_reward = avg_reward
                torch.save(policy_net.state_dict(), "dqn_pong_best.pth")

    return episode_rewards

def plot_training_curves(rewards, window=20):
    """Grafica la recompensa por episodio y su media móvil."""
    fig, ax = plt.subplots(1, 1, figsize=(10, 4))

    episodes = range(1, len(rewards) + 1)

    # Reward crudo (transparente)
    ax.plot(episodes, rewards, alpha=0.3, color='#00b894', linewidth=0.8,
            label='Reward por episodio')

    # Media móvil
    if len(rewards) >= window:
        moving_avg = np.convolve(
            rewards, np.ones(window) / window, mode='valid'
        )
        ax.plot(
            range(window, len(rewards) + 1), moving_avg,
            color='#00b894', linewidth=2,
            label=f'Media móvil ({window} episodios)'
        )

    # Líneas de referencia
    ax.axhline(y=0, color='#74b9ff', linestyle='--', alpha=0.5, label='Empate')
    ax.axhline(y=21, color='#55efc4', linestyle=':', alpha=0.5, label='Victoria perfecta')
    ax.axhline(y=-21, color='#ff7675', linestyle=':', alpha=0.5, label='Derrota total')

    ax.set_xlabel('Episodio')
    ax.set_ylabel('Reward total')
    ax.set_title('Curva de aprendizaje — DQN en Pong')
    ax.legend(loc='lower right', fontsize=8)
    ax.grid(True, alpha=0.2)
    fig.tight_layout()
    plt.savefig('training_curve.png', dpi=150, bbox_inches='tight')
    plt.show()

plot_training_curves(rewards_history)

def evaluate(model, env_name="PongNoFrameskip-v4", n_episodes=10):
    """
    Evalúa el modelo entrenado sin exploración (ε=0).
    Devuelve la recompensa media y un resumen por episodio.
    """
    eval_env = make_env(env_name)
    model.eval()
    rewards = []

    for ep in range(n_episodes):
        state, _ = eval_env.reset()
        episode_reward = 0
        done = False

        while not done:
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(DEVICE)
            with torch.no_grad():
                q_values = model(state_tensor)
            action = q_values.argmax(dim=1).item()

            state, reward, terminated, truncated, _ = eval_env.step(action)
            done = terminated or truncated
            episode_reward += reward

        rewards.append(episode_reward)
        print(f"  Evaluación ep {ep+1}: reward = {episode_reward:.0f}")

    eval_env.close()

    mean_r = np.mean(rewards)
    std_r = np.std(rewards)
    print(f"\n📊 Resultado: {mean_r:.1f} ± {std_r:.1f} (sobre {n_episodes} episodios)")
    return rewards

# Cargar el mejor modelo y evaluar
policy_net.load_state_dict(torch.load("dqn_pong_best.pth", weights_only=True))
eval_rewards = evaluate(policy_net)

def record_game(model, env_name="PongNoFrameskip-v4", filename="pong_agent.mp4"):
    """Graba un episodio completo del agente jugando."""
    import imageio

    # Crear entorno con render_mode para capturar frames
    raw_env = gym.make(env_name, render_mode="rgb_array")
    eval_env = MaxAndSkip(raw_env, skip=4)
    eval_env = PreprocessFrame(eval_env)
    eval_env = FrameStack(eval_env, k=4)
    eval_env = ScaleObservation(eval_env)

    model.eval()
    state, _ = eval_env.reset()
    frames = []
    done = False

    while not done:
        # Capturar frame RGB del juego original
        frame = raw_env.render()
        frames.append(frame)

        # Decidir acción
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            action = model(state_tensor).argmax(dim=1).item()

        state, _, terminated, truncated, _ = eval_env.step(action)
        done = terminated or truncated

    eval_env.close()

    # Guardar como video MP4
    imageio.mimsave(filename, frames, fps=30)
    print(f"🎬 Video guardado: {filename} ({len(frames)} frames)")

# Grabar un partido
# record_game(policy_net)  # Descomentar para ejecutar