[詳解] Deepmind如何用強化學習控制核聚變

Deepmind是個很有意思的公司，在Alpha go之後他們一度沉寂了一小段時間，一度被負面新聞籠罩，和google鬧不愉快啦，不盈利被嫌棄啦，等等等等。但是最近的一系列機器引人矚目的研究，包括將深度學習/深度強化學習應用在醫藥（Alpha Fold［1］/Alpha Fold2［2］）， 應用在數學研究［3］上，應用在語言研究上［4］，用在寫程式碼上［8］，甚至用在考古學上［5］，和核聚變上［6，7］。

Deepmind的研究方向可謂非常寬廣

Deepmind最近又在Nature上發表了使用強化學習訓練控制器來進行核聚變控制的文章，我仔細閱讀了這篇文章後，覺得這個研究非常有趣，於是寫一篇文章專門來聊一下。

問題-Deepmind到底想控制啥

那麼到底Deepmind想要控制的是什麼呢？官方部落格也非常貼心，深怕我們看不懂，給出了一個實際控制的實時影片：

DeepMind控制核聚變等離子體形狀

https：//www。zhihu。com/video/1487396063706968064

上面的這個影片中，左邊是就是聚變反應器內實時控制聚變反應中的等離子體的形狀的錄影，右邊是與其對應的等離子體形狀的截面圖，Deepmind想要控制的就是鉅變容器內等離子體的形狀。

眾所周知（也許並不是那麼眾所周知），上圖的這個裝置叫做

託卡馬克

裝置，是一種用強大的電磁力約束核聚時超高溫度等離子體的裝置，如果有一天人類找到了能夠長期維持和控制這種等離子體的方法，那麼我們就會擁有無窮無盡的能量，不用再擔心電費了。

而為什麼要控制等離子體的形狀呢？因為這可是託卡馬克裝置的一個主要研究方向。比如中科院等離子體物理研究所也有一個

託卡馬克

裝置，叫做

EAST

（實驗“Experimental”、先進“Advanced”、超導“Superconducting”、託卡馬克“Tokamak”）。在它的官方介紹裡，有這麼一段：

。。。它不僅是一個全超導託卡馬克（左圖為託卡馬克示意圖），而且

具有會改善等離子體約束狀況的大拉長非圓截面的等離子體位形

，它的建成使我國成為世界上少數幾個擁有這種型別超導託卡馬克裝置的國家，使我國磁約束核聚變研究進入世界前沿。。。

中科院的EAST裝置，是一個比較小巧的託卡馬克裝置，大小大約是國際熱核聚變實驗堆（ITER）的四分之一

介紹中還特別註明了，中科院等離子體物理研究所的這個託卡馬克，是可以拉出

非圓形的等離子體

的哦。在本來就不長的介紹文章裡，還專門分出一段介紹它的這個能力，足以說明其重要性。

不同形狀的等離子體具有不同的散熱，穩定等屬性，如果我們想用核聚變發電（俗稱燒開水），那麼我們就需要熱量能夠傳匯出來，用來燒開水發電，不能完全的約束在等離子體中，但是又不能太熱，太熱容器就融化了，所以很大程度上需要一種動態控制的手段，而核聚變的容器本身是固定的沒法改變形狀，那麼就只能透過控制核聚變離子體的形狀等來控制散熱，穩定等等屬性。

隨意控制等離子體形狀的能力無疑是朝可控核聚變徹底被攻克方向邁的堅實的一步。

方法-Deepmind怎麼利用強化學習做控制？

輸入輸出，以及問題建模

託卡馬克裝置的控制系統，包括16個極向線圈控制器，2個電阻線圈控制器，和1個快速線圈控制器

想要控制託卡馬克裝置，必須先要看看它的輸入輸出有哪些，與Deepmind合作的是瑞士等離子體中心（SPC），他們的託卡馬克裝置有19個

輸入

控制線圈，包括16個極向線圈控制器，2個電阻線圈控制器，和1個快速線圈控制器，Deepmind將會透過直接控制這19個線圈中的電流來控制託卡馬克裝置。

那麼

輸出

呢？當然是透過感測器啦，瑞士等離子體中心（SPC）的託卡馬克的感測器分為兩種，實時的和非實時的。實時的感測器包括一些電磁傳場感器等等，非實時的感測器包括攝像頭等更加複雜的器械。Deepmind的演算法中，僅僅使用了實時感測器，準確的說，他們使用了34個線圈感測器來測量磁通量的迴路，38個prob來檢測區域性測場，加上19個控制線圈中的19個電流感測器。透過這些感測器輸出，他們可以實時還原出託卡馬克裝置中流體的性質，比如大致形狀等等。

攝像頭感測器

而要定義控制結果的好壞，就需要使用一個“獎勵函式”，來量化最後的結果，Deepmind使用的獎勵函式包括好幾個項，有磁流體和目標的形狀差別，磁流體的X-point（一個特徵點）和目標的差別，等等，相對簡單但構成比較複雜，考慮到應該沒有多少人能去復現這篇論文，這塊就不做過多展開了。

