【DL輪読会】Stabilizing Reinforcement Learning in Differentiable Multiphysics Simulation

DEEP LEARNING JP [DL Papers] Stabilizing Reinforcement Learning in Differentiable Multiphysics Simulation Ku Onoda, Matsuo-Iwasawa Lab, M1 http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 • Stabilizing Reinforcement Learning in Differentiable Multiphysics Simulation – ICLR 2025 (Spotlight) • 著者 – Eliot Xing, Vernon Luk, Jean Oh • リンク – https://rewarped.github.io/ 2

• https://rewarped.github.io/

概要 • 課題 – 変形体（ロープ・布・流体・粘土など）は剛体より計算負荷が高く、従来の RLではサンプル効率が致命的に低い • 本研究の貢献 1．Rewarped：GPU並列・微分可能なマルチフィジックス環境のシミュレータ 2．SAPO：最大エントロピー×解析勾配の FO-MBRL 3．様々なタスクで従来の手法（PPO/SAC/APG/SHAC）を上回る 3

背景① RL × 物理シミュレーション • RL 累積期待報酬を最大化する方策𝜋𝜃 のパラメータ𝜃を最適化 – 目的関数 – 歩行・把持・マニピュレーションなどのロボット制御への応用がある • シミュレータの役割 – 実機でのデータ収集は高コスト → シミュレーションで大量のデータ生成 – GPU並列環境により剛体タスクは数時間で学習が可能 • 現状の課題 – 変形体（deformable objects）では計算が2-3桁遅い – モデルフリーRLでは10^8 step以上要求され現実的でない 4

背景② 微分可能シミュレータ • 普通のシミュレータ – 遷移関数や報酬関数の内部がブラックボックス – 勾配の0次推定量 • 微分可能シミュレータ – シミュレータ内部の遷移関数・報酬関数内の計算を微分可能な演算のみに限定 – 方策のパラメータに対して累積報酬の勾配を解析的に利用可能 → より効率的で安定した学習を可能にする 5

背景③ 既存シミュレータ比較 • 「変形体対応」「微分可能」「並列化可能」全てを満たすものはない 6

論文の貢献 1．シミュレータ Rewarped：WarpベースでGPU上に全環境をバッチ展開 逆モード自動微分+CUDA Graphによる高速勾配計算 2．アルゴリズム SAPO：最大エントロピー FO-MBRL（First-Order Model Base RL） →探索と安定性を両立 3．実験的検証 剛体 2 タスク + 変形体 4 タスク 全タスクで既存手法を上回る（最大で7倍） 7

Rewarped コンセプト • Rewarped: Warp（NVIDIAが公開するGPU向けDSL）をもとに実装さ れたシミュレーションプラットフォーム • Warpとの差分 – 既存のWarpの実装は”単一環境”前提 – Rewarpedは N環境 × Tステップ をバッチ並列化 • 対応マテリアル – 剛体・関節・弾性体 – 粘土 – 流体 8

Rewarped 実装の工夫① • Material Point Method（MPM） – 変形体を粒子+格子で解く方法 – 既存のMLS-MPM手法（Ma et al., 2023）は単一環境向け – 粒子データをバッチ化することでN環境を同時に計算 • 剛体との相互作用 – 一方向結合：関節剛体の衝突形状をkinematics（外力に影響されない）とし てMPM粒子に力を与える – 逆方向の力は無視することで安定性を確保し、実用上の多くのマニピュレー ションタスクをカバー 9

Rewarped 実装の工夫② • CUDA Graph – forward/backward をテープに記録 – 後から”adjointカーネル”を自動生成し、逆伝播を実行 – ソースコード変換（Griewank & Walther, 2008）で解析的勾配を取得 • 有限差分法よりも高速かつ高精度 • Gradient Check-pointing – 順伝播では中間値を保存せず、逆伝播時に該当ステップを再計算 • PyTorch Autograd ブリッジ – Warp 側テンソルを PyTorch Tensor と共有 – RLアルゴリズムは通常のPyTorchコードで記述可 – シミュレータ側はGPUネイティブで高速 10

