[RL] Actor Critic

import gym
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

# AttributeError: module 'numpy' has no attribute 'bool8' 방지 
if not hasattr(np, 'bool8'):
    np.bool8 = np.bool_     

## ActorCriticNetwork
class ActorCriticNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        # actor params 
        self.actor_fc1 = nn.Linear(input_size, hidden_size)
        self.actor_fc2 = nn.Linear(hidden_size, output_size)

        # critic params 
        self.critic_fc1 = nn.Linear(input_size, hidden_size)
        self.critic_fc2 = nn.Linear(hidden_size, hidden_size)
        self.critic_fc3 = nn.Linear(hidden_size, 1)

    def forward(self, x):
        actor_x = F.tanh(self.actor_fc1(x))
        policy = F.softmax(self.actor_fc2(actor_x), dim=-1)

        critic_x = F.tanh(self.critic_fc1(x))
        critic_x = F.tanh(self.critic_fc2(critic_x))
        value = self.critic_fc3(critic_x)

        return policy, value

## Agent
class Agent:
    def __init__(self, env, hyper_parameters:dict):
        self.env = env 
        self.action_size = env.action_space.n
        self.state_shape = env.observation_space.shape[0]

        # hyper params 
        self.gamma = hyper_parameters['gamma']
        self.lr = hyper_parameters['lr']

        # actor critic network
        self.a2c = ActorCriticNetwork(self.state_shape, 32, self.action_size)
        self.optimizer = optim.Adam(self.a2c.parameters(), lr=self.lr)

    def get_action(self, state:torch.Tensor):
        policy, _ = self.a2c(state)
        policy = policy.detach().numpy()
        action = np.random.choice(self.action_size, p=policy)
        return action 

    def train(self, state:torch.Tensor, action:int, reward:int, next_state:torch.Tensor, done:bool):

        self.a2c.train()

        policy, value = self.a2c(state)
        _, next_value = self.a2c(next_state)

        target = reward + (1 - done) * self.gamma * next_value
        advantage = target - value 

        # actor(policy) nn 
        action_prob = policy[action]
        log_prob = torch.log(action_prob + 1e-8)
        actor_loss = -log_prob * advantage.detach() # policy update 과정이니, value 신경망이 개입하면 안된다. 

        # critic nn 
        critic_loss = advantage.pow(2) # 제곱한다는 의미 

        # total loss
        total_loss = actor_loss + critic_loss

        # back propagation 
        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()

        return total_loss.item()


## Hyper Params
HYPER_PARAMETERS = {
    'gamma' : 0.99, 
    'lr' : 1e-3
}

N_EPISODES = 1000
MAX_STEPS = 500
PRINT_INTERVAL = 20  


## Env Setting
seed = 2000
env_name = "CartPole-v1"
env = gym.make(env_name)
env.seed(seed)

## main 
env = gym.make(env_name)
agent = Agent(env, HYPER_PARAMETERS)

reward_history = []
loss_history = []

for e in range(N_EPISODES):
    done = False
    state = env.reset()
    state = torch.tensor(state, dtype=torch.float32)
    
    total_reward = 0
    total_n_steps = 0
    loss_list = []

    while not done:
        action = agent.get_action(state)
        next_state, reward, done, _ = env.step(action)
        next_state = torch.tensor(next_state, dtype=torch.float32)

        total_reward += reward
        loss = agent.train(state, action, reward, next_state, done)
        if not np.isnan(loss):
            loss_list.append(loss)

        total_n_steps += 1
        if total_n_steps >= MAX_STEPS:
            break

        state = next_state

    reward_history.append(total_reward)
    loss_history.append(np.mean(loss_list))

    # PRINT_INTERVAL 마다 평균 출력
    if (e + 1) % PRINT_INTERVAL == 0:
        avg_reward = np.mean(reward_history[-PRINT_INTERVAL:])
        avg_loss = np.mean(loss_history[-PRINT_INTERVAL:])
        print(f"[Episode {e+1}] Average Reward: {avg_reward:.2f}, Average Loss: {avg_loss:.4f}")