要点
学习如何使用PyTorch和Lightning Fabric(lightning.ai/pages/open-…
强化学习简介
强化学习（RL）是一种机器学习算法，通过让智能代理与环境互动并随时间调整行为以实现某个目标的最大化，从而训练智能代理做出决策。它的灵感来源于人类和动物如何从经验中学习并相应地调整行动。
强化学习在各种应用中取得了极大的成功，包括机器人、自动驾驶汽车、推荐系统和游戏玩家等。其中最著名的例子是DeepMind开发的AI系统AlphaGo。它将强化学习与深度神经网络相结合，击败了世界冠军围棋选手。围棋是一种策略性双人棋盘游戏，起源于2500多年前的中国。共有近2×10^170种可能的合法棋盘位置，游戏在19×19的棋盘上进行。目标是围住比对手更多的领地。
强化学习的组成部分
代理：代理是根据与环境的互动来学习和做决策的实体（例如，AI算法或机器人）。
环境：环境代表代理操作的外部背景或世界。它可以简单到一个二维的井字游戏网格，也可以复杂到真实世界。
状态：状态是环境在某个时间点的快照，代表代理所感知的内容。它为代理提供了做决策所需的必要信息。
动作：动作是代理在给定状态下可以进行的一系列可能的移动或选择。代理的目标是根据对环境的当前理解选择最合适的动作。
奖励：奖励是代理在执行动作后从环境中获得的反馈。它们表明代理在实现目标方面的表现如何。代理的目标是学习一种策略，使得随着时间的推移，累积奖励最大化。**
**
高层次表示如下图所示：
在上图中，时间被离散化并用t表示，代理与环境进行互动。代理从环境中接收到一个观察值，该观察值表示在那个时间点的环境状态。然后，代理根据该状态执行动作，并获得相应的奖励。奖励是一个标量值，表示相对于代理试图实现的特定目标或任务，该动作的好坏程度。
与PyTorch一起进行强化学习之旅
意大利的一家AI公司Orobix与一家视频游戏公司合作，开发了一个强化学习（RL）框架。其目标是提高游戏中非玩家角色（NPC）的赛车表现。这一合作旨在创造更具竞争力和沉浸感的体验。
该框架从零开始构建，以便我们可以完全控制训练循环和分布式训练基础设施。
我们需要对分布式训练、半精度以及代码的每个部分进行手动控制，以使其更具灵活性。Lightning AI（前身为PyTorch Lightning）团队推出的一款新库Fabric，帮助我们实现了这一目标。它让我们在自定义训练循环上具有完全的灵活性，同时抽象了多个设备、分布式和半精度训练。
Fabric-加速强化学习**
**
现在我们将构建并训练一个强化学习代理在CartPole环境中玩耍，其中一个杆通过一个非驱动关节连接到一个沿着无摩擦轨道移动的手推车。这个代理基于近端策略优化（PPO）算法。目标是通过在手推车的左右方向施加力来平衡杆子：
需要做的是
安装以下几个库:**
**
Gymnasium：用于强化学习的标准API，包含各种参考环境。
Fabric：用于加速和分布我们的训练。
所需库列表在这里:github.com/Lightning-A…
环境与智能体相结合
让我们首先了解环境何时与代理程序耦合。主要思想如下图所示：
我们将会创建N+1个进程,称为rank-0, …,rank-N;每个进程都包含环境(可能有多个,上述图中为M+1个,并有多个副本)和实体:它们在同一进程中并发地连接在一起。
让我们首先定义我们的main(…)函数，在其中使用Fabric初始化分布式训练设置。
from lightning.fabric import Fabric
def main(args):
# Initialize Fabric
fabric = Fabric()
rank = fabric.global_rank # The rank of the current process
world_size = fabric.world_size # Number of processes spawned
device = fabric.device
fabric.seed_everything(42) # We seed everything for reproduciability purpose
接下来，我们使用gymnasium创建环境。首先，我们定义一个名为make_env的辅助函数，用于创建单个环境。
def make_env(env_id: str, seed: int, idx: int, capture_video: bool, run_name: Optional[str] = None, prefix: str = ""):
def thunk():
env = gym.make(env_id, render_mode="rgb_array")
env = gym.wrappers.RecordEpisodeStatistics(env)
if capture_video:
if idx == 0 and run_name is not None:
env = gym.wrappers.RecordVideo(
env, os.path.join(run_name, prefix + "_videos" if prefix else "videos"), disable_logger=True
)
env.action_space.seed(seed)
env.observation_space.seed(seed)
return env
return thunk
现在，我们将使用刚刚创建的make_env函数，通过SyncVectorEnv对象创建一个并行同步环境池。
import gymnasium as gym
# given an initial seed of 42 and 4 environments per rank, then
# rank-0 will seed the environments with --> 42, 43, 44, 45
# rank-1 will seed the environments with --> 46, 47, 48, 49
# and so on
rl_environment = gym.vector.SyncVectorEnv([
make_env(
args.env_id,
args.seed + rank * args.num_envs + i,
rank,
args.capture_video,
logger.log_dir,
"train"
)
for i in range(args.num_envs)
])
在最后一步中，我们创建代理程序、优化器，并将其与Fabric集成，以实现更快的训练。
agent = PPOLightningAgent(
rl_environment,
act_fun=args.activation_function,
vf_coef=args.vf_coef,
ent_coef=args.ent_coef,
clip_coef=args.clip_coef,
clip_vloss=args.clip_vloss,
ortho_init=args.ortho_init,
normalize_advantages=args.normalize_advantages,
)
optimizer = agent.configure_optimizers(args.learning_rate)
# accelerated training with Fabric
agent, optimizer = fabric.setup(agent, optimizer)
我们定义了PPOLightningAgent，它是一个LightningModule，也是一个Actor-Critic代理程序。在Actor-Critic代理程序中，演员（actor）在给定状态下提出一组可能的动作，评论家（critic）评估演员所采取的动作。
现在我们需要创建一个“无限”循环，在其中执行以下操作：
1. 代理程序与环境进行交互，收集经验。一个单独的经验由以下内容组成：
done是一个标识游戏是否完成的布尔值
代理程序会收集经验，直到游戏终止或达到预定义的步数。
2. 根据收集到的经验，训练代理程序以改进其行为。
3. 重复步骤1，直到达到收敛或与环境的最大交互次数。
经验收集循环:
import torch
with fabric.device:
# with fabric.device is only supported in PyTorch 2.x+
obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape)
actions = torch.zeros((args.num_steps, args.num_envs) + envs.single_action_space.shape)
rewards = torch.zeros((args.num_steps, args.num_envs))
dones = torch.zeros((args.num_steps, args.num_envs))
# Log-probabilities of the action played are needed later on during the training phase
logprobs = torch.zeros((args.num_steps, args.num_envs))
# The same happens for the critic values
values = torch.zeros((args.num_steps, args.num_envs))
# Global variables
global_step = 0
single_global_rollout = int(args.num_envs * args.num_steps * world_size)
num_updates = args.total_timesteps // single_global_rollout
with fabric.device:
# Get the first environment observation and start the optimization
next_obs = torch.tensor(envs.reset(seed=args.seed)[0])
next_done = torch.zeros(args.num_envs)
# Collect `num_steps` experiences `num_updates` times
for update in range(1, num_updates + 1):
# Learning rate annealing
if args.anneal_lr:
linear_annealing(optimizer, update, num_updates, args.learning_rate)
for step in range(0, args.num_steps):
global_step += args.num_envs * world_size
obs[step] = next_obs
dones[step] = next_done
# Sample an action given the observation received by the environment
with torch.no_grad():
action, logprob, _, value = agent.get_action_and_value(next_obs)
values[step] = value.flatten()
actions[step] = action
logprobs[step] = logprob
# Single environment step
next_obs, reward, done, truncated, info = envs.step(action.cpu().numpy())
# Check whether the game has finished or not
done = torch.logical_or(torch.tensor(done), torch.tensor(truncated))
with fabric.device:
rewards[step] = torch.tensor(reward).view(-1)
next_obs, next_done = torch.tensor(next_obs), done
为了训练演员和评论家，我们需要估计回报（returns）和优势（advantages）。
- 优势描述了在某个状态下，采取特定行动相对于根据行动者的随机选择行动的优势有多大。
- 回报是环境所获得的折扣奖励的总和。
是折扣银子。直观上，回报意味着现在的奖励比以后的奖励更有价值。
# Estimate advantages and returns with GAE ()
returns, advantages = agent.estimate_returns_and_advantages(
rewards, values, dones, next_obs, next_done, args.num_steps, args.gamma, args.gae_lambda
)
我们现在终于能够训练代理了。
# Flatten the batch
local_data = {
"obs": obs.reshape((-1,) + envs.single_observation_space.shape),
"logprobs": logprobs.reshape(-1),
"actions": actions.reshape((-1,) + envs.single_action_space.shape),
"advantages": advantages.reshape(-1),
"returns": returns.reshape(-1),
"values": values.reshape(-1),
}
# Train the agent
train(fabric, agent, optimizer, local_data, global_step, args)
from torch.utils.data import BatchSampler, RandomSampler
def train(
fabric: Fabric,
agent: PPOLightningAgent,
optimizer: torch.optim.Optimizer,
data: Dict[str, Tensor],
global_step: int,
args: argparse.Namespace,
):
sampler = RandomSampler(list(range(data["obs"].shape[0])))
sampler = BatchSampler(sampler, batch_size=args.per_rank_batch_size, drop_last=False)
for _ in range(args.update_epochs):
for batch_idxes in sampler:
loss = agent.training_step({k: v[batch_idxes] for k, v in data.items()})
optimizer.zero_grad(set_to_none=True)
fabric.backward(loss)
fabric.clip_gradients(agent, optimizer, max_norm=args.max_grad_norm)
optimizer.step()
agent.on_train_epoch_end(global_step)
有关代理的完整训练步骤的更详细信息，请参考此链接。(github.com/Lightning-A…)
正如我们所见，分布式训练不需要任何样板代码；Fabric为我们抽象了这个过程。要以分布式方式训练我们的代理，只需执行以下命令：
lightning run model \
--accelerator=gpu \
--strategy=ddp \
--devices=2 \
train_fabric.py \
--capture-video \
--env-id CartPole-v1 \
--total-timesteps 100000 \
--num-envs 2 \
--num-steps 512
训练完的代理应该会像下面这样玩游戏：
结论
强化学习是一种强大的机器学习技术，使代理能够通过经验学习并随着时间推移改进其决策能力。它有潜力革新各个行业，并为更智能和适应性强的人工智能系统的发展做出贡献。**
**
在这篇博文中，我们简要介绍了强化学习的高级概念，并展示了如何训练一个代理以在Cart-Pole游戏中实现最佳表现，感谢Fabric的帮助，我们能够加速训练，而无需编写样板代码。
