import gym
import torch
from torch import nn
from torch.distributions import Categorical

class Agent(nn.Module):
    def __init__(self):
        super(Agent, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(4, 32),
            nn.ReLU(),
            nn.Linear(32, 32),
            nn.ReLU(),
            nn.Linear(32, 2),
            nn.Softmax(dim=0)
        )

    def forward(self, state):
        return self.net(state)


class Trainer:
    def __init__(self):
        self.net = Agent()
        self.env = gym.make('CartPole-v1')
        self.opt = torch.optim.Adam(self.net.parameters())

    def __call__(self):
        # 采样
        for epoch in range(1000):
            state = self.env.reset()
            memory = []  # 存放得分函数和回报
            gamma = 0.9  # 折扣系数
            while True:
                if epoch>300:
                    self.env.render()
                action, log_prob = self.action_select(state)
                state, reward, done ,info= self.env.step(action)
                memory.append([log_prob, reward])
                if done:
                    break
            # 计算总回报
            G = 0
            GS = []  # 存放每个状态的总回报
            rewards = 0  # 所有回报的和
            for _, reward in memory[::-1]:
                G = reward + G * gamma
                GS.insert(0, G)
                rewards += reward
            print(rewards)
            # 数据标准化
            GS = torch.tensor(GS)
            GS = (GS - GS.mean()) / (GS.std() + torch.finfo(torch.float))

            # 计算损失
            loss = 0
            for G, (log_prob, _) in zip(GS, memory):
                loss += -G * log_prob
            self.opt.zero_grad()
            loss.backward()
            self.opt.step()

    def action_select(self, state):
        state = torch.from_numpy(state).float()
        probs = self.net(state[None])
        prob = probs[0]
        m = Categorical(prob)  # 探索机制同时按照动作概率进行采样
        action = m.sample()  # 根据采样选择动作
        return action.item(), m.log_prob(action)


if __name__ == '__main__':
    train = Trainer()
    train()