【SSL-RL】自监督强化学习：随机网络蒸馏（RND）方法

pip install gym torch torchvision numpy

"""《 RND 算法简单实现》
    时间：2024.10.28
    环境：cartpole
    作者：不去幼儿园
"""
import numpy as np  # 导入NumPy库，用于数组和数学操作
import gym  # 导入OpenAI Gym库，用于环境模拟
import torch  # 导入PyTorch库，用于深度学习
import torch.nn as nn  # 导入PyTorch的神经网络模块
import torch.optim as optim  # 导入PyTorch的优化器模块
import torch.nn.functional as F  # 导入PyTorch的功能模块（激活函数等）
from torch.distributions import Categorical  # 导入Categorical分布，用于策略选择

# RND网络（随机网络）
class RandomNetwork(nn.Module):  # 定义一个随机网络类，继承自nn.Module
    def __init__(self, input_size, output_size):  # 初始化方法
        super(RandomNetwork, self).__init__()  # 调用父类初始化方法
        self.fc1 = nn.Linear(input_size, 128)  # 第一层全连接层，输入大小为input_size，输出128
        self.fc2 = nn.Linear(128, output_size)  # 第二层全连接层，输入128，输出output_size
        self.random_weights = nn.Parameter(torch.randn(1, output_size))  # 随机权重参数

    def forward(self, x):  # 前向传播方法
        x = F.relu(self.fc1(x))  # 通过第一层并应用ReLU激活
        return self.fc2(x) + self.random_weights  # 通过第二层并加上随机权重

class PredictorNetwork(nn.Module):  # 定义一个预测网络类
    def __init__(self, input_size, output_size):  # 初始化方法
        super(PredictorNetwork, self).__init__()  # 调用父类初始化方法
        self.fc1 = nn.Linear(input_size, 128)  # 第一层全连接层
        self.fc2 = nn.Linear(128, output_size)  # 第二层全连接层

    def forward(self, x):  # 前向传播方法
        x = F.relu(self.fc1(x))  # 通过第一层并应用ReLU激活
        return self.fc2(x)  # 通过第二层

# A3C代理
class A3CAgent:  # 定义A3C代理类
    def __init__(self, input_size, action_size):  # 初始化方法
        self.policy_net = nn.Sequential(  # 策略网络
            nn.Linear(input_size, 128),  # 输入层
            nn.ReLU(),  # 激活函数
            nn.Linear(128, action_size),  # 输出层
            nn.Softmax(dim=-1)  # Softmax函数
        )
        self.value_net = nn.Sequential(  # 价值网络
            nn.Linear(input_size, 128),  # 输入层
            nn.ReLU(),  # 激活函数
            nn.Linear(128, 1)  # 输出层，值为一个标量
        )
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=1e-3)  # 策略网络优化器
        self.optimizer_value = optim.Adam(self.value_net.parameters(), lr=1e-3)  # 价值网络优化器

        self.rnd = RandomNetwork(input_size, output_size=1)  # 随机网络实例
        self.predictor = PredictorNetwork(input_size, output_size=1)  # 预测网络实例
        self.optimizer_rnd = optim.Adam(list(self.rnd.parameters()) + list(self.predictor.parameters()), lr=1e-3)  # 随机网络优化器

    def select_action(self, state):  # 选择动作方法
        state = torch.FloatTensor(state).unsqueeze(0)  # 将状态转换为张量并增加一个维度
        probs = self.policy_net(state)  # 通过策略网络计算动作概率
        distribution = Categorical(probs)  # 创建Categorical分布
        action = distribution.sample()  # 根据分布采样动作
        return action.item(), distribution.log_prob(action)  # 返回动作及其对数概率

    def update(self, rewards, log_probs, values, next_value):  # 更新方法
        # 计算回报
        returns = []  # 初始化回报列表
        R = next_value  # 初始化R为下一个值
        for r in rewards[::-1]:  # 反向遍历奖励
            R = r + 0.99 * R  # 计算当前回报
            returns.insert(0, R)  # 将回报插入列表开头

        # 计算损失
        returns = torch.FloatTensor(returns).view(-1, 1)  # 转换回报为张量并调整形状
        log_probs = torch.stack(log_probs)  # 堆叠对数概率
        values = torch.stack(values)  # 堆叠价值

        advantage = returns - values.detach()  # 计算优势

        policy_loss = -log_probs * advantage.detach()  # 策略损失
        value_loss = F.mse_loss(values.view(-1, 1), returns)  # 价值损失，确保形状一致

        self.optimizer.zero_grad()  # 清空优化器梯度
        policy_loss.mean().backward()  # 反向传播策略损失
        self.optimizer.step()  # 更新策略网络

        self.optimizer_value.zero_grad()  # 清空价值优化器梯度
        value_loss.backward()  # 反向传播价值损失
        self.optimizer_value.step()  # 更新价值网络

    def compute_intrinsic_reward(self, state):  # 计算内在奖励
        state_tensor = torch.FloatTensor(state).unsqueeze(0)  # 状态转换为张量
        with torch.no_grad():  # 在不计算梯度的情况下进行前向传播
            rnd_output = self.rnd(state_tensor)  # 获取随机网络输出
            pred_output = self.predictor(state_tensor)  # 获取预测网络输出

        intrinsic_reward = F.mse_loss(rnd_output, pred_output)  # 计算内在奖励
        return intrinsic_reward.item()  # 返回内在奖励的数值

