首页 > 其他分享 >GPU释放威力:在Gymnasium环境中使用稳定基线3在AMD GPU上训练强化学习代理

GPU释放威力:在Gymnasium环境中使用稳定基线3在AMD GPU上训练强化学习代理

时间:2024-11-07 16:49:24浏览次数:7  
标签:Gymnasium AMD 学习 env time frames GPU model reward

GPU Unleashed: Training Reinforcement Learning Agents with Stable Baselines3 on an AMD GPU in Gymnasium Environment — ROCm Blogs

2024年4月11日 作者: Douglas Jia.

本博客将深入探讨深度强化学习的基本原理,通过一个实用的代码示例,指导您如何利用AMD GPU在Gymnasium环境中训练深度Q网络(DQN)策略。

您可以在 Github repo中找到本博客提到的所有文件。 

介绍

什么是强化学习

强化学习(Reinforcement Learning,简称RL)是一种强大的机器学习范式,其核心思想是训练智能体在动态环境中做出连续决策。与监督学习不同,RL智能体通过与环境的互动,根据所采取的行动接收奖励或惩罚形式的反馈。强化学习的根本目标是让智能体学习一种策略,将状态映射到行动的方式,以增加预期奖励。学习过程包括智能体探索环境、做出决策、观察结果,并根据所获得的反馈调整其策略。这种探索与利用的迭代循环使RL智能体能适应变化的条件,学习复杂的决策策略,最终优化其行为以实现长期目标。强化学习在各个领域有广泛的应用,从机器人技术和游戏到金融和自主系统,智能体可以在复杂和动态的场景中自主学习并改善其决策能力。

深度学习如何发挥作用

在强化学习 (RL) 中,深度学习能够解决由高维状态空间带来的挑战,并提升函数(即状态-行动值函数或策略)近似。这在具有广泛或连续状态空间的环境中尤为重要,尤其是当这些环境由高维输入(如图像)表示时。神经网络凭借其相互连接的神经元层,擅长从原始数据中自动学习层次表示。这种能力对于捕捉决策所需的复杂模式至关重要。深度 RL 算法,如 DQN 或策略梯度方法,利用深度学习的力量来泛化各种状态,使其非常适合传统方法可能难以应对的真实世界应用。深度学习和RL的协同,使得训练出的代理能够在动态和多样化的环境中掌握复杂任务。具体来说,深度强化学习系统的输入通常是来自环境的观测值,目标是学习一个策略(由神经网络表示),该策略输出动作以在一段时间内最大化累积奖励。学习过程涉及根据从环境中收到的奖励形式的反馈来优化策略的参数。

如果你想深入了解深度强化学习,你可以访问 Hugging Face 提供的深度强化学习课程 .

什么是Gymnasium和Stable Baselines3?

想象一下一个供人工智能运动员使用的虚拟操场——这就是Gymnasium!Gymnasium是OpenAI的Gym库的一个维护分支。这个开源工具包提供了各种虚拟环境,从平衡Cartpole机器人到导航Lunar Lander挑战。这里是你的AI代理通过实际操作学习的地方。

但训练这些运动员需要一个有经验且技巧娴熟的教练。这时,Stable Baselines3就登场了。这个包为你提供了现成的强化学习算法,从DQN到策略梯度。你可以把它看作一个装满了经过验证的训练技术的工具箱,随时可以在Gym的多样环境中释放你的AI代理的全部潜力。

本质上,Gymnasium是应用Stable Baselines3提供的深度学习算法来学习和优化策略的环境。为了提高训练过程的效率,我们利用了AMD GPU的强大性能。在下面的代码示例中,我们将展示通过这种方法可以实现的加速效果。

代码示例:训练一个赛车代理

实现环境

我们使用一块 AMD GPU,并安装了 PyTorch 2.0.1 和 ROCm 5.7.0 来实现代码示例(具体使用的 Docker 镜像)。关于支持的操作系统和硬件列表,请参阅此 ROCm 文档页面。我们建议您在具有 AMD GPU 的机器上使用下面的代码块,拉取预装有 ROCm 和 PyTorch 的 Docker 镜像。

