在本文中，我们通过使用 PyTorch 和 TorchRL 库实现近端策略优化（PPO）算法，对倒立摆任务进行深入学习。从环境创建、模型设计、数据收集、损失函数实现、训练循环至模型评估，每一步均详尽介绍。首先，定义环境配置，创建倒立摆环境，并在 GPU 上运行以提升训练效率。接着，构建策略网络和价值网络，设计数据收集器，定义损失函数与优化器。最后，执行训练循环，评估模型性能，并讨论应用示例与结果。通过本文，您将深入了解 PPO 算法及其在强化学习领域的实际应用。
1. 环境与转换创建首先，我们需要准备好运行环境。我们将使用 MuJoCo 库中的倒立摆环境。确保已经安装了 gymnasium 库，并根据可用资源选择在 GPU 上运行训练。
!pip install gymnasium定义超参数和配置环境：
import gymnasium as gymnasium
# 环境配置
device = 'cuda' if torch.cuda.is_available() else 'cpu'
total_frames = 50_000
frames_per_batch = 1000
sub_batch_size = 64
num_epochs = 10
clip_epsilon = 0.2
gamma = 0.99
lmbda = 0.95
entropy_eps = 1e-4
# 创建倒立摆环境
env = gymnasium.make('CartPole-v1', render_mode='human')
env = gymnasium.wrappers.TransformObservation(env, gymnasium.wrappers.NormalizeObservation())接下来，我们将构建策略网络和价值网络。使用 TensorDictModule 和 ProbabilisticActor。
from torchrl.modules import (SequentialModule, TensorDictModule,
ProbabilisticActor, TanhNormal)
class PolicyModule(nn.Module):
def __init__(self, input_size, output_size):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_size, 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, output_size * 2)
)
self.distractor = NormalParamExtractor()
def forward(self, x):
out, distr_params = self.net(x).chunk(2, dim=-1)
distr_params = out.view(out.size(0), 2, -1)
return distr_params, out
policy_module = PolicyModule(env.observation_space.shape[0], env.action_space.n)
value_module = ValueOperator(
module=nn.Sequential(
nn.Linear(env.observation_space.shape[0], 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, 1)
),
in_keys=['observation']
)
actor_module = ProbabilisticActor(
module=policy_module,
spec=env.action_space,
distribution_class=TanhNormal,
distribution_kwargs={'min': -1.0, 'max': 1.0},
return_log_prob=True
)定义数据收集器，用于收集训练数据：
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
replay_buffer = ReplayBuffer(
storage=LazyTensorStorage(max_size=frames_per_batch),
sampler=SamplerWithoutReplacement()
)
collector = SyncDataCollector(
env,
policy_module,
frames_per_batch=frames_per_batch,
total_frames=total_frames,
split_trajs=False,
device=device
)构建损失函数和优化器：
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE
advantage_module = GAE(gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True)
loss_module = ClipPPOLoss(
actor_network=actor_module,
critic_network=value_module,
clip_epsilon=clip_epsilon,
entropy_bonus=bool(entropy_eps),
entropy_coef=entropy_eps,
loss_critic_type="smooth_l1"
)
optimizer = torch.optim.Adam(loss_module.parameters(), lr=1e-3)执行训练循环，并在每 10 批数据之后评估模型：
import tqdm
rewards = []
step_counts = []
lrs = []
for i in tqdm.trange(total_frames // frames_per_batch):
tensordict_data = collector.collect()
for _ in range(num_epochs):
advantage_module(tensordict_data)
data_view = tensordict_data.reshape(-1)
replay_buffer.extend(data_view.to(device))
for _ in range(frames_per_batch // sub_batch_size):
subdata = replay_buffer.sample(sub_batch_size)
loss = loss_module(subdata)
loss.backward()
torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 1)
optimizer.step()
optimizer.zero_grad()
rewards.append(tensordict_data["next", "reward"].mean().item())
step_counts.append(tensordict_data["step_count"].max().item())
lrs.append(optimizer.param_groups[0]['lr'])
# 评估模型
if i % 10 == 0:
with torch.no_grad():
eval_reward, eval_step_count = evaluate_policy(actor_module, env, num_steps=1000)
rewards.append(eval_reward)
step_counts.append(eval_step_count)在执行完上述步骤后，我们可以观察到模型在训练过程中的性能提升和策略的改进。结果讨论部分可以包括训练期间奖励和步数的可视化，以及在环境中的表现示例。此外，可以讨论模型的局限性、可能的优化方向以及与其他强化学习算法（如 A3C、DQN 等）的对比分析。
结论通过以上步骤，我们构建了一个完整的 PPO 系统，用于解决倒立摆任务，并在每一步中详细解释了每个组件的实现和功能。这不仅提供了一个实际的代码示例，还涵盖了强化学习中关键概念的深入理解，如策略优化、模型设计和训练过程优化。