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

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') # 创建测试环境 te