Theseus-教程4

运动规划 第 1 部分：将运动规划建模为非线性最小二乘优化

在本教程中，我们将学习如何实现 GPMP2（Mukadam 等人，2018）运动规划算法，应用于一个在二维平面环境中的机器人。

目标是：在给定起点位姿和目标位姿以及环境表示的情况下，找到机器人的轨迹（位姿和速度）。这个问题可以建模为一个优化问题，其优化变量为机器人在一段总时间步长轨迹上的二维位姿和二维速度（在固定的时间间隔上）。

在本例中，我们将优化的目标函数定义为以下代价项的加权组合：

• 边界条件：轨迹应从起始位姿以零速度出发，并在目标位姿以零速度结束。
• 避障约束：轨迹应避免与环境中的障碍物发生碰撞（我们使用带符号距离场 signed distance fields）。
• 平滑性约束：轨迹应当平滑（我们使用零加速度先验 zero acceleration prior）。

import random

import matplotlib as mpl
import numpy as np
import torch
import torch.utils.data

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"
seed = 0
torch.random.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)

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

1. 加载并可视化轨迹数据

首先，让我们从数据集中加载一些运动规划问题，这些数据集包含使用 dgpmp2 中的代码生成的地图和轨迹。

dataset_dir = "data/motion_planning_2d"
dataset = theg.TrajectoryDataset(True, 2, dataset_dir, map_type="tarpit")
data_loader = torch.utils.data.DataLoader(dataset, 2)

batch = next(iter(data_loader))

该批次（batch）是一个从字符串到 torch.Tensor 的字典，包含以下键：

for k, v in batch.items():
    if k != "file_id":
        print(f"{k:20s}: {v.shape}")

输出得到结果

map_tensor          : torch.Size([2, 128, 128])
sdf_origin          : torch.Size([2, 2])
cell_size           : torch.Size([2, 1])
sdf_data            : torch.Size([2, 128, 128])
expert_trajectory   : torch.Size([2, 4, 101])

让我们绘制加载的地图和轨迹。th.eb.SignedDistanceField2D 是一个带符号距离场对象，它包含一个将 $x, y$ 坐标转换为地图网格单元的函数，我们在这里用于绘图。为了完整性，我们展示加载的专家轨迹（expert trajectories），尽管在本示例中不会使用它们——将在本教程的第 2 部分使用。我们还将展示每个地图的带符号距离场。

sdf = th.eb.SignedDistanceField2D(
    th.Point2(batch["sdf_origin"]),
    th.Variable(batch["cell_size"]),
    th.Variable(batch["sdf_data"]),
)
figs = theg.generate_trajectory_figs(
    batch["map_tensor"], 
    sdf, 
    [batch["expert_trajectory"]], 
    robot_radius=0.4, 
    labels=["expert trajectory"], 
    fig_idx_robot=0,
    figsize=(10, 4),
    plot_sdf=True,
)
figs[0].show()
figs[1].show()

以下是示例中我们将使用的一些常数变量

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 = 20.0
boundary_w = 100.0

2. 建模问题

2.1 定义变量对象

在本示例中，我们的目标是使用 Theseus 为上述加载的地图生成规划。如前言所述，我们需要为轨迹上的每个点优化一个二维位姿（2D pose）和一个二维速度（2D velocity）。为此，我们将创建一组带有独立名称的 th.Point2 变量，并将它们分别存储在两个列表中，以便稍后传递给相应的代价函数。

# 创建优化变量
poses = []
velocities = []
for i in range(trajectory_len):
    poses.append(th.Point2(name=f"pose_{i}", dtype=torch.double))
    velocities.append(th.Point2(name=f"vel_{i}", dtype=torch.double))

除了优化变量之外，我们还需要一组辅助变量，用于封装代价函数计算中依赖地图的量，但在整个优化过程中保持不变。这包括起点/目标的参考值，以及碰撞代价和动力学代价函数所需的参数。

# 位姿边界代价函数的目标值
start_point = th.Point2(name="start")
goal_point = th.Point2(name="goal")