在Alpha go［10］和Alpha Star［9］ 之後，大家都知道，強化學習可以用於尋找一些環境中的最優策略，來實現獎勵的最大化，而且運用得當的話，找出來的最優策略會還不錯。但是強化學習有個毛病，就是一個強化學習的ai需要少則幾千多則百萬的學習迭代才能達到一個比較好的效果，而核聚變的研究裝置（

託卡馬克

裝置非常貴，而且少）大家都是排隊來的，不可能讓你一家霸佔著幾年慢慢做實驗，而且某次實驗要是ai沒控制好，輕則熄火重則爆炸，於是open ai只能使用託卡馬克裝置的

模擬器

［12］來進行強化學習ai的訓練。

託卡馬克模擬器的一個結果，從圖中可以看出其可以模擬等離子體溫度，形狀等

trick1 - 用數量彌補速度

這個時候deepmind就遇到了第一個問題，託卡馬克裝置模擬器

非常慢

，因為需要模擬很多物理作用之類的。Deepmind透過堆疊算力的方法解決這個問題，直接利用自己充裕的計算資源並行啟動了5000個模擬器，要知道託卡馬克模擬器是個很複雜的出程式，每一個例項都需要大量的運算資源，5000個並行環境要消耗的資源已經遠遠不是學校或者小公司可以觸及的，同時執行這麼多的託卡馬克模擬器，一股豪氣撲面而來，直接用數量彌補了效率上的問題。

trick2 - “非對稱”Actor Critic

解決了第一個問題之後馬上deepmind就遇上了第二個問題，實際核聚變控制中，對於控制器的效能要求非常高，一次狀態訊號輸入到控制訊號輸出最多隻能有50 μs的時間，這個時間是一毫秒的20分之一。這麼短的時間直接對深度強化學習中的神經網路結構提出了很大的挑戰，有做過深度強化學習的應該都知道，就算很簡單的網路，在很好的cpu/gpu上，也很少能跑到1毫秒以內的。

現在流行的強化學習框架中一般有兩個角色，actor和critic，actor負責給出行動建議，critic負責評估當前的局面。

舉例：在託卡馬克裝置中，Actor會負責判斷應該怎麼輸出控制訊號，Critic會負責判斷當前的狀態是否符合預期

Actor和Critic在訓練的時候兩者都是需要的，而在部署的時候，只需要actor，原因也很簡單，因為ai在部署的時候一般不關心當前局勢好不好，只要給出行動建議就行。

於是，deepmind為了解決這個控制器的效能問題，提出了一種方案，原本在強化學習中，一般來說Actor和Critic的網路都是差不多大小的，而由於應用場景要求極快的運算速度，deepmind提出使用“非對稱”的actor-critic結構，讓Actor擁有更簡單的網路結構，而Critic擁有更復雜的結構，這樣就可以在訓練過程中仍然可以吃到模型較為複雜的Critic的紅利，而在部署時又擁有比較快的速度。

Deepmind使用較小的Actor和較大的Critic，來獲得更適合核聚變的AI模型

透過這種方法，Deepmind很大程度上解決了速度的問題。

trick3 - 繞過模擬器不準的問題

模擬器畢竟是模擬器，很難完全還原真實世界的所有屬性，這個問題不僅在核聚變時會遇到，在很多其他場景也會遇到，比如Deepmind 的另一個應用，在模擬器中學會轉魔方，然後遷移到現實的機械臂［13］，也遇到了這個問題。

仿生手轉魔方，在模擬器中

https：//www。zhihu。com/video/1489239592260304897

仿生手轉魔方-現實

https：//www。zhihu。com/video/1489241223747174400

這類從模擬器到現實轉換時的問題統稱為sim2real問題，要想解決這種問題，方法有很多，但是基本都不適合這裡的核聚變場景。因為Deepmind需要實現Zero-shot的sim2real，也就是說，在模擬器上訓練之後，不進行實際環境的訓練或者微調，而直接應用在實際場景，並且由於核聚變的特殊性，在實際場景中失敗代價極大，所以方案要儘可能保守。

不僅如此，Deepmind還發現，在某些條件下，模擬器的模擬結果會和真實結果相差甚遠（幸好大部分情況下，模擬器都是準確的），這近一步增加了強化學習訓練的難度。

不過Deepmind還是找到了解決方法，他們管這種方法叫做

learned-region avoidance(學習區域規避)

，簡單來說，就是主動人為的去識別模擬器中模擬不好的部分，並且設計算法去規避它們。

舉一個例子，比如我們想讓強化學習在模擬器中學習轉魔方，然後應用到現實當中，但是我們發現模擬器中當機械手把魔方扔到空中的時候，模擬結果會和現實差別大：

假設：把魔方扔到空中，模擬器和現實差別比較大

那麼deepmind的方法就是，設計一個獎勵函式，鼓勵機械臂不要把魔方扔到空中去。

設計獎勵函式，讓ai主動不去把魔方拋起來

透過這樣的方法，就可以讓模擬器中的環境對於強化學習ai來說盡可能貼近現實。

trick4-工程最佳化，極限加速

