增广拉格朗日方法在iLQR算法中的应用
背景
在之前的文章《iLQR算法公式推导》中，我们推导了iLQR算法在无约束情况下的迭代公式，但是在实际应用中，我们往往需要考虑约束条件。本文将介绍如何使用增广拉格朗日法Augmented Lagrangian Methods来处理约束条件。
约束优化问题
具有约束的优化问题具有如下形式：
\[\begin{aligned} \min_{x_{0:N},u_{0:N-1}} \quad & \ell_N(x_N)+\sum_{k=0}^{N-1}\ell_k(x_k,u_k)\\ s.t. \quad & x_{k+1}=f(x_k,u_k),\quad k=0,\cdots,N-1\\ & g_k(x_k,u_k)\leq 0,\\ & h_k(x_k,u_k)=0, \end{aligned} \tag{1}\]其中$k$是时间步，$x_k\in\mathbb{R}^{n_x}$是状态，$u_k\in\mathbb{R}^{n_u}$是控制量，$f$是状态转移函数，$\ell_f$是终端损失函数，$\ell_k$是中间损失函数，$g_k$是不等式约束，$h_k$是等式约束。
增广拉格朗日法
将约束优化问题转化为增广拉格朗日函数的形式：
\[\begin{aligned} \mathcal{L}(x_{0:N},u_{0:N-1},\lambda_{0:N},\mu_{0:N-1})&=\ell_N(x_N)+(\lambda_N + \frac{1}{2}I_{\mu_N}c_N(x_N))^Tc_N(x_N)\\ &+\sum_{k=0}^{N-1}\left[\ell_k(x_k,u_k)+(\lambda_k + \frac{1}{2}I_{\mu_k}c_k(x_k,u_k))^Tc_k(x_k,u_k)\right]\\ &=\mathcal{L}_N(x_N,\lambda_N,\mu_N)+\sum_{k=0}^{N-1}\mathcal{L}_k(x_k,u_k,\lambda_k,\mu_k), \end{aligned} \tag{2}\]其中$\lambda_k\in\mathbb{R}^{p_k}$是拉格朗日乘子，$\mu_k\in\mathbb{R}^{p_k}$是惩罚系数，$c_k=(g_k,h_k)\in\mathbb{R}^{p_k}$是不等式约束和等式约束的组合，相应的，不等式约束的序号和等式约束的序号的集合分别是$\mathcal{I}k$和$\mathcal{E}_k$，$I{\mu_k}$是对角矩阵，它的定义如下：
\[I_{\mu_k,ii}= \begin{cases} 0, & \text{if $c_{k_i}<0$ $\land$ $\lambda_{k_i} = 0$},i\in\mathcal{I}_k\\ \mu_{k_i}, & \text{otherwise}, \end{cases} \tag{3}\]其中$k_i$表示第$k$个时间步的第$i$个约束，$c_{k_i}$表示第$k$个时间步的第$i$个约束的值，$\lambda_{k_i}$表示第$k$个时间步的第$i$个约束的拉格朗日乘子，$\mu_{k_i}$表示第$k$个时间步的第$i$个约束的惩罚系数。
系统的动力学约束可以通过如下方式显式得处理：使用系统的初始状态$x_0$和控制量$u_{0:N-1}$，通过状态转移函数$f$得到状态序列$x_{1:N}$。
增广拉格朗日法的迭代公式
根据iLQR算法的反向传播过程，我们可以通过固定拉格朗日乘子和惩罚系数来定义cost-to-go函数：
\[\begin{aligned} \hat J_N(x_N)|_{\lambda,\mu}&=\mathcal{L}_N(x_N,\lambda_N,\mu_N)\\ &=\ell_N(x_N)+(\lambda_N + \frac{1}{2}I_{\mu_N}c_N(x_N))^Tc_N(x_N)\\ \hat J_k(x_k)|_{\lambda,\mu}&=\min_{u_k}\left[\mathcal{L}_k(x_k,u_k,\lambda_k,\mu_k)+\hat{J}_{k+1}(f(x_k,u_k))|_{\lambda,\mu}\right]\\&=\min_{u_k}\tilde{J}_k(x_k,u_k)|_{\lambda,\mu}, \end{aligned} \tag{4}\]根据上述新的 $\hat J_k(x_k)\vert_{\lambda,\mu}$ 和 $\tilde J_k(x_k)\vert_{\lambda,\mu}$ 的定义对其进行泰勒展开：
\[\begin{aligned} \delta \hat J_k(x_k)\vert_{\lambda,\mu}&\approx p_k^T\delta x_k+\frac{1}{2}\delta x_k^TP_k\delta x_k\\ \end{aligned} \tag{5}\]根据定义，我们可以得到终末（terminal）的最优cost-to-go函数的二阶展开之后的系数：
\[\begin{aligned} p_N&=\nabla_x\mathcal{L}_N(x_N,\lambda_N,\mu_N)\\ &=(\ell_N)_x+(c_N)^T_x(\lambda+I_{\mu_N}c_N)\\ P_N&=\nabla_{xx}\mathcal{L}_N(x_N,\lambda_N,\mu_N)\\ &=(\ell_N)_{xx}+(c_N)^T_{x}I_{\mu_N}(c_N)_x\\ \end{aligned} \tag{6}\]其中：
\[\begin{aligned} (\ell_N)_x&=\frac{\partial \ell_N}{\partial x}\\ (c_N)_x&=\frac{\partial c_N}{\partial x}\\ (\ell_N)_{xx}&=\frac{\partial^2 \ell_N}{\partial x^2}\\ (c_N)_{xx}&=\frac{\partial^2 c_N}{\partial x^2}\\ \end{aligned} \tag{7}\]接下来我们看action-value函数的二阶展开之后的系数，根据iLQR中的推导：
\[\begin{aligned} \delta \tilde{J}_i(x_i,u_i)&=\frac{1}{2} \begin{bmatrix} \delta x_i \\ \delta u_i \end{bmatrix}^T \begin{bmatrix} Q_{x_i^2} & Q_{x_iu_i}\\ Q_{u_ix_i} & Q_{u_i^2} \end{bmatrix} \begin{bmatrix} \delta x_i \\ \delta u_i \end{bmatrix}+ \begin{bmatrix} Q_{x_i} \\ Q_{u_i} \end{bmatrix}^T \begin{bmatrix} \delta x_i \\ \delta u_i \end{bmatrix}\\ \end{aligned} \tag{8}\]其中$Q_{x_i}$，$Q_{u_i}$，$Q_{x_i^2}$，$Q_{x_iu_i}$，$Q_{u_ix_i}$，$Q_{u_i^2}$的形式和iLQR中的有所不同，它们的定义如下：
\[\begin{aligned} Q_{x_i}&=\ell_{x_i}+A_i^Tp_{i+1}+(c_i)^T_{x_i}(\lambda_i+I_{\mu_i}c_i)\\ Q_{u_i}&=\ell_{u_i}+B_i^Tp_{i+1}+(c_i)^T_{u_i}(\lambda_i+I_{\mu_i}c_i)\\ Q_{x_i^2}&=\ell_{x_ix_i}+A_i^TP_{i+1}A_i+(c_i)^T_{x_i}I_{\mu_i}(c_i)_{x_i}\\ Q_{x_iu_i}&=\ell_{x_iu_i}+A_i^TP_{i+1}B_i+(c_i)^T_{x_i}I_{\mu_i}(c_i)_{u_i}\\ Q_{u_ix_i}&=\ell_{u_ix_i}+B_i^TP_{i+1}A_i+(c_i)^T_{u_i}I_{\mu_i}(c_i)_{x_i}\\ Q_{u_i^2}&=\ell_{u_iu_i}+B_i^TP_{i+1}B_i+(c_i)^T_{u_i}I_{\mu_i}(c_i)_{u_i}\\ \end{aligned} \tag{9}\]其中：
\[\begin{aligned} A_i&=\frac{\partial f}{\partial x}|_{x_i,u_i}\\ B_i&=\frac{\partial f}{\partial u}|_{x_i,u_i}\\ \end{aligned} \tag{10}\]接下来就和iLQR中的过程一样了，其中反向传播的过程如下：
\[\begin{aligned} \delta \hat u_i &= -Q_{u_i^2}^{-1}(Q_{u_i}+Q_{u_ix_i}^T\delta x_i)\\ &\triangleq K_i\delta x_i + d_i\\ K_i &= -Q_{u_i^2}^{-1}Q_{u_ix_i}^T\\ d_i &= -Q_{u_i^2}^{-1}Q_{u_i}\\ p_i&=Q_{x_i}+K_i^TQ_{u_i^2}d_i+Q_{x_iu_i}d_i+K_i^TQ_u\\ P_i&=Q_{x_i^2}+K_i^TQ_{u_i^2}K_i+Q_{x_iu_i}K_i+K_i^TQ_{u_ix_i}\\ \Delta \hat J_i&=\frac{1}{2}d_i^TQ_{u_i^2}d_i+Q_{u_i}^Td_i\\ \end{aligned} \tag{11}\]当然，如果考虑正则化，$\delta \hat u_i$的计算方式如下：
方法1：
\[\begin{aligned} \overline{Q}_{u_i^2} &= Q_{u_i^2}+\rho I\\ &= \ell_{u_i^2}+B_i^TP_{i+1}B_i+\rho I\\ \\ d_i &= -\overline{Q}_{u_i^2}^{-1}Q_{u_i}\\ K_i &= -\overline{Q}_{u_i^2}^{-1}Q_{u_ix_i}\\ \\ \Delta \hat J_i &= \frac{1}{2}d_i^TQ_{u_i^2}d_i+Q_{u_i}^Td_i\\ p_i &= Q_{x_i}+K_i^TQ_{u_i^2}d_i+Q_{x_iu_i}d_i+K_i^TQ_{u_i}\\ P_i &= Q_{x_i^2}+K_i^TQ_{u_i^2}K_i+Q_{x_iu_i}K_i+K_i^TQ_{u_ix_i}\\ \end{aligned} \tag{12}\]方法2：
\[\begin{aligned} \overline{Q}_{u_i^2} &= \ell_{u_i^2}+B_i^T(P_{i+1}+\rho I)B_i+(c_i)^T_{u_i}I_{\mu_i}(c_i)_{u_i}\\ \overline{Q}_{u_ix_i} &= \ell_{u_ix_i}+B_i^T(P_{i+1}+\rho I)A_i+(c_i)^T_{u_i}I_{\mu_i}(c_i)_{x_i}\\ \\ d_i &= -\overline{Q}_{u_i^2}^{-1}Q_{u_i}\\ K_i &= -\overline{Q}_{u_i^2}^{-1}\overline{Q}_{u_ix_i}\\ \\ \Delta \hat J_i &= \frac{1}{2}d_i^TQ_{u_i^2}d_i+Q_{u_i}^Td_i\\ p_i &= Q_{x_i}+K_i^TQ_{u_i^2}d_i+Q_{x_iu_i}d_i+K_i^TQ_{u_i}\\ P_i &= Q_{x_i^2}+K_i^TQ_{u_i^2}K_i+Q_{x_iu_i}K_i+K_i^TQ_{u_ix_i}\\ \end{aligned} \tag{13}\]更新增广拉格朗日乘子和惩罚系数
在保持$\lambda$和$\mu$不变的情况下进行完iLQR操作之后（如上），我们可以得到新的控制量序列$\delta \hat u_{0:N-1}$以及对应的状态序列$\delta \hat x_{0:N}$，接下来我们在这个基础上更新$\lambda$和$\mu$：
\[\begin{aligned} \lambda_{k_i}^+&=\begin{cases} \lambda_{k_i}+\mu_{k_i}c_{k_i}(\hat x_k,\hat u_k) & i \in \mathcal{E}_k\\ \max(0,\lambda_{k_i}+\mu_{k_i}c_{k_i}(\hat x_k,\hat u_k)) & i \in \mathcal{I}_k\\ \end{cases}\\ \end{aligned} \tag{14}\] \[\begin{aligned} \mu_{k_i}^+&=\phi\mu_{k_i} \end{aligned} \tag{15}\]其中$\phi$是一个大于1的常数，$\lambda_{k_i}^+$和$\mu_{k_i}^+$表示第$k$个时间步的第$i$个约束的更新后的拉格朗日乘子和惩罚系数。
AL-iLQR算法整体流程
- 初始化$\lambda$，$\mu$和$\phi$；
- 根据当前的$\lambda$和$\mu$，使用iLQR算法计算$\delta \hat u_{0:N-1}$和$\delta \hat x_{0:N}$；
- 根据$\delta \hat u_{0:N-1}$和$\delta \hat x_{0:N}$，更新$u_{0:N-1}$和$x_{0:N}$。
- 根据$\delta \hat u_{0:N-1}$和$\delta \hat x_{0:N}$，更新$\lambda$和$\mu$；
- 检查$\max(c) > tol$如果是真的，转到步骤2，否则结束。
- 输出$u_{0:N-1}$和$x_{0:N}$和$\lambda$。
至此AL-iLQR算法的推导就完成了。
未经允许，禁止转载。