# 训练过程
def train(agent, env, num_episodes):  # 训练函数
    for episode in range(num_episodes):  # 遍历每个回合
        state, _ = env.reset()  # 重置环境并获取初始状态
        log_probs = []  # 初始化对数概率列表
        values = []  # 初始化价值列表
        rewards = []  # 初始化奖励列表
        intrinsic_rewards = []  # 初始化内在奖励列表

        done = False  # 初始化done标志
        while not done:  # 在未完成时循环
            action, log_prob = agent.select_action(state)  # 选择动作
            next_state, reward, done, _, _ = env.step(action)  # 执行动作并获取下一个状态及奖励

            intrinsic_reward = agent.compute_intrinsic_reward(state)  # 计算内在奖励
            intrinsic_rewards.append(intrinsic_reward)  # 将内在奖励添加到列表中

            log_probs.append(log_prob)  # 添加对数概率
            value = agent.value_net(torch.FloatTensor(state).unsqueeze(0))  # 计算当前状态的价值
            values.append(value)  # 添加价值

            rewards.append(reward + intrinsic_reward)  # 添加总奖励（外在 + 内在）
            state = next_state  # 更新状态

        next_value = agent.value_net(torch.FloatTensor(state).unsqueeze(0))  # 计算下一个状态的价值
        agent.update(rewards, log_probs, values, next_value)  # 更新代理
        print(f"Episode {episode}: Total Reward = {sum(rewards)}")  # 打印回合总奖励


# 主程序
env = gym.make("CartPole-v1")  # 创建CartPole环境
agent = A3CAgent(input_size=env.observation_space.shape[0], action_size=env.action_space.n)  # 实例化代理
train(agent, env, num_episodes=100)  # 训练代理
env.close()  # 关闭环境

# 测试阶段显示动画
def test_agent(agent, env, num_episodes=5):  # 测试函数
    for episode in range(num_episodes):  # 遍历每个测试回合
        state, _ = env.reset()  # 重置环境
        total_reward = 0  # 初始化总奖励
        done = False  # 初始化done标志
        while not done:  # 在未完成时循环
            env.render()  # 显示动画
            action, log_prob = agent.select_action(state)  # 选择动作
            state, reward, done, _, _ = env.step(action)  # 执行动作并获取下一个状态及奖励
            intrinsic_reward = agent.compute_intrinsic_reward(state)  # 计算内在奖励
            total_reward += reward + intrinsic_reward  # 更新总奖励
        print(f"Test Episode {episode}: Total Reward = {total_reward}")  # 打印测试回合总奖励
    env.close()  # 关闭环境


# 测试模型
env_test = gym.make('CartPole-v1', render_mode='human')  # 创建测试环境
test_agent(agent, env_test)  # 测试代理

env = gym.make("Taxi-v3", render_mode="human")  # 使用人类可视化的方式