运动规划 第 2 部分：可微分运动规划

在本教程中，我们将在第一部分的基础上，展示如何对使用 Theseus 实现的运动规划器进行求导。特别地，我们将演示如何在 PyTorch 中设置一个模仿学习（imitation learning）循环，为 TheseusLayer 提供初始化值，从而使其更快收敛到高质量轨迹。如果你还没有看过第一部分的运动规划教程，建议先回顾第一部分再继续本部分内容。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import random 

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.data
from IPython.display import clear_output

import theseus as th
import theseus.utils.examples as theg

%load_ext autoreload
%autoreload 2

torch.set_default_dtype(torch.double)

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch.random.manual_seed(1)
random.seed(1)
np.random.seed(1)

mpl.rcParams["figure.facecolor"] = "white"
mpl.rcParams["font.size"] = 16

1. 初始设置

和运动规划教程第 1 部分一样，第一步是从数据集中加载几个规划问题，并设置一些在整个实验中使用的常量。在本示例中，我们将使用从加载器获得的 2 个问题的批次（batch）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
dataset_dir = "data/motion_planning_2d"
num_prob = 2
dataset = theg.TrajectoryDataset(True, num_prob, dataset_dir, "tarpit")
data_loader = torch.utils.data.DataLoader(dataset, num_prob, shuffle=False)

batch = next(iter(data_loader))
map_size = batch["map_tensor"].shape[1]
trajectory_len = batch["expert_trajectory"].shape[2]
num_time_steps = trajectory_len - 1
map_size = batch["map_tensor"].shape[1]
safety_distance = 0.4
robot_radius = 0.4
total_time = 10.0
dt_val = total_time / num_time_steps
Qc_inv = [[1.0, 0.0], [0.0, 1.0]]
collision_w = 5.0
boundary_w = 100.0

接下来我们创建运动规划器。theg.MotionPlanner 类存储了一个 TheseusLayer，该 Layer 是按照第 1 部分描述的步骤构建的，同时还提供了一些有用的工具函数，用于从优化器的当前变量中获取轨迹。

1
2
3
4
5
6
7
8
9
10
planner = theg.MotionPlanner(
    optimizer_config=("LevenbergMarquardt", {"max_optim_iters": 2, "step_size": 0.3}),
    map_size=map_size,
    epsilon_dist=safety_distance + robot_radius,
    total_time=total_time,
    collision_weight=collision_w,
    Qc_inv=Qc_inv,
    num_time_steps=num_time_steps,
    device=device,
)

由于我们处理的是单个数据批次，我们可以用一些张量初始化运动规划器的输入字典，这些张量将在整个示例中使用。需要提醒的是，输入字典将 TheseusLayer 中的 th.Variable 名称与对应的张量值关联起来。

1
2
3
4
5
6
7
8
9
start = batch["expert_trajectory"][:, :2, 0]
goal = batch["expert_trajectory"][:, :2, -1]
planner_inputs = {
    "sdf_origin": batch["sdf_origin"].to(device),
    "start": start.to(device),
    "goal": goal.to(device),
    "cell_size": batch["cell_size"].to(device),
    "sdf_data": batch["sdf_data"].to(device),
}

翻译如下：

2. 模仿学习循环

概览

在本示例中，我们考虑以下模仿学习流程（参见第 2.2 节）：

对若干训练轮（epochs）重复以下步骤：

  1. 使用一个以地图信息为输入的 nn.Module 生成初始变量值（即轨迹）。
  2. 使用 TheseusLayer 规划轨迹，并用第 1 步的结果初始化优化变量。
  3. 计算损失，使第 2 步的输出轨迹接近专家轨迹且质量较高。
  4. 使用反向传播更新第 1 步中模块的参数。

2.1 基本的初始轨迹模型

下面的单元格创建了一个用于生成初始轨迹的基本模型。该模型以地图 ID 的 one-hot 表示作为输入，并生成从地图起点到终点的轨迹。输出是一个字典，其中键对应变量名称，值对应每个变量的初始值（张量），表示生成的轨迹。

注意：为了保持重点在学习部分，我们没有包含模型代码，但感兴趣的读者可以在这里查看。该模型利用 GPMP2 的概率解释，能够轻松生成从起点到终点的多样化平滑轨迹。其核心思路如下：轨迹首先通过生成一个抛物线得到，其焦点位于起点和终点的中点，焦点与顶点之间的距离由一个 MLP 输出；该抛物线使模型更容易学习较宽的曲线。然后，模型可以通过在抛物线周围构建一个“采样”的轨迹来向轨迹添加高阶曲率，使用运动规划问题的概率形式。在这种情况下，采样是另一个 MLP 的输出，用作在抛物线周围“采样”轨迹的种子。模型最终将该采样轨迹作为初始轨迹返回。

1
2
3
init_trajectory_model = theg.InitialTrajectoryModel(planner)
init_trajectory_model.to(device)
model_optimizer = torch.optim.Adam(init_trajectory_model.parameters(), lr=0.04)

2.2 模仿学习循环

有了模型之后，我们现在可以将所有部分结合起来，通过运动规划器进行求导，并为两个地图找到优化所需的良好初始轨迹。该循环基本上遵循概览小节中的步骤 1-4。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
initial_trajectory_dicts = []
best_loss = float("inf")
losses = []
num_epochs = 100
best_epoch = None

