ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

基于模型的强化学习工具箱:从动力学模型到连续动作规划

基于模型的强化学习工具箱:从动力学模型到连续动作规划 1. 为什么我要自己搭一套基于模型的强化学习工具箱搞强化学习的人大概都有过这种体验策略梯度跑一晚上第二天早上打开终端一看reward 曲线跟心电图似的要么原地踏步要么直接崩掉。尤其是做连续动作控制的任务比如机械臂抓取、四足机器人步态、无人机悬停无模型方法动辄需要几百万步的交互仿真环境跑得冒烟实际机器人根本扛不住这种采样量。我最早接触强化学习是从 DDPG 和 SAC 入手的这两个算法在连续控制里算是标杆了但它们的样本效率实在让人头疼。后来开始转向基于模型的强化学习Model-Based RL核心思路很朴素先学一个环境动力学模型然后用这个模型来生成虚拟经验或者直接做规划。这样一来真实环境里的交互次数可以大幅减少对于实际机器人系统来说这个优势是决定性的。这个系列我打算分几篇来写第一篇先把工具箱的骨架搭起来重点解决三个问题环境交互数据的采集与组织、动力学模型的训练、以及如何用学到的模型做连续动作的预测和控制。代码全部用 Python 实现依赖尽量精简PyTorch 做网络训练NumPy 做数值计算Gym 做环境接口。整套东西我会放在一个统一的工具箱结构里方便后续扩展不同算法。如果你正在做机器人控制、自动驾驶仿真、或者任何需要连续动作输出的决策任务这套东西可以直接拿去改。哪怕你只是刚入门强化学习跟着走一遍也能把基于模型的方法论理清楚。我尽量把每个模块的设计理由讲透不只是贴代码而是让你知道为什么这么写、换一种写法会踩什么坑。2. 工具箱整体架构与模块拆解2.1 核心模块划分与职责边界一个完整的基于模型强化学习工具箱我把它拆成五个核心模块每个模块各司其职模块之间通过明确定义的接口通信。这种拆分方式的好处是你换掉任何一个模块都不会影响其他部分比如想把动力学模型从神经网络换成高斯过程只需要实现同样的接口就行。第一个模块是环境封装层负责对接 Gym 或者自定义的仿真环境统一观测空间和动作空间的接口同时记录每一步的转移数据。第二个模块是数据缓冲区专门存储和管理交互数据支持按批次采样、按时间窗口切片、以及数据归一化。第三个模块是动力学模型也就是整个工具箱的核心输入当前状态和动作输出下一状态的预测。第四个模块是规划与控制利用学到的模型做动作序列优化常见的方法有随机打靶法、交叉熵方法、模型预测控制等。第五个模块是训练循环与评估把前面几个模块串起来管理模型训练和策略评估的流程。这五个模块的依赖关系是单向的环境封装层产生数据数据缓冲区存储数据动力学模型消费数据规划与控制调用动力学模型训练循环协调全局。这种单向依赖保证了代码的可测试性每个模块都可以单独写单元测试。2.2 为什么选择这样的架构而不是端到端有人可能会问为什么不直接把环境交互、模型学习、策略优化全部塞进一个类里搞成端到端的我早期确实这么干过结果就是代码越写越乱调试的时候根本定位不到问题出在哪。比如 reward 不涨你没法判断是动力学模型学得不准还是规划算法有问题还是数据归一化没做好。拆成模块之后每个环节都可以独立验证。动力学模型的预测误差可以单独画曲线规划算法的优化效果可以单独测试数据缓冲区的采样逻辑可以单独跑单元测试。这种可验证性是工程化的基础尤其是做研究的时候你需要快速定位哪个环节是瓶颈。另一个考虑是复用性。动力学模型这块我可能今天用神经网络明天想试试高斯过程或者集成模型如果架构拆得干净换起来就是改一个配置文件的事。规划算法也是同理随机打靶法和交叉熵方法的接口可以统一上层训练循环完全不用改。2.3 目录结构与依赖管理工具箱的目录结构我按照功能划分顶层是mbrl_toolbox包下面分envs、buffers、models、planners、trainers五个子包外加一个utils放通用工具函数。每个子包下面再按具体实现分文件比如models下面有base_model.py、neural_dynamics.py、ensemble_dynamics.py。依赖管理我用的是requirements.txt核心依赖就四个gym做环境接口torch做神经网络训练numpy做数值计算matplotlib做可视化。版本方面Gym 建议用 0.26 以上因为 0.26 之后 API 有较大变动老版本的step返回值和新版本不一样。PyTorch 用 2.0 以上主要是为了torch.compile的加速虽然这个工具箱里模型不大但编译之后训练速度还是有提升的。注意Gym 0.26 之后的版本reset()返回的是(obs, info)元组step()返回的是(obs, reward, terminated, truncated, info)五元组。如果你用的是老版本代码直接迁移会报错需要把done拆成terminated和truncated两个变量。3. 环境交互与数据采集的实操细节3.1 连续动作环境的接口统一连续动作环境和离散动作环境最大的区别在于动作空间的类型。离散环境用Discrete空间连续环境用Box空间。Gym 里Box空间的每个维度都有上下界比如Box(low-1.0, high1.0, shape(3,))表示三维连续动作每一维都在 -1 到 1 之间。我在环境封装层做了一个统一的EnvWrapper类核心功能有三个一是把观测和动作都转成float32的 NumPy 数组避免后续训练时类型不匹配二是记录每一步的转移数据包括观测、动作、奖励、下一观测、终止标志三是支持设置随机种子保证实验可复现。import gym import numpy as np class EnvWrapper: def __init__(self, env_name, seed0): self.env gym.make(env_name) self.env.reset(seedseed) self.obs_dim self.env.observation_space.shape[0] self.act_dim self.env.action_space.shape[0] self.act_low self.env.action_space.low self.act_high self.env.action_space.high def step(self, action): action np.clip(action, self.act_low, self.act_high) obs, reward, terminated, truncated, info self.env.step(action) done terminated or truncated return obs.astype(np.float32), reward, done, info def reset(self): obs, info self.env.reset() return obs.astype(np.float32)这里有个细节值得展开动作裁剪。神经网络输出的动作可能超出环境允许的范围如果不裁剪有些环境会直接报错有些环境会静默截断但行为不可预期。我选择在step里显式裁剪这样至少能保证行为一致。裁剪之后动作的实际值和网络输出值会有偏差这个偏差在训练动力学模型时需要注意因为模型学的是“实际执行的动作”到“下一状态”的映射而不是“网络输出动作”到“下一状态”的映射。3.2 数据缓冲区的设计与采样策略数据缓冲区我实现了一个ReplayBuffer类底层用 NumPy 数组做循环存储避免用列表导致内存碎片化。缓冲区的大小根据任务复杂度来定简单的倒立摆任务几万条就够了复杂的机械臂任务可能需要几十万条。class ReplayBuffer: def __init__(self, capacity, obs_dim, act_dim): self.capacity capacity self.obs np.zeros((capacity, obs_dim), dtypenp.float32) self.act np.zeros((capacity, act_dim), dtypenp.float32) self.rew np.zeros((capacity, 1), dtypenp.float32) self.next_obs np.zeros((capacity, obs_dim), dtypenp.float32) self.done np.zeros((capacity, 1), dtypenp.float32) self.ptr 0 self.size 0 def add(self, obs, act, rew, next_obs, done): self.obs[self.ptr] obs self.act[self.ptr] act self.rew[self.ptr] rew self.next_obs[self.ptr] next_obs self.done[self.ptr] done self.ptr (self.ptr 1) % self.capacity self.size min(self.size 1, self.capacity) def sample(self, batch_size): idx np.random.randint(0, self.size, sizebatch_size) return (self.obs[idx], self.act[idx], self.rew[idx], self.next_obs[idx], self.done[idx])采样策略上我默认用均匀随机采样但实际用下来发现对于动力学模型训练近期数据比早期数据更重要。因为随着策略更新智能体会访问新的状态区域早期数据可能集中在初始状态附近对模型在新区域的预测帮助不大。所以我加了一个recent_ratio参数每次采样时按比例从最近的数据里抽一部分剩下的从全量数据里抽。实测下来这个比例设在 0.3 到 0.5 之间效果比较好。实操心得数据归一化对动力学模型的训练影响非常大。观测的各个维度量纲可能差好几个数量级比如位置是米级速度是米每秒角度是弧度。如果不做归一化神经网络会偏向大量纲的维度小量纲的维度学不准。我的做法是在缓冲区里维护观测和动作的均值和方差每次采样后做标准化预测时再反标准化。3.3 数据采集策略随机策略还是已有策略初始数据怎么来最朴素的做法是用随机策略跑一批数据。但随机策略在连续动作空间里效率很低尤其是高维动作空间随机探索几乎不可能覆盖到有意义的状态区域。我的做法是分两阶段第一阶段用随机策略采集少量种子数据大概几千条用来训练一个初始的动力学模型第二阶段用基于模型的规划器生成动作同时加入高斯噪声做探索采集新数据并持续更新模型。这个思路本质上就是基于模型强化学习的标准流程模型指导探索探索改进模型。噪声的方差也需要调节。初始阶段噪声大一点保证探索范围随着模型越来越准噪声逐渐减小让策略更精细。我一般用线性衰减从 0.3 衰减到 0.05衰减步数根据任务难度调整。4. 动力学模型的核心实现与训练技巧4.1 神经网络动力学模型的结构选择动力学模型要学的是 $s_{t1} f(s_t, a_t)$ 这个映射。输入是当前状态和动作的拼接输出是下一状态。对于连续控制任务状态维度通常在 10 到 50 之间动作维度在 1 到 10 之间所以输入维度不大不需要太深的网络。我用的结构是三层全连接网络每层 256 个隐藏单元激活函数用 SiLU也叫 Swish。为什么选 SiLU 而不是 ReLU因为 ReLU 在负半轴梯度为零对于动力学模型这种需要精细预测的任务负半轴的梯度信息也很重要。SiLU 是平滑的梯度处处存在实测下来预测误差比 ReLU 低 10% 到 15%。import torch import torch.nn as nn class NeuralDynamics(nn.Module): def __init__(self, obs_dim, act_dim, hidden_dim256): super().__init__() self.net nn.Sequential( nn.Linear(obs_dim act_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, obs_dim) ) def forward(self, obs, act): x torch.cat([obs, act], dim-1) return self.net(x)输出层没有加激活函数因为下一状态的预测值可以是任意实数。损失函数用均方误差MSE但这里有个细节不同维度的预测误差量纲不同直接求 MSE 会让大量纲维度主导梯度。我的做法是在归一化后的空间里计算 MSE这样每个维度的权重是均衡的。4.2 集成模型与不确定性估计单个神经网络做动力学预测有个致命问题在训练数据覆盖不到的区域预测值可能完全离谱但网络本身不知道。这就是所谓的模型不确定性问题。解决方案是用集成模型训练多个通常 5 到 7 个结构相同但初始化不同的网络预测时取均值和方差。均值作为最终预测方差作为不确定性度量。在规划时如果某个动作序列的预测方差很大说明模型对这个区域不熟悉应该降低它的优先级。这个思路在 PETS 和 MBPO 这些经典工作里都有体现实测下来对样本效率的提升非常明显。class EnsembleDynamics(nn.Module): def __init__(self, obs_dim, act_dim, num_models5, hidden_dim256): super().__init__() self.models nn.ModuleList([ NeuralDynamics(obs_dim, act_dim, hidden_dim) for _ in range(num_models) ]) def forward(self, obs, act): preds torch.stack([m(obs, act) for m in self.models], dim0) mean preds.mean(dim0) var preds.var(dim0) return mean, var集成模型的训练方式是每个网络独立训练但共享同一批数据。每个网络用不同的随机种子初始化训练时的批次顺序也打乱保证网络之间的差异性。如果所有网络都收敛到同一个解那集成就失去意义了。注意集成模型的推理开销是单个模型的 N 倍。如果规划算法需要大量前向传播比如随机打靶法采样几千条动作序列集成模型的计算量会很大。我的优化方案是先用单个模型做粗筛选出 top-k 候选序列再用集成模型做精细评估。这样既保证了不确定性估计的准确性又控制了计算量。4.3 训练循环与早停策略动力学模型的训练循环比较直接从缓冲区采样一批数据前向传播计算预测值和真实下一状态算 MSE反向传播更新参数。但有几个细节需要处理。第一是验证集划分。我通常从缓冲区里留出 10% 的数据做验证不参与训练。每个 epoch 结束后在验证集上算预测误差如果连续几个 epoch 验证误差不降反升就触发早停。这个策略能有效防止过拟合尤其是数据量不大的时候。第二是学习率调度。初始学习率设 1e-3用余弦退火降到 1e-5。动力学模型的训练不需要太激进的学习率因为目标函数比较平滑学习率太大会导致预测值震荡。第三是梯度裁剪。虽然动力学模型的梯度通常不会爆炸但在数据分布变化剧烈的时候偶尔会出现梯度异常。我设置梯度范数上限为 10超过就按比例缩放。def train_dynamics(model, buffer, epochs100, batch_size256, lr1e-3): optimizer torch.optim.Adam(model.parameters(), lrlr) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs) best_val_loss float(inf) patience 10 wait 0 for epoch in range(epochs): model.train() for _ in range(len(buffer) // batch_size): obs, act, _, next_obs, _ buffer.sample(batch_size) obs torch.tensor(obs) act torch.tensor(act) next_obs torch.tensor(next_obs) pred model(obs, act) loss nn.functional.mse_loss(pred, next_obs) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 10.0) optimizer.step() scheduler.step() # 验证集评估逻辑省略实测下来动力学模型在几万条数据上训练 100 个 epoch 大概需要几分钟单卡预测误差能降到观测标准差的 5% 以内。这个精度对于后续的规划算法来说已经够用了。5. 连续动作规划与控制的落地实现5.1 随机打靶法的原理与参数选择有了动力学模型之后怎么用它来选动作最简单的方法是随机打靶法Random Shooting。具体做法是在当前状态 $s_t$ 下随机采样 K 条动作序列每条序列长度 H用动力学模型预测每条序列的累积奖励选奖励最高的那条序列只执行第一个动作然后进入下一状态重新规划。这个方法听起来很暴力但在低维动作空间和短规划时域下效果出奇地好。K 一般取 1000 到 5000H 取 10 到 30。动作序列的采样分布用高斯分布均值为零方差根据动作空间的上下界来定。def random_shooting(model, obs, act_dim, act_low, act_high, num_samples1000, horizon20): obs_tensor torch.tensor(obs, dtypetorch.float32).unsqueeze(0) best_return -float(inf) best_action None for _ in range(num_samples): action_seq np.random.uniform(act_low, act_high, size(horizon, act_dim)) total_return 0.0 current_obs obs_tensor.clone() for t in range(horizon): act_tensor torch.tensor(action_seq[t], dtypetorch.float32).unsqueeze(0) with torch.no_grad(): next_obs model(current_obs, act_tensor) reward compute_reward(current_obs, act_tensor, next_obs) total_return reward current_obs next_obs if total_return best_return: best_return total_return best_action action_seq[0] return best_action这里有个关键点奖励函数的设计。在基于模型的规划里奖励函数可以是环境给出的真实奖励也可以是你自己定义的代理奖励。如果环境奖励是稀疏的规划算法很难找到有效动作序列这时候需要设计稠密的代理奖励比如距离目标的负距离、动作平滑度惩罚等。5.2 交叉熵方法CEM的改进与调参随机打靶法的问题是采样效率低大部分随机序列都很差浪费计算资源。交叉熵方法Cross-Entropy Method, CEM做了改进先采样一批动作序列选出 top 20% 的精英序列用这些精英序列的均值和方差更新采样分布迭代几轮之后采样分布会集中到高奖励区域。CEM 的关键参数是精英比例和迭代次数。精英比例一般取 0.1 到 0.2太小会导致分布更新不稳定太大则收敛慢。迭代次数取 3 到 5 轮就够了再多边际收益递减。初始方差设大一点保证探索范围后续轮次方差自然收缩。def cem_planner(model, obs, act_dim, act_low, act_high, num_samples500, horizon20, elite_ratio0.1, num_iters5): mean np.zeros((horizon, act_dim)) std np.ones((horizon, act_dim)) * (act_high - act_low) / 2 for _ in range(num_iters): samples np.random.normal(mean, std, size(num_samples, horizon, act_dim)) samples np.clip(samples, act_low, act_high) returns evaluate_sequences(model, obs, samples) elite_idx np.argsort(returns)[-int(num_samples * elite_ratio):] elite_samples samples[elite_idx] mean elite_samples.mean(axis0) std elite_samples.std(axis0) 1e-6 return mean[0]实测下来CEM 比随机打靶法在相同计算量下能找到更好的动作序列尤其是在规划时域较长的时候。但 CEM 的迭代特性意味着它不能完全并行化每轮迭代依赖上一轮的结果所以 GPU 利用率不如随机打靶法。如果计算资源充足随机打靶法加大采样量也能达到类似效果。5.3 模型预测控制MPC的滚动时域策略不管是随机打靶还是 CEM核心逻辑都是滚动时域控制在当前状态规划一条动作序列只执行第一个动作然后在新状态下重新规划。这个策略的好处是能及时纠正模型预测误差因为每次只执行一步下一步的规划会基于真实观测而不是模型预测。滚动时域的代价是计算量大每个时间步都要重新规划。对于实时控制任务比如机器人关节控制规划时间必须小于控制周期。我的优化方案是用较短的规划时域H10减少采样数量K200同时用 GPU 并行加速。在 RTX 3060 上单次规划大概 5 到 10 毫秒对于 100Hz 的控制频率来说勉强够用。实操心得MPC 的规划时域和采样数量需要根据任务动态调整。如果任务变化缓慢比如机械臂慢速移动可以用长时域少采样如果任务变化剧烈比如四足机器人奔跑需要用短时域多采样。我一般会先跑几组对比实验找到精度和速度的平衡点。6. 常见问题排查与避坑经验实录6.1 动力学模型预测误差大的排查思路动力学模型预测不准是最常见的问题排查思路按优先级排列如下。第一检查数据归一化。这是最容易忽略也最容易出问题的地方。如果观测的某个维度方差特别大不归一化的话MSE 损失会被这个维度主导其他维度学不好。验证方法是把预测值和真实值都反归一化后画散点图看是否在某些维度上系统性偏离。第二检查数据覆盖范围。如果训练数据只覆盖了状态空间的一小部分模型在其他区域的预测必然不准。验证方法是统计训练数据的观测分布和当前策略访问的观测分布做对比看是否有明显偏移。如果有需要增加探索噪声或者用集成模型的不确定性来指导探索。第三检查网络容量。如果状态维度和动作维度都很高比如 50 维状态加 10 维动作三层 256 单元的网络可能不够。可以尝试增加到四层或者每层 512 单元但要注意过拟合风险。第四检查训练轮数。动力学模型通常需要比策略网络更多的训练轮数因为它的目标函数更复杂。我一般至少训练 100 个 epoch数据量大的话训练 500 个 epoch 也不奇怪。6.2 规划算法选不出有效动作的解决方案规划算法选不出好动作通常有三个原因。一是奖励函数设计不合理。如果奖励全是负的规划算法会倾向于选动作幅度最小的序列导致智能体不动。解决方案是给奖励加一个基准值或者用相对奖励而不是绝对奖励。二是规划时域太短。如果任务需要多步才能获得奖励短时域规划看不到长期收益会选短视动作。解决方案是增加规划时域但要注意计算量会线性增长。三是动作空间采样范围不对。如果动作空间的上下界设置得和环境不匹配采样出来的动作大部分都被裁剪了有效探索范围很小。解决方案是检查环境返回的action_space.low和action_space.high确保采样范围一致。6.3 训练不稳定与性能波动的调优技巧基于模型的强化学习训练不稳定通常是因为模型误差和策略优化之间的耦合。模型不准导致策略选错动作选错动作导致采集的数据分布偏移数据偏移又导致模型更不准形成恶性循环。打破这个循环的关键是控制模型更新的节奏。我的做法是每采集一批新数据后不立即更新模型而是等积累了几批数据后一起更新。这样模型的变化更平滑策略不会因为模型突变而崩溃。具体来说我设置一个更新间隔比如每 5 轮数据采集更新一次模型每次更新用全量数据训练。另一个技巧是限制策略更新的幅度。在基于模型的规划里策略更新体现为规划时域和采样分布的变化。如果突然把规划时域从 10 增加到 30策略行为会剧烈变化。我的做法是逐步增加规划时域每次增加 5等性能稳定后再继续增加。问题现象可能原因排查方法解决方案预测误差大数据未归一化检查各维度方差标准化观测和动作预测误差大数据覆盖不足对比训练与当前分布增加探索噪声规划动作无效奖励设计不合理检查奖励取值范围调整奖励基准值规划动作无效规划时域太短增加时域看效果逐步增加时域训练不稳定模型更新太快检查更新频率降低更新频率训练不稳定策略变化剧烈检查规划参数平滑调整参数6.4 计算资源优化与加速方案基于模型的强化学习计算量主要花在三个地方动力学模型训练、规划算法采样、环境交互。环境交互通常不是瓶颈因为基于模型的方法本身就是为了减少交互。瓶颈在模型训练和规划采样。模型训练的加速方案用 GPU 训练批次大小设大一点512 或 1024用混合精度训练torch.cuda.amp这些都能显著提升速度。如果模型不大可以用torch.compile编译实测有 20% 到 30% 的加速。规划采样的加速方案把动作序列的评估向量化不要用 Python 循环逐条评估。具体做法是把 K 条序列拼成一个批次一次性输入动力学模型输出也是批量的。这样能充分利用 GPU 的并行能力速度提升几十倍。def evaluate_sequences_vectorized(model, obs, action_seqs): # action_seqs: (num_samples, horizon, act_dim) num_samples, horizon, act_dim action_seqs.shape obs_dim obs.shape[0] obs_batch np.tile(obs, (num_samples, 1)) total_returns np.zeros(num_samples) for t in range(horizon): act_batch action_seqs[:, t, :] obs_tensor torch.tensor(obs_batch, dtypetorch.float32) act_tensor torch.tensor(act_batch, dtypetorch.float32) with torch.no_grad(): next_obs model(obs_tensor, act_tensor).numpy() rewards compute_reward_batch(obs_batch, act_batch, next_obs) total_returns rewards obs_batch next_obs return total_returns这个向量化版本比逐条循环快 50 倍以上是我在实际项目里最重要的优化之一。早期我没做向量化的时候规划一次要几百毫秒根本没法做实时控制。向量化之后降到几毫秒整个系统才真正可用。7. 从数据到动作的完整闭环串联把前面几个模块串起来整个训练流程是这样的初始化环境和缓冲区用随机策略采集种子数据训练初始动力学模型然后进入主循环。主循环里用规划算法选动作执行动作采集新数据存入缓冲区定期更新动力学模型同时评估策略性能。这个循环持续到策略性能收敛或者达到预设的训练步数。评估环节我通常每 1000 步跑一次测试用确定性策略不加探索噪声跑几个 episode记录平均回报。测试环境和训练环境用同一个但随机种子不同避免过拟合到特定初始状态。整个工具箱的代码量大概在 2000 行左右核心逻辑集中在动力学模型和规划算法这两块。后续我会继续写第二篇讲怎么把这套东西扩展到多任务场景以及怎么用真实机器人数据做微调。如果你跟着这篇把基础版本跑通了扩展起来会顺利很多。最后分享一个我在调试过程中总结的小技巧每次修改代码后先用一个极简任务比如倒立摆跑 10 分钟确认整个流程能跑通再去跑复杂任务。复杂任务动辄几小时如果代码有 bug浪费的时间成本太高。倒立摆虽然简单但能验证数据采集、模型训练、规划控制、评估这条完整链路是最高效的冒烟测试。
返回列表