SAPO コアの考え • FO-MBRL – シミュレータが微分可能なら勾配をREINFORCEではなく解析勾配として 直接取得可能 • 課題：解析勾配は、目的関数の非連続性やカオス的挙動により学習が 不安定となる – 特に、接触が伴うタスク、長いホライズン • 方向性：報酬にエントロピー項を加える – LandscapeをSmoothing – 探索も促進 → 局所最適を回避 11

既存手法① FO-MBRL • Analytic Policy Gradient（APG, Freeman et al. (2021)） – 直接割引報酬を最大化する – horizonがエピソード全ての場合はBPTTと呼ばれるもの • Short Horizon Actor Critic （SHAC, Xu et al. (2021)） – 方策𝜋𝜃 と ホライズンの終端の 価値関数𝑉𝜓 を学習 ෨ – 𝑉はTD(𝜆)を使って計算 – 1次勾配 12

既存手法② エントロピー最大化RL • SAC（Soft Actor-Critic） – モデルフリーの連続行動空間に対応するRLアルゴリズム 割引率を導入 この時のQ値と価値関数V ↓ 最終的な目的関数 13

SAPO • SHACの枠組みにエントロピー最大化を加える – H stepのリターン – 1次推定量を利用した勾配 – 価値関数 – 価値関数は以下の式を最小化するように学習 14

設計選択① エントロピーの利用 • 𝛼の自動チューニング – Lagrange Dualにより学習中に𝛼を最適化 – 目標エントロピー に近づくように更新 – これにより、異なるタスクの異なる行動次元でも手動でのチューニングが不要 • エントロピー正規化 – エントロピーを [0, +1]に正規化し、報酬依存性を低減 15

設計選択② SHACとの違い • 状態依存分散 – Actorは平均𝜇のみでなく分散σ2 も出力（SHACは𝜇のみ） – 不確実ならσ↑で探索、確実ならσ↓ • Critic ensemble, no target networks – clipped double critic • 2つのCriticを学習、TD targetを作るとき二つの最小値を用いる （worst caseを見積り、過大評価を防ぐ） • Actor更新時には平均を利用し、学習スピードは保つ – Target networkは使用しない • 最適化安定化 – 活性化関数：ELU → SiLUを使用 – 最適化手法：Adam → AdamW – 勾配クリッピング：1.0 → 0.5 16

実験 タスク設定 • 剛体タスク – AntRun：4脚の剛体の歩行 – HandReorient：立方体を連続姿勢制御 • 変形体タスク – Rolling Flat：粘土をめん棒でならす – SoftJumper：弾性4足 ジャンプ移動 – HandFlip：粘土を手のひらで半回転 – FluidMove：コップ内の流体をこぼさず移送 17

実験 共通条件 • 並列環境数 N=32 or 64 • エピソード長 T=200-400 step • 総学習ステップ 4M-6M • ベースライン – モデルフリー：PPO・SAC – FO-MBRL：APG・SHAC – TrajOpt：開ループ最適化 18

実験結果 定量結果 • 全てのタスクで提案手法（SAPO）が最も良い性能 • HandReorientの大幅改善が顕著 19

実験結果 学習曲線 • 他のFO-MBRLの手法に比べて、学習の安定性が高い 20

実験結果 アブレーション • 検討パターン – 状態価値関数のエントロピー項を除いた場合 – Returnのエントロピー項を除いた場合 – 両方除いた場合 • エントロピー項が最重要だった • V_softによりターゲットの変動を抑制し、安定化 21

実験結果 ランドスケープ • SAPOでは他のアルゴリズムに比べてloss関数のランドスケープが滑らかに 22

考察と今後の方向性 • 観測問題 – 今回は粒子状態を直接入力 → 実機では取得困難 • sim2realギャップ – マテリアルパラメータのずれ、センサノイズ • 世界モデル × SAPO – Differentiable Simulatorを教師データとしてLatent Dynamicsを学習 – モデル予測でHorizon拡大 23

まとめ • Rewarped – “並列 × 微分可能 × マルチフィジックス” を初めて実現 • SAPO – 解析勾配 + 最大エントロピー • サンプル効率の向上 • 変形体でも安定学習 • 全てのタスクで最も良い性能 24