聲明:本文來自于微信公眾號 新智元，作者:新智元，授權微新創想轉載發布。

【新智元導讀】AI可控核聚變新突破!DeepMind最新研究用RL算法，將等離子體形狀精度提升65%，為「人造太陽」精準放電指明道路。

AI可控核聚變，指日可待。

秘密研發3年，DeepMind去年宣稱，首次成功用AI控制「托卡馬克」內部等離子體。其重磅成果登上Nature。

時隔一年，谷歌AI團隊在這一領域再次取得突破。

最新實驗模擬中，將等離子體形狀精度提高了65%。

DeepMind團隊基于上次的研究，對智能體架構和訓練過程提出了算法改進。

研究發現，等離子形狀精度提高的同時，還降低了電流的穩態誤差。

甚至，學習新任務所需的訓練時間減少了3倍還要多。

論文地址:https://arxiv.org/pdf/2307.11546.pdf

從「星際爭霸」AI碾壓人類，到AlphaGo大戰李世石、AI預測蛋白質折疊，DeepMind已經將人工智能算法深入到了足以改世界的不同領域。

這次，DeepMind最細實驗模擬結果，為RL實現精確放電指明了道路。

這一里程碑式的成果，標志著「人造太陽」可控放電離人類終極能源的未來又進了一步。

RL揭開核聚變奧秘

一旦人類掌握了可控核聚變能，將可擁有無窮不盡的清潔能源。

要知道，反饋控制對于「托卡馬克裝置」的運行至關重要。

而控制系統會主動管理磁線圈，以控制拉長離子體的不穩定性，防止破壞性的垂直事件發生。

此外，人類若能實現對等離子體電流、位置和形狀的精確控制，還可以實現熱排放，甚至對其能量的管理。

一直以來，科學家們致力于研究等離子體配置變化對這些相關量的影響。因此就需要能夠用于新配置，以及圍繞標稱場景快速變化的系統。

傳統上，等離子體的精確控制是通過等離子體電流、形狀和位置的連續閉環來實現的。

在這種模式下，控制設計者預先計算出一組前饋線圈電流，然后為每個受控量建立反饋回路。等離子體形狀和位置無法直接測量，必須通過磁場測量實時間接估算。

尤其是等離子體的形狀，必須使用平衡重構代碼進行實時估算。

雖然這類系統已成功穩定了大范圍的放電，但其設計不僅具有挑戰性，還耗時，特別是針對新型等離子體情況。

值得一體的是，強化學習（RL）已成為構建實時控制系統的另一種全新范式。

2022年，DeepMind團隊登上Nature的一篇論文表明，RL設計的系統能夠成功實現「托卡馬克磁控制」的主要功能。

論文地址:https://www.nature.com/articles/s41586-021-04301-9

這項工作提出了一個系統，RL智能體通過與FGE 托卡馬克模擬器交互，學習控制托卡馬克配置變量（TCV）。

智能體學習的控制策略隨后被集成到TCV控制系統中，通過觀察TCV的磁場測量，并為所有19個磁控線圈輸出控制指令。

尤其，研究人員展示了RL智能體控制各種情況的能力，包括高度拉長的等離子體、雪花。

甚至還展示了同時在真空室中，使用兩個獨立等離子體穩定「液滴 」配置的新方法。

AI控制下生成的幾種不同等離子幾何形狀

但是，RL方法有許多缺點，限制了其作為控制托卡馬克等離子體的實用解決方案的應用。

最新研究中，DeepMind決定要解決三個挑戰:

– 指定一個既可學習又能激發精確控制器性能的標量獎勵函數

– 追蹤誤差的穩態誤差

– 較長的訓練時間

首先，團隊提出了「獎勵塑形」的方法，以提高控制精度。

然后，通過向智能體提供明確的錯誤信號，和集成錯誤信號來解決積分器反饋中的穩態誤差問題。這縮小了經典控制器和強化學習控制器之間的精度差距。

最后，在片段分塊和遷移學習中，解決了生成控制策略所需的訓練時間問題。