docker run -it --ipc=host --network=host --device=/dev/kfd --device=/dev/dri \
           --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
           --name=rl-gym rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1 /bin/bash

要开始这个示例,我们需要安装运行此代码示例所需的 Python 包。

%%capture

!pip install stable-baselines3 gymnasium[all]

Gymnasium 环境

我们将使用 Gymnasium 中具有离散动作空间的 CarRacing-v2 环境。有关此环境的详细信息,请参阅 官方文档. 。Gymnasium 中的 Car Racing 环境是一种模拟环境,旨在训练强化学习代理进行汽车赛车。该环境的主要目标是让代理有效地导航赛道,考虑速度、转向和刹车等因素,同时最大化累积奖励。CarRacing 中的状态空间包括以 RGB 图像表示的视觉观测,为代理提供有关当前环境状态的信息。离散动作空间由 5 个动作组成:`0: 什么都不做`, 1: 左转向2: 右转向3: 加速4: 制动,允许代理控制汽车的运动。奖励函数可以从官方文档 here中找到。
因此,挑战在于学习一种策略,使代理能够有效地驾驶汽车通过赛道,避开障碍物,并根据环境中定义的奖励结构实现最佳性能。
通过下面的代码块,您可以输出动作空间、观测空间和其他环境规格。

import gymnasium as gym
env = gym.make("CarRacing-v2", domain_randomize=True, continuous=False)

print("Action Space:", env.action_space)

# 注意,在输出的观测空间 Box(0, 255, (96, 96, 3), uint8) 中,
# 0 和 255 分别表示该 (96, 96, 3) 数组中数值的下限和上限。

print("Observation Space: ", env.observation_space)

print("Env Spec: ", env.spec)

输出:

Action Space: Discrete(5)
    Observation Space:  Box(0, 255, (96, 96, 3), uint8)
    Env Spec:  EnvSpec(id='CarRacing-v2', entry_point='gymnasium.envs.box2d.car_racing:CarRacing', reward_threshold=900, nondeterministic=False, max_episode_steps=1000, order_enforce=True, autoreset=False, disable_env_checker=False, apply_api_compatibility=False, kwargs={'domain_randomize': True, 'continuous': False}, namespace=None, name='CarRacing', version=2, additional_wrappers=(), vector_entry_point=None)

接下来,我们将通过一个随机做出动作决策的代理向您展示 Gymnasium 环境是如何工作的。代理将做 10 个连续动作。实现每个动作后,环境将返回:

  • 下一次观测(或状态)

  • 执行该动作的奖励

  • 环境是否由于满足预定义条件而终止或截断

  • 有关步骤的信息,例如度量和调试信息。

import gymnasium as gym

# 首先,我们创建一个名为 CarRacing-v2 的环境。
env = gym.make("CarRacing-v2", domain_randomize=True, continuous=False)

# 然后我们重置这个环境
observation, info = env.reset()

for _ in range(10):
    # 从动作空间中随机选择一个动作。
    action = env.action_space.sample()

    # 在环境中执行这个动作并获取
    # next_state, reward, terminated, truncated and info
    observation, reward, terminated, truncated, info = env.step(action)
    print(f"Action taken: {action}, reward: {reward}.")

    # 如果游戏终止(在我们的情况下是着陆、碰撞)或截断(超时),
    # 重置环境。
    if terminated or truncated:
        # 重置环境
        print("Environment is reset")
        observation, info = env.reset()

env.close()

输出:

Action taken: 0, reward: 7.1463768115942035.
Action taken: 4, reward: -0.09999999999999964.
Action taken: 1, reward: -0.09999999999999964.
Action taken: 0, reward: -0.09999999999999964.
Action taken: 4, reward: -0.09999999999999964.
Action taken: 1, reward: -0.09999999999999964.
Action taken: 4, reward: -0.09999999999999964.
Action taken: 4, reward: -0.09999999999999964.
Action taken: 2, reward: -0.09999999999999964.
Action taken: 0, reward: -0.09999999999999964.