# 碰撞避免代价函数所需变量
sdf_origin = th.Point2(name="sdf_origin")  # 带符号距离场原点
cell_size = th.Variable(torch.empty(1, 1), name="cell_size")  # 地图网格大小
sdf_data = th.Variable(torch.empty(1, map_size, map_size), name="sdf_data")  # 带符号距离场数据
cost_eps = th.Variable(torch.tensor(robot_radius + safety_distance).view(1, 1), name="cost_eps")  # 安全距离加机器人半径

# GP 动力学代价函数所需变量
dt = th.Variable(torch.tensor(dt_val).view(1, 1), name="dt")  # 时间步长

2.2 代价权重

接下来，我们将为优化问题中涉及的每个代价函数创建一系列代价权重（cost weights）。

# 所有 GP 动力学代价函数使用的代价权重
gp_cost_weight = th.eb.GPCostWeight(torch.tensor(Qc_inv), dt)

# 所有碰撞避免代价函数使用的代价权重
collision_cost_weight = th.ScaleCostWeight(th.Variable(torch.tensor(collision_w)))

# 对于所有硬约束（终点位置/速度），使用单一标量权重，并设置为较大值
boundary_cost_weight = th.ScaleCostWeight(boundary_w)

2.3 代价函数

在本节中，我们将创建一个 Theseus 目标函数（Objective），并添加用于运动规划的 GPMP2 代价函数。首先，我们创建目标函数：

objective = th.Objective(dtype=torch.double)

边界代价函数

在这里，我们为边界条件创建代价函数，为它们分配名称，并将其添加到 Objective 中。对于边界条件，我们需要四个代价函数，每个代价函数都使用类型为 th.Difference 的代价函数。该类型的代价函数输入包括一个优化变量、一个代价权重、一个目标辅助变量以及名称。其误差函数是优化变量与目标之间的局部差值。

例如，考虑下面添加的第一个 Difference（名称为 pose_0）。该代价函数会告诉优化器尝试使 poses[0] 的变量值接近辅助变量 start_point 的值（它也是一个有名称的变量，如第 2.1 节所述）。另一方面，对于速度约束（如 vel_0），我们不需要传递命名的辅助变量作为目标，因为我们知道希望其值为 torch.zeros(1, 2)，无论地图数据如何（机器人初始速度为零）。

最后，所有这些代价函数共用相同的 boundary_cost_weight，如你可能记得的，它是一个 ScaleCostWeight(100.0)。

# 固定起始位置
objective.add(
    th.Difference(poses[0], start_point, boundary_cost_weight, name="pose_0")
)
# 固定初始速度
objective.add(
    th.Difference(
        velocities[0],
        th.Point2(tensor=torch.zeros(1, 2)),
        boundary_cost_weight,
        name="vel_0",
    )
)
# 固定终点位置
objective.add(
    th.Difference(
        poses[-1], goal_point, boundary_cost_weight, name="pose_N"
    )
)
# 固定终点速度
objective.add(
    th.Difference(
        velocities[-1],
        th.Point2(tensor=torch.zeros(1, 2)),
        boundary_cost_weight,
        name="vel_N",
    )
)

碰撞代价函数

为了实现碰撞避免，我们使用 th.eb.Collision2D 代价函数类型，其输入包括：

• 一个 th.Point2 优化变量。
• 辅助变量：
  • 三个表示带符号距离场的数据（sdf_origin、sdf_data、cell_size）。
  • 引发碰撞代价的距离阈值（cost_eps）。
• 一个代价权重。

由于轨迹中的每个内部点都需要一个这样的代价函数，我们在循环中创建这些代价函数，并传入上面定义的对应位姿变量。

for i in range(1, trajectory_len - 1):
    objective.add(
        th.eb.Collision2D(
            poses[i],
            sdf_origin,
            sdf_data,
            cell_size,
            cost_eps,
            collision_cost_weight,
            name=f"collision_{i}",
        )
    )

GP 动力学代价函数

为了确保轨迹平滑，我们使用 th.eb.GPMotionModel 代价函数，其输入包括：