研究人員針對復雜的放電情況采用了多重啟動方法，使得訓練時間大幅縮減。

此外，研究還表明，當相關新情景與之前的情景接近時，使用現有控制策略進行熱啟動訓練，是一種非常有效的工具。

總之，這些技術大大縮短了訓練時間，提高了精確度，從而使RL成為等離子體控制的常規可用技術取得了長足進步。

強化學習控制等離子體

最新論文中，研究人員采用與Nature那篇論文相同的基本實驗。

RL通過與模擬環境的交互，學習特定實驗的控制策略?，然后TCV上部署由此產生的放電策略。

具體來講，使用自由邊界模擬器FGE進行動態建模，并添加了額外隨機性，以模擬傳感器值和電源的噪聲，并改變等離子體的參數。

傳感器噪聲適用于每個環境步驟，而等離子體參數變化（等離子體電阻率??、歸一化等離子體壓力??、等離子體軸安全系數

）則經過簡化，因此其值在一個事件內是恒定的，但在兩個事件之間隨機取樣。

然后，研究人員使用最大后驗優化（MPO）算法來制定控制策略。

MPO依靠兩個神經網絡:一個是輸出當前策略?的actor網絡，另一個是近似該策略預期累積獎勵的critic網絡。

智能體與1000份FGE環境進行交互，收集看到的觀察結果、采取的行動，以及獲得的獎勵。

每一步獲得的獎勵，都是根據等離子體狀態與參考值中包含的目標值的接近程度來計算的，并輔以其他因素，如避免不良等離子體狀態。

從最優控制范式到強化學習的直接轉換是，為每個要最小化的誤差項設置一個獎勵分量，其中每個分量?都被映射為一個標量值??。

然后將這些值合并為一個標量獎勵值。

根據觀察、行動和獎勵的記錄序列，智能體使用正則化損失函數上的梯度下降交替更新策略和critic網絡。更新后的actor網絡參數將用于未來與環境的交互。

對于等離子體放電，actor網絡被限制在一個能以10kHz頻率執行的小型架構中，但critic網絡只在訓練過程中使用，因此可以足夠復雜地學習環境動態。

面向實用的的強化學習控制器

在具體任務實操中，研究人員演示了智能體具體訓練過程。

首先討論了通過獎勵塑形來提高控制精度。然后介紹了通過積分觀測來減少穩態誤差的工作，討論了使用「episode chunking」來改善現實的訓練時間。最后探討了遷移學習作為提高訓練效率的手段。

獎勵塑形（reward shaping）

傳統控制算法用各種辦法來最小化主動測量（或估計）的數量誤差，而強化學習(RL)算法則旨在最大化一個通用定義的獎勵信號。

在訓練過程中，這種獎勵最大化目標能推動智能體行為的演化，但是在部署時不會計算獎勵值。

在經典控制算法中，控制器的性能可以通過顯式調整控制增益（例如，修改響應性或干擾抑制）和調整多項輸入多項輸出(MIMO)系統的權衡權重來進行調整。

相比之下，在強化學習中，獎勵函數對于被學習的控制器行為至關重要。

因此，需要仔細設計獎勵函數來調整控制器行為。

在本節中，研究人員探討了如何修改獎勵的設計，以引發最終訓練得到的智能體去進行我們所期望的行為。

研究人員發現，通過調整獎勵函數的設計，他們可以快速適應智能體的行為，并權衡目標的不同方面。

此外，研究人員證明了塑形獎勵函數對于創建準確的強化學習控制策略是必不可少的。

而且他們進一步展示了通過使用更新后的獎勵函數繼續訓練，可以將智能體應用到新的目標上。

獎勵設計介紹

研究人員在先前研究的基礎上修改了為磁控而設計的獎勵函數。

研究人員使用加權的SmoothMax函數來組合獎勵組件的值。

在某些情況下，一個單獨的獎勵組件由多個相關的誤差量構成，比如在多個控制點處的形狀誤差。

研究人員還利用SmoothMax函數將這些誤差組合成一個單一的標量獎勵組件。

SmoothMax函數的定義如下所示:

許多喂給SmoothMax函數的單獨組件的構建方式與經典控制器類似（例如，將等離子體電流保持接近期望值）。

然而，獎勵組件并不受限于從傳感器測量中獲得，這在構建中就能提供了額外的靈活性。

獎勵組件還可以是多模態的，例如鼓勵智能體遠離狀態空間中不理想或模擬器建模較差的區域。

研究人員使用用SoftPlus轉換來獲得標量獎勵組件:

理論上，許多參數的選擇應該是近似等效的，因為它們是獎勵的單調調整，不應該對最優策略產生很大影響。

然而，在實踐中，研究者依賴于梯度下降（gradient descent），并沒有一個完美的全局優化器(global optimizer)。

研究人員需要在面對隨機回報的情況下探索全局空間。

很好和很差的緊密值使得很難找到任何可觀的獎勵區域（或者在如何改進方面有明顯的梯度）。

另一方面，較寬松的很差值使得更容易找到獎勵信號，但更難以發現精確的控制，因為改進時獎勵變化較小。

直觀上，因此，「緊密」獎勵參數可能更適用于初始條件接近目標狀態的情況，因此獎勵不需要塑造目標發現，而應更注重精確性。

在簡單環境中的獎勵塑形

在研究人員的初始實驗中，考慮了三種訓練方法，重點是通過修改「shape_70166」任務中形狀誤差的獎勵組件的超參數來最小化形狀誤差。

1. 基準線:采用之前研究的默認獎勵參數 – good =0.005，bad =0.05。

參考值產生了一個較為寬松的獎勵函數，該設置使獎勵信號集中在較高的誤差值，對于較小的誤差值也提供了引導信號，激勵增加形狀控制的準確性。

2. 窄化獎勵:將參數更新為good =0和bad =0.025。

這些參考值產生了一個更為嚴格的獎勵函數。該設置將獎勵信號集中在較低的誤差值，甚至對于小的誤差值也提供了引導信號，鼓勵在控制形狀時提高準確性。

3. 獎勵調度（reward schedule）:將good和bad的值在訓練過程中逐漸調整為更加尖峰(more Peaked)，good =0，bad從0.1逐漸減少到0.025，共進行600萬次策略更新步驟。

該調度在訓練開始時提供了一個較寬的獎勵區域來幫助探索，隨著訓練的進行逐漸收緊獎勵函數，以鼓勵準確性。

歷史數據在獎勵函數演變過程中不會重新被標記，但過時的數據最終會從學習智能體的回放緩沖區中消失。

這一系列的實驗結果如下圖所示。該研究證明了用于訓練的獎勵選擇對最終訓練的智能體的性能有著顯著影響。

通過對形狀誤差的關注，研究人員注意到對最終智能體性能影響最大的是采用了高度嚴格的靜態獎勵函數的「窄化獎勵」。

在這個簡單的任務中，更精確的獎勵函數為控制器提供了強烈的準確性激勵。

盡管如上所述，這樣尖銳的獎勵信號可能會對策略發現造成影響，但該任務的目標是保持交接位置，因此在這個任務中探索并不是一個主要的挑戰。

由于幾乎不需要探索來找到高度獎勵的狀態，智能體可以專注于滿足嚴格的獎勵信號。

此外，任務的簡單性意味著在獎勵組件之間準確控制很少或幾乎不需要權衡取舍（trade off）。

復雜任務的獎勵塑形

研究人員轉向「snowflake_to_perfect」任務，這個任務訓練成本更高，獎勵調整更為復雜，因為涉及到時變目標和更多的關注指標。

而且他們試圖通過獎勵塑形來提高X點位置的準確性。

以下是針對X點位置準確性的獎勵塑形方法:

1. 基準線:使用從Degrave等人先前的采取的默認參數進行訓練 good =0.005，bad =0.05。

2. X點微調（X-Point Fine Tuned）:首先使用默認參數進行訓練，然后進行第二階段的訓練，使用更為嚴格的獎勵，強調X點位置的準確性 — good =0，bad =0.025。

3. 窄化X點獎勵（Narrow X-Point Reward）:從訓練開始就使用更為嚴格的獎勵函數 — good =0，bad =0.025。

4. 額外訓練:在不更新獎勵函數的情況下進行額外的訓練。這樣使得研究人員能區分更多訓練和改變獎勵函數所帶來的影響。

研究人員比較了上述四種不同的訓練配置的性能，結果總結在下表中。

積分器（integrator）反饋

積分誤差的近似可以通過遞歸神經網絡來計算，然而，它們更容易過度擬合仿真動態。

在這項工作中，研究人員采用了一種更簡單的解決方案:沒有讓策略（policy）學習積分誤差，而是手動計算它，并將其附加到前饋策略所觀察到的觀測集中。

他們特別關注了減少等離子體電流（??）的穩態誤差，之前研究的的訓練策略表現出明顯的偏差，并且該誤差可以很容易地計算。

與傳統方法稍有不同，研究人員向網絡提供了時間?的平均等離子體電流誤差定義如下:

研究人員在「shape_70166」任務中評估了將平均誤差信號納入考慮的好處。

在該任務中，等離子體電流和形狀的參考值是恒定的，環境初始化后實際值接近參考值。

因此，智能體的主要目標是控制穩態誤差（steady-state）。

下圖顯示了使用積分器反饋訓練和未使用積分器反饋訓練的策略的模擬等離子體電流誤差軌跡，每種情況下進行了三次隨機運行。

研究人員發現，積分器反饋顯著降低了等離子體電流偏差，正如預期的那樣。

Episode Chunking

在TCV上的實驗持續1-2秒，相當于以10kHz的控制頻率進行10，000-20，000個時間步。

FGE模擬器（如上所述用于訓練智能體）在訓練過程中使用一顆AMD EPYC7B12CPU核心，每個典型的模擬步驟大約需要2秒鐘，使用隨機動作。

因此，FGE生成包含10，000個步驟的一次完整episode大約需要5小時的時間。

這意味著在最理想的情況下，即智能體在第一次嘗試之前已經知道最佳策略，訓練時間仍然會約為5小時（以觀察高質量的結果）。

實際上，強化學習智能體需要探索動作空間以找到最佳策略。因此，根據任務復雜性，訓練時間可能從幾天到幾周不等。

此外，研究人員的任務結構使得智能體需要按順序學習相對獨立的「技能」。例如，在「showcase_xpoint」任務中，智能體必須先使等離子體變形，然后移動其垂直位置，然后改變其流向，最后恢復原始形狀（參見下圖1）。研究人員觀察到該任務的學習過程發生在兩個明顯的階段(見下圖2a)。

首先，智能體學會操作有限的等離子體，理解如何延展、移動和保持等離子體，這對應于獎勵曲線，就是從0平滑上升至約80。

在此階段，智能體嘗試（但失敗了）生成一個轉向形狀，取而代之的是獲得具有非活動X點的圓形LCFS，如上圖b所示。

獎勵在此水平上保持穩定，直到最后，智能體發現如何成功地將等離子體轉向，這時獎勵值從80突變至接近1。

將分塊（chunking）技術應用于展示_x點(showcase_xpoint)任務，并分別使用兩個/三個塊(如下圖一所示)，可以顯著縮短訓練時間，如下圖2所示。

兩個塊的設置（橙色曲線）已經比基準線(藍色曲線)更快。三個塊的設置(3_chunks和3_chunks_eq_weights)不僅提供進一步的訓練加速，而且學習曲線更加平滑。

智能體在約10小時內就能達到96（滿分100）的獎勵，而基準線需要40小時。

在這里，研究人員嘗試了兩種不同的三塊設置:所有參與者（actor）被平均分為相同大小的組(3_chunks_eq_weights);與每個其他塊相比，整個episode使用三倍更多的參與者。這兩種設置給出了類似的結果。

轉移學習

在試圖減少訓練時間時，一個自然的問題是問是否可以重用之前放電時訓練的模型，也就是說，智能體在解決一個初始任務時積累的知識在多大程度上可以轉移到一個相關的目標任務上。

研究人員以兩種形式考察遷移學習的性能:

1.零樣本（Zero-shot）:研究人員在目標任務上運行在初始任務上學習的策略，而無需進行任何額外的數據收集或策略參數更新。

2.微調（Fine tuning）:研究人員使用在初始任務上學習的模型的權重來初始化策略和值函數，然后使用這些權重在新的目標任務上通過與環境交互進行訓練，其中目標任務作為獎勵。需要注意的是，這要求在兩個任務中使用相同的架構(actor和critic網絡)。

在兩種情況下，研究人員使用在showcase_xpoint任務上訓練的智能體參數作為遷移的初始參數。

在第一個實驗中，研究人員考察當參考等離子體電流調整到新的參考水平時的遷移學習。

具體而言，研究人員選擇了三種變化，其中目標??從基準線-150kA調整到-160kA，然后-170kA，最后-100kA（具體而言，在圖1中除了初始交接水平和最終降溫水平外的所有時間片中調整參考電流）。

研究人員測試了在showcase_xpoint上訓練的策略，首先在目標任務上沒有任何額外訓練，然后允許在目標任務上進行新的訓練。

零樣本結果的獎勵和??誤差如下表所示，在小的??變化情況下，智能體表現良好，但在較大的變化情況下，尤其是對于較大的??變化，智能體表現較差。

微調的結果如下圖a、b、c所示，微調智能體在所有情況下比從頭開始訓練的智能體更快地收斂到近乎最優的策略，盡管在最大的50??變化情況下差異較小。

第二個實驗考察了等離子體目標位置的變化。

具體而言，研究人員沿著z軸向下調整目標形狀，分別平移2厘米、10厘米和20厘米。對于這個實驗，研究人員觀察到以下結果:

1. 零樣本（Zero-shot）:結果如下表所示。研究人員發現對于最小的平移(2厘米)，零樣本遷移效果非常好，任務的表現達到了最佳可實現性能的97%以上(滿分100分)，形狀誤差也很小。

對于較大的10厘米平移，表現較為一般，只獲得了85的獎勵，并且形狀位置誤差更大。對于最大的20厘米平移，表現較差，只獲得了35的獎勵，由于未能成功轉向等離子體。

2. 微調（Fine tuning）:微調的結果如上圖d、e、f所示，表明對于2厘米的平移，遷移學習效果顯著，對于10厘米平移，三個不同的種子中有兩個種子的效果有效。而對于較大的20厘米平移，遷移學習似乎對性能產生了不利影響。

總體而言，結果表明遷移學習在當前形式下是有用的，但也有一定的局限性。

正如預期的那樣，目標任務與初始任務之間的差距越大，遷移學習的性能就會降低，尤其是在零樣本學習的情況下。

然而，值得注意的是，在運行硬件實驗之前，通過模擬進行零樣本評估的成本相對較低（以CPU小時為單位）。

研究人員還發現，某些類型的任務變化比其他任務更容易進行遷移學習，在他們的實驗中，相對較大的等離子體電流變化似乎更適合于遷移學習，而不是大的位置變化，這在考慮到任務的相對復雜性時是可以理解的。

需要進一步研究來了解哪些任務適合于遷移學習，并如何擴展有效遷移的范圍，包括零樣本和微調學習。

TCV上的托卡馬克放電實驗

之前的部分僅關注使用FGE模擬器進行仿真、訓練和評估控制策略。

考慮到托卡馬克建模（Tokamak modeling）的復雜性和挑戰，重要的是不能盲目地認為仿真中的性能改進與實際放電中的性能改進完全相同。

雖然更好的仿真結果可能對實際托卡馬克的改進結果是必要的，但往往是不夠的。

如果沒有額外明確的工作來減小仿真與實際之間的差距，模型不匹配誤差可能會變成一個很主要的問題。

對于使用強化學習獲得的策略，已知會過度擬合到不完美的模擬器，這種情況尤為明顯。

因此，研究人員在TCV托卡馬克上對一些上述的仿真改進進行了測試。

通過這種方式，研究人員可以評估當前工作的優勢和局限性，并為下一步的改進提供方向。

等離子體形狀精度的獎勵塑形

研究人員檢查了獎勵塑形在兩種不同配置和目標上所帶來的精度改進:減少形狀穩定任務中的LCFS誤差和提高「snowflake_to_perfect」任務配置中的X點精度。

研究人員將模擬結果與TCV上的實驗結果以及來自Degrave等人（2022）的可比實驗進行了比較。與先前的研究一樣，研究人員通過將演員網絡(由JAX圖定義)創建為共享庫對象來部署控制策略，其中命令的動作是輸出高斯分布的均值。

研究人員首先測試了一個控制策略，該策略通過在獎勵塑形部分中討論的獎勵塑形方法來減少shape_70166穩定任務中的LCFS誤差。

對于這個穩定任務，研究人員使用了TCV的標準擊穿過程和初始等離子體控制器。在0.45秒時，控制權移交給學習的控制策略，然后它試圖在1秒的持續時間內維持固定的等離子體電流和形狀。

放電后，研究人員使用LIUQE代碼計算重構的平衡態。在1秒的放電過程中的每個0.1毫秒時間片內，研究人員計算等離子體形狀的誤差。研究人員比較了三個實驗的精度，分別從模擬放電和TCV放電中測量形狀誤差:

（a） 一種在本研究之前已經存在的基線RL控制器(「Previous」)， (b) 一種使用本研究中更新的訓練基礎設施的更新的基線代理(「Updated」)， (c) 一種使用獎勵塑形訓練的代理，就像在獎勵塑形部分描述的Fixed Reward一樣。

這些運行的結果在下表中。

X點位置精度的獎勵塑形

接下來，研究人員將比較獎勵塑形對更復雜的「snowflake」配置的影響，如下圖所示。

該策略的訓練獎勵被塑形以增加X點控制的準確性。

與穩定實驗中一樣，等離子體是通過標準的TCV程序創建和初始控制的，在0.45秒時將控制權移交給強化學習控制器。

在這個實驗中，RL訓練的策略成功地建立了一個兩個X點距離為34厘米的「snowflake」。

然后，該策略成功將兩個X點帶到了目標距離6.7厘米的位置，接近建立一個所謂的「完美snowflake」。

然而，在1.0278秒（即交接后的0.5778秒），等離子體因垂直不穩定性而發生破裂。

經檢查，發現控制器在保持一致形狀方面存在困難，其中垂直振蕩增加，活動的X點在兩個X點之間切換，導致失控。

下表顯示了在等離子體成功控制期間對X點追蹤的準確性。

通過 「Episode Chunking 」來加速訓練

最后，研究人員驗證了使用「Episode Chunking」來減少訓練時間，特別是驗證在TCV放電中是否出現可能的「不連續性」。

研究人員進行了一個在showcase配置下使用3個塊進行訓練的實驗。這個實驗的重建平衡態的時間軌跡可以在下圖中看到。

研究人員發現實驗按預期進行，沒有因為「episode chunking」而產生明顯的偽影。

這證明了這種訓練加速方法沒有損失質量。

20世紀50年代起，眾多科學家們致力于探索、攻克可控核聚變這一難題。

DeepMind最新研究，用強化學習算法大幅提升了等離子體的精度，極大縮短了學習新任務的訓練時間。

這為可控核聚變在未來實現「精準放電」，能量管理鋪平了道路。

在為人類獲取海量清潔能源，以改變未來的能源路線圖上，DeepMind再次點亮了一盞明燈。

參考資料:

https://arxiv.org/abs/2307.11546

https://twitter.com/GoogleDeepMind/status/1684217852289601541