根据输出,可以看到大多数随机动作都会导致负奖励。 

使用Stable Baselines3训练智能代理

显然,依赖随机动作不会在这款赛车游戏中取得成功。我们的目标是通过从环境中学习知识来训练深度神经网络,使其能够做出最大化累积奖励的动作选择。

在这个特定的环境中,我们将使用带有卷积神经网络(CNN)的DQN来训练智能代理。通过最小化预测的Q值与目标Q值之间的时间差误差来训练DQN。CNN架构用于处理原始像素输入(游戏帧或观察状态)并提取有用的特征来进行预测。

import gymnasium as gym
from stable_baselines3 import DQN
from stable_baselines3.common.callbacks import EvalCallback

from typing import Callable

# 在我们的实验中,使用逐渐减小学习率的线性调度比使用固定学习率表现更好。
def linear_schedule(initial_value: float) -> Callable[[float], float]:
    """
    线性学习率调度

    :param initial_value: 初始学习率
    :return: 根据剩余进度计算当前学习率的调度
    """
    def func(progress_remaining: float) -> float:
        """
        进度将从1(开始)减少到

        :param progress_remaining:剩余进度
        :return: 当前学习率
        """
        return progress_remaining * initial_value

    return func

要了解以下单元格输出中的每个度量意味着什么,你可以访问 this page 。请注意,我们已经简化了输出以节省页面空间。然而,你仍然可以看到代理逐渐获得环境知识,并对应地在以下显示的输出中增加奖励。

import os

env = gym.make("CarRacing-v2", domain_randomize=True, continuous=False)

# 设置种子以确保可复现性
env.reset(seed=1234)
env.action_space.seed(123)

# 实例化代理并指定Tensorboard日志的目录。
# 过Tensorboard日志,我们可以监控模型的训练进度。
model = DQN('CnnPolicy', env, verbose=1, device="cuda", learning_rate=linear_schedule(0.0001), tensorboard_log="./dqn_carrace_tensorboard_lr0001_1e7/")
# 指定保存最佳模型的目录。
dnq_path = os.path.join('./Saved_Models/dqn_best_model_lr0001_1e7')

# EvalCallback定期评估代理的性能,使用一个独立的测试环境。如果
# 指定了best_model_save_path文件夹,它将保存最佳模型。
eval_env = model.get_env()
eval_callback = EvalCallback(eval_env=eval_env, best_model_save_path=dnq_path,
                             n_eval_episodes=50,
                             eval_freq=100000,verbose=1,
                             deterministic=True, render=False)

# 训练代理
model.learn(total_timesteps=10000000,callback=eval_callback)

# 指定保存最终模型的路径
final_path = os.path.join('./Saved_Models/dqn_final_model_lr0001_1e7')
model.save(final_path)

Using cuda device
Wrapping the env with a `Monitor` wrapper
Wrapping the env in a DummyVecEnv.
Wrapping the env in a VecTransposeImage.
Logging to ./dqn_carrace_tensorboard_lr0001_1e7/DQN_1
----------------------------------
| rollout/            |          |
|    ep_len_mean      | 1e+03    |
|    ep_rew_mean      | -52.4    |
|    exploration_rate | 0.996    |
| time/               |          |
|    episodes         | 4        |
|    fps              | 167      |
|    time_elapsed     | 23       |
|    total_timesteps  | 4000     |
----------------------------------
----------------------------------
| rollout/            |          |
|    ep_len_mean      | 1e+03    |
|    ep_rew_mean      | -52.7    |
|    exploration_rate | 0.992    |
| time/               |          |
|    episodes         | 8        |
|    fps              | 167      |
|    time_elapsed     | 47       |
|    total_timesteps  | 8000     |
----------------------------------
----------------------------------
| rollout/            |          |
|    ep_len_mean      | 1e+03    |
|    ep_rew_mean      | -53.9    |
|    exploration_rate | 0.989    |
| time/               |          |
|    episodes         | 12       |
|    fps              | 167      |
|    time_elapsed     | 71       |
|    total_timesteps  | 12000    |
----------------------------------

...
...
...