• 四个 th.Point2 优化变量：时间 $t-1$ 的位姿、时间 $t-1$ 的速度、时间 $t$ 的位姿、时间 $t$ 的速度。
• 一个辅助变量，描述连续时间步之间的时间差。
• 一个代价权重（通常为 th.eb.GPCostWeight 类型）。

我们需要为每对连续状态（位姿和速度）创建一个这样的代价函数，因此在循环中添加这些代价函数，并传入上面创建的优化变量列表中的对应变量。

for i in range(1, trajectory_len):
    objective.add(
        (
            th.eb.GPMotionModel(
                poses[i - 1],
                velocities[i - 1],
                poses[i],
                velocities[i],
                dt,
                gp_cost_weight,
                name=f"gp_{i}",
            )
        )
    )

翻译如下：

---

为运动规划创建 TheseusLayer

在本示例中，我们将使用 Levenberg-Marquardt 作为非线性优化器，并结合基于 Cholesky 分解 的密集线性求解器。

optimizer = th.LevenbergMarquardt(
    objective,
    th.CholeskyDenseSolver,
    max_iterations=50,
    step_size=1.0,
)
motion_planner = th.TheseusLayer(optimizer)
motion_planner.to(device=device, dtype=torch.double)

3. 运行优化器

最后，我们可以生成一些最优规划了。我们首先初始化所有依赖于地图的辅助变量（例如起点和终点位置，或 SDF 数据）。同时，我们为优化变量提供一些合理的初始值；在本示例中，我们将优化变量初始化为从起点到终点的直线路径。下面的辅助函数对此非常有用：

def get_straight_line_inputs(start, goal):
    # 返回一个字典，将位姿（pose）和速度（velocity）变量名与
    # 从起点到终点的直线路径对应起来
    start_goal_dist = goal - start
    avg_vel = start_goal_dist / total_time  # 平均速度
    unit_trajectory_len = start_goal_dist / (trajectory_len - 1)  # 每个时间步的位移增量
    input_dict = {}
    for i in range(trajectory_len):
        input_dict[f"pose_{i}"] = start + unit_trajectory_len * i  # 线性插值计算位姿
        if i == 0 or i == trajectory_len - 1:
            input_dict[f"vel_{i}"] = torch.zeros_like(avg_vel)  # 起点和终点速度为零
        else:
            input_dict[f"vel_{i}"] = avg_vel  # 中间点速度为平均速度
    return input_dict

现在，让我们将运动规划数据传入 TheseusLayer 并开始生成轨迹；请注意，我们可以利用 Theseus 的批处理支持，同时求解多个轨迹。为了初始化变量，我们创建一个字典，将字符串映射到 torch.Tensor，其中键为 th.Variable 的名称，值为应使用的初始张量。

start = batch["expert_trajectory"][:, :2, 0].to(device)
goal = batch["expert_trajectory"][:, :2, -1].to(device)
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),
}
planner_inputs.update(get_straight_line_inputs(start, goal))    
with torch.no_grad():        
    final_values, info = motion_planner.forward(
        planner_inputs,
        optimizer_kwargs={
            "track_best_solution": True,
            "verbose": True,
            "damping": 0.1,
        }
    )

优化结果如下：

