Theseus-教程3

创建自定义代价函数

在本教程中，我们将展示如何创建一个自定义的代价函数，以满足特定应用的需求。虽然我们始终可以通过编写误差函数来使用 AutoDiffCostFunction，但对于计算量较大的应用，派生一个新的 CostFunction 子类并使用解析形式的雅可比矩阵通常会更高效。

在本教程中，我们将演示如何编写一个自定义的 VectorDifference 代价函数。该代价函数的误差是两个向量之间的差值。

注意：VectorDifference 是 Theseus 库中已提供的 Difference 代价函数的简化版本（见教程 0）。Difference 可以用于任意 LieGroup，而 VectorDifference 只能用于向量。

初始化

任何 CostFunction 子类在初始化时，都应当接收一个 CostWeight 以及计算代价函数所需的所有参数。

在这个示例中，我们为 VectorDifference 设置了 __init__ 函数，它需要输入两个向量来计算差值：

需要被优化的向量 var，
作为比较基准的向量 target。

此外，__init__ 函数还需要注册优化变量和所有辅助变量。

在这个例子中，优化变量 var 使用 register_optim_vars 进行注册。
计算代价所需的另一个输入 target 使用 register_aux_vars 进行注册。

这是非线性优化器能够正确工作的必要条件：这些函数会将优化变量和辅助变量注册到内部列表中，随后相关的 Objective 就能方便地使用它们，将其添加进去，确保变量名不冲突，并在需要时更新它们的值。

CostWeight 用于对误差和雅可比矩阵进行加权，它是每个 CostFunction 子类必须提供的（误差和雅可比加权函数由父类 CostFunction 继承）。

from typing import List, Optional, Tuple
import theseus as th

class VectorDifference(th.CostFunction):
    def __init__(
        self,
        cost_weight: th.CostWeight,
        var: th.Vector,
        target: th.Vector,
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name) 

        # 添加检查，确保输入参数属于相同的类，并且自由度 (dof) 一致
        if not isinstance(var, target.__class__):
            raise ValueError(
                "VectorDifference 的变量类型与给定 target 不一致。"
            )
        if not var.dof() == target.dof():
            raise ValueError(
                "VectorDifference 中的变量和 target 必须具有相同的自由度。"
            )

        self.var = var
        self.target = target

        # 注册优化变量和目标变量
        self.register_optim_vars(["var"])
        self.register_aux_vars(["target"])

实现抽象函数

接下来，我们需要实现 CostFunction 的抽象函数：dim、error、jacobians 和 _copy_impl：

dim：返回误差的自由度 (dof)；在本例中，就是优化变量 var 的自由度。
error：返回向量的差值，即 var - target。
jacobians：返回误差相对于 var 的雅可比矩阵。
\_copy\_impl：创建内部类成员的深拷贝。

下面我们给出实现示例（包括上文的 __init__ 函数，以便完整定义这个类）。

class VectorDifference(th.CostFunction):
    def __init__(
        self,
        cost_weight: th.CostWeight,
        var: th.Vector,
        target: th.Vector,
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name) 
        self.var = var
        self.target = target
        # 为了提高可读性，这里省略了前面示例中的数据检查
        self.register_optim_vars(["var"])
        self.register_aux_vars(["target"])

    def error(self) -> torch.Tensor:
        # 误差定义为 var - target
        return (self.var - self.target).tensor

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        return [
            # 误差函数对 var 的雅可比是单位矩阵 I
            torch.eye(self.dim(), dtype=self.var.dtype)  
            # 在 batch 的每个元素上复制雅可比
            .repeat(self.var.shape[0], 1, 1)  
            # 将张量放到与 var 相同的设备上
            .to(self.var.device)  
        ], self.error()

    def dim(self) -> int:
        # 返回误差的自由度，这里等于 var 的自由度
        return self.var.dof()

    def _copy_impl(self, new_name: Optional[str] = None) -> "VectorDifference":
        # 创建当前类的深拷贝
        return VectorDifference(  # type: ignore
            self.var.copy(),      # 拷贝 var
            self.weight.copy(),   # 拷贝权重
            self.target.copy(),   # 拷贝 target
            name=new_name         # 可选的新名字
        )

用法（Usage）

现在我们展示 VectorDifference 代价函数的效果是否符合预期。

为此，我们创建一组 VectorDifference 代价函数，每个代价函数作用在一对向量 $a_i$ 和 $b_i$ 上，并将它们加入一个 Objective。随后，我们为每个 VectorDifference 代价函数的向量 $a_i$ 和 $b_i$ 构造数据，并用这些数据更新 Objective。下面的代码片段展示了 Objective 的误差能够被正确计算。

在这里，我们使用 ScaleCostWeight 作为输入的 CostWeight：它是一个标量实数权重，用于对 CostFunction 进行加权。在这个示例中，为了简化，我们固定其值为 1.。

cost_weight = th.ScaleCostWeight(1.0)

# 构建代价函数并添加到目标函数中
objective = th.Objective()
num_test_fns = 10
for i in range(num_test_fns):
    a = th.Vector(2, name=f"a_{i}")
    b = th.Vector(2, name=f"b_{i}")
    cost_fn = VectorDifference(cost_weight, a, b)
    objective.add(cost_fn)
    
# 创建数据并添加到目标函数
theseus_inputs = {}
for i in range(num_test_fns):
    # 每对 var/target 的差值为 [1, 1]
    theseus_inputs.update({
        f"a_{i}": torch.ones((1,2)), 
        f"b_{i}": 2 * torch.ones((1,2))
    })

objective.update(theseus_inputs)
# 10 个代价函数中，每个误差都是 [1, 1]，平方和结果应该是 20
error_sq = objective.error_metric()
print(f"示例误差平方范数: {error_sq.item()}")

输出：

1	示例误差平方范数: 20