离散控制Matlab代码：当HJB方程撞上Tensor Train，为什么状态维度一高就集体翻车

简介这份资源面向从事最优控制、科学计算与数值分析的研究生、工程师及科研人员提供离散控制场景下Hamilton-Jacobi-Bellman方程牛顿策略迭代的Tensor Train实现属于系统开源类Matlab代码。压缩包共15个文件以13个.m脚本为核心辅以1个md说明与1个license整体约24KB体积轻量却覆盖求解器、谱微分、Legendre递推、控制函数与交叉近似等模块。数值测试脚本可复现一维与二维Allen-Cahn方程及Fokker-Planck方程实验并支持通过umax参数开启控制约束二维测试对CPU耗时较高需留意算力安排。每个文件开头附详细说明也可通过Matlab help查询hjb_leg等求解器语法参数由parse_parameter.m统一解析默认值便于快速起步。目前已有1194人学习适合希望将张量列车方法落地到高维HJB求解的读者参考与二次开发。1. 离散控制Matlab代码当HJB方程撞上Tensor Train为什么状态维度一高就集体翻车做离散时间最优控制的人迟早会撞上一堵墙Hamilton-Jacobi-Bellman方程。一维二维状态还能靠网格硬扛状态维度一到四维五维网格点数按指数爆炸内存直接吃满算到天亮也出不来结果。我最早做的一个倒立摆加多关节的仿真状态六维网格法跑一次要几个小时参数稍微一改又得重来那种绝望感做过的人都懂。这个标题讲的就是一条绕开维度灾难的路子用Tensor Train张量列分解去压缩HJB方程里的值函数再用牛顿策略迭代去求解。核心思路是把一个高维值函数拆成一串低维张量的乘积形式参数量从指数级降到线性级然后每一轮策略迭代里用牛顿法加速收敛。它适合谁适合已经会写基础Matlab控制代码、被维度问题卡住、想找一个能真正跑起来的高维离散控制求解方案的工程师。下面我把这套东西从原理到代码到踩坑一层层拆开。2. 从HJB到Tensor Train值函数为什么能被拆成一串小张量2.1 离散HJB方程到底在解什么离散时间最优控制的标准形式是这样的系统状态 $x_{k1} f(x_k, u_k)$代价函数是累积的 $\sum \gamma^k g(x_k, u_k)$其中 $\gamma$ 是折扣因子。我们要找的是最优值函数 $V^(x)$它满足Bellman最优性方程$$V^(x) \min_u \left[ g(x,u) \gamma V^*(f(x,u)) \right]$$这就是离散HJB方程。注意它是一个不动点方程不是微分方程——离散时间下没有对时间的偏导只有对状态的依赖。这一点很关键因为很多从连续HJB转过来的人会习惯性去找微分算子其实离散版本的核心难点在于值函数 $V$ 是一个定义在高维状态空间上的函数你得把它表示出来、存下来、算出来。网格法的做法是把每个状态维度离散成 $n$ 个点总网格数 $n^d$$d$ 是状态维度。$d6$、$n50$ 的时候就是 $50^6 1.56 \times 10^{10}$ 个点每个点存一个double就是上百GB。这就是维度灾难不是算法不够快是根本存不下。2.2 Tensor Train分解把指数级参数量压成线性级Tensor Train的核心思想很朴素一个 $d$ 维张量 $\mathcal{V}(i_1, i_2, \ldots, i_d)$ 可以写成 $d$ 个三阶张量叫core的缩并$$\mathcal{V}(i_1, \ldots, i_d) \sum_{\alpha_1, \ldots, \alpha_{d-1}} G_1(1, i_1, \alpha_1) G_2(\alpha_1, i_2, \alpha_2) \cdots G_d(\alpha_{d-1}, i_d, 1)$$每个core $G_k$ 的尺寸是 $r_{k-1} \times n_k \times r_k$其中 $r_k$ 叫TT-rank。如果所有rank都限制在 $r$总参数量是 $O(d \cdot n \cdot r^2)$对维度 $d$ 是线性的。原来存不下的东西现在几百MB甚至几十MB就能搞定。但这里有个前提值函数必须“低秩可近似”。好消息是最优控制里的值函数通常有光滑结构TT-rank不会太大坏消息是如果你硬压到很小的rank精度会崩。我一般会从rank4开始试逐步加到8、16看残差什么时候进入可接受范围。2.3 牛顿策略迭代比值迭代快在哪标准的值迭代是 $V_{k1} \mathcal{T}(V_k)$每次只做一次Bellman备份收敛是线性的折扣因子 $\gamma$ 接近1的时候慢得让人想砸键盘。策略迭代是给定当前策略 $\pi$精确求值 $V^\pi$然后贪心改进策略。它收敛是二次的但每步求值代价大。牛顿策略迭代也叫Newton-PI是在策略迭代的基础上用牛顿法去解那个非线性不动点方程。具体来说把Bellman方程写成残差形式 $F(V) V - \mathcal{T}(V) 0$然后牛顿步是 $V_{new} V - J_F^{-1} F(V)$其中 $J_F$ 是残差对值函数的雅可比。在TT格式下这个雅可比和求逆都要用TT代数来做复杂度可控。实际效果是值迭代可能要几百上千轮牛顿策略迭代通常十几轮就收敛到机器精度。3. Matlab里把TT-HJB跑起来核心函数与参数设置3.1 环境准备与TT格式的初始化Matlab本身没有内置TT工具你需要自己写核心的TT代数函数或者用一些开源的TT-Toolbox注意版本兼容性老版本在R2020之后有些函数会报错。我一般会自己写一套精简版只保留本方案需要的几个操作TT乘法、TT加法、TT-SVD截断、TT内积。先定义网格和TT结构。假设状态维度 $d4$每维网格数 $n32$rank上限 $r8$% 参数设置 d 4; % 状态维度 n 32; % 每维网格点数 r_max 8; % TT-rank上限 gamma 0.95; % 折扣因子 tol 1e-8; % 收敛容差 max_iter 50; % 最大迭代轮数 % 生成每维网格以[-1,1]区间为例 grids cell(1, d); for k 1:d grids{k} linspace(-1, 1, n); end % 初始化值函数为全零TTrank1 V_tt tt_zeros(n * ones(1, d));这里tt_zeros返回一个所有core都是零的TT表示rank为1。注意网格范围要根据你的实际系统调整不要照搬[-1,1]状态超出网格范围会导致插值外推结果完全不可信。3.2 构建代价函数和动力学函数的TT表示Bellman方程里需要计算 $g(x,u) \gamma V(f(x,u))$。在TT框架下$g$ 和 $f$ 也要能作用在TT格式上。最简单的情况是 $g$ 和 $f$ 可以分离变量那它们本身就能写成低rank的TT。如果不可分离就需要先做函数插值再转TT。% 代价函数 g(x,u) x*Q*x u*R*u 的TT表示 % 假设Q和R是对角矩阵g可以分离为各维之和 Q_diag [1, 1, 0.5, 0.5]; % 各维权重 g_tt tt_zeros(n * ones(1, d)); for k 1:d % 构造第k维的二次代价core xk grids{k}; gk Q_diag(k) * xk.^2; % 将一维函数嵌入TT结构 g_tt tt_add(g_tt, tt_from_1d(gk, k, d, n)); end % 动力学 f(x,u) A*x B*u线性系统 % 这里以A为对角、B为列向量为例 A_diag [0.9, 0.85, 0.8, 0.75]; B_vec [0.1; 0.1; 0.05; 0.05];tt_from_1d是把一个一维函数放到第 $k$ 个维度上、其余维度为常数的辅助函数。tt_add做TT加法注意加法之后rank会相加必须紧跟一次截断。3.3 牛顿策略迭代主循环主循环的结构是外层策略迭代内层牛顿步。每一轮里先固定策略求值用TT线性求解器再做策略改进在TT格式下逐点求min。for iter 1:max_iter % 步骤1策略评估——解 (I - gamma*P_pi) V g_pi % P_pi是当前策略下的转移算子TT表示 % 用TT-GMRES或TT-DMRG求解线性系统 V_tt tt_solve_linear(P_pi_tt, g_pi_tt, V_tt, tol); % 步骤2策略改进——逐点求min % 对每个状态网格点计算所有控制候选下的Q值 u_candidates linspace(-1, 1, 11); % 控制离散化 V_new_tt tt_zeros(n * ones(1, d)); for iu 1:length(u_candidates) u u_candidates(iu); % 计算 Q(x,u) g(x,u) gamma * V(f(x,u)) Q_tt tt_evaluate_Q(V_tt, g_tt, A_diag, B_vec, u, gamma, grids); V_new_tt tt_min(V_new_tt, Q_tt); % 逐点取min end % 步骤3TT截断控制rank增长 V_new_tt tt_truncate(V_new_tt, r_max, tol); % 步骤4检查收敛 diff_norm tt_norm(tt_add(V_new_tt, tt_scale(V_tt, -1))); if diff_norm tol fprintf(Converged at iteration %d, residual %.2e\n, iter, diff_norm); break; end V_tt V_new_tt; end这段代码里几个关键点tt_solve_linear是TT格式的线性求解器我一般用TT-GMRES因为转移算子通常不是对称正定的tt_min是逐点取最小值在TT格式下需要先转成张量再逐点操作再转回TT这一步是rank增长的主要来源tt_truncate用TT-SVD做截断把rank压回 $r_{max}$。参数方面r_max是最关键的太小精度不够太大内存和计算量飙升。我的经验是从4开始每轮翻倍直到残差不降。tol一般设1e-8到1e-10再小就是数值噪声了。u_candidates的数量影响策略改进的精度11个点对大多数问题够用控制维度高的时候要考虑用TT表示控制空间。4. 避坑指南TT-HJB实现里最容易翻车的五个地方4.1 现象迭代几轮后残差不降反升原因TT截断的误差累积。每轮截断都会引入误差如果截断太狠rank太小误差会像滚雪球一样越滚越大最终发散。解决监控每轮截断前后的相对误差如果截断误差超过当前残差的10%就自动提高rank上限。我一般会设一个rank自适应策略初始rank4如果连续两轮残差下降不到5%rank翻倍直到rank32封顶。4.2 现象策略改进后值函数出现剧烈振荡原因控制离散化太粗或者网格边界处插值外推。当最优控制落在候选点之间时逐点min会选到一个次优但相邻的点导致策略在相邻网格点之间跳变。解决控制候选点至少11个边界处加密同时在网格边界加惩罚项把状态拉回网格内部。另一个办法是用TT表示控制策略本身做连续优化而不是离散枚举但实现复杂度高不少。4.3 现象内存爆了明明rank设得不大原因TT加法和逐点操作会让中间结果的rank临时膨胀。比如两个rank8的TT相加结果rank16如果不及时截断下一步再乘一下rank就32、64很快就爆了。解决每次TT加法、TT乘法、逐点操作之后立刻截断不要等到一轮结束。截断的容差可以比最终容差松一个量级比如最终要1e-8中间截断用1e-6。4.4 现象收敛速度远慢于预期牛顿法退化成梯度下降原因牛顿步里的雅可比矩阵近似太粗糙。在TT格式下精确求雅可比代价很高很多人会用对角近似或者低rank近似近似太差的时候牛顿方向就不准了。解决雅可比近似至少保留到rank2*V_rank并且用TT-GMRES解牛顿方程时残差容差设紧一点1e-10。如果还是慢检查一下是不是折扣因子太接近1$\gamma 0.99$ 的时候牛顿法也会退化这时候要考虑换用其他加速策略。4.5 现象结果和低维暴力网格法对不上原因TT-SVD截断把一些“看起来不重要”的奇异值扔了但这些奇异值对应的方向可能恰好是控制关键方向。低秩近似在值函数平坦区域没问题在最优控制切换面附近会出大误差。解决在切换面附近加密网格或者用自适应网格细化。另一个办法是先用低rank跑出一个粗略解然后在这个解的基础上做局部高rank修正。我一般会拿二维或三维的暴力解做基准验证TT解在关键区域的一致性确认没问题再上高维。5. 进阶技巧用TT交叉插值绕过显式网格把维度推到10以上前面讲的方案本质上还是在网格上做TT每维32个点$d10$ 的时候总网格数 $32^{10} \approx 1.1 \times 10^{15}$虽然TT存得下但构建代价函数和动力学函数的TT表示时如果需要遍历所有网格点照样跑不动。这时候要用TT交叉插值TT-cross approximation它不需要遍历全部网格点只需要在 $O(d \cdot n \cdot r^2)$ 个采样点上求值就能重构出TT表示。具体做法是把Bellman方程的右端 $\mathcal{T}(V)$ 看成一个黑箱函数输入是状态索引 $(i_1, \ldots, i_d)$输出是标量。用TT-cross算法通过自适应采样最大元素来逐步构建TT cores。Matlab里实现TT-cross的核心是maxvol算法找最大体积子矩阵这个在数值线性代数里有标准实现。% TT-cross主循环简化版 function V_tt tt_cross_bellman(V_tt, f_handle, n, d, r_max, tol) % f_handle: 输入索引向量返回Bellman右端值 % 初始化索引集 idx cell(1, d); for k 1:d idx{k} randi(n, r_max, 1); % 随机初始采样 end for sweep 1:2*d % 左右各扫一遍 for k 1:d % 构建采样矩阵 sample_mat zeros(r_max * n, r_max); for i 1:r_max for j 1:n idx_temp idx; idx_temp{k}(i) j; sample_mat((i-1)*n j, :) f_handle(idx_temp); end end % maxvol选主元 [~, piv] maxvol(sample_mat); idx{k} piv(1:r_max); % 更新core % ...具体更新公式略核心是SVD截断 end end V_tt assemble_tt_from_cores(cores); end这个方案的验证方法是先用 $d4$ 的问题对比TT-cross解和全网格TT解确认误差在1e-6以内再逐步加维度。我自己的经验是 $d10$、每维64点、rank16的时候TT-cross构建一次大约几秒牛顿策略迭代十几轮总共几分钟内存峰值不到1GB。同样的精度全网格法需要的内存是天文数字。最后一个习惯每次改完参数先跑一个 $d2$ 或 $d3$ 的小规模版本和暴力解对一遍确认没问题再上高维。这个后悔药我吃过太多次亏才养成希望帮到你。本文还有配套的精品资源点击获取