Neural Network by Pytorch.py

import random
from collections import deque, namedtuple
from typing import Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from gymnasium.spaces import Discrete


torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")


class DQN(nn.Module):
    def __init__(
        self,
        state_size: int,
        action_size: int,
        hidden_size: int = 128,
        dropout_rate: float = 0.2,
    ):
        super().__init__()
        self.fc1 = nn.Linear(state_size, hidden_size)
        self.bn1 = nn.BatchNorm1d(hidden_size)
        self.dropout1 = nn.Dropout(dropout_rate)

        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.bn2 = nn.BatchNorm1d(hidden_size)
        self.dropout2 = nn.Dropout(dropout_rate)

        self.fc_out = nn.Linear(hidden_size, action_size)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        x = self.fc1(state)
        if x.size(0) > 1:
            x = self.bn1(x)
        x = F.relu(x)
        x = self.dropout1(x)

        x = self.fc2(x)
        if x.size(0) > 1:
            x = self.bn2(x)
        x = F.relu(x)
        x = self.dropout2(x)

        return self.fc_out(x)


Experience = namedtuple("Experience", ["state", "action", "reward", "next_state", "done"])


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

    def push(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        next_state: np.ndarray,
        done: bool,
    ) -> None:
        self.buffer.append(Experience(state, action, reward, next_state, done))

    def sample(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states_t = torch.as_tensor(np.array(states), dtype=torch.float32, device=device)
        actions_t = torch.as_tensor(actions, dtype=torch.long, device=device)
        rewards_t = torch.as_tensor(rewards, dtype=torch.float32, device=device)
        next_states_t = torch.as_tensor(np.array(next_states), dtype=torch.float32, device=device)
        dones_t = torch.as_tensor(dones, dtype=torch.float32, device=device)

        return states_t, actions_t, rewards_t, next_states_t, dones_t

    def __len__(self) -> int:
        return len(self.buffer)


class DQNAgent:
    def __init__(
        self,
        state_size: int,
        action_size: int,
        learning_rate: float = 1e-4,
        gamma: float = 0.99,
        epsilon_start: float = 1.0,
        epsilon_end: float = 0.01,
        epsilon_decay: float = 0.995,
        batch_size: int = 64,
        buffer_size: int = 100000,
        target_update_freq: int = 10,
        tau: float = 0.001,
    ):
        self.state_size = state_size
        self.action_size = action_size
        self.gamma = gamma
        self.batch_size = batch_size
        self.target_update_freq = target_update_freq
        self.tau = tau

        self.epsilon = epsilon_start
        self.epsilon_end = epsilon_end
        self.epsilon_decay = epsilon_decay

        self.policy_net = DQN(state_size, action_size).to(device)
        self.target_net = DQN(state_size, action_size).to(device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=learning_rate)
        self.criterion = nn.SmoothL1Loss()

        self.memory = ReplayBuffer(capacity=buffer_size)

        self.steps_done = 0
        self.episodes_done = 0

        self.losses = []
        self.episode_rewards = []
        self.episode_lengths = []

    def select_action(self, state: np.ndarray) -> int:
        if random.random() < self.epsilon:
            return int(np.random.randint(0, self.action_size))

        with torch.no_grad():
            state_tensor = torch.as_tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
            q_values = self.policy_net(state_tensor)
            return int(q_values.argmax(dim=1).item())

    def store_experience(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        next_state: np.ndarray,
        done: bool,
    ) -> None:
        self.memory.push(state, action, reward, next_state, done)

    def train_step(self) -> float:
        if len(self.memory) < self.batch_size:
            return 0.0

        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)

        current_q_values = self.policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)

        with torch.no_grad():
            next_q_values = self.target_net(next_states).max(1)[0]
            target_q_values = rewards + self.gamma * next_q_values * (1.0 - dones)

        loss = self.criterion(current_q_values, target_q_values)

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

        self.steps_done += 1
        self.losses.append(float(loss.item()))
        return float(loss.item())

    def update_target_network(self, soft_update: bool = True) -> None:
        if soft_update:
            for target_param, policy_param in zip(self.target_net.parameters(), self.policy_net.parameters()):
                target_param.data.copy_(self.tau * policy_param.data + (1.0 - self.tau) * target_param.data)
            return

        self.target_net.load_state_dict(self.policy_net.state_dict())

    def decay_epsilon(self) -> None:
        self.epsilon = max(self.epsilon_end, self.epsilon * self.epsilon_decay)

    def save(self, filepath: str) -> None:
        checkpoint = {
            "policy_net_state_dict": self.policy_net.state_dict(),
            "target_net_state_dict": self.target_net.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "epsilon": self.epsilon,
            "steps_done": self.steps_done,
            "episodes_done": self.episodes_done,
            "losses": self.losses,
            "episode_rewards": self.episode_rewards,
            "episode_lengths": self.episode_lengths,
            "config": {
                "state_size": self.state_size,
                "action_size": self.action_size,
                "gamma": self.gamma,
                "batch_size": self.batch_size,
                "target_update_freq": self.target_update_freq,
                "tau": self.tau,
            },
        }
        torch.save(checkpoint, filepath)
        print(f"Model saved to {filepath}")

    def load(self, filepath: str) -> None:
        checkpoint = torch.load(filepath, map_location=device)
        self.policy_net.load_state_dict(checkpoint["policy_net_state_dict"])
        self.target_net.load_state_dict(checkpoint["target_net_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

        self.epsilon = float(checkpoint.get("epsilon", self.epsilon))
        self.steps_done = int(checkpoint.get("steps_done", 0))
        self.episodes_done = int(checkpoint.get("episodes_done", 0))
        self.losses = checkpoint.get("losses", [])
        self.episode_rewards = checkpoint.get("episode_rewards", [])
        self.episode_lengths = checkpoint.get("episode_lengths", [])

        print(f"Model loaded from {filepath}")


def train_dqn(
    env_name: str = "CartPole-v1",
    episodes: int = 300,
    max_steps: int = 500,
    model_path: str = "dqn_cartpole.pt",
    log_every: int = 10,
) -> DQNAgent:
    try:
        import gymnasium as gym
    except ImportError as error:
        raise ImportError("Please install gymnasium first: pip install gymnasium") from error

    env = gym.make(env_name)

    if hasattr(env.observation_space, 'shape') and env.observation_space.shape is not None:
        state_size = env.observation_space.shape[0]
    else:
        raise ValueError(f"Environment {env_name} does not have a valid observation space shape")
    
    if isinstance(env.action_space, Discrete):
        action_size = int(env.action_space.n)
    else:
        raise ValueError(f"Environment {env_name} does not have a valid discrete action space")
    agent = DQNAgent(state_size=int(state_size), action_size=action_size)

    for episode in range(1, episodes + 1):
        state, _ = env.reset(seed=42 + episode)
        state = np.asarray(state, dtype=np.float32)

        total_reward = 0.0
        episode_loss_sum = 0.0
        updates = 0

        for step in range(1, max_steps + 1):
            action = agent.select_action(state)
            next_state, reward, terminated, truncated, _ = env.step(action)

            done = bool(terminated or truncated)
            next_state = np.asarray(next_state, dtype=np.float32)
            reward = float(reward)

            agent.store_experience(state, action, reward, next_state, done)
            loss = agent.train_step()
            if loss > 0.0:
                episode_loss_sum += loss
                updates += 1

            if agent.steps_done % agent.target_update_freq == 0:
                agent.update_target_network(soft_update=True)

            state = next_state
            total_reward += reward

            if done:
                break

        avg_loss = episode_loss_sum / updates if updates > 0 else 0.0

        agent.episodes_done += 1
        agent.episode_rewards.append(total_reward)
        agent.episode_lengths.append(step)
        agent.decay_epsilon()

        if episode % log_every == 0 or episode == 1:
            recent_rewards = agent.episode_rewards[-10:]
            recent_avg_reward = float(np.mean(recent_rewards))
            print(
                f"Episode {episode}/{episodes} | "
                f"Reward: {total_reward:.1f} | "
                f"Avg(10): {recent_avg_reward:.1f} | "
                f"Loss: {avg_loss:.4f} | "
                f"Epsilon: {agent.epsilon:.4f}"
            )

    env.close()
    agent.save(model_path)
    return agent


def evaluate_dqn(
    agent: DQNAgent,
    env_name: str = "CartPole-v1",
    episodes: int = 5,
    max_steps: int = 500,
) -> float:
    try:
        import gymnasium as gym
    except ImportError as error:
        raise ImportError("Please install gymnasium first: pip install gymnasium") from error

    env = gym.make(env_name)
    original_epsilon = agent.epsilon
    agent.epsilon = 0.0

    rewards = []
    for episode in range(1, episodes + 1):
        state, _ = env.reset(seed=1000 + episode)
        state = np.asarray(state, dtype=np.float32)

        total_reward = 0.0
        for _ in range(max_steps):
            action = agent.select_action(state)
            next_state, reward, terminated, truncated, _ = env.step(action)
            state = np.asarray(next_state, dtype=np.float32)
            total_reward += float(reward)
            if terminated or truncated:
                break

        rewards.append(total_reward)
        print(f"Eval Episode {episode}/{episodes} | Reward: {total_reward:.1f}")

    env.close()
    agent.epsilon = original_epsilon

    avg_reward = float(np.mean(rewards)) if rewards else 0.0
    print(f"Evaluation complete | Avg Reward: {avg_reward:.2f}")
    return avg_reward


def print_final_summary(agent: DQNAgent, eval_avg_reward: float) -> None:
    best_reward = float(max(agent.episode_rewards)) if agent.episode_rewards else 0.0
    last_10_avg = float(np.mean(agent.episode_rewards[-10:])) if agent.episode_rewards else 0.0
    avg_loss_last_100 = float(np.mean(agent.losses[-100:])) if agent.losses else 0.0

    print("\n" + "=" * 50)
    print("FINAL RESULT")
    print("=" * 50)
    print(f"Episodes trained     : {agent.episodes_done}")
    print(f"Best train reward    : {best_reward:.2f}")
    print(f"Train avg reward(10) : {last_10_avg:.2f}")
    print(f"Avg loss(last 100)   : {avg_loss_last_100:.6f}")
    print(f"Final epsilon        : {agent.epsilon:.4f}")
    print(f"Evaluation avg reward: {eval_avg_reward:.2f}")
    print("=" * 50)


if __name__ == "__main__":
    trained_agent = train_dqn()
    print(f"Training complete. Episodes: {trained_agent.episodes_done}")
    eval_reward = evaluate_dqn(trained_agent, episodes=5)
    print_final_summary(trained_agent, eval_reward)