----------------------------------
| rollout/            |          |
|    ep_len_mean      | 960      |
|    ep_rew_mean      | 844      |
|    exploration_rate | 0.05     |
| time/               |          |
|    episodes         | 4060     |
|    fps              | 91       |
|    time_elapsed     | 43427    |
|    total_timesteps  | 3966401  |
| train/              |          |
|    learning_rate    | 6.03e-05 |
|    loss             | 0.669    |
|    n_updates        | 979100   |
----------------------------------
----------------------------------
| rollout/            |          |
|    ep_len_mean      | 960      |
|    ep_rew_mean      | 846      |
|    exploration_rate | 0.05     |
| time/               |          |
|    episodes         | 4064     |
|    fps              | 91       |
|    time_elapsed     | 43456    |
|    total_timesteps  | 3970095  |
| train/              |          |
|    learning_rate    | 6.03e-05 |
|    loss             | 0.838    |
|    n_updates        | 980023   |
----------------------------------

在加载TensorBoard文件后 (请参考 this tutorial 进行指导), 查看`eval/mean_reward`下的奖励,你会注意到在完成大约350万次训练步数后,奖励收敛到约870分。这是一个值得注意的成就,尤其是考虑到理论上的最高奖励在1000以下。请注意,黑线和亮线实际上基于同一组数据。黑线是亮线的平滑版本(基于原始数据),可以更好地反映整体趋势。

PNG

你可以在`src`文件夹中找到我们保存的最佳模型,文件名为`best_model.zip`,以及TensorBoard文件。使用这个模型文件,你可以加载并使用它来以最高累计奖励玩游戏。

可视化赛车

现在,让我们深入到令人兴奋的部分! 我们可以在卡通设置中直观地比较训练好的智能体与随机选择动作的效果。

随机动作的智能体

在本节中,我们将让智能体从动作空间中随机选择动作。正如卡通中所展示的,当随机选择动作时,汽车移动缓慢且始终获取负奖励(卡通左下角显示累积奖励)。

%%capture
# 通过以下命令来安装依赖项。如果你已经安装了 ffmpeg,可以跳过这一步。
!sudo apt update -y && sudo apt upgrade -y
!sudo apt install ffmpeg -y
import gymnasium as gym
from matplotlib import rcParams
from matplotlib import pyplot as plt, animation
from IPython.display import HTML
from stable_baselines3 import DQN
rcParams['animation.embed_limit'] = 2**128

# 定义一些用于生成和保存卡通的实用函数。
def create_anim(frames, dpi, fps):
    plt.figure(figsize=(frames[0].shape[1] / dpi, frames[0].shape[0] / dpi), dpi=dpi)
    patch = plt.imshow(frames[0])
    def setup():
        plt.axis('off')
    def animate(i):
        patch.set_data(frames[i])
    anim = animation.FuncAnimation(plt.gcf(), animate, init_func=setup, frames=len(frames), interval=fps)
    return anim

def display_anim(frames, dpi=72, fps=60):
    anim = create_anim(frames, dpi, fps)
    return anim.to_jshtml()

def save_anim(frames, filename, dpi=72, fps=50):
    anim = create_anim(frames, dpi, fps)
    anim.save(filename)

env = gym.make("CarRacing-v2", domain_randomize=True, continuous=False, render_mode='rgb_array')

# 使用随机动作生成一个回合的观察结果。
frames = []
episodes = 1

for episode in range(1, episodes+1):
    state = env.reset()
    terminated = truncated = False
    score = 0
    
    while not (terminated or truncated):
        frames.append(env.render())
        # 注意: 这里我们随机选择一个动作。
        action = env.action_space.sample()
        observation, reward, terminated, truncated, info = env.step(action)
        score += reward
    # 在回合结束时打印出回合数和累积奖励。
    print("Episode:{} Score:{}".format(episode,score))
