模仿学习.py

"""
假设存在一个专家智能体，其策略可以看成最优策略，我们就可以通过直接模仿这个专家在环境中交互的状态动作数据来训练一个策略，并且不需要用到环境提供的奖励信号

模仿学习的方法基本上可以分为3类：
* 行为克隆（behavior cloning,BC）
    * 用监督学习的方法将专家数据中的s看做样本输入，a看做标签，大多数实际场景作为策略预训练，但存在克隆的复合误差
* 逆强化学习（inverse RL）
* 生成对抗模仿学习（generative adversarial imitation learning）

"""
import random

import torch
import torch.nn.functional as F
from mistune import toc
from tqdm import tqdm

import rl_utils
import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt

"""生成专家内容,PPO作为专家"""
class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim,hidden_dim,action_dim):
        super().__init__()
        self.l1 = torch.nn.Linear(state_dim, hidden_dim)
        self.l2 = torch.nn.Linear(hidden_dim, action_dim)
    def forward(self, x):
        x = F.relu(self.l1(x))
        return F.softmax(self.l2(x))

class ValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super().__init__()
        self.l1 = torch.nn.Linear(state_dim, hidden_dim)
        self.l2 = torch.nn.Linear(hidden_dim, 1)
    def forward(self, x):
        x = F.relu(self.l1(x))
        return self.l2(x)

class PPO:
    """PPO算法，采用截断式"""
    def __init__(self,state_dim,hidden_dim,action_dim,actor_lr,critic_lr,lmbda,epochs,eps,gamma,device):
        self.actor = PolicyNet(state_dim,hidden_dim,action_dim).to(device)
        self.critic = ValueNet(state_dim,hidden_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),lr=actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),lr=critic_lr)
        self.gamma = gamma
        self.lmbda = lmbda
        self.epochs = epochs#一条序列的数据用来训练轮数
        self.eps = eps #PPO中截断范围的参数
        self.device = device

    def take_action(self,state):
        state = torch.tensor([state],dtype=torch.float).to(self.device)
        probs = self.actor(state)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()#脱下tensor还给环境

    def update(self,transition_dict):
        states = torch.tensor(transition_dict['states'],dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions'],dtype=torch.long).view(-1,1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],dtype=torch.float).view(-1,1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],dtype=torch.float).view(-1,1).to(self.device)

        td_target = rewards+self.gamma*self.critic(next_states)*(1-dones)
        td_delta = td_target - self.critic(states)
        # 引入lmbda把走多步的优势加权平均，把短视和远见平衡起来，降低训练的方差
        advantage = rl_utils.compute_advantage(self.gamma, self.lmbda, td_delta.cpu()).to(self.device)
        # 旧策略的baseline
        old_log_probs = torch.log(self.actor(states).gather(1,actions)).detach()#旧策略的log概率，detach()脱离计算图不更新

        for _ in range(self.epochs):
            log_probs = torch.log(self.actor(states).gather(1,actions))
            # 新旧比例，好的情况就是新旧策略的概率差不多，坏的情况就是新策略的概率比旧策略大很多（过拟合了当前数据），或者小很多（更新太保守了）
            ratio = torch.exp(log_probs - old_log_probs)
            # 如果这个动作好（advantage>0）网络希望ratio越大越好
            surr1 = ratio * advantage
            # 动作再好也不能太好，把ratio锁在【1-eps,1+eps]之间
            surr2 = torch.clamp(ratio,1-self.eps,1+self.eps)*advantage # 截断
            # 取surr1和surr2的最小值，保证了更新的稳定性
            actor_loss = -torch.min(surr1,surr2).mean() #PPO损失函数
            # 教练也要自我反省，但是要训成常数
            critic_loss = F.mse_loss(self.critic(states),td_target.detach())

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()

env_name = 'CartPole-v1'
seed = 0 # 提取出一个统一的种子变量，方便以后修改
rl_utils.set_all_seeds(seed)
env = rl_utils.create_env(env_name,seed)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n  # CartPole 是离散动作空间，所以用 .n
hidden_dim = 128
actor_lr = 1e-3
critic_lr = 1e-2
lmbda = 0.95
epochs = 10
eps = 0.2
gamma = 0.98
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
n_episode = 100 # 设定采样的专家轨迹数量
ppo_agent = PPO(state_dim,hidden_dim,action_dim,actor_lr,critic_lr,lmbda,epochs,eps,gamma,device)

# ========== 先训练专家！ ==========
print("开始训练专家 PPO 智能体...")
# 让 PPO 在环境里先训个几百回合，直到它能稳定拿满分
return_list = rl_utils.train_on_policy_agent(env, ppo_agent, num_episodes=500)
print("专家训练完成！")#此时的专家就是ppo_agent
# # 将训练好的 Actor 策略网络保存到本地文件
# torch.save(ppo_agent.actor.state_dict(), "ppo_expert_cartpole.pth")
# print("专家模型已保存到 ppo_expert_cartpole.pth")
# ===========================================


def sample_expert_data(n_episode):
    states = []
    actions = []
    for episode in range(n_episode):
        state,info = env.reset()
        done = False
        while not done:
            action = ppo_agent.take_action(state)
            states.append(state)
            actions.append(action)
            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            state = next_state
    return np.array(states), np.array(actions)

print("开始采集专家数据...")
expert_s, expert_a = sample_expert_data(n_episode)
n_samples = min(1000, expert_s.shape[0]) # 比如采样 1000 条
random_index = random.sample(range(expert_s.shape[0]), n_samples)
expert_s = expert_s[random_index]
expert_a = expert_a[random_index]
print(f"成功采集并截取了 {n_samples} 条专家数据准备进行 BC 训练。")


class BehaviorClone:
    def __init__(self, state_dim, hidden_dim, action_dim,lr):
        self.policy = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.optimizer = torch.optim.Adam(self.policy.parameters(), lr=lr)

    def learn(self,states,actions):
        states = torch.tensor(states,dtype=torch.float).to(device)
        actions = torch.tensor(actions).view(-1,1).to(device)
        log_probs = torch.log(self.policy(states).gather(1,actions))
        bc_loss = -torch.mean(log_probs)#最大似然估计

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

    def take_action(self,state):
        state = torch.tensor([state],dtype=torch.float).to(device)
        probs = self.policy(state)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()

def test_agent(agent,env,n_episode):
    return_list = []
    for episode in range(n_episode):
        episode_return = 0
        state,info = env.reset()
        done = False
        while not done:
            action = agent.take_action(state)
            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            state = next_state
            episode_return += reward
        return_list.append(episode_return)
    return np.mean(return_list)
seed = 0
rl_utils.set_all_seeds(seed)
lr = 1e-3
bc_agent = BehaviorClone(state_dim, hidden_dim, action_dim,lr)
n_iterations = 1000
batch_size = 64
test_returns =[]

with tqdm(total=n_iterations,desc="进度条") as pbar:
    for i in range(n_iterations):
        sample_indices = np.random.randint(low=0,high=expert_s.shape[0],size=batch_size)
        bc_agent.learn(expert_s[sample_indices],expert_a[sample_indices])
        current_return = test_agent(bc_agent,env,5)
        test_returns.append(current_return)
        if(i+1)%10==0:
            pbar.set_postfix({"return":'%.3f'%np.mean(test_returns[-10:])})
        pbar.update(1)

iteration_list = list(range(len(test_returns)))
plt.plot(iteration_list,test_returns)
plt.xlabel("iterations")
plt.ylabel("returns")
plt.title(f"BC on{env_name}")
plt.show()