INFO:__main__:PPOAgent initialized successfully.
c:\Users\heman\Desktop\Notes\2024notes\RLNotes\example\ppocodegen.py:109: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\torch\csrc\utils\tensor_new.cpp:281.)
  state = torch.FloatTensor(state).to(self.device)
ERROR:__main__:Error selecting action: expected sequence of length 4 at dim 1 (got 0)
ERROR:__main__:Error during training: expected sequence of length 4 at dim 1 (got 0)
ERROR:__main__:Training terminated due to an error: expected sequence of length 4 at dim 1 (got 0)
import gym
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
from typing import List, Tuple, Optional
import numpy as np
import logging
import sys

# Configure logging
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger(__name__)


class ActorCritic(nn.Module):
    """
    Combined Actor-Critic Network for PPO Algorithm.
    """
    def __init__(self, state_dim: int, action_dim: int, hidden_sizes: List[int] = [64, 64]) -> None:
        super(ActorCritic, self).__init__()
        try:
            layers = []
            input_dim = state_dim
            for hidden_size in hidden_sizes:
                layers.append(nn.Linear(input_dim, hidden_size))
                layers.append(nn.ReLU())
                input_dim = hidden_size
            # Actor Head
            self.actor = nn.Sequential(*layers, nn.Linear(input_dim, action_dim), nn.Softmax(dim=-1))
            # Critic Head
            self.critic = nn.Sequential(*layers, nn.Linear(input_dim, 1))
        except Exception as e:
            logger.error(f"Error initializing ActorCritic network: {e}")
            raise

    def forward(self, state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass to compute action probabilities and state value.
        """
        try:
            action_probs = self.actor(state)
            state_value = self.critic(state)
            return action_probs, state_value
        except Exception as e:
            logger.error(f"Error during forward pass: {e}")
            raise


class Memory:
    """
    Memory buffer for storing trajectories.
    """
    def __init__(self) -> None:
        self.states: List[np.ndarray] = []
        self.actions: List[int] = []
        self.rewards: List[float] = []
        self.is_terminals: List[bool] = []
        self.log_probs: List[float] = []

    def clear_memory(self) -> None:
        """
        Clears the memory buffer.
        """
        try:
            self.states = []
            self.actions = []
            self.rewards = []
            self.is_terminals = []
            self.log_probs = []
        except Exception as e:
            logger.error(f"Error clearing memory: {e}")
            raise


class PPOAgent:
    """
    Proximal Policy Optimization (PPO) Agent.
    """
    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        lr: float = 3e-4,
        gamma: float = 0.99,
        K_epochs: int = 80,
        eps_clip: float = 0.2,
        hidden_sizes: List[int] = [64, 64],
        device: Optional[torch.device] = None
    ) -> None:
        self.device = device if device else torch.device("cpu")
        self.policy = ActorCritic(state_dim, action_dim, hidden_sizes).to(self.device)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.policy_old = ActorCritic(state_dim, action_dim, hidden_sizes).to(self.device)
        self.policy_old.load_state_dict(self.policy.state_dict())
        self.MseLoss = nn.MSELoss()

        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs
        self.memory = Memory()
        logger.info("PPOAgent initialized successfully.")

    def select_action(self, state: np.ndarray) -> Tuple[int, float]:
        """
        Selects an action based on the current policy.
        """
        try:
            state = torch.FloatTensor(state).to(self.device)
            with torch.no_grad():
                action_probs, _ = self.policy_old(state)
            dist = Categorical(action_probs)
            action = dist.sample()
            self.memory.states.append(state.cpu().numpy())
            self.memory.actions.append(action.item())
            self.memory.log_probs.append(dist.log_prob(action).item())
            return action.item(), dist.log_prob(action).item()
        except Exception as e:
            logger.error(f"Error selecting action: {e}")
            raise

    def compute_returns_and_advantages(
        self,
        rewards: List[float],
        is_terminals: List[bool],
        state_values: List[float],
        next_state_value: float
    ) -> Tuple[List[float], List[float]]:
        """
        Computes returns and advantage estimates using GAE (Generalized Advantage Estimation).
        """
        try:
            returns = []
            advantages = []
            G = next_state_value
            A = 0
            for step in reversed(range(len(rewards))):
                mask = 1.0 - float(is_terminals[step])
                G = rewards[step] + self.gamma * G * mask
                returns.insert(0, G)
            returns = torch.tensor(returns, dtype=torch.float32).to(self.device)
            state_values = torch.tensor(state_values, dtype=torch.float32).squeeze().to(self.device)
            advantages = returns - state_values
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
            return returns.tolist(), advantages.tolist()
        except Exception as e:
            logger.error(f"Error computing returns and advantages: {e}")
            raise

    def update(self) -> None:
        """
        Updates the policy and value networks.
        """
        try:
            # Convert lists to tensors
            old_states = torch.tensor(self.memory.states, dtype=torch.float32).to(self.device)
            old_actions = torch.tensor(self.memory.actions, dtype=torch.int64).to(self.device)
            old_log_probs = torch.tensor(self.memory.log_probs, dtype=torch.float32).to(self.device)

            # Compute state values
            with torch.no_grad():
                _, state_values = self.policy_old(old_states)
                state_values = state_values.squeeze().cpu().numpy()

            # Compute returns and advantages
            returns, advantages = self.compute_returns_and_advantages(
                rewards=self.memory.rewards,
                is_terminals=self.memory.is_terminals,
                state_values=state_values,
                next_state_value=0 if self.memory.is_terminals[-1] else self.policy_old(torch.FloatTensor(self.memory.states[-1])).detach().cpu().numpy()[1].squeeze()
            )

            # Optimize policy for K epochs
            for _ in range(self.K_epochs):
                # Evaluate current policy
                action_probs, state_values = self.policy(old_states)
                dist = Categorical(action_probs)
                entropy = dist.entropy().mean()
                new_log_probs = dist.log_prob(old_actions)

                # Calculate ratios
                ratios = torch.exp(new_log_probs - old_log_probs)

                # Calculate surrogate losses
                surr1 = ratios * torch.tensor(advantages, dtype=torch.float32).to(self.device)
                surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * torch.tensor(advantages, dtype=torch.float32).to(self.device)
                loss = -torch.min(surr1, surr2) + 0.5 * self.MseLoss(state_values.squeeze(), torch.tensor(returns, dtype=torch.float32).to(self.device)) - 0.01 * entropy

                # Take gradient step
                self.optimizer.zero_grad()
                loss.mean().backward()
                self.optimizer.step()

            # Update old policy
            self.policy_old.load_state_dict(self.policy.state_dict())
            self.memory.clear_memory()
            logger.info("Policy and value networks updated successfully.")
        except Exception as e:
            logger.error(f"Error during policy update: {e}")
            raise

    def save(self, checkpoint_path: str) -> None:
        """
        Saves the current policy network state.
        """
        try:
            torch.save(self.policy.state_dict(), checkpoint_path)
            logger.info(f"Policy network saved to {checkpoint_path}")
        except Exception as e:
            logger.error(f"Error saving policy network: {e}")
            raise

    def load(self, checkpoint_path: str) -> None:
        """
        Loads the policy network state.
        """
        try:
            self.policy.load_state_dict(torch.load(checkpoint_path, map_location=self.device))
            self.policy_old.load_state_dict(self.policy.state_dict())
            logger.info(f"Policy network loaded from {checkpoint_path}")
        except Exception as e:
            logger.error(f"Error loading policy network: {e}")
            raise


def train_ppo(
    env_name: str = "CartPole-v1",
    max_episodes: int = 1000,
    max_timesteps: int = 300,
    update_timestep: int = 2000,
    lr: float = 3e-4,
    gamma: float = 0.99,
    K_epochs: int = 80,
    eps_clip: float = 0.2,
    hidden_sizes: List[int] = [64, 64],
    log_interval: int = 10,
    save_model: bool = False,
    checkpoint_path: str = "ppo_cartpole.pth"
) -> None:
    """
    Trains the PPO agent on the specified environment.
    """
    try:
        env = gym.make(env_name)
        state_dim = env.observation_space.shape[0]
        action_dim = env.action_space.n

        agent = PPOAgent(
            state_dim=state_dim,
            action_dim=action_dim,
            lr=lr,
            gamma=gamma,
            K_epochs=K_epochs,
            eps_clip=eps_clip,
            hidden_sizes=hidden_sizes
        )

        timestep = 0
        running_reward = 0
        avg_length = 0
        for episode in range(1, max_episodes + 1):
            state = env.reset()
            for t in range(max_timesteps):
                timestep += 1
                action, _ = agent.select_action(state)
                state, reward, done, _ = env.step(action)

                # Store rewards and dones
                agent.memory.rewards.append(reward)
                agent.memory.is_terminals.append(done)

                running_reward += reward
                if timestep % update_timestep == 0:
                    agent.update()
                if done:
                    break

            avg_length += t

            # Logging
            if episode % log_interval == 0:
                avg_length = int(avg_length / log_interval)
                running_reward = int((running_reward) / log_interval)
                logger.info(f"Episode {episode}\tAverage Length: {avg_length}\tAverage Reward: {running_reward}")
                running_reward = 0
                avg_length = 0

        if save_model:
            agent.save(checkpoint_path)
        env.close()
        logger.info("Training completed successfully.")
    except Exception as e:
        logger.error(f"Error during training: {e}")
        raise


if __name__ == "__main__":
    try:
        train_ppo()
    except Exception as e:
        logger.error(f"Training terminated due to an error: {e}")