env.close()
Episode:1 Score:-56.20437956204454
# 内嵌显示卡通。这个智能体采取的是随机动作。
# 由于视频的大小太大,我们删除了输出。
# 你可以查看'videos'文件夹以查看根据下面步骤保存的生成视频。
HTML(display_anim(frames))
# 卡通中的动作基于随机动作。
# 保存卡通。我们在 /videos 文件夹中提供了生成的卡通。
filename = 'CarRacing-v2_random.mp4'
save_anim(frames, filename=filename)

生成的视频如下所示:

https://github.com/ROCm/rocm-blogs/assets/13548555/c836a82e-1ad8-44cf-90dd-fba86f648d45

基于最佳模型预测动作的智能体 

在本节中,我们将让智能体根据训练模型预测的最佳动作进行操作。正如你在卡通动画中看到的那样,使用训练后的智能体,汽车快速加速并巧妙地穿越各种路况。这很了不起,不是吗?

# 根据评估结果加载最佳模型。您可能需要更改模型路径,具体取决于您的最佳模型存储位置。
model_path = "./Saved_Models/dqn_best_model_lr0001_1e7/best_model"
model = DQN.load(model_path, env=env)
Wrapping the env with a `Monitor` wrapper
Wrapping the env in a DummyVecEnv.
Wrapping the env in a VecTransposeImage.
# 生成一集由最佳模型预测的动作的观察结果。

frames = []
episodes = 1

for episode in range(1, episodes+1):
    observation, _ = env.reset()
    terminated = truncated = False
    score = 0
    
    while not (terminated or truncated):
        frames.append(env.render())
        action , _ = model.predict(observation) 
        observation, reward, terminated, truncated, info = env.step(action)
        score += reward
    print("Episode:{} Score:{}".format(episode,score))
env.close()
Episode:1 Score:917.1999999999867
# 内嵌显示卡通动画。此智能体采取随机动作。
# 由于视频尺寸过大,我们删除了输出内容。
# 您可以参考'videos'文件夹查看生成的视频,
# 并通过以下步骤保存。
HTML(display_anim(frames))
# 卡通动画中的动作是使用训练后的最佳模型选择的。
# 保存卡通动画。我们将生成的卡通动画保存在 /videos 文件夹中。
filename = 'CarRacing-v2_best.mp4'
save_anim(frames, filename=filename)

这是生成的视频:

https://github.com/ROCm/rocm-blogs/assets/13548555/1db55b04-afde-4ab5-9bbc-291d160dbcf1

CPU与GPU性能对比

鉴于它们在深度学习模型训练中的并行处理能力,GPU通常显著优于CPU。在下面的代码块中,我们展示了通过利用AMD GPU,性能能够远超CPU。具体的数值差异可能会根据你所使用的硬件而有所不同。

import time
from stable_baselines3 import DQN
import gymnasium as gym

env = gym.make("CarRacing-v2", domain_randomize=True,continuous=False)
env.reset(seed=1234)
env.action_space.seed(123)
t1 = time.time()
model = DQN("CnnPolicy", env, verbose=0, device="cuda")
model.learn(total_timesteps=100000)
print(f"Time with gpu : {time.time()-t1:.2f}s")

env.reset(seed=1234)
env.action_space.seed(123)
t1 = time.time()
model = DQN("CnnPolicy", env, verbose=0, device="cpu")
model.learn(total_timesteps=100000)
print(f"Time with cpu : {time.time()-t1:.2f}s")

env.close()
Time with gpu : 652.19s
Time with cpu : 5837.96s

致谢

我们对这个笔记本的作者Manthan Bhagat表示感谢,他的工作为我们的博客提供了宝贵的参考,尤其是在创建和保存卡通动画的函数代码方面。

标签:Gymnasium,AMD,学习,env,time,frames,GPU,model,reward
From: https://blog.csdn.net/eidolon_foot/article/details/143581395

