PyTorch MPPI 实现

本仓库实现了基于近似动力学的模型预测路径积分（MPPI），使用 PyTorch 框架。通常，MPPI 需要真实的轨迹样本，但这篇论文表明，通过重要性采样，即使使用近似动力学模型（例如神经网络），也可以实现 MPPI。

因此，它可以替代其他轨迹优化方法，如交叉熵法（CEM）或随机射击法。

---

自 2024 年 8 月起新增平滑方法，包括我们自己的 KMPPI，请参阅下方的[平滑部分]。

安装

pip install pytorch-mppi

若需自动调优超参数，请使用以下命令安装：

pip install pytorch-mppi[tune]

若需运行测试，请使用以下命令安装：

pip install pytorch-mppi[test]

若需进行开发，请克隆仓库并以可编辑模式安装：

pip install -e .

使用

请参阅 tests/pendulum_approximate.py，其中展示了如何使用神经网络近似摆的动力学模型。此外，not_batch 分支提供了更易读的算法实现。基本用法如下：

from pytorch_mppi import MPPI

# 创建控制器并设置相关参数
ctrl = MPPI(dynamics, running_cost, nx, noise_sigma, num_samples=N_SAMPLES, horizon=TIMESTEPS,
            lambda_=lambda_, device=d,
            u_min=torch.tensor(ACTION_LOW, dtype=torch.double, device=d),
            u_max=torch.tensor(ACTION_HIGH, dtype=torch.double, device=d))

# 假设你有一个类似 Gym 的环境
obs = env.reset()
for i in range(100):
    action = ctrl.command(obs)
    obs, reward, done, _ = env.step(action.cpu().numpy())

要求

• PyTorch (>= 1.0)
• 动力学函数：next state <- dynamics(state, action)（不一定是真实动力学）
  • state 是 K x nx，action 是 K x nu
• 运行代价函数：cost <- running_cost(state, action)
  • cost 是 K x 1，state 是 K x nx，action 是 K x nu

特性

• 基于重要性采样的近似动力学 MPPI
• 并行/批处理的 PyTorch 实现，加速采样
• 通过从修正高斯分布中采样控制噪声来施加控制约束
• 支持随机动力学模型（假设每次调用都是一个样本），可通过为同一动作轨迹采样多条状态轨迹来处理，使用 rollout_samples 参数

参数调优与提示

terminal_state_cost - 函数(state (K x T x nx)) -> cost (K x 1)。默认情况下没有终端代价，但如果发现轨迹接近目标却始终无法到达，可以加入终端代价来帮助收敛。该函数应与时间步长（T）成比例，以匹配运行代价的缩放。

lambda_ - 较高的值会增加控制噪声的成本，从而使更多样本集中在均值附近；通常较低的值效果更好（可尝试 1e-2）。

num_samples - 采样轨迹的数量；一般来说，数量越多越好。 运行性能对 num_samples 的依赖远大于对 horizon 的依赖，尤其是在使用 GPU 设备时（记得将其传递进去！）。

noise_mu - 默认值为 0，适用于所有控制维度。然而，如果存在控制约束且允许范围并非以 0 为中心，则可能效果不佳。对于非对称控制维度，务必将其调整为合适的值。

平滑

从版本 0.8.0 开始，可以使用对控制信号进行平滑处理的 MPPI 变体。我们实现了 SMPPI 以及我们自己的核插值 MPPI（KMPPI）。在基础算法中，可以通过增大 lambda_ 来获得较为平滑的轨迹，但这会牺牲一定的最优性。而显式的平滑算法则可以在不损失最优性的前提下实现平滑效果。

我们在最近的一篇论文中使用并描述了这些方法（arXiv），在我们发布专门针对 KMPPI 的工作之前，您可以引用该论文。下面展示了 MPPI、SMPPI 和 KMPPI 在一个二维导航问题上的差异，其中控制是受约束的位置增量。您可以在 tests/smooth_mppi.py 中查看详细内容。

API 大致相同，仅增加了部分构造函数选项：