# 为了提高速度，我们将优化器的最大迭代次数设置为较低的值（2）。
# 这也会促使初始轨迹模型生成更高质量的轨迹。
planner.layer.optimizer.set_params(max_iterations=2)
for epoch in range(num_epochs):
    clear_output(wait=True)
    model_optimizer.zero_grad()
    
    # 步骤 1：通过模型生成初始轨迹
    initial_traj_dict = init_trajectory_model.forward(batch)
    # 将上述生成的轨迹更新到运动规划器的输入字典中
    planner_inputs.update(initial_traj_dict)

    # 步骤 2：优化以改进模型生成的初始轨迹
    planner.layer.forward(
        planner_inputs,
        optimizer_kwargs={
            "verbose": False,
            "damping": 0.1,
        }
    )
    # 保存初始轨迹的副本
    initial_trajectory_dicts.append(
        dict([(k, v.detach().clone()) for k, v in initial_traj_dict.items()])
    )

    # 步骤 3：计算评估轨迹质量的损失
    # 损失包含两个部分。第一部分是让最终轨迹尽量匹配该地图的专家轨迹（imitation_loss）。
    # 第二部分使用轨迹规划器的总平方误差，也鼓励轨迹平滑且避开障碍。
    # 我们将第二部分缩放一个较小的因子，否则它会完全压倒仿真损失。
    error_loss = planner.objective.error_metric().mean() / planner.objective.dim()

    solution_trajectory = planner.get_trajectory()
    imitation_loss = F.mse_loss(
        batch["expert_trajectory"].to(device), solution_trajectory)
    loss = imitation_loss + 0.001 * error_loss
    
    # 步骤 4：通过 TheseusLayer 进行反向传播并更新模型参数
    loss.backward()
    model_optimizer.step()
    
    if loss.item() < best_loss:
        best_loss = loss.item()
        best_epoch = epoch
    losses.append(loss.item())
    print("------------------------------------")
    print(f"             Epoch {epoch}")
    print("------------------------------------")
    print(f"{'Imitation loss':20s}: {imitation_loss.item():.3f}")
    print(f"{'Error loss':20s}: {error_loss.item():.3f}")
    print(f"{'Total loss':20s}: {loss.item():.3f}")
    print("------------------------------------")
    print("------------------------------------")

输出结果如下：

1
2
3
4
5
6
7
8
------------------------------------
             Epoch 99
------------------------------------
Imitation loss      : 0.009
Error loss          : 0.063
Total loss          : 0.009
------------------------------------
------------------------------------

可视化图像

1
2
3
4
5
plt.figure(figsize=(12, 6))
plt.plot(losses)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.show()

3. 结果

现在让我们可视化使用学习到的初始值生成的轨迹，并对优化器再运行几次迭代。下面的函数对于从变量值字典中绘制轨迹非常有用。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def get_trajectory(values_dict):
    trajectory = torch.empty(values_dict[f"pose_0"].shape[0], 4, trajectory_len, device=device)
    for i in range(trajectory_len):
        trajectory[:, :2, i] = values_dict[f"pose_{i}"]
        trajectory[:, 2:, i] = values_dict[f"vel_{i}"]
    return trajectory

def plot_trajectories(initial_traj_dict, solution_traj_dict, include_expert=False):
    initial_traj = get_trajectory(initial_traj_dict).cpu()
    sol_traj = get_trajectory(solution_traj_dict).detach().clone().cpu()

    sdf = th.eb.SignedDistanceField2D(
        th.Point2(batch["sdf_origin"]),
        th.Variable(batch["cell_size"]),
        th.Variable(batch["sdf_data"]),
    )
    trajectories = [initial_traj, sol_traj]
    if include_expert:
        trajectories.append(batch["expert_trajectory"])
    figs = theg.generate_trajectory_figs(
        batch["map_tensor"], 
        sdf, 
        trajectories,
        robot_radius=0.4, 
        labels=["initial trajectory", "solution trajectory", "expert"], 
        fig_idx_robot=1,
        figsize=(6, 6)
    )
    for fig in figs:
        fig.show()

3.1 从直线初始化的轨迹

作为参考，下面展示了从直线初始化后，经过 10 次优化器迭代获得的轨迹质量。正如图中所示，从直线生成的轨迹质量较差；需要超过 10 次迭代才能生成高质量的轨迹（在第 1 部分中，我们使用了 50 次迭代）。

1
2
3
4
5
6
7
8
9
10
11
12
13
straight_traj_dict = planner.get_variable_values_from_straight_line(
    planner_inputs["start"], planner_inputs["goal"])

planner_inputs.update(straight_traj_dict)
planner.layer.optimizer.set_params(max_iterations=10)
solution_dict, info = planner.layer.forward(
    planner_inputs,
    optimizer_kwargs={
        "verbose": False,
        "damping": 0.1,
    }
)
plot_trajectories(straight_traj_dict, solution_dict)

3.2 学习得到的初始轨迹

另一方面，使用学习得到的初始轨迹，下面的图显示经过 10 次迭代就足以生成平滑且避开所有障碍物的轨迹，这展示了通过轨迹 planner.trj 进行求导的潜力。

1
2
3
4
5
6
7
8
9
10
11
planner_inputs.update(initial_trajectory_dicts[best_epoch])
planner.layer.optimizer.set_params(max_iterations=10)
solution_dict, info = planner.layer.forward(
    planner_inputs,
    optimizer_kwargs={
        "verbose": False,
        "damping": 0.1,
    }
)
plot_trajectories(
    initial_trajectory_dicts[best_epoch], solution_dict, include_expert=True)