Nonlinear optimizer. Iteration: 0. Error: 3905.6714736579306
Nonlinear optimizer. Iteration: 1. Error: 2282.251394473349
Nonlinear optimizer. Iteration: 2. Error: 136.22542880573795
Nonlinear optimizer. Iteration: 3. Error: 50.78312661084182
Nonlinear optimizer. Iteration: 4. Error: 4.300118887813326
Nonlinear optimizer. Iteration: 5. Error: 14.84557244071468
Nonlinear optimizer. Iteration: 6. Error: 2.3698518778509152
Nonlinear optimizer. Iteration: 7. Error: 2.1764857434368388
Nonlinear optimizer. Iteration: 8. Error: 163.46637040887495
Nonlinear optimizer. Iteration: 9. Error: 2.213071348805293
Nonlinear optimizer. Iteration: 10. Error: 181.50298247241238
Nonlinear optimizer. Iteration: 11. Error: 2.2187748890968124
Nonlinear optimizer. Iteration: 12. Error: 169.89215894334885
Nonlinear optimizer. Iteration: 13. Error: 3.805443232691481
Nonlinear optimizer. Iteration: 14. Error: 132.1010533389907
Nonlinear optimizer. Iteration: 15. Error: 1.830013377706033
Nonlinear optimizer. Iteration: 16. Error: 137.40931841459334
Nonlinear optimizer. Iteration: 17. Error: 2.3630850988848158
Nonlinear optimizer. Iteration: 18. Error: 127.7326186066636
Nonlinear optimizer. Iteration: 19. Error: 1.6962584333307351
Nonlinear optimizer. Iteration: 20. Error: 133.13349342246815
Nonlinear optimizer. Iteration: 21. Error: 2.3779312169651745
Nonlinear optimizer. Iteration: 22. Error: 100.65554701995993
Nonlinear optimizer. Iteration: 23. Error: 1.5484209548141115
Nonlinear optimizer. Iteration: 24. Error: 79.11945618450362
Nonlinear optimizer. Iteration: 25. Error: 1.5620953068549461
Nonlinear optimizer. Iteration: 26. Error: 60.766726980963455
Nonlinear optimizer. Iteration: 27. Error: 1.4306538819862062
Nonlinear optimizer. Iteration: 28. Error: 1.3989819422747745
Nonlinear optimizer. Iteration: 29. Error: 1.1621472075446577
Nonlinear optimizer. Iteration: 30. Error: 1.1368880251022904
Nonlinear optimizer. Iteration: 31. Error: 1.1085143552454548
Nonlinear optimizer. Iteration: 32. Error: 1.0931716928138968
Nonlinear optimizer. Iteration: 33. Error: 1.0802935750315232
Nonlinear optimizer. Iteration: 34. Error: 1.0712316967724411
Nonlinear optimizer. Iteration: 35. Error: 1.0644195614703251
Nonlinear optimizer. Iteration: 36. Error: 1.0603135136138264
Nonlinear optimizer. Iteration: 37. Error: 1.0569988398213055
Nonlinear optimizer. Iteration: 38. Error: 1.0549630384559077
Nonlinear optimizer. Iteration: 39. Error: 1.0530582611647743
Nonlinear optimizer. Iteration: 40. Error: 1.0515998330960024
Nonlinear optimizer. Iteration: 41. Error: 1.050424648528516
Nonlinear optimizer. Iteration: 42. Error: 1.049765287190003
Nonlinear optimizer. Iteration: 43. Error: 1.0490200351376027
Nonlinear optimizer. Iteration: 44. Error: 1.048551093917665
Nonlinear optimizer. Iteration: 45. Error: 1.048095388794331
Nonlinear optimizer. Iteration: 46. Error: 1.04779239077781
Nonlinear optimizer. Iteration: 47. Error: 1.0475756086931332
Nonlinear optimizer. Iteration: 48. Error: 1.0474201617339547
Nonlinear optimizer. Iteration: 49. Error: 1.0473066061954055
Nonlinear optimizer. Iteration: 50. Error: 1.0472228497233835

4. 结果

在优化完成后，我们可以查询优化变量以获取最终轨迹，并可视化结果。下面的函数将 TheseusLayer 输出的字典转换为轨迹张量（trajectory tensor）。

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

现在让我们绘制最终的轨迹。

trajectory = get_trajectory(info.best_solution).cpu()

sdf = th.eb.SignedDistanceField2D(
    th.Point2(batch["sdf_origin"]),
    th.Variable(batch["cell_size"]),
    th.Variable(batch["sdf_data"]),
)
figs = theg.generate_trajectory_figs(
    batch["map_tensor"], 
    sdf, 
    [trajectory], 
    robot_radius=robot_radius, 
    labels=["solution trajectory"], 
    fig_idx_robot=0,
    figsize=(6, 6)
)
figs[0].show()
figs[1].show()