import pytorch_mppi as mppi
ctrl = mppi.KMPPI(args, 
                 kernel=mppi.RBFKernel(sigma=2), # 轨迹时间空间中的核（一维）
                 num_support_pts=5,              # 采样控制点的数量，不超过时间步长
                 **kwargs)

核可以是 mppi.TimeKernel 的任何子类。它作用于轨迹的时间维度（一维空间）。需要注意的是，使用 B 样条核可以实现 B 样条平滑。支持点的数量即为采样的控制点数，其余轨迹点将通过核进行插值。例如，如果轨迹时间步长为 20，而 num_support_pts 为 5，则会在整个时间步长上均匀选取 5 个控制点（首尾分别对应轨迹的起点和终点），其余部分则由核插值生成。核作用于控制信号，而非状态信号。

未平滑的 MPPI

SMPPI，通过在动作导数空间中采样噪声进行平滑，在此问题中效果不佳。

采用 RBF 核的 KMPPI 平滑效果良好。

自动调优

从版本 0.5.0 开始，可以自动调优超参数。我们实现了一个兼容流行库 ray tune 的便捷调优器。用户可以选择多种前沿的黑盒优化算法，如 CMA-ES、HyperOpt、fmfn/BayesianOptimization 等。示例请参阅 tests/auto_tune_parameters.py，后续将提供基于该示例的教程。

该调优器也可用于其他控制器，但需要定义相应的 TunableParameter 子类。

首先，我们创建一个用于控制的二维玩具环境，并使用一些默认参数创建控制器。

import torch
from pytorch_mppi import MPPI

device = "cpu"
dtype = torch.double

# 创建用于控制的玩具环境（默认起点和终点）
env = Toy2DEnvironment(visualize=True, terminal_scale=10)

# 使用一些初始参数创建 MPPI
mppi = MPPI(env.dynamics, env.running_cost, 2,
            terminal_state_cost=env.terminal_cost,
            noise_sigma=torch.diag(torch.tensor([5., 5.], dtype=dtype, device=device)),
            num_samples=500,
            horizon=20, device=device,
            u_max=torch.tensor([2., 2.], dtype=dtype, device=device),
            lambda_=1)

然后我们需要创建一个评估函数，供调优器进行调优。该函数不应接受任何参数，并应返回一个至少包含成本的 EvaluationResult 对象。如果在评估成本时不需要进行轨迹回放，可以在返回值中将其设置为 None。以下注释中介绍了如何创建评估函数的提示：

from pytorch_mppi import autotune
# 为了公平起见，所有评估都使用相同的标称轨迹作为起点
nominal_trajectory = mppi.U.clone()
# 我们示例评估函数的参数 - 评估函数有很多选择
evaluate_running_cost = True
num_refinement_steps = 10
num_trajectories = 5

def evaluate():
    costs = []
    rollouts = []
    # 我们针对同一个起始到目标的问题采样多条轨迹，但在实际应用中，您应该考虑在多样化的轨迹数据集上进行评估
    for j in range(num_trajectories):
        mppi.U = nominal_trajectory.clone()
        # 如果时间 horizon 发生变化，起始时的标称轨迹也会不同
        mppi.change_horizon(mppi.T)
        # 通常情况下，MPPI 的标称轨迹会从前一次迭代中热启动
        # 为了公平地测试调优效果，我们将每次的标称轨迹重置为相同的随机轨迹
        # 然后我们手动通过几次细化步骤来对其进行预热
        for k in range(num_refinement_steps):
            mppi.command(env.start, shift_nominal_trajectory=False)

        rollout = mppi.get_rollouts(env.start)

        this_cost = 0
        rollout = rollout[0]
        # 在这里我们评估 MPPI 生成的轨迹的总成本
        // 其他用于调参的成本函数也是可行的，比如仅考虑终端成本
        if evaluate_running_cost:
            for t in range(len(rollout) - 1):
                this_cost += env.running_cost(rollout[t], mppi.U[t])
        this_cost += env.terminal_cost(rollout, mppi.U)

        rollouts.append(rollout)
        costs.append(this_cost)
    // 如果不需要计算轨迹回放，可以返回 None
    return autotune.EvaluationResult(torch.stack(costs), torch.stack(rollouts))

