设置 TorchRL
本节将展示如何安装并开始使用 TorchRL。
先决条件
在安装和使用 TorchRL 之前,您需要以下依赖软件包:
- PyTorch:TorchRL 构建于 PyTorch 之上,因此需要预先安装。
- Gymnasium:您需要 Gymnasium 包来导入 RL 环境。截至 2025 年 1 月,最新版本的 Gymnasium 与 TorchRL 不兼容,详见此Git Discussions 页面。因此请安装较早版本
0.29.1。 - PyGame:一个用于电子游戏的 Python 包。TorchRL 需要它来模拟类似 CartPole 的类游戏 RL 环境。
- TensorDict:在 TorchRL 中使用类似字典的数据结构来存储神经网络的输入与输出,可更便捷地操作张量。TensorDict 提供将张量以键值对形式存储的容器。
安装上述依赖包:
!pip install torch tensordict gymnasium==0.29.1 pygame
安装 TorchRL
使用 pip 安装 TorchRL。如果您在个人电脑或服务器上工作,强烈建议在 Conda 环境中安装这些包。
!pip install torchrl
验证安装
安装完成后,测试 TorchRL 是否成功安装。尝试在 Python 交互环境(或笔记本)中导入 torchrl。使用 check_env_specs() 方法检查一个标准环境(例如 CartPole)是否符合 torchrl 的规范:
import torchrl
from torchrl.envs import GymEnv
from torchrl.envs.utils import check_env_specs
check_env_specs(GymEnv("CartPole-v1"))
输出应表明环境已成功创建,且可与 TorchRL 正常工作。
[torchrl][INFO] check_env_specs succeeded!
TorchRL 的关键组件
在构建您的第一个 RL 智能体之前,先来了解 TorchRL 的核心积木块。
环境
TorchRL 提供统一的 API 来对接不同环境。其方法是将各环境的特定功能封装到一组标准化的封装类与函数中。您向封装器传入相应参数,内部便会将命令映射为对具体环境的函数调用。尤其是:
- TorchRL 将环境状态(观测)、动作与奖励转换为 PyTorch 张量对象,可被实现 RL 算法的各模块直接使用。
- 它支持应用预处理与后处理步骤,例如对输入张量进行归一化与缩放,或以特定格式输出张量。
例如,要从 Gymnasium 创建一个环境,可使用 GymEnv 模块:
env = GymEnv("CartPole-v1")
变换
在基础环境上扩展附加功能与变换是常见做法,否则您需要自行实现。例如,您可以向环境添加步数计数模块,而无需自己编写计数器。步数计数器会记录每个回合的步数。transformedEnv 模块可以帮助实现:
from torchrl.envs import GymEnv, StepCounter, TransformedEnv
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())
同样,操作前对张量进行归一化是一种常见的预处理步骤。使用 ObservationNorm 模块可完成归一化。TransformedEnv 函数文档介绍了多种变换方式。
您还可以使用 compose() 参数组合多个变换:
base_env = GymEnv('CartPole-v1', device=device)
env = TransformedEnv(
base_env,
Compose(
ObservationNorm(in_keys=["observation"]),
StepCounter()
)
)
智能体与策略
在 RL 中,智能体依据策略并基于对环境状态的观测来决定动作。智能体的目标是最大化来自环境的累计奖励。当它选择了正确动作(例如在红灯前停车)时会获得奖励。
例如,在自动驾驶场景中,智能体负责驾驶。它的决策(如将车朝某方向转向)基于对环境状态的观测(如交通状况、车辆位置等)与其策略(如避让行人和障碍物、在红灯前停车等)。
最简单的策略是随机选择动作。执行者会从可能动作空间中随机挑选一个。有时会先使用随机策略生成一批初始交互数据,再开始训练模型。
使用 RandomPolicy 模块可创建随机策略。该随机策略接收 action_spec 参数,用于指定动作空间。在下例中,动作空间为 -1 到 +1 之间的(连续)数值。因此,随机策略会选择一个介于 -1 与 +1 之间的随机数来表示动作。
import torchrl
import torch
from tensordict import TensorDict
from torchrl.data.tensor_specs import Bounded
action_spec = Bounded(-torch.ones(1), torch.ones(1))
actor = torchrl.envs.utils.RandomPolicy(action_spec=action_spec)
td = actor(TensorDict({}, batch_size=[]))
print(td.get("action"))
输出应为一个张量,如下所示:
tensor([0.9258])
再次运行执行者,检查策略是否确为随机:
td = actor(TensorDict({}, batch_size=[]))
print(td.get("action"))
应会输出不同的随机张量。
如果您需要 RL 概念速成,可在 DataCamp 查看 Python 强化学习技能路径!
用 TorchRL 构建您的第一个 RL 智能体
本节将演示如何使用 TorchRL 实现一个简单的强化学习智能体。
开始之前,在 Python 中导入所需软件包:
time用于计量训练时长。GymEnv、StepCounter与TransformedEnv用于操作 Gymnasium 环境。MLP用于创建简单的 MLP(多层感知机)神经网络。EGreedyModule用于平衡对环境的探索与对最优已知策略的利用。QValueModule与DQNLoss用于实现深度 Q 学习算法。SoftUpdate用于更新神经网络。SyncDataCollector用于收集智能体交互数据。ReplayBuffer用于存储智能体交互数据。Adam用于反向传播。matplotlib用于可视化训练进度。torchrl_logger用于记录训练过程。
import time
import matplotlib.pyplot as plt
from torchrl.envs import GymEnv, StepCounter, TransformedEnv
from tensordict.nn import TensorDictModule as TensorDict, TensorDictSequential as Seq
from torchrl.modules import EGreedyModule, MLP, QValueModule
from torchrl.objectives import DQNLoss, SoftUpdate
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer
from torch.optim import Adam
from torchrl._utils import logger as torchrl_logger
步骤 1:定义环境
在本示例中,我们要解决 CartPole 环境。将该环境与计步器一同从 Gymnasium 导入,以记录训练步数:
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())
为 Python 与 RL 环境设定随机种子,以便在多次训练中复现实验结果。
torch.manual_seed(0)
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())
env.set_seed(0)
定义用于训练的参数与超参数:
INIT_RAND_STEPS:在采用策略前,智能体以随机方式行动的步数。这些初始步用于收集初始数据以启动策略训练。FRAMES_PER_BATCH:每个训练批次中的数据点数量(每次交互/时间步一个)。OPTIM_STEPS:在执行一次反向传播前累计损失的步数。EPS_0:初始 epsilon(探索系数)。BUFFER_LEN:回放缓冲区大小。ALPHA:学习率。TARGET_UPDATE_EPS:使用软更新模块更新目标网络时的衰减因子。REPLAY_BUFFER_SAMPLE:每次训练迭代从回放缓冲区中随机采样的大小。回放缓冲区存储智能体在各时间步与环境交互的结果。LOG_EVERY:每隔多少步打印一次训练进度。MLP_SIZE:神经网络规模。
INIT_RAND_STEPS = 5000
FRAMES_PER_BATCH = 100
OPTIM_STEPS = 10
EPS_0 = 0.5
BUFFER_LEN = 100_000
ALPHA = 0.05
TARGET_UPDATE_EPS = 0.95
REPLAY_BUFFER_SAMPLE = 128
LOG_EVERY = 1000
MLP_SIZE = 64
步骤 2:创建策略
定义一个简单的神经网络以实现策略:
- 定义 MLP(神经网络)。给定一个观测,MLP 输出动作值,实质上扮演 Q 函数的角色。
- 使用
TensorDict模块基于上述 MLP 定义张量字典。该字典将环境的状态观测(字典键)映射到对应状态的动作概率值(字典值)。 - 使用QValueModule 实现 Q 函数。给定包含动作值的张量,它返回对应最高动作值的动作,实现贪心(利用)策略。
- 将 MLP 的张量字典与
QValueModule进行组合(使用 TensorDictSequential 模块),得到策略。
value_mlp = MLP(out_features=env.action_spec.shape[-1], num_cells=[MLP_SIZE, MLP_SIZE])
value_net = TensorDict(value_mlp, in_keys=["observation"], out_keys=["action_value"])
policy = Seq(value_net, QValueModule(spec=env.action_spec))
- 使用 EGreedyModule 定义探索函数——其输入为环境的动作空间、回放缓冲区长度与探索系数 epsilon。将该探索函数与上面定义的策略(通过 TensorDictSequential 模块)串联,得到最终策略:
exploration_module = EGreedyModule(
env.action_spec, annealing_num_steps=BUFFER_LEN, eps_init=EPS_0
)
policy_explore = Seq(policy, exploration_module)
步骤 3:训练智能体
训练智能体的第一步是收集其与环境交互产生的数据。使用 SyncDataCollector 构建一个收集器,用于执行策略并收集交互结果:
collector = SyncDataCollector(
env,
policy_explore,
frames_per_batch=FRAMES_PER_BATCH,
total_frames=-1,
init_random_frames=INIT_RAND_STEPS,
)
创建回放缓冲区以存储交互结果:
rb = ReplayBuffer(storage=LazyTensorStorage(BUFFER_LEN))
您还需要声明与训练相关的模块,例如损失函数(使用基于 DQN 的损失计算)、优化器(经典 Adam 算法)与更新器函数(用于更新神经网络)。这些均基于 TorchRL 的预定义模块:
loss = DQNLoss(value_network=policy, action_space=env.action_spec, delay_value=True)
optim = Adam(loss.parameters(), lr=ALPHA)
updater = SoftUpdate(loss, eps=TARGET_UPDATE_EPS)
初始化用于记录总步数与回合数、每回合成功步数以及执行时间的计数器:
total_count = 0
total_episodes = 0
t0 = time.time()
success_steps = []
训练函数由 2 层 for 循环组成:
- 最外层循环执行策略,并将智能体的交互结果追加到回放缓冲区。
- 内层循环将训练样本拆分为批次,并按以下方式处理每个批次:
- 从回放缓冲区中随机抽样。
- 使用损失函数计算损失。
- 运行优化器与更新器。
- 更新计数器。
每个步骤都调用 TorchRL 的预构建模块,无需从零编码。以下代码展示了如何实现这些步骤:
for i, data in enumerate(collector):
rb.extend(data)
max_length = rb[:]["next", "step_count"].max()
if len(rb) > INIT_RAND_STEPS:
for _ in range(OPTIM_STEPS):
sample = rb.sample(REPLAY_BUFFER_SAMPLE)
loss_vals = loss(sample)
loss_vals["loss"].backward()
optim.step()
optim.zero_grad()
# Update exploration factor
exploration_module.step(data.numel())
# Update target params
updater.step()
total_count += data.numel()
total_episodes += data["next", "done"].sum()
success_steps.append(max_length)
步骤 4:评估智能体
上一节的循环会持续训练策略。我们需要设定性能评估标准,并决定在何时认为训练成功。同时也希望输出训练进度。
使用 TorchRL 的 logger 周期性打印训练进度:
if total_count > 0 and total_count % LOG_EVERY == 0:
torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")
我们使用智能体在最近一个回合中达到的最大步数来判定训练是否成功。CartPole-v1 环境将每回合的最大步数与总奖励上限设为 500。通常当策略达到超过 475 步时即可视为成功:
if max_length > 475:
print("TRAINING COMPLETE")
break
以上两段代码应追加到(前文所示)训练循环的末尾。下方代码展示了包含评估逻辑的完整训练循环:
for i, data in enumerate(collector):
# Write data in replay buffer
rb.extend(data)
max_length = rb[:]["next", "step_count"].max()
if len(rb) > INIT_RAND_STEPS:
for _ in range(OPTIM_STEPS):
sample = rb.sample(REPLAY_BUFFER_SAMPLE)
loss_vals = loss(sample)
loss_vals["loss"].backward()
optim.step()
optim.zero_grad()
# Update exploration factor
exploration_module.step(data.numel())
# Update target params
updater.step()
total_count += data.numel()
total_episodes += data["next", "done"].sum()
success_steps.append(max_length)
if total_count > 0 and total_count % LOG_EVERY == 0:
torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")
if max_length > 475:
print("TRAINING COMPLETE")
break
最后,在训练与评估完成后,打印总训练时间并绘制训练进度:
t1 = time.time()
torchrl_logger.info(
f"solved after {total_count} steps, {total_episodes} episodes and in {t1-t0}s."
)
def plot_steps():
plt.plot(success_steps)
plt.title('Successful steps over training episodes')
plt.xlabel('Training episodes')
plt.ylabel('Steps')
plt.show()
plot_steps()
您可以在这本 DataLab 工作簿中找到并运行使用 TorchRL 实现 DQN 的完整代码。
在本节中,我们使用简单的 DQN 算法构建并训练了一个 RL 智能体。下一节将探讨如何利用 TorchRL 的预构建模块实现更复杂的算法,如 PPO。
探索 TorchRL 中的预构建算法
如前文所述,TorchRL 自带一些预构建算法。下面看看它们以及其工作方式。
支持的算法
TorchRL 为许多常见的深度强化学习算法提供了预构建模块,例如:
- 深度 Q 网络(DQN)
- 深度确定性策略梯度(DDPG)
- 软行动者-评论家(SAC)
- 随机集成双重 Q 学习(REDQ)
- CrossQ
- 隐式 Q 学习(IQL)
- 连续 Q 学习(CQL)
- 生成式对抗模仿学习(GAIL)
- 决策 Transformer(DT)
- 双延迟 DDPG(TD3)
- 行动者-优势评论家(A2C)
- 近端策略优化(PPO)
- REINFORCE
- 等等
这让您能高效地尝试不同类型的算法,并比较它们在解决特定问题时的表现。
示例:TorchRL 中的 PPO
近端策略优化使用裁剪的替代目标函数,以实现兼顾探索与利用的平滑训练过程。
本节演示如何使用 TorchRL 中预构建的 PPO 模块。与先前的实现类似,先导入所需模块。除此前实现所需的模块(如 ReplayBuffer、SyncDataCollector 等)外,PPO 还需要以下额外包:
- ProbabilisticActor,用于以随机方式选择动作。随机选择动作(而非总是遵循策略最大化奖励)有助于探索环境,发现在训练中更优的路径。
- OneHotCategorical,用于对表示动作对数概率的张量生成独热编码。
- ValueOperator,用于基于实现价值函数的神经网络构建 TorchRL 模块。
- GAE,用于实现广义优势估计。PPO 使用优势函数作为价值函数的代理。上面的价值模块是 GAE 的输入。
- ClipPPOLoss,用于构建实现 PPO 的裁剪目标函数。
import torch
from torch import nn
from torchrl.envs import Compose, ObservationNorm, DoubleToFloat, StepCounter, TransformedEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.utils import check_env_specs
from torchrl.modules import ProbabilisticActor, OneHotCategorical, ValueOperator
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.storages import LazyTensorStorage
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.objectives.value import GAE
from torchrl.objectives import ClipPPOLoss
from tensordict.nn import TensorDictModule
声明训练所需的参数与超参数:
FRAMES_PER_BATCH:每个批次中包含的帧数(时间步),对应于智能体与环境交互产生的数据。SUB_BATCH_SIZE:每个子批次的时间步数。PPO 训练中,每个批次会被划分为多个子批次。TOTAL_FRAMES:训练运行的总时间步数。GAMMA:用于折扣未来时间步奖励价值的折扣因子。CLIP_EPSILON:定义每次训练迭代中策略可变动的信任域。通过裁剪更新后的策略,使新旧概率之比落在一定范围内实现。ALPHA:学习率。ENTROPY_EPS:用于控制探索与利用比例的超参数。OPTIM_STEPS:对每个子批数据运行优化器的次数。LOG_EVERY:每隔多少步评估策略并打印奖励。
FRAMES_PER_BATCH = 1024
TOTAL_FRAMES = 1048576
GAMMA = 0.99
LAMBDA = 0.95
CLIP_EPSILON = 0.2
ALPHA = 1e-4
ENTROPY_EPS = 5e-4
SUB_BATCH_SIZE = 64
OPTIM_STEPS = 8
LOG_EVERY = 16
从 Gymnasium 导入基础的 CartPole 环境:
device="cpu"
base_env = GymEnv('CartPole-v1', device=device)
声明 TorchRL 环境:导入基础环境并添加对观测张量归一化与计步的模块:
env = TransformedEnv(
base_env,
Compose(
ObservationNorm(in_keys=["observation"]),
DoubleToFloat(),
StepCounter()
)
)
初始化环境并设置随机数种子:
env.transform[0].init_stats(1024)
torch.manual_seed(0)
env.set_seed(0)
check_env_specs(env)
将行动者声明为一个含 1 个隐藏层、16 个权重的神经网络。其输入为环境观测,输出为动作空间中各动作的可能性。
actor_net = nn.Sequential(
nn.Linear(env.observation_spec["observation"].shape[-1], 32, device=device),
nn.ReLU(),
nn.Linear(32, 32, device=device),
nn.ReLU(),
nn.Linear(32, env.action_spec.shape[-1], device=device),
nn.ReLU()
)
基于上述神经网络创建一个张量字典,其键值将状态映射到动作概率。
actor_module = TensorDictModule(actor_net, in_keys=["observation"], out_keys=["logits"])
在训练 PPO(及许多其他 RL 算法)时,您并不总是选择能在下一步带来最高奖励的动作。引入一定的随机性有助于智能体探索环境,发现可带来更高长期回报的路径。创建一个概率行动者来实现这一点:
actor = ProbabilisticActor(
module = actor_module,
spec = env.action_spec,
in_keys = ["logits"],
distribution_class = OneHotCategorical,
return_log_prob = True
)
创建实现价值函数的网络。该网络含有 16 个权重。输入为环境状态(观测),输出为该状态的期望价值。
value_net = nn.Sequential(
nn.Linear(env.observation_spec["observation"].shape[-1], 16, device=device),
nn.ReLU(),
nn.Linear(16, 1, device=device),
nn.ReLU()
)
使用 ValueOperator() 将上述价值函数(实现于神经网络)封装为 TorchRL 模块,使其可与其他 TorchRL 组件对接。
value_module = ValueOperator(
module = value_net,
in_keys = ["observation"]
)
创建回放缓冲区,用于存储智能体与环境交互的结果:
replay_buffer = ReplayBuffer(
storage = LazyTensorStorage(max_size=FRAMES_PER_BATCH),
sampler = SamplerWithoutReplacement()
)
创建数据收集器,在环境中运行策略并收集每个时间步(帧)的交互结果:
collector = SyncDataCollector(
env,
actor,
frames_per_batch = FRAMES_PER_BATCH,
total_frames = TOTAL_FRAMES,
split_trajs = True,
reset_at_each_iter = True,
device=device
)
使用 GAE 模块实现 PPO 的优势函数。优势函数基于价值函数。
advantage_module = GAE(
gamma = GAMMA,
lmbda = GAMMA,
value_network = value_module,
average_gae = True
)
使用内置的 ClipPPOLoss 模块按 PPO 算法实现损失函数:
loss_module = ClipPPOLoss(
actor_network = actor,
critic_network = value_module,
clip_epsilon = CLIP_EPSILON,
entropy_bonus = bool(ENTROPY_EPS),
entropy_coef = ENTROPY_EPS
)
声明 Adam 优化器:
optim = torch.optim.Adam(loss_module.parameters(), lr=ALPHA)
创建调度器,以在训练过程中逐步降低学习率:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, TOTAL_FRAMES // FRAMES_PER_BATCH)
使用上述组件构建训练循环。训练函数包含三层 for 循环:
- 最外层循环遍历智能体与环境交互所收集的数据。参数
TOTAL_FRAMES决定智能体运行的时间步数。 - 中间层循环执行以下操作:
- 按固定次数(通常 5 到 10 次,由
OPTIM_STEPS超参数定义)运行最内层循环。 - 在每次训练迭代后更新学习率。
- 为每个时间步计算优势函数的值。
- 周期性评估策略并打印奖励。
- 最内层循环运行训练过程:
- 从回放缓冲区中采样一个训练批次。
- 使用损失模块计算损失。
- 对损失执行反向传播。
- 使用优化器进行梯度下降。
以下代码实现了训练循环:
for i, tensordict_data in enumerate(collector):
for _ in range(OPTIM_STEPS):
advantage_module(tensordict_data)
replay_buffer.extend(tensordict_data.reshape(-1).cpu())
for _ in range(FRAMES_PER_BATCH // SUB_BATCH_SIZE):
data = replay_buffer.sample(SUB_BATCH_SIZE)
loss = loss_module(data.to(device))
loss_value = loss["loss_objective"] + loss["loss_critic"] + loss["loss_entropy"]
loss_value.backward()
optim.step()
optim.zero_grad()
scheduler.step()
if i % LOG_EVERY == 0:
with torch.no_grad():
rollout = env.rollout(FRAMES_PER_BATCH, actor)
reward_eval = rollout["next","reward"].sum()
#print(rollout["next","reward"].sum())
print(reward_eval)
rewards.append(reward_eval)
del rollout
这本DataLab 工作簿包含了上述使用 TorchRL 中的 PPO 算法训练 RL 智能体的代码。您可将其作为起点微调参数,提升智能体性能。
定制算法
TorchRL 采用模块化设计,便于定制 RL 解决方案。框架负责不同组件之间的对接。您可以随意组合环境、不同策略、回放缓冲区与 RL 算法,以测试并比较其性能。得益于这种模块化:
- 您可以在无需重写大量程序的情况下替换不同组件。
- 还可独立修改各组件以创建定制化解决方案。
- 您可以尝试不同模块,构建优化的解决方案。
TorchRL 中的一些定制示例包括:
- 可在训练循环中加入学习率调度器、日志记录器与自定义指标。
- 通过调整回放缓冲区长度来控制训练时长。
- 可对环境进行定制,例如添加缩放与归一化函数、计步器等。
可视化与调试 RL 训练
前文展示了如何使用 TorchRL 模块来训练 RL 智能体。本节讨论如何监控与可视化训练进度。
监控训练进度
在训练过程中记录各类指标有助于监控进展。像 TensorBoard 这样的工具使得在训练过程中直接记录结果变得简便。例如,TensorBoard 的 SummaryWriter 允许在训练循环中调用以记录指标。以下伪代码展示了方法:
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir="training_logs")
for _ in training_data (): # training loop
loss = …
loss.backprop()
writer.add_scalar(“Loss - “, loss.item(), step)
…
writer.close()
在完成训练与记录后,您可以绘图并可视化结果。例如,使用 tensorboard 命令来可视化训练进度。
tensorboard --logdir=”training_logs”
使用 TorchRL 进行调试
调试对于验证智能体与环境的交互至关重要。第一步是检查动作空间与观测空间。环境文档通常包含其规范说明。例如,CartPole 的文档显示:
- 动作空间包含两个离散值——0(向左推)与 1(向右推)
- 观测空间有 4 个值:
- 小车位置,取值范围
- -4.8 到 +4.8
- 小车速度,任意实数值
- 杆角度,取值范围 -24 ° 到 +24 °
- 杆角速度,任意实数值
检查观测与动作空间是否符合预期值。例如:
print("Observation space: ", env.observation_spec)
print("Action space: ", env.action_spec)
此外,您可以从观测与动作空间中随机采样,并检查其数值:
print("Sample Observation:", base_env.observation_spec.sample().get("observation"))
print("Sample Action:", base_env.action_spec.sample())
输出类似如下示例:
Sample Observation: tensor([-4.5960e+00, 3.4028e+38, 2.2261e-02, 3.4028e+38])
Sample Action: tensor([1, 0])
注意,上述 4 个状态值落在之前指定的范围内。多次重复上述命令,可以看到不同(随机采样的)输出值仍落在规定范围内。
上述两条命令基于直接从 Gymnasium 导入的基础环境。在实践中,我们使用经变换的环境,它会对原始张量值施加变换,如归一化。因此,当您使用经变换的 TorchRL 环境来采样观测时,其数值可能不再与原始 Gymnasium 环境中的范围一致。
例如,从经变换的环境中采样:
print("Sample Observation:", env.observation_spec.sample().get("observation"))
输出类似:
Sample Observation: tensor([-57.5869, nan, 4.5276, nan])
可见小车位置与杆角度的值已超出之前的范围。这是因为张量值已被归一化。此外,由于智能体尚未开始与环境交互,小车速度与杆角速度为 nan。
可视化智能体表现
除绘制训练进度外,渲染环境并直观观察智能体的交互也很有帮助,这有助于理解智能体的实际表现。
最务实的做法是渲染视频。您需要安装以下附加包:
torchvision,用于处理多媒体文件av,用于封装ffmpeg,以便与 PyTorch 对接。
pip install torchvision
pip install av==12.0.0
在训练完成后,按以下准备步骤渲染视频:
- 声明用于存储视频的(相对)路径。
- 初始化记录器,以 CSV 格式保存智能体的交互结果。
- 初始化视频记录器,使用 CSV 日志生成视频。
- 将 TorchRL 环境变换为包含视频记录器的环境。
…
# training loop
# for _ in …
…
…
from torchrl.record import CSVLogger, VideoRecorder
path = "./training_loop"
logger = CSVLogger(exp_name="dqn", log_dir=path, video_format="mp4")
video_recorder = VideoRecorder(logger, tag="video")
record_env = TransformedEnv(
GymEnv("CartPole-v1", from_pixels=True, pixels_only=False), video_recorder
)
在环境中运行(已训练的)策略,并将渲染结果写入视频文件:
record_env.rollout(max_steps=1000, policy=policy)
video_recorder.dump()
程序运行结束后,您将在此前设置的目录路径(上述示例中的 ./training_loop)中找到视频。请注意,使用 TorchRL 实现 DQN 的 DataLab 工作簿并未包含生成视频的代码,因为在在线笔记本中导出视频文件较为困难。请在本地电脑或服务器上运行程序(并添加上述步骤)以录制视频。
使用 TorchRL 的最佳实践
最后,让我们简要讨论一些最佳实践。以下是我的建议。
从简单环境开始
尽管 TorchRL 简化了开发流程,但要在复杂环境中训练表现良好的 RL 智能体仍然具有挑战性。因此,在尝试解决困难问题前,务必先解决诸如 CartPole 等简单环境。
TorchRL 的模块化设计便于尝试不同算法与参数。在实际使用替代算法与定制选项前,您需要先对其工作机理形成直观认识。最好的方式是在简单环境中进行尝试,在那里个别改动的影响更容易通过可视化观察到。
尝试不同超参数
RL 模型的训练表现对学习率、探索率、折扣率以及其他算法特定的超参数(如 PPO 的裁剪比)都很敏感。例如,学习率过高会导致训练不稳定,而过低则收敛太慢。
同样,探索率过高会妨碍训练收敛,而过低又可能使智能体无法发现最优路径。
并不存在确定合适超参数的公式。尽管标准环境有一些经验指南与推荐值,但并不适用于所有问题。因此,必须尝试诸如网格搜索或随机搜索等方法,来确定最佳超参数。
也可以使用 Weights & Biases Sweeps 或 Optuna 等自动化库,测试多种超参数组合下的训练表现。
充分利用预构建算法
如您在前述示例所见,使用 TorchRL 的预构建模块可以节省大量开发工作。否则从零实现全部功能将耗费显著更多的时间与成本来开发与测试。
由于有许多开发者在使用 TorchRL 模块,它们也在多种场景下被反复验证。因此,这些模块通常比自研模块更少缺陷。对于大多数标准用例,强烈建议优先使用预构建模块。
结语
本文介绍了 TorchRL 的基础概念——一个基于 PyTorch 的 RL 算法实现框架。我们也通过实例演示了如何用 TorchRL 实现较为简单的深度 Q 学习与更复杂的 PPO 算法。
下一步,您可以基于这些程序去尝试其他环境与算法。
如果您希望进一步加深对 RL 基础与实践应用的理解,欢迎查看 Reinforcement Learning with Gymnasium in Python,获取更多基于 Gymnasium 环境的实操体验。

