DeepMindの深層強化学習による核融合炉制御論文を読んだので論文内容と論文を理解するために調べた技術背景をまとめます。
Accelerating fusion science through learned plasma control
※筆者は原子物理/核融合理論について完全な素人であり当該分野の記述の正しさについてなんら保証がありません。間違いがあればコメントにて指摘お願いします
要約
- トカマク核融合炉の磁気コイル群は複雑にネストしたPIDコントローラー群によって制御されている
- この制御方法で性能的な問題はない
- しかし制御系設計・構築コストの重さが研究開発のボトルネックになっている
- 深層強化学習でPID制御を置き換えることで制御システムの構築をとても簡単かつ低コストにした
- これからは研究者の新たなアイデアをすぐに実験投入することができるから核融合炉研究が捗る
技術背景：核融合炉の仕組み
核分裂エネルギーと核融合エネルギー
一般に、原子力発電とは核分裂反応によってエネルギーを取り出す方法を示します。この核分裂反応エネルギーとは大きな原子が分裂する際に生じるエネルギーです。核分裂反応による原子力発電は、CO2排出ほぼ無しで莫大なエネルギーを得られるという大きなメリットがある一方で、連鎖反応による暴走の危険性や使用済み燃料の廃棄および廃炉問題など多くの課題を抱えています。
一方、核融合炉では小さな原子核が融合し重原子核を生成する反応によってエネルギーを取り出します。CO２排出無しで莫大なエネルギーを得られるという従来の原子力発電のメリットはそのままに、連鎖反応による暴走の危険性が低く、高レベル放射性廃棄物を伴わず、かつ水素という普遍的な資源が燃料ということでしばしば「２１世紀の夢のエネルギー」という表現を与えられ、活発に研究開発が行われています。身近な例では太陽のエネルギー源もまた核融合反応です。
核融合反応のおこしかた
原子核と原子核を十分な速度（1000km/sec 以上）でぶつければ核融合反応を起こすことができます。そこでまずは原子を超高温に加熱することで原子運動を加速しプラズマ状態とすします。ここで、プラズマ相とは原子核と電子が自由に動きまわっている状態です。
しかしながら、原子核は正荷電しているために反発しあってすぐに散逸してしまうので単純に加熱しただけでは核融合に必要な原子核の衝突がほぼ起こりません。
実用的な核融合反応を起こすためにはプラズマを狭い空間に閉じ込めることで高い密度を保つ必要があります。お互いに反発しあうたくさんの原子核を核融合しないと出られない部屋に閉じ込めてやりましょう。また、できるだけ長く閉じ込めることができるほど原子核衝突≒核融合反応発生の可能性が上がります。
トカマク型磁気閉じ込め方式
ここまで、プラズマを狭い空間に閉じ込めることができれば核融合を実現できることがわかりました。しかしどうやって閉じ込めましょう？
太陽では超重力によるプラズマ圧縮で核融合を起こしているようですが、地球上では非現実的です。また、超合金的な素材で密閉容器を作るのは一見よさそうなアイデアですが、これをやると物凄い勢いで容器が削れる ＆ プラズマが冷えてしまうのでダメっぽいです。なんとかプラズマに直接触れないで圧縮する必要があります。
そこで、ドーナッツ状の強力な磁力線を発生させることでプラズマに直接触らず閉じ込めを実現するのがトカマク型磁気閉じ込め方式です。
大雑把かつ乱暴な表現をすると、ドーナッツ環に沿った磁力線（トロイダル方向磁力線）と ドーナッツ断面に沿った環状磁力線（ポロイダル方向磁力線）の組み合わせによってドーナッツ表面に沿ったらせん状の磁力線を発生させることにより、プラズマはドーナッツ環内を移動しつづけるためにプラズマの散逸を防ぐことができるようです。
深層強化学習によるトカマク型核融合炉の制御
研究者はさまざまなプラズマ形状の特性を探索したい
要約するとトカマク方式の核融合炉とは、プラズマを強力な磁力線によってドーナッツ型に圧縮することにより核融合反応を発生させる方式のようです。このトカマク方式の重要な研究課題の一つは閉じ込めの安定性やエネルギー取り出しを最適化するために、より良いプラズマ分布形状（ドーナッツの断面）を探索することです。
課題：目標形状ごとに磁気コイル制御システム構築するのがつらい
論文がターゲットとしているスイス/ローザンヌのトカマク炉ではプラズマの周囲に配置された19個の磁気コイルの精密制御によってプラズマ形状をコントロールしているとのことです*1。問題は、目標とするプラズマ形状ごとに磁気コイルの制御システムを実装する必要があることです。
制御系の実装には相当なエンジニアリング/設計作業および専門知識を必要とするため、研究者のアイデアを容易に実験で確認できず試行錯誤のコストが研究領域の進歩のボトルネックになってしまっています。
提案：深層強化学習で任意形状の制御を実現する
論文ではセンサー測定値と達成すべきプラズマ形状を入力として、19個の磁気コイルの操作値を出力する方策ネットワークの訓練が深層強化学習で実現できることを示しました。この系においては実験者が新たなプラズマ形状を実験したいときに行うべきことは目標とするプラズマ形状を指示するだけであるため、制御系の実装コストが大きく低減され、トカマク核融合炉の研究が大きく加速することが期待されます。
制御ポリシーのトレーニング
Sim to Real
トカマク炉は物理モデルに基づくシミュレータが利用可能ということで、まずはシミュレータ環境で方策の事前訓練を行った後、実機から収集されたサンプルでファインチューニングを行います。ただし、実機データからのファインチューニングについては厳しいリアルタイム制約からかオンライントレーニングを行ったわけではないようです。
学習回避領域の設定：
トカマクシミュレータ環境について、シミュレーションがうまく現実に一致しないことが事前に分かっている領域があるようです。そこで、指定された条件が発生したときにシミュレーションを停止することで、このような領域をエージェントが学習してしまうことを回避する仕組みを導入しています。
分散並列強化学習
多数の並列actorによって収集される多様な遷移情報が学習の安定性を向上させることが経験的に知られており、この論文でもオフポリシー分散並列強化学習アーキテクチャを採用しています。論文によると環境と相互作用して遷移情報を収集しReplayBufferに送信する多数のActorとひたすらネットワーク更新を繰り返すLearnerで構成されるApe-Xっぽいアーキテクチャとなっているようです。
MPOアルゴリズム
Maximum a Posteriori Policy Optimisation | OpenReview
Actor-Critic系のMaximum a Posteriori Policy Optimization (MPO)を採用して方策ネットワークを訓練します。MPOはオフポリシーゆえの高いサンプル効率とTRPOのような更新安定性を兼ね備えた使いやすいアルゴリズムです。*2 ざっくりとは方策関数がQ値のボルツマン分布を近似するように更新する手法です。
Actor-Criticネットワーク：
大きめかつRNNつきのCriticネットワークとたった４層MLPのPolicyネットワークという非対称な構造を採用しています。これはPolicyネットワークは実機環境におけるリアルタイム推論が必要なために十分に高速に動作しなければならないという実用上の制約のためであるようです。Criticネットワークはネットワーク更新時にしか使われないので動作が遅くともまったく問題ありません。方策ネットワークは92次元のセンサー測定値と132次元で表現される目標状態を入力され、19個のコイルそれぞれの電圧値を出力します。
目標状態tを方策に入力する：
132次元で表現される目的状態tがセンサ測定値mとともに方策関数に入力されます。この研究の目的は研究者の指定する任意のプラズマ形状を実現する制御ポリシーを訓練することなので実現すべきゴール状態ｔを方策に知らせる必要があるためです。階層強化学習でサブゴールをpolicyに入力するのと同様です。
即時報酬rの設計
即時報酬rは目標状態tと現在のセンサー観測値mに基づいて決定されるスカラ値です。とても大雑把にはセンサー観測mから推定される現在のプラズマ状態と目標プラズマ状態tが似ていれば高い即時報酬が与えられるような即時報酬関数になっているのですが、ドメイン知識に基づいたさまざまな指標（下表）が即時報酬要素として使われており、さらにベースとするプラズマ形状の種別に応じて即時報酬要素を取捨選択しているようなので報酬設計にはかなりの力が入っていることが推察されます。
実機へのデプロイ
実機への方策デプロイにおいては厳しいリアルタイム推論性能（50μs以内の応答）が求められるため、tfcompile（https://www.tensorflow.org/xla/tfcompile?hl=ja） で高速化を行っているとのことです。
性能検証
所感
PID制御でうまくいっている系を強化学習で置き換えることにパフォーマンス上の意味はないが、R&Dにおいて試行錯誤コストが大きく低減される意義は大きい、という着眼点がさすがだなと感じます。やってることはただの効率化でも、その効率化のケタが違うとゲームチェンジをもたらすというあたりは研究開発デジタルトランスフォーメーションの見本といった印象。