有了这些内容，我们就可以开始调优了。例如，我们可以使用 CMA-ES 优化器进行迭代调优。

# 这些是我们想要调优的 TunableParameter 子类（具体来说是 MPPIParameter）
params_to_tune = [autotune.SigmaParameter(mppi), autotune.HorizonParameter(mppi), autotune.LambdaParameter(mppi)]
// 创建一个使用 CMA-ES 优化器的调优器
tuner = autotune.Autotune(params_to_tune, evaluate_fn=evaluate, optimizer=autotune.CMAESOpt(sigma=1.0))
// 调优若干次
iterations = 30
for i in range(iterations):
  // 返回本次优化步骤的结果
  res = tuner.optimize_step()
  // 我们可以在环境中绘制轨迹回放
  env.draw_rollouts(res.rollouts)
// 获取最佳结果并将其应用到控制器上
// （默认情况下，控制器会采用最新调优的参数，但这可能不是最佳的）
res = tuner.get_best_result()
tuner.apply_parameters(res.param_values)

这是一种从初始定义的参数开始进行优化的局部搜索方法。对于全局搜索，我们可以使用与 Ray Tune 兼容的搜索算法。请注意，您可以修改每个参数的搜索空间，但系统已经提供了合理的默认值。

// 也可以使用 Ray Tune 优化器，详情请参阅：
// https://docs.ray.io/en/latest/tune/api_docs/suggestion.html#search-algorithms-tune-search
// 与调整现有参数不同，这些优化器允许您为每个参数定义搜索空间，并在此空间内进行搜索。
from pytorch_mppi import autotune_global
from ray.tune.search.hyperopt import HyperOptSearch
from ray.tune.search.bayesopt import BayesOptSearch

// 全局版本的参数为每个参数定义了一个合理的搜索空间
params_to_tune = [autotune_global.SigmaGlobalParameter(mppi),
                  autotune_global.HorizonGlobalParameter(mppi),
                  autotune_global.LambdaGlobalParameter(mppi)]

// 在进行 Ray Tune 优化之前，请务必关闭所有图表，否则它们会被重复显示
env.visualize = False
plt.close('all')
tuner = autotune_global.AutotuneGlobal(params_to_tune, evaluate_fn=evaluate,
                                       optimizer=autotune_global.RayOptimizer(HyperOptSearch))
// Ray 调优器不能进行迭代调优，但您可以指定要调优的次数
res = tuner.optimize_all(100)
res = tuner.get_best_result()
tuner.apply_parameters(res.params)

例如，仅对玩具问题使用 CMA-ES 调整超参数（每次都会重置标称轨迹，因此是从噪声中采样）：

如果您不仅想要找到最佳解决方案，还希望获得超参数值的多样性，或者您的评估函数存在较大的不确定性，那么可以直接查询历史结果：

for res in tuner.optim.all_res:
    // 成本
    print(res.metrics['cost'])
    // 提取参数
    params = tuner.config_to_params(res.config)
    print(params)
    // 将参数应用到控制器
    tuner.apply_parameters(params)

或者，您也可以尝试使用 CMA-ME 优化器 进行质量多样性优化。该优化器会在确保参数多样性的同时，努力优化高质量的参数。不过，它的速度非常慢，您可能更适合使用 RayOptimizer 并在检查多样性的同时选择最佳结果。

要使用它，您需要安装：

pip install ribs

然后您可以这样使用：

import pytorch_mppi.autotune_qd

optim = pytorch_mppi.autotune_qd.CMAMEOpt()
tuner = autotune_global.AutotuneGlobal(params_to_tune, evaluate_fn=evaluate,
                                       optimizer=optim)

iterations = 10
for i in range(iterations):
  // 返回本次优化步骤的结果
  res = tuner.optimize_step()
  // 我们可以在环境中绘制轨迹回放
  best_params = optim.get_diverse_top_parameters(5)
  for res in best_params:
    print(res)

测试