强化学习(RL)被用于解决诸多复杂问题,从训练自动驾驶汽车到训练大语言模型(LLM)生成类人回答。
RL 智能体常借助人类反馈进行训练,即基于人类反馈的强化学习(RLHF)。尽管像 Keras 和 TensorFlow 这样的Python 框架广泛用于企业级深度学习应用,但如今大多数新项目都基于 PyTorch 和 PyTorch Lightning。
TorchRL 是一个基于 PyTorch 的开源库,用于构建 RL 解决方案。本文教程将带您完成 TorchRL 的环境搭建,理解其核心组件,并用它构建一个简单的 RL 智能体。我们还将讨论如何使用 TorchRL 来实现诸如近端策略优化(PPO)等预构建 RL 算法版本。最后,我们会介绍记录与监控 RL 智能体的基本原则。
什么是 TorchRL,为什么要使用它?
RL 算法往往很复杂。您需要创建智能体、计算回报与损失、定义前向与反向传播,并评估智能体的性能。
TorchRL 将许多常用 RL 功能打包为可直接调用的模块。这既方便实现与试验多种算法以解决实际问题,也让研究者在无需从零构建生态组件的前提下更容易提出并实现新算法。
TorchRL 自带大量预构建模块,使 RL 开发更高效。例如:
基于以上功能,TorchRL 已被证明可高效简化多种用例下 RL 解决方案的构建,例如: