Reinforcement learning agents for games and navigation

import random
from collections import deque

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

# --- 1. Q-Network ---
class DQN(nn.Module):
    def __init__(self, state_dim: int, action_dim: int):
        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: torch.Tensor) -> torch.Tensor:
        return self.net(x)

# --- 2. Hyperparameters ---
ENV_NAME    = "CartPole-v1"   # swap for Atari env + wrappers
GAMMA       = 0.99
LR          = 1e-3
BATCH_SIZE  = 64
BUFFER_SIZE = 10_000
SYNC_EVERY  = 500             # steps between target network updates
EPSILON_START, EPSILON_END, EPSILON_DECAY = 1.0, 0.05, 5_000

env = gym.make(ENV_NAME)
state_dim  = env.observation_space.shape[0]
action_dim = env.action_space.n

online_net = DQN(state_dim, action_dim)
target_net = DQN(state_dim, action_dim)
target_net.load_state_dict(online_net.state_dict())
target_net.eval()

optimizer  = optim.Adam(online_net.parameters(), lr=LR)
replay     = deque(maxlen=BUFFER_SIZE)
total_steps = 0

# --- 3. Training loop ---
for episode in range(500):
    state, _ = env.reset()
    done = False
    ep_reward = 0.0

    while not done:
        # Epsilon-greedy action selection
        epsilon = EPSILON_END + (EPSILON_START - EPSILON_END) * \
                  np.exp(-total_steps / EPSILON_DECAY)
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            with torch.no_grad():
                q = online_net(torch.tensor(state, dtype=torch.float32))
                action = q.argmax().item()

        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        replay.append((state, action, reward, next_state, float(done)))
        state = next_state
        ep_reward += reward
        total_steps += 1

        # Learn once the buffer has enough samples
        if len(replay) >= BATCH_SIZE:
            batch = random.sample(replay, BATCH_SIZE)
            s, a, r, ns, d = zip(*batch)

            s  = torch.tensor(np.array(s),  dtype=torch.float32)
            a  = torch.tensor(a,             dtype=torch.long).unsqueeze(1)
            r  = torch.tensor(r,             dtype=torch.float32).unsqueeze(1)
            ns = torch.tensor(np.array(ns), dtype=torch.float32)
            d  = torch.tensor(d,             dtype=torch.float32).unsqueeze(1)

            # Bellman target (Double DQN variant: select action with online, evaluate with target)
            with torch.no_grad():
                best_actions = online_net(ns).argmax(1, keepdim=True)
                target_q     = r + GAMMA * (1 - d) * target_net(ns).gather(1, best_actions)

            current_q = online_net(s).gather(1, a)
            loss = nn.functional.mse_loss(current_q, target_q)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # Sync target network
        if total_steps % SYNC_EVERY == 0:
            target_net.load_state_dict(online_net.state_dict())

    print(f"Episode {episode+1:4d} | reward: {ep_reward:7.1f} | ε: {epsilon:.3f}")

env.close()