在 tests 目录下，你可以找到将 MPPI 方法应用于已知的摆系统动力学以及近似摆系统动力学（使用一个两层前馈神经网络来估计状态残差）的示例。采用连续角度表示法（即输入 cos(θ), sin(θ) 而不是直接输入 θ）会带来显著差异。尽管两种方法都能工作，但连续表示法对控制器参数和随机种子的鲁棒性要强得多。此外，过度摆动后持续旋转的问题也不会出现。

使用 100 步随机策略数据初始化动力学模型，在近似动力学上的示例结果如下：

相关项目

• pytorch CEM - 一种具有与此项目相似 API 的替代 MPC 射线法实现
• pytorch iCEM - 基于采样的另一种 MPC 实现

PyTorch MPPI 快速上手指南

本指南帮助中国开发者快速掌握 pytorch_mppi，一个基于 PyTorch 实现的模型预测路径积分（MPPI）控制库。该工具支持使用近似动力学模型（如神经网络）进行轨迹优化，适用于替代 CEM 或随机射击等方法。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.8+
• 核心依赖：
  • pytorch (版本 >= 1.0)
  • 可选依赖（用于自动调参）：ray[tune], cma, hyperopt 等
• 硬件加速：推荐使用 NVIDIA GPU 以利用其并行采样加速特性（需在代码中指定 device）

安装步骤

您可以直接通过 pip 安装稳定版。为了获得更好的下载速度，建议使用国内镜像源（如清华源或阿里源）。

1. 基础安装

仅安装核心功能：

pip install pytorch-mppi -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装自动调参功能

如果您需要使用内置的超参数自动调优功能（基于 Ray Tune 等优化器）：

pip install "pytorch-mppi[tune]" -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 开发模式安装

如果您需要修改源码或运行测试：

git clone https://github.com/your-repo/pytorch_mppi.git
cd pytorch_mppi
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下是使用 pytorch_mppi 进行控制的最简流程。您需要定义动力学模型和代价函数，然后实例化控制器并在循环中调用。

1. 定义必要函数

MPPI 需要两个核心函数，输入需支持 Batch 维度（K 为样本数）：

• 动力学模型：next_state = dynamics(state, action)
  • 输入：state (K x nx), action (K x nu)
  • 输出：next_state
• 运行代价：cost = running_cost(state, action)
  • 输入：同上
  • 输出：cost (K x 1)

2. 控制器初始化与运行

import torch
from pytorch_mppi import MPPI

# 假设已定义好 dynamics, running_cost, nx (状态维度), nu (动作维度)
# 设置设备
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.double

# 创建控制器
ctrl = MPPI(
    dynamics, 
    running_cost, 
    nx=nx, 
    noise_sigma=torch.ones(nu, dtype=dtype, device=device), # 噪声标准差
    num_samples=500,          # 采样轨迹数量
    horizon=20,               # 预测时域
    lambda_=0.1,              # 温度参数，控制探索程度
    device=device,
    u_min=torch.tensor([-1.0], dtype=dtype, device=device), # 动作下限
    u_max=torch.tensor([1.0], dtype=dtype, device=device)   # 动作上限
)

# 模拟控制循环 (假设有一个 gym 风格的环境 env)
obs = env.reset()
# 将观察值转换为 tensor 并添加到 batch 维度
current_state = torch.tensor(obs, dtype=dtype, device=device).unsqueeze(0)

for i in range(100):
    # 获取最优动作 (返回值为 tensor)
    action = ctrl.command(current_state)
    
    # 执行动作 (转换回 numpy 供环境使用)
    next_obs, reward, done, _ = env.step(action.cpu().numpy())
    
    if done:
        obs = env.reset()
    
    # 更新当前状态
    current_state = torch.tensor(next_obs, dtype=dtype, device=device).unsqueeze(0)

关键参数提示

• num_samples：采样轨迹数。越多效果通常越好，且在 GPU 上扩展性极佳。
• lambda_：控制噪声成本。值越小探索性越强（通常尝试 1e-2 到 0.1），值越大动作越接近均值。
• noise_mu：如果动作空间不是以 0 为中心（例如 [0, 1]），请务必将此参数设置为范围中心，否则采样效率会极低。