基于TT分解与牛顿策略迭代的MATLAB高维HJB求解器

简介这份面向离散控制与科学计算学习者的Matlab代码包围绕Hamilton-Jacobi-Bellman方程的牛顿策略迭代展开采用Tensor Train张量分解降低高维状态空间的计算与存储开销适合具备最优控制、数值迭代和Matlab基础的研究生或工程人员复现实验、理解算法框架。压缩包共15个文件以13个.m函数与脚本为主辅以1个README说明和1个许可证文件整包约24KB结构轻量便于按求解器、测试脚本和辅助函数分层查阅。内容包含一维与二维Allen-Cahn方程、Fokker-Planck方程等数值测试入口并提供参数解析、谱Chebyshev微分、Legendre递推与控制函数等支撑实现测试支持键盘输入模型与近似参数默认值可直接开始实验也可设置有限umax开启控制约束。已有1202人学习下载适合需要对照论文实验、复现HJB求解流程并扩展离散控制算例的读者。1. 从一次内存报错说起TT-HJB 在做什么41 的 6 次方大约是 47 亿。一个 6 维状态、每维 41 个格点的离散控制问题如果把值函数老老实实存成稠密数组MATLAB 会先抛出内存不足你才会意识到 HJB 方程的维数灾难并非论文里的修辞。这份基于 MATLAB 的离散控制代码走了另一条路用 Tensor TrainTT分解把值函数压成 d 个三阶核用牛顿策略迭代把 Hamilton-Jacobi-Bellman 方程拆成「策略评估解线性系统 策略改进逐点取最小」两步让 6 维乃至 10 维的 HJB 方程在一台普通工作站上能收敛到可用精度。它面向的是已经懂 Bellman 方程和动态规划、手头有 MATLAB、想做高维最优控制数值实验的人状态维度只有 23 维时稠密网格反而更省事用不上这套东西。2. 工程结构一个 MATLAB 求解器的分层与数据流TT 格式的代码最容易写崩的地方就是把「与问题无关的张量代数」和「与问题强相关的动力学、代价」搅在一起。这套工程是系统开源的 MATLAB 实现目录的划分基本遵循一条线core 不知道 HJB 的存在hjb 只描述问题solver 只调用两边的接口。2.1 三块分层TT 核库、HJB 模型、牛顿迭代驱动目录大体是这个形状文件名可能和线上版本有出入但层次和职责基本一致tt-hjb/ ├── start.m % 入口加路径、跑 demo、出图 ├── core/ % 与问题无关的 TT 线性代数 │ ├── tt_from_cores.m % 用核 cell 组装 TT 对象 │ ├── tt_round.m % QR SVD 截断控制秩 │ ├── tt_add.m / tt_scale.m % 线性运算 │ ├── tt_eval.m % 单点与批量求值 │ ├── tt_cross.m % 采样交叉插值把光滑函数做成 TT │ └── tt_min.m % 逐点最小 argmin ├── hjb/ % 具体问题 │ ├── model_*.m % 动力学 f(x,u)、单步代价 g(x,u) │ ├── grid_setup.m % 每维网格、动作集 │ ├── compose_dynamics.m % 计算 V(f(x,u))TT 插值 截断 │ └── policy_improve.m ├── solver/ │ ├── newton_pi.m % 外层牛顿策略迭代主循环 │ ├── dmrg_solve.m % 内层TT 格式下解线性系统 │ └── value_iteration.m % 对照用的值迭代 └── demo/ ├── demo_lqr.m % 线性二次有解析解先验正确性 ├── demo_pendulum.m % 2 维摆能画策略图 └── demo_highdim.m % 610 维看 TT 的实际收益分层带来的直接好处是换问题不动核心。想加一个新的动力学模型只需要在hjb/下补一个model_xxx.m提供f(x,u)的批量求值和每个动作下的gcore/一行都不用改。反过来想把 DMRG 换成 AMEn 或者 TT-GMRES只动solver/dmrg_solve.m外层主循环的接口签名保持不变。二次开发时这条边界值钱很多同学一上手就把动力学硬编码进求解器改到第三个模型就开始后悔。有一点值得单独说demo_lqr.m是这套代码的校准用例并非示例。线性二次问题的值函数是二次型TT 秩很小而且有 Riccati 方程给出的解析解可以逐点比对。任何对 TT 核库的改动先在这个 demo 上跑通再说。2.2 状态网格与 TT 的内存账稠密网格的存储量是 $N n^d$其中 n 是每维格点数、d 是状态维数。TT 在均匀秩 r 下的存储量是 $d \cdot n \cdot r^2$。两者的差距不是常数倍是指数对多项式配置稠密条目数内存doubleTT 条目数r32TT 内存d2, n411.7e313 KB8.4e40.7 MBd4, n412.8e622 MB1.7e51.3 MBd6, n414.7e938 GB2.5e52 MBd10, n211.7e13133 TB2.2e51.7 MBd2 的时候 TT 反而更占内存这是正常的核的元数据开销摊不平。真正的分水岭在 d≥5稠密方案直接出局。比存储更关键的是求值方式。稠密实现里习惯ndgrid一次性把 N×d 的坐标矩阵建出来TT 方案里这么做等于又把内存吃回去了。常见做法是不建全网格需要哪个点的值就用线性索引换算成多维下标逐维收缩核function v tt_eval(tt, idx) % idx: 1×d 的整数下标返回该点的标量值 d numel(tt.cores); v 1; for k 1:d C tt.cores{k}; % r_{k-1} × n_k × r_k v v * squeeze(C(:, idx(k), :)); % 逐维收缩全程只有 r×r 的小矩阵 end end这里的idx是多维下标并非线性下标。squeeze那一步在 r1 时会退化成标量写通用代码时得用reshape兜住否则第一次跑秩 1 的初值就会报维度错误。批量求值时把循环换成pagefun或者手工展开成矩阵乘速度能差一个数量级这一点在后面对比图里很直观。2.3 主循环数据流一次迭代里张量怎么流动外层一次迭代的数据流是固定的五步看图不如看这条链路当前值函数V是一个 TT秩为rV。对动作集里的每个 $u_i$算出 $Q_i g_i \gamma \cdot V(f(x,u_i))$其中 $V \circ f$ 通过 TT 交叉插值得到算完立刻截断。对所有 $Q_i$ 逐点取最小得到新的V和策略policy每个格点上的最优动作下标。用新策略组装算子 $A I - \gamma P_\pi$调用 DMRG 解 $A V g_\pi$这一步是内层迭代。算外层残差 $|V_{\text{new}} - V| / |V|$小于阈值就停否则回到第 1 步。第 2 步和第 3 步是廉价操作第 4 步吃掉 90% 的时间。参数调优的重心全部在第 4 步的内层精度上这一点在第 5 章会展开。3. 核心模块拆解一次牛顿迭代的完整链路离散时间、有限动作集下的策略迭代和连续时间的 Newton-HJB 不是一回事得先把这个概念掰清楚。有限动作集时策略空间是有限的策略迭代理论上有限步终止——每一步要么严格改进要么已经最优。代码里之所以还叫牛顿策略迭代是因为它对应的是「精确策略评估 贪心改进」这个牛顿型结构而连续动作集或者用函数逼近时它退化成的就是标准牛顿法。理解这一点很重要因为它决定了收敛判据怎么写外层看的是策略是否稳定而非残差是否单调下降。3.1 策略评估把 Bellman 方程写成线性方程组固定策略 $\pi$ 之后Bellman 方程 $V g_\pi \gamma P_\pi V$ 就是一个线性方程$$(I - \gamma P_\pi) V g_\pi$$其中 $P_\pi$ 是沿策略 $\pi$ 的确定性转移算子。代码里从不显式构造 $N \times N$ 的矩阵只提供算子作用在 TT 上的函数function y apply_A(V, model, policy, gamma) % 计算 (I - γ P_π) VV 是 TT返回值也是 TT PV compose_dynamics(V, model, policy); % V(f(x, π(x)))TT 交叉插值 y tt_add(V, tt_scale(PV, -gamma)); % V - γ·P V y tt_round(y, 1e-12, 256); % 算子作用后秩会涨必须截断 endcompose_dynamics是整个项目里最贵的单点操作。它的输入是 TT 格式的 V 和一个由策略决定的位移场输出是 $V(f(x,\pi(x)))$ 的 TT 表示。一般的动力学不满足按维可分离所以这里只能走采样交叉插值在每对相邻维上取若干指标集用 TT-cross 把复合函数重建出来再做一次截断。插值的容差1e-12是刻意设小的因为这里引入的误差会直接进内层线性系统和内层容差同量级时会把 DMRG 的收敛判断搅乱。还有个容易被忽略的细节tt_round的上限256是硬约束。复合之后秩经常从 32 涨到几百不设上限的话第一次迭代就能把内存吃光而且涨上去的秩里绝大部分是数值噪声截掉不掉精度。3.2 DMRG 单核扫描把大线性系统切成小系统内层要解的是 $A V b$A 和 b 都以 TT 或 TT 算子的形式给出未知量 V 也要求是 TT。直接对整个 TT 做最小二乘没有闭式解标准做法是交替最小二乘固定除第 k 个核以外的所有核只解第 k 个核解完右移一位扫完一遍再反向扫一遍。固定其它核之后第 k 个核的未知量是一个 $r_{k-1} \times n_k \times r_k$ 的张量展平成 $\mathbb{R}^{(r_{k-1} n_k) \times r_k}$ 的矩阵。局部系统的规模是 $(r_{k-1} n_k r_k)^2$r32、n33 时大概是 $3.4 \times 10^6$ 个元素稠密求解完全扛得住。function V dmrg_solve(V, b, model, policy, gamma, opts) % 交替最小二乘解 (I - γP)V bV 与 b 均为 TT d numel(V.cores); for sweep 1:opts.max_sweeps for k 1:d % 左扫 L env_left(V, k); % 左环境k 左边所有核的收缩 R env_right(V, k); % 右环境k 右边所有核的收缩 Hk local_operator(L, R, model, policy, gamma); % 局部算子 bk local_rhs(L, R, b, k); % 局部右端项 vk Hk \ bk; % 小规模稠密解 [U, S, Vr] svd(reshape(vk, [], size(V.cores{k},3)), econ); r trunc_rank(diag(S), opts.tol_in); % 按能量比截断 V.cores{k} reshape(U(:,1:r), size(V.cores{k},1), [], r); V.cores{k1} fold_left(S(1:r,1:r)*Vr(:,1:r), V.cores{k1}); end for k d:-1:2 % 右扫对称操作 % 同上方向反过来 end res tt_norm(apply_A(V, model, policy, gamma) - b) / tt_norm(b); if res opts.tol_in, break; end end end几个参数值得单独解释。opts.max_sweeps控制单次策略评估里最多扫几遍常见取值 1020opts.tol_in是内层残差阈值一般比外层阈值小两到三个数量级。trunc_rank用的是能量截断 $r \min{k : \sum_{ik} s_i^2 \le \epsilon^2 \sum_i s_i^2}$和tt_round里保持一致的判据。还有个隐含前提Hk必须是对称正定的这样\走 Cholesky 才既快又稳。$I - \gamma P_\pi$ 在 $\gamma 1$ 时确实正定但数值上接近奇异——$\gamma 0.99$ 时条件数能到 $10^2$ 量级$\gamma 0.999$ 就上四位数了。所以求解器里通常会给Hk加一个极小的对角正则项或者在条件数超阈值时切到pinv。3.3 策略改进离散动作集上的逐点最小离散控制的动作集是有限集合 $U {u_1, \dots, u_m}$「最小化」这一步变成逐个动作算 Q 值再逐点比较不需要任何优化器。这是离散控制相对连续控制的一个便宜之处代价是动作数多时内存线性增长。function [Vnew, policy] policy_improve(V, model, gamma, opts) m numel(model.U); % 离散控制动作是有限集合 Q cell(1, m); for i 1:m Vi compose_dynamics(V, model, i); % V(f(x, u_i))TT gi model.g(:,:,i); % 单步代价TT 或可采样函数 Q{i} tt_round(tt_add(gi, tt_scale(Vi, gamma)), opts.tol_round, opts.rmax); end [Vnew, policy] tt_min(Q); % 逐点最小同时记录 argmin 下标 endtt_min的实现是这套代码里比较有意思的一块。两个 TT 逐点取最小结果一般不是低秩的直接做会炸。工程上的做法是在每对相邻维上取交叉插值的采样点在采样点上求最小再用 TT-cross 把结果重建出来最后截断。误差不会为零但在光滑问题上控制得住。policy输出的是每个格点上的动作下标尺寸本身是 $n^d$这张表存不下来——实际存法是把它也做成 TT秩低的时候几个核就够需要查某个点的策略时用tt_eval取。3.4 秩控制tt_round 的 QRSVD 与截断阈值前面几块都在调tt_round这里单独看它的实现。整个算法里它的调用次数最多写法稍微拖沓一点整体运行时间就会翻倍。function tt tt_round(tt, tol, rmax) % 逐核 QR SVD 截断把 TT 秩压到 tol 精度内且不超过 rmax d numel(tt.cores); for k 1:d-1 C tt.cores{k}; [rl, nk, rr] size(C); M reshape(C, rl*nk, rr); % (r_{k-1}·n_k) × r_k [Q, R] qr(M, 0); % 先 QR避免 SVD 直接吃大矩阵 [U, S, V] svd(R, econ); s diag(S); r find(cumsum(s.^2, reverse) tol^2 * sum(s.^2), 1); if isempty(r), r numel(s); end r min([r, rmax, numel(s)]); tt.cores{k} reshape(Q*U(:,1:r), rl, nk, r); tt.cores{k1} fold_left(S(1:r,1:r)*V(:,1:r), tt.cores{k1}); end end先 QR 再 SVD 是老套路目的是把 $r_{k-1} n_k \times r_k$ 的矩阵乘缩到 $r_k \times r_k$ 上n 越大收益越明显。cumsum(..., reverse)那种写法是算尾部能量占比比从头累加再比较更直观。tol的选取有个经验外层迭代初期可以设大一点1e-6让秩快速降下来跑通流程接近收敛时收紧到1e-10以上否则残差曲线会卡在一个平台上再也下不去看起来像发散。rmax是安全阀宁小勿大超出预期就说明问题本身不适合低秩表示继续硬撑只会越跑越慢。4. 跑起来环境、参数与结果验证4.1 环境与启动MATLAB 版本上R2018b 之后的版本都行用到的都是基础算子没有工具箱依赖pagefun需要 Parallel Computing Toolbox代码里做了 fallback没装也能跑只是慢一点。把工程根目录加进路径直接跑start.m或者进demo/里挑一个。% demo_highdim.m 的关键几行 addpath(genpath(pwd)); cfg.d 6; % 状态维数 cfg.n 33; % 每维格点数取奇数让 0 落在网格上 cfg.rmax 32; % TT 秩上限先小后大 cfg.gamma 0.95; % 折扣因子越接近 1 内层越难解 cfg.tol_out 1e-6; % 外层策略稳定判据 cfg.tol_in 1e-8; % 内层 DMRG 残差阈值 cfg.max_sweeps 20; % 单次策略评估的最大扫描次数 model model_highdim(cfg); % 返回 f、g、动作集 U V tt_zero(cfg.n, cfg.d); % 秩 1 的零初值 [V, info] newton_pi(V, model, cfg); plot_convergence(info); % 外层残差与每轮耗时第一次运行时建议把cfg.d改成 2、cfg.n改成 21先跑demo_lqr.m。它几十秒内就能收敛还能和 Riccati 解逐点比对这是排查「跑出来但结果是错的」这类问题最快的手段。4.2 参数表改哪个会有什么后果参数常见取值调大之后调小之后cfg.d210指数级变难TT 收益才显现稠密方案更合适cfg.n2165精度提高核矩阵线性变大截断误差被网格误差掩盖cfg.rmax1664精度上限变高内存平方增长收敛到某个残差平台后就下不去cfg.gamma0.90.99内层条件数变大扫描次数上升问题变「近视」可能退化成一步最优化cfg.tol_in1e-101e-6内层求解时间明显增加外层残差曲线抖动、策略来回切换cfg.tol_out1e-81e-5迭代轮数增加提前停止值函数精度不够gamma是最容易踩的一脚。$\gamma$ 趋近 1 时 $I - \gamma P_\pi$ 的条件数按 $1/(1-\gamma)$ 增长内层迭代次数跟着涨配合不当就会出现「内层没解到位 → 外层策略乱换 → 看着像发散」的连锁反应。4.3 结果验证三条线跑出图不等于跑对。这套代码的自检一般走三条线。第一条是残差曲线。外层看的是 $|V_{k1} - V_k|/|V_k|$它应该单调下降到阈值以下如果是锯齿状震荡说明内层精度不够。内层看的是 DMRG 的相对残差正常形态是先快速下降再缓慢收敛如果扫了两轮不动八成是秩上限卡住了。第二条是降维比对。把cfg.d设成 2同一组参数下再跑一遍稠密值迭代两张值函数切面图的差应该在网格分辨率量级以内。这条能抓住绝大部分索引错位和张量模式顺序写反的 bug。第三条是策略可视化。2 维问题上直接imagesc画出策略分区物理上应该是一片片连通区域如果出现大量棋盘格噪点说明tt_min的采样精度不够或者tt_round的容差太松把决定性信息截掉了。5. 排查清单TT 秩爆炸、牛顿反跳与索引错位5.1 现象跑几轮之后内存一路涨到爆最后 MATLAB 卡死原因基本只有一个某个地方漏了截断或者截断的容差设得太松。最常见的位置是compose_dynamics的返回值——复合之后秩从 32 涨到几百是常态只要这个结果没有立刻tt_round下一步tt_add就会把两个高秩张量拼在一起秩相乘。另一个隐蔽的来源是策略改进里Q{i}的累加如果动作数 m 是 21每个 Q 都涨一点最后tt_min的输入就已经失控了。解决给所有张量运算的出口都挂上截断宁可多调一次也别省。在newton_pi的主循环里加一行日志打印每轮结束后每个核的尺寸秩的增长趋势一眼就能看出来。真要定位就把rmax强行设成 8 跑一遍看什么时候报维度不匹配。5.2 现象外层残差先降后升策略在两个状态之间来回切换这不是 TT 的问题是内层没解干净。策略评估给出的 V 不够准贪心改进就会指向一个「看起来更好但其实更差」的动作下一轮再切回来形成循环。$\gamma$ 大于 0.98 时尤其明显。原因有两层一是 DMRG 的max_sweeps太小内层残差还在 $10^{-4}$ 量级就返回了二是I - γP本身病态局部系统Hk的条件数已经大到 Cholesky 分解掉精度。解决先把cfg.tol_in收紧到1e-10、max_sweeps提到 40看循环是否消失。如果还在循环给Hk加1e-12的对角正则或者换Hk \ bk为带列主元的 QR 求解。实在不行就对策略改进加阻尼只接受让残差严格下降的更新否则退回原策略继续把内层多扫几轮。5.3 现象2 维能跑出漂亮结果3 维以上切面图明显不对称但问题本身是对称的这类「小维数对、大维数错」的 bug 九成是张量的模式顺序错位。MATLAB 是列优先存储reshape的顺序和 TT 的 mode 约定一旦不一致2 维时刚好因为对称性看不出来3 维就露馅。原因通常出在核的重排上tt_round里reshape(C, rl*nk, rr)假定第一维是左秩、第二维是格点、第三维是右秩而某处permute之后顺序变成了rl, rr, nk两个函数对同一份数据的理解就不一样了。解决写一个单元测试构造一个已知的秩 1 张量比如 $V(x) \prod_k \sin(x_k)$跑一遍所有核心函数逐点比对tt_eval和直接乘出来的值。这个测试十几行能挡掉后面 80% 的调试时间。另外统一约定所有核一律r_{k-1} × n_k × r_k任何中间态都用注释标清楚。5.4 现象边界附近值函数发散或者策略在边界上完全乱掉原因在于动力学的像点跑出了网格定义域。compose_dynamics做插值时落在网格外的点只能外插而 TT 插值在外插区的行为是不可控的误差随距离指数增长再被 $\gamma$ 一轮轮放大。解决三件事一起做。网格区间留裕量按动力学的最大可达范围往外扩 10%20%在插值函数里加裁剪超界的点强制拉回最近的网格点单步代价g在边界处加一个二次惩罚项让最优策略主动远离边界。第三条最有效因为它从问题定义上就把解推离了不可靠区域。6. 二次开发换动力学、换代价、换内层求解器想加一个新的动力学模型改动面比想象中小。只需要在hjb/下新建一个model_xxx.m返回三样东西批量求值的f(x, u)、按动作组织的单步代价g、离散动作集U。f的签名建议统一成「输入 N×d 的状态矩阵和标量动作输出 N×d 的像点矩阵」这样compose_dynamics里不需要为每个模型写特例。如果动力学对状态是可分离的比如每个维度只依赖自己把它写成若干个秩 1 项的和TT 交叉插值的采样量能降一个数量级这个优化在 d≥6 时非常值得做。换代价函数要小心一点。g如果不能在网格上按维分离就得走tt_cross采样重建而采样点选得不好会直接把最优策略带偏。经验做法是先在 d2 的稠密网格上把g和它对应的最优策略算出来作为基准换成 TT 版本之后比对策略分区重合度低于 95% 就回去调采样点密度。内层求解器是最值得动的地方。DMRG 的收敛速度强依赖于扫描顺序和初始秩如果发现某些问题上扫 40 轮还不如别人扫 5 轮可以换成 AMEn——它在每次局部求解后做一次全局的秩自适应调整对病态问题更稳。接口不用变只要新写的函数满足「输入 TT 的 V 和 b、输出 TT 的 V」这个契约主循环一行不改。我自己的习惯是保留两套实现用一个cfg.inner_solver开关切换跑对比实验的时候直接横着比。验证方法我固定用一套任何改动之后先跑demo_lqr把值函数和 Riccati 解的最大相对误差打出来超过1e-4就停下来查再跑demo_pendulum看策略图有没有出现棋盘噪点最后才上高维。这三步加起来不到五分钟但省下来的时间远超投入。从那以后我每次改完core/里的任何一个函数都强制走一遍这个三步验证哪怕只改了一个reshape——血泪经验是TT 代码里一个reshape写反能让结果看起来「差不多对」而实际上完全错。希望帮到你。本文还有配套的精品资源点击获取