前文提到過，在核聚變中決策的時間只有50 μs，一毫秒的1/20，這個時間限制對工程能力也有極高的要求，如果按照tensorflow一般的部署流程，匯出graph，使用tf serving，肯定是達不到效能要求的。Deepmind為了解決這個問題，使用tfcompile ［14］ 來將網路轉化成可執行的二進位制c程式碼，並且他們小心的掉一些無用的依賴來減小程式碼量和執行時間，並且定製了網路結構來讓推理的時候能夠充分利用上cpu的cache，達到效率的最大化。

雖然Deepmind在這部分沒有特別展開來講，但是光看這些文字就能感覺到，他們這裡也沒少踩坑。

Deepmind的狼子野心

不得不說Deepmind這幾年的發展真的令人刮目相看。似乎他們什麼都做，每個領域上都能做出有影響力，非常前沿的東西，真的是非常令人羨慕。

其實國內也有一家公司，他什麼都做，在每個領域都能賺錢

“有什麼業務是騰訊不做的嗎？”美團網CEO王興的語氣中難掩鬱悶。 —— 計算機世界2010《“XX的”騰訊——攪局者還是終結者》

這家公司也做過很多獨創性的前沿研究，比如。。。額。。。。嗯。。比如。。。。。嗯，有點難度，確實編不出來。

哈哈，扯遠了。Deepmind這幾年的戰略越來越明顯，就是“由虛向實，由簡單向複雜”。

最早Deepmind起家靠的是遊戲，AlphaGo戰勝了最強的人類，隨後的AlphaZero / Muzero / AlphaStar等一些列研究解決了非常多的遊戲問題。如果Deepmind是國內一些大網際網路公司的子公司，那麼可以想像，Deepmind會搖身一變成為很多遊戲的ai演算法提供商，目標也會是“世界一流的遊戲演算法團隊”，如果母公司是阿里或類似的公司的話，很快就會要求Deepmind透過提供遊戲演算法盈利，一切不盈利的專案都會被砍掉，估計過個三五年就泯然眾人矣。

但是Deepmind明顯並沒有止步於遊戲，他們繼續向前，在研發出AlphaFold ［1］這種劃時代的生物演算法的時候，他們沒有說 “我們要大力賦能醫藥領域”，在將強化學習應用到核聚變［6］之後，他們也沒有說 “我們要加大力量賦能核聚變”。

而是繼續向前，在他們走過的路上，除了腳印，還留下了一座又一座的金礦。

Deepmind公司CEO Demis Hassabis

另一個令人震驚的數字，deepmind目前只有1000多個員工，與之對比的是據我所知，無論是國內的商湯還是曠視，光光做人臉識別的人數就已經遠遠超出這個數量。

要給國內的這些公司找理由的話，我張口就能來很多：環境不同啦，員工質量不同啦，公司耐心不同啦，不賺錢啦，等等等等。但是，不做一件事情有一萬個理由，做卻只有一個：這件事情值得做，應該做。

理論上來說，Deepmind並沒有花很多錢（每年也就幾個億刀），這些錢對於國內的很多網際網路公司根本是九牛一毛，根本是在財報的統計誤差之內，所以，為什麼不能做？為什麼不去做呢？

我並非這些公司的CEO，也沒有什麼特別遠大的抱負，沒法回答這些問題。2010年高考北京卷的作為題目叫“仰望星空與腳踏實地”，也許時不時的，我們確實應該抬起頭看看星星，如果你看得夠仔細，也許會發現，其中的一些，離我們也沒那麼遠。

引用

［1］ Senior A W， Evans R， Jumper J， et al。 Improved protein structure prediction using potentials from deep learning［J］。 Nature， 2020， 577（7792）： 706-710。

［2］ Jumper J， Evans R， Pritzel A， et al。 Highly accurate protein structure prediction with AlphaFold［J］。 Nature， 2021， 596（7873）： 583-589。

［3］ Exploring the beauty of pure mathematics in novel ways

［4］ Language modelling at scale

［5］ Predicting the past with Ithaca

［6］ Accelerating fusion science through learned plasma control

［7］ Degrave J， Felici F， Buchli J， et al。 Magnetic control of tokamak plasmas through deep reinforcement learning［J］。 Nature， 2022， 602（7897）： 414-419。

［8］ Competitive programming with AlphaCode

［9］ AlphaStar： Grandmaster level in StarCraft II using multi-agent reinforcement learning

［10］ Silver D， Huang A， Maddison C J， et al。 Mastering the game of Go with deep neural networks and tree search［J］。 nature， 2016， 529（7587）： 484-489。

［11］ 中國科學院等離子體物理研究所

［12］ Carpanese， F。 Development of Free-boundary Equilibrium and Transport Solvers for Simulation and Real-time Interpretation of Tokamak Experiments。 PhD thesis， EPFL

［13］

https：//

openai。com/blog/solving

-rubiks-cube/

［14］ Abadi， M。 et al。 TensorFlow： a system for large-scale machine learning。 In Proc。 12th USENIX Symposium on Operating Systems Design and Implementation （OSDI ’16） 265–283 （2016）。