相关文章

  • 使用AMD GPU进行图像分类的ResNet模型
    ResNetforimageclassificationusingAMDGPUs—ROCmBlogs2024年4月9日,作者:LoganGrado。在这篇博客中,我们演示了如何使用ROCm在AMDGPU上训练一个简单的ResNet模型来进行CIFAR10数据集的图像分类。在AMDGPU上训练ResNet模型非常简单,仅需安装ROCm和适当的PyTorch库,无......
  • CUDA开始的GPU编程 - 第四章:C++封装GPU上的数组
    第四章:C++封装GPU上的数组std::vector的秘密:第二模板参数**你知道吗?**std::vector作为模板类,其实有两个模板参数:std::vector<T,AllocatorT>那为什么我们平时只用了std::vector呢?因为第二个参数默认是std::allocator。也就是std::vector等价于std::vector<T,s......
  • GPU 环境搭建指南:如何在裸机、Docker、K8s 等环境中使用 GPU
    本文主要分享在不同环境,例如裸机、Docker和Kubernetes等环境中如何使用GPU。跳转阅读原文:GPU环境搭建指南:如何在裸机、Docker、K8s等环境中使用GPU1.概述仅以比较常见的NVIDIAGPU举例,系统为Linux,对于其他厂家的GPU设备理论上流程都是一样的。省流:对于裸......
  • **AI的三大支柱:神经网络、大数据与GPU计算的崛起之路**
      每周跟踪AI热点新闻动向和震撼发展想要探索生成式人工智能的前沿进展吗?订阅我们的简报,深入解析最新的技术突破、实际应用案例和未来的趋势。与全球数同行一同,从行业内部的深度分析和实用指南中受益。不要错过这个机会,成为AI领域的领跑者。点击订阅,与未来同行!订阅:https://......
  • [使用整理-1] bamdst:统计覆盖度和测序深度
    BAMDSTBamdst是一个轻量级工具,用于统计bam文件中目标区域的测序深度、覆盖度。必须使用排序后的bam文件,bed文件和输出目录必须优先指定。1.下载安装gitclonehttps://github.com/shiquan/bamdst.gitcdbamdst/make或者使用conda:condainstallxdgene::bamdst......
  • cuda、cudnn、zlib 深度学习GPU必配三件套(Ubuntu)
    跨大版本不推荐,到处是坑、坑、坑~。tensorrt10、cuda12、cudnn9是目前最新的大版本,但是对于一般的老显卡(1050等),太新可能提醒一些错误(主要是tensorrt太新导致的)。为了不折腾,使用如下版本:tensorrt8.6.1、cuda11.8、cudnn8.9.7默认已经安装了英伟达显卡的最新版本驱动。系统Ubuntu22......
  • Ollama简介,以及中文使用指南和AMD图形显卡驱动下载
    Ollama是一个开源的大型语言模型(LLM)服务工具,它旨在简化在本地运行大语言模型的过程,降低使用大语言模型的门槛。Ollama使得开发者、研究人员和爱好者能够在本地环境快速实验、管理和部署最新的大语言模型,包括但不限于如Qwen2、Llama3、Phi3、Gemma2等开源的大型语言模型。Oll......
  • 基于AMD显卡安装Pytorch(小白攻略)
    安装的时候看了很多博客,踩了一些雷,现在把成功下载的流程汇总下来。假设这个时候已经安装好了ubuntu,我安装的是ubuntu22.04.安装rocmLinux®DriversforAMDRadeon™andRadeonPRO™Graphics可以点击上面这个链接,点击ubuntux8664-bit.我选的是带rocm的版本复制这......
  • 深入解析Amdahl定律与Gustafson定律:并行计算的加速之道
    如果你觉得这篇文章对你有帮助,请不要吝惜你的“关注”、“点赞”、“评价”、“收藏”,你的支持永远是我前进的动力~~~个人收藏的技术大会分享PDF文档,欢迎点击下载查看!!!本文将探讨并行计算中的两个重要定律——Amdahl定律和Gustafson定律。通过分析这两个定律的原理、区别及实......
  • Termux安装GUI应用并启用虚拟GPU
    目录1.你知道吗,即使不用容器Termux也能运行桌面应用1.使用VNC远程桌面2.使用Termux-x112.没有桌面总感觉怪怪的,不是吗1.安装桌面环境2.VNC配置3.Termux-x11配置3.使用虚拟GPU驱动1.安装virglrenderer-android2.启动virglrenderer-android3.virglrenderer-android的......