ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

从单线程到多进程:我用C++写了个能跑在超算上的PPO训练系统

2026/8/30 22:20:58 拓冰建站 浏览量
从单线程到多进程:我用C++写了个能跑在超算上的PPO训练系统 一、先聊聊为啥要用C写强化学习说起强化学习Reinforcement Learning简称RL大家第一反应可能是Python——OpenAI Gym、Stable Baselines3这些耳熟能详的工具确实都是Python写的。Python写起来快、库多、调试方便这是它的优势。但问题也很现实当环境模拟本身就很重的时候Python的性能瓶颈就暴露出来了。比如你要模拟一个复杂的物理系统、机器人动力学、或者流体力学仿真每一步环境推进可能要算好一会儿。这时候如果还用单线程Python跑训练一个模型可能得等到地老天荒。所以有人就想了能不能把环境模拟这部分交给C来跑同时让多个环境并行推进一个主进程负责学策略多个子进程负责跑模拟这样既能利用C的高性能又能通过并行化大幅提速。这就是本文要聊的这个框架的核心思路。二、这玩意儿到底是干嘛的简单来说这是一个用C写的、支持多进程并行的强化学习训练框架。目前它只实现了PPOProximal Policy Optimization近端策略优化这一个算法但架构上已经预留了扩展空间。它的工作模式可以这么理解想象你有一个教练训练节点同时带了好几个运动员环境节点在训练。每个运动员在自己的场地上做动作、拿反馈然后把经验汇报给教练。教练汇总所有人的经验更新策略再告诉每个人下一步该怎么改。大家各司其职互不耽误。具体来说框架分成两大类进程训练节点Rank 0只有一个负责跑神经网络、做PPO策略更新。它用LibtorchPyTorch的C版本来训练和推理。环境节点Rank 1~N可以有多个每个节点独立跑一个环境实例比如CartPole倒立摆。它们通过MPIMessage Passing Interface和训练节点通信。三、核心设计思路一个大脑管多个身体整个框架的设计可以概括为**“集中学习、分布执行”**3.1 进程分工进程角色数量职责训练节点1个维护Actor-Critic网络接收观测数据输出动作执行PPO更新环境节点N个独立运行环境模拟执行动作返回新的观测奖励3.2 MPI通信流程环境节点和训练节点之间的通信靠MPI完成大致流程是环境节点每走一步把当前状态observation、奖励reward、是否结束terminate/timeup打包通过MPI_Send发给训练节点。训练节点用MPI_Irecv非阻塞接收多个环境节点的数据收到后喂给神经网络做推理得到动作。训练节点通过MPI_Isend把动作发回给对应的环境节点。环境节点收到动作执行下一步模拟循环往复。这里用了非阻塞接收MPI_Irecv好处是训练节点不会傻等某一个环境而是可以同时监听所有环境节点的消息谁先来就处理谁效率更高。3.3 经验收集与合并每个环境节点在跑的过程中会把自己的经验状态、动作、log概率、奖励、是否结束存在本地的一个MemoryNN对象里。当某个环境的一轮episode结束或者总步数达到更新阈值时训练节点会把所有环境节点的经验合并到一起形成一个大的经验池然后做一次PPO策略更新。这种设计的妙处在于多个环境同时采样相当于用并行化增加了样本多样性而且训练节点可以一次性用更多数据来更新策略样本利用率更高。四、PPO算法通俗解读PPO是目前强化学习领域用得最广泛的算法之一OpenAI的很多项目包括当年击败Dota2职业玩家的OpenAI Five都用的是它。4.1 传统策略梯度的问题在PPO之前策略梯度Policy Gradient算法有个毛病策略更新的时候步子迈太大容易把模型搞崩。比如上一轮策略还不错结果一次更新后策略突然变差奖励断崖式下跌而且很难恢复。4.2 PPO的解决思路给更新加个安全带PPO的核心思想很简单每次更新策略时新策略和旧策略的差别不能太大。它引入了一个概率比ratioratio π_new(a|s) / π_old(a|s)这个ratio表示在同样的状态下新策略选这个动作的概率是旧策略的多少倍。然后PPO做了一个裁剪操作Clipsurr1 ratio * advantage surr2 clip(ratio, 1-ε, 1ε) * advantage loss -min(surr1, surr2)用人话说就是如果新策略把某个动作的概率调得太高或太低ratio超出[1-ε, 1ε]范围就直接截断不让它继续变。这样策略每次只微调一小步训练就稳定多了。4.3 Actor-Critic结构PPO通常配合Actor-Critic架构使用Actor演员负责输出动作。输入是当前状态输出是动作的均值这里假设动作服从正态分布。Critic评论家负责评估状态价值。输入也是当前状态输出是这个状态值多少钱预期未来总奖励。训练的时候Critic帮Actor算优势函数AdvantageAdvantage 实际回报 - Critic估计的价值如果Advantage为正说明这个动作比平均水平好Actor就多学学如果为负就少学甚至反向调整。五、代码实现原理拆解5.1 网络结构框架里的Actor和Critic都是简单的全连接网络结构如下输入层state_dim比如CartPole是6维 ↓ Linear(6 → 64) Tanh ↓ Linear(64 → 32) Tanh ↓ Actor输出层Linear(32 → action_dim) Tanh → 动作均值 Critic输出层Linear(32 → 1) → 状态价值动作不是直接输出一个确定值而是从以动作均值为中心的正态分布里采样得到的。这样策略就有了探索性不会永远只做一个动作。5.2 经验存储MemoryNN每个环境节点都有一个MemoryNN对象记录每一步的states观测到的状态actions执行的动作logprobs动作的对数概率用于PPO的ratio计算rewards获得的奖励is_terminals是否因为失败而结束is_timeups是否因为达到最大步数而结束当训练节点决定更新时会调用merge()把所有环境节点的经验合并然后清空各自的环境记忆。5.3 折扣回报计算PPO更新前需要先算每个状态的折扣回报Discounted Reward。代码里是从后往前算的for(从最后一步往前遍历){if(这一步是因为timeup结束的){discounted_rewardCritic(最后状态的价值);// 用Critic估计剩余价值}elseif(这一步是因为terminate结束的){discounted_reward0;// 失败了后面没有奖励了}discounted_rewardrewardgamma*discounted_reward;}这个gamma是折扣因子通常设为0.99意思是未来的奖励要打点折越远的奖励越不重要。5.4 PPO更新流程合并完经验后PPO会重复更新K个epoch代码里默认80次。每次更新用当前策略重新计算旧经验上的log概率、状态价值、分布熵。算ratio exp(新logprob - 旧logprob)。算Advantage 折扣回报 - 状态价值。做Clip裁剪取min。总Loss 策略损失 价值函数MSE损失 - 熵奖励。反向传播Adam优化器更新网络参数。六、相关领域知识点知识点作用PPO策略优化算法通过Clip限制更新幅度训练稳定Actor-Critic结合策略梯度Actor和价值估计Critic的混合架构MPI多进程通信标准用于跨进程发送观测和动作数据LibtorchPyTorch的C前端提供张量运算和自动求导正态分布采样动作从分布中采样保证策略的探索性折扣回报把未来奖励折算到当前衡量长期收益优势函数衡量某个动作比平均水平好多少七、这框架能用来干啥目前这个框架已经在CartPole环境上验证通过了。CartPole是个经典控制问题一根杆立在小车上你要左右移动小车让杆不倒。框架训练出来的策略能很快学会保持平衡。但它的潜力不止于此机器人控制把CartPole替换成你自己的机器人动力学模型就能训练机器人走路、抓取等任务。工程仿真比如流体力学控制、结构优化这些仿真本身计算量大用C并行跑很合适。算法研究框架结构清晰方便在此基础上扩展其他RL算法比如SAC、TD3。八、怎么跑起来手把手教你编译运行8.1 环境依赖需要准备这几样东西LibtorchPyTorch的C库OpenMPI多进程通信CMake或Make编译工具clang-format可选用于代码格式化安装Libtorchwgethttps://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zipunziplibtorch-shared-with-deps-latest.zipexportTorch_DIR/absolute/path/to/libtorchechoexport Torch_DIR/absolute/path/to/libtorch~/.bashrc安装OpenMPI和CMakeUbuntu/Debiansudoaptupdatesudoaptinstallcmake libopenmpi-dev openmpi-bin8.2 编译项目方式一用Makefilemakeall# 编译 格式化# 或者makebuild# 只编译# 或者makedebug# 编译调试版本方式二用CMakemkdirbuildcdbuild cmake..cmake--build.--configRelease# 或者 Debug 版本cmake--build.--configDebug如果在集群环境比如CARC上编译module load gcc/8.3.0 module load openmpi/4.0.2 module load cmakeexportLD_PRELOAD/spack/apps/gcc/8.3.0/lib64/libstdc.so.6 cmake-DCMAKE_C_COMPILERgcc-DCMAKE_CXX_COMPILERg..8.3 配置文件项目根目录下有个config文件内容大概长这样1 # control_vars动作维度CartPole是1维力 6 # obs_vars观测维度CartPole是6维 10000 # Nepisodes总训练回合数 100 # N_timestep每回合最大步数 1000000 # Max_timestep总训练步数上限 2000 # updateTimestep每隔多少步更新一次策略你可以根据任务调整这些参数。8.4 运行编译成功后会生成cpp-rl-training可执行文件。# 用4个进程运行1个训练节点 3个环境节点mpirun-n4./cpp-rl-training /path/to/configout.txt# 如果不指定配置文件默认用 ../configmpirun-n4./cpp-rl-trainingout.txt注意进程数至少要有2个1个训练节点 1个环境节点否则跑不起来。运行过程中每个环境节点会把自己的episode奖励写到procX_log.txt文件里你可以用这些日志画学习曲线。8.5 在集群上用SLURM提交如果在有SLURM调度系统的集群上跑可以写个.sl脚本#!/bin/bash#SBATCH --nodes1#SBATCH --ntasks-per-node4#SBATCH --time01:00:00module load gcc/8.3.0 module load openmpi/4.0.2 mpirun-n4./cpp-rl-training../configout然后提交sbatch your_script.sl九、从实验结果看并行效果从CartPole的实验来看单环境节点大概需要2000个episode左右收敛到满分奖励稳定在100。3个环境节点并行收敛速度明显加快大概500个episode左右就稳定了。这说明并行环境采样确实能加速训练而且多个环境同时探索样本更多样策略学得更稳。If you need the complete source code, please add the WeChat number (c17865354792)十、总结这个框架虽然还处于早期版本但已经搭好了一个不错的骨架C高性能环境模拟和神经网络推理都在C层完成避免了Python GIL的限制。MPI多进程并行一个训练节点带多个环境节点通信简洁高效。PPO算法实现完整从Actor-Critic网络到Clip裁剪核心逻辑都到位了。可扩展性强环境接口清晰换成你自己的仿真模型也不难。如果你正在做一个需要大量环境交互的强化学习任务又嫌Python太慢不妨参考这个思路用C搭一个自己的并行训练 pipeline。毕竟训练快一点迭代就能多一轮好策略就早一点出来。Welcome to follow WeChat official account【程序猿编码】