RL强化学习从小白到老鸟(一)——速通贪吃蛇游戏
简介
很久没有更新了,刚好最近在带新人学习强化学习相关的内容,就写一个教学贴吧。
游戏用的是经典的贪吃蛇游戏,因为是为了教学就没有设计特别花哨的背景音乐了,纯python代码手撸出来的(想要加音乐和图片的可以自己改进)
废话不多说,直接整活!
环境
- python3.9+ (都2024了,至少也要3.9了吧)
- 依赖库 pygame gym pytorch numpy (pytorch带不带cuda,本项目都能跑,代码中自动适配了cpu和cuda)
- Windows 系统,不然可视化不方便哈(代码支持在linux训练)
介绍一下用到的算法–SAC
用过SAC(soft-actor-critic)算法的同学,可以直接跳过本节,看后续。
SAC(Soft Actor-Critic)算法是一种基于策略的强化学习算法,属于Actor-Critic方法的一种变体。它由Tuomas Haarnoja等人于2018年提出,旨在解决连续动作空间中的强化学习问题。SAC算法结合了最大熵强化学习(Maximum Entropy Reinforcement Learning)和确定性策略梯度(Deterministic Policy Gradient)的思想,通过最大化策略的期望回报和策略的熵来实现探索和利用的平衡。
算法简介
SAC算法的核心思想是在传统的Actor-Critic框架中引入最大熵目标,即在优化策略的同时,鼓励策略产生多样化的动作,以增加探索。这通过在目标函数中添加策略的熵项来实现,使得策略不仅追求高回报,还追求高不确定性。
SAC算法的主要组成部分包括:
-
Actor(策略网络):负责根据当前状态选择动作。在SAC中,Actor输出动作的概率分布,而不是确定性动作。
-
Critic(价值网络):评估在给定状态下采取某个动作的价值。SAC通常使用两个独立的Critic网络来减少估计误差。
-
目标函数:SAC的目标函数是期望回报和策略熵的加权和,即最大化以下目标:
[ J(\pi) = \mathbb{E}{s_t \sim \rho\pi, a_t \sim \pi}[r(s_t, a_t) + \alpha \pazocal{H}(\pi(\cdot|s_t))] ]
其中,( r(s_t, a_t) ) 是奖励函数,(\rho_\pi) 是策略(\pi)下的状态分布,(\pazocal{H}(\pi(\cdot|s_t))) 是策略在状态(s_t)下的熵,(\alpha) 是温度参数,控制熵的相对重要性。 -
温度参数(\alpha):用于平衡期望回报和熵的权重。在训练过程中,SAC会自动调整这个参数,以确保策略既有足够的探索性,又能获得高回报。
算法作用
SAC算法的主要作用和优势包括:
-
连续动作空间:SAC专门设计用于处理连续动作空间,这在机器人控制、游戏AI等领域非常常见。
-
稳定性和样本效率:SAC算法通常表现出较高的稳定性和样本效率,这得益于其基于策略梯度的方法和最大熵目标的结合。
-
探索与利用的平衡:通过最大化策略的熵,SAC能够在探索未知状态和利用已知信息之间找到平衡,这有助于避免局部最优并提高学习效率。
-
自适应温度调节:SAC算法中的温度参数(\alpha)可以自适应调节,这有助于算法在不同的环境中自动找到合适的探索-利用平衡点。
-
多目标优化:SAC通过同时优化期望回报和策略熵,实现了多目标优化,这有助于提高策略的鲁棒性和泛化能力。
-
易于并行化:SAC算法可以很容易地并行化,通过在多个环境中同时运行策略来收集数据,加快学习速度。
应用领域
SAC算法在多个领域都有应用,包括但不限于:
- 机器人控制:用于机器人的运动规划和控制,如四足机器人行走、机械臂操作等。
- 游戏AI:在需要连续动作空间决策的游戏中,如模拟赛车、模拟飞行等。
- 自动驾驶:用于车辆的轨迹规划和控制。
- 金融交易:在自动交易系统中,用于生成交易策略。
总之,SAC算法是一种高效、稳定的强化学习算法,特别适合处理连续动作空间的复杂决策问题。通过结合最大熵强化学习和Actor-Critic框架,SAC在多个应用领域都展现出了优异的性能。
项目逻辑
代码逻辑主要分为三个部分,网络模型,环境设计,训练参数
游戏在二维平面上,我们可以将游戏环境视为一个图像信息(环境信息二维矩阵,加上实体类型可以视为一个灰度图信息),于是可以锁定模型为卷积网络,更为适合处理空间上的变化情况。以下是网络模型的核心代码
import torch
import torch.nn as nn
class ConvActorCritic(nn.Module):
def __init__(self, input_channels, output_dim, grid_size, lr=1e-3, weight_decay=1e-5):
super(ConvActorCritic, self).__init__()
self.conv1 = 8
self.conv2 = 16
self.feature_extractor = nn.Sequential(
nn.Conv2d(input_channels, self.conv1, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Flatten()
)
reduced_grid_size = grid_size // 1 # 下采样后,特征边缘缩小2倍
self.actor = nn.Sequential(
nn.Linear(self.conv1 * reduced_grid_size * reduced_grid_size, output_dim),
nn.Softmax(dim=-1)
)
self.critic = nn.Linear(self.conv1 * reduced_grid_size * reduced_grid_size, 1)
self.optimizer = torch.optim.Adam(self.parameters(), lr, weight_decay=weight_decay)
def forward(self, x):
features = self.feature_extractor(x)
action_probs = self.actor(features)
value = self.critic(features).squeeze(-1)
return action_probs, value
可以看到这里使用了一个非常小的卷积参数,只用了一层卷积网络,这将会训练收敛得比较快(因为初始任务在一个10X10的网格中,任务比较简单)
这里采用了Relu函数作为激活函数,帮助网络快速收敛
这里有两个网络,一个是actor,一个是critic,顾名思义,一个是用于行动的网络,一个是用于评估的网络,这里采用了forward,向前传播的算法,返回动作和评估价值。
环境的设计主要䢍两部,一部分是强化学习训练需要用到的奖励函数设计逻辑,一部分是游戏可视化的逻辑。先讲一下游戏可视化吧,这一部分代码也是为了方便直观得感受训练的效果。
def render(self, mode='human', fps=1000):
if not hasattr(self, 'screen'):
pygame.init()
self.cell_size = 20
self.screen_width = self.grid_size * self.cell_size
self.screen_height = self.grid_size * self.cell_size
self.screen = pygame.display.set_mode((self.screen_width, self.screen_height))
self.clock = pygame.time.Clock()
for event in pygame.event.get():
if event.type == pygame.QUIT:
pygame.quit()
return
self.screen.fill((0, 0, 0)) # 背景设为黑色
# 绘制网格
for x in range(0, self.screen_width, self.cell_size):
for y in range(0, self.screen_height, self.cell_size):
rect = pygame.Rect(x, y, self.cell_size, self.cell_size)
pygame.draw.rect(self.screen, (50, 50, 50), rect, 1) # 网格颜色较深
# 绘制食物
food_x, food_y = self.food_pos
pygame.draw.rect(self.screen, (0, 255, 0),
(food_x * self.cell_size, food_y * self.cell_size, self.cell_size, self.cell_size))
# 绘制蛇
# 蛇头
head_color = (255, 0, 0) # 蛇头颜色,例如红色
head_x, head_y = self.snake[0]
head_rect = (head_x * self.cell_size, head_y * self.cell_size, self.cell_size, self.cell_size)
self._draw_head(self.screen, head_rect, head_color)
# 蛇身
for segment in self.snake[1:]:
x, y = segment
pygame.draw.rect(self.screen, (255, 0, 0),
(x * self.cell_size, y * self.cell_size, self.cell_size, self.cell_size))
# 结束标志
if self.game_over:
if len(self.snake) >= self.end_score:
string = "YOU ARE WIN!"
else:
string = "YOU ARE LOST!"
if mode == 'human':
print(string)
# 显示"WIN!"
font = pygame.font.Font(None, 36)
text = font.render(string, True, (255, 255, 255))
text_rect = text.get_rect(center=(self.screen_width // 2, self.screen_height // 2))
self.screen.blit(text, text_rect)
# 停止游戏更新,但保持渲染
while True:
for event in pygame.event.get():
if event.type == pygame.QUIT:
pygame.quit()
return
self.clock.tick(fps) # 控制帧率
pygame.display.flip()
else:
pygame.display.flip()
self.clock.tick(fps) # 控制帧率
def _draw_head(self, screen, rect, color):
# 绘制圆形蛇头
pygame.draw.circle(screen, color, (rect[0] + self.cell_size // 2, rect[1] + self.cell_size // 2),
self.cell_size // 2)
# 根据蛇头方向绘制眼睛
eye_color = (0, 0, 0) # 眼睛颜色,例如黑色
if self.current_direction[0] > 0: # 向右
eye_pos = (rect[0] + self.cell_size // 4, rect[1] + self.cell_size // 4)
elif self.current_direction[0] < 0: # 向左
eye_pos = (rect[0] + 3 * self.cell_size // 4, rect[1] + self.cell_size // 4)
elif self.current_direction[1] > 0: # 向下
eye_pos = (rect[0] + self.cell_size // 4, rect[1] + 3 * self.cell_size // 4)
else: # 向上
eye_pos = (rect[0] + 3 * self.cell_size // 4, rect[1] + 3 * self.cell_size // 4)
# 绘制两个眼睛
pygame.draw.circle(screen, eye_color, eye_pos, self.cell_size // 8)
pygame.draw.circle(screen, eye_color, (eye_pos[0], eye_pos[1] + self.cell_size // 2), self.cell_size // 8)
代码中绘制了游戏的网格背景和食物以及我们的代理蛇,专门画了一个蛇头,方便观察运动中的方向选择情况。也支持手动玩,不过需要添加手动输入的判断逻辑。
接下来讲一下强化学习训练需要用到的一个重要函数step,这是强化学习中用来迭代agent与环境互动的重要步骤。代码如下:
def step(self, action):
self.reward = 0
info = {}
# 尝试所有可能的动作并选择一个最佳的
best_action = action
best_reward = -float('inf')
#---------------------------------预训练结束后注释本区代码------------------------
for trial_action in range(3):
trial_direction = self._get_direction(trial_action)
trial_head = self.snake[0] + trial_direction
if self._is_safe(trial_head) and self._bfs_safe_path(trial_head):
trial_reward = self._calculate_potential_reward(trial_head)
if trial_reward > best_reward:
best_reward = trial_reward
best_action = trial_action
# 执行最佳动作
self.current_direction = self._get_direction(best_action)
new_head = self.snake[0] + self.current_direction
# 检查新的蛇头位置是否超出边界或者撞到自己
if self._is_collision(new_head) or any(new_head < 0) or any(new_head) >= self.grid_size:
self.game_over = True
self.reward -= 5
return self._get_observation(), self.reward, self.game_over, info
# --------------------------------------------------------------------------------
self._calculate_reward(new_head)
self.last_action = action
done = self._update_snake_and_food(new_head)
return self._get_observation(), self.reward, done, info
这段代码是预训练的代码,预训练成功后,正式训练的时候注释掉备注中的代码即可。
代码中设置了一个最佳动作和最佳奖励,这是为了通过动作空间选择评估在预训练的时候获得一些比较不错的动作策略,后面的训练可以直接基于动作策略学习,能加快训练收敛。所以预训练的设计非常关键,然后是在训练过程中,奖励函数的设置会引导AI学习我们期望的动作,最后通过验证就完成这个学习的过程了。
下面就是训练的代码了
if __name__ == '__main__':
# 检查是否有可用的CUDA设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = r"./save_model"
log_dir = r"./logs"
grid_size = 20 # 或其他你的环境设置的值
end_score = 60 # 游戏结束条件
env = SnakeEnv(grid_size=grid_size, end_score=end_score)
max_episodes = 10000000
max_steps = 2000
eval_freq = 1000 # 评估评率
num_eval_episodes = 50 # 评估周期
save_interval = 500 # 固定保存频率
lr = 1e-4
weight_decay = 1e-5
gamma = 0.99
alpha = 0.2
beta = 1
input_channels = 3 # Assuming the observation space is a single-channel image
output_dim = env.action_space.n
agent = ConvActorCritic(input_channels, output_dim, grid_size, lr, weight_decay).to(device)
if not os.path.exists(model_path):
os.makedirs(model_path)
load_model_path = f"{model_path}\\model_2500.pth"
if os.path.exists(load_model_path):
agent.load_model(agent.to(device), filename=load_model_path)
print(f"Loaded model from {load_model_path}")
else:
print("No model to load. Starting a new training session.")
if not os.path.exists(log_dir):
os.makedirs(log_dir)
writer = SummaryWriter(log_dir)
best_reward = float('-inf') # 初始化最佳奖励值为负无穷
for episode in range(2501, max_episodes):
state = env.reset()
state = torch.tensor(state, dtype=torch.float).unsqueeze(0).to(device) # # 为了输入卷积,增加一个维度(卷积4维)
# print(state)
step_count = 0
episode_reward = 0
done = False
while not done and step_count < max_steps:
action_probs, value = agent(state)
dist = torch.distributions.Categorical(action_probs)
action = dist.sample()
next_state, reward, done, _ = env.step(action.item())
next_state = torch.tensor(next_state, dtype=torch.float).unsqueeze(0).to(device)
# 智能体根据奖励函数更新策略
agent.update(state, action, reward, next_state, done, gamma)
episode_reward += reward
state = next_state.to(device)
if done:
break
step_count += 1
# env.render(mode="train",fps=100) # 运行可视化
writer.add_scalar('Reward', round(episode_reward, 5), episode)
writer.add_scalar('Steps', step_count, episode) # 记录步长
print(f"Episode {episode}, Reward: {round(episode_reward, 3)}, step_count: {step_count}")
if episode % eval_freq == 0 and episode != 0:
avg_reward, rewards = evaluate_model(agent, env, device, num_eval_episodes)
# 更新最佳模型
if avg_reward > best_reward:
best_reward = avg_reward
agent.save_model(agent.to(device), filename=f"{model_path}\\best_model.pth")
# 早停逻辑
# if early_stopping_criteria(rewards, patience=30): # patience 连续多个周期没有表现提升
# print("Early stopping triggered.")
# break
# 保存周期性模型
if episode % save_interval == 0:
agent.save_model(agent.to(device), filename=f"{model_path}\\model_{episode}.pth")
env.close() # 结束游戏运行
writer.close()
在这里可以设置不同的网格大小,以及任务结束的标志,这里默认是吃到50个食物就认为通关了,这里已经加入了tensorboard,这个训练过程都是可视化的,可以通过tensorboard监视奖励的变化情况。
命令:
tensorboard --logdir=.\logs
验证模型(可视化游戏)
为了查看训练的模型效果,这里提供了一个加载模型和运行游戏的代码
def play_game(model_path):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 初始化环境和模型
grid_size = 20
end_score = 80
env = SnakeEnv(grid_size=grid_size, end_score=end_score)
input_channels = 3
output_dim = env.action_space.n
model = ConvActorCritic(input_channels, output_dim, grid_size).to(device)
# 加载训练好的模型
model = load_model(model, model_path)
state = env.reset()
done = False
while not done:
state = torch.tensor(state, dtype=torch.float).unsqueeze(0).to(device)
with torch.no_grad():
action_probs, _ = model(state)
action = torch.argmax(action_probs).item() # 选择概率最高的动作
state, _, done, _ = env.step(action)
env.render(fps=1000) # 展示游戏动画
if __name__ == "__main__":
model_path = 'save_model/best_model.pth' # 修改为你的模型路径
# model_path = 'save_model/model_500.pth' # 修改为你的模型路径
play_game(model_path)
这里可以选择自己训练的模型,加载测试一下效果,我训练了几千轮后,基本就打通关了游戏。
附图:
这蛇头蛮可爱的吧,哈哈哈哈
这个项目的代码已经放到github了,需要完整项目代码的同学,可以直接跳转去完整项目源码 记得给我star哈,如果star够多,后面会更新更多的学习源码喲!
遇到环境安装问题,可以看我这篇文章
python从爬虫开始(一)——Python3的安装与环境配置以及网络爬虫的手把手教学
第二篇出炉了,欢迎大家阅读
RL强化学习从小白到老鸟(二)——手撕GPT(零基础保姆级教学)