[論文レビュー] Low-Precision Reinforcement Learning: Running Soft Actor-Critic in Half Precision

Johan Björck, Xiangyu Chen+3 more|arXiv (Cornell University)|Feb 26, 2021

Reinforcement Learning in Robotics参考文献 56被引用数 7

ひとこと要約

この論文は、6つの数値安定化技術を導入することで、Soft Actor-Critic (SAC)エージェントを半精度（bfloat16/fp16）で訓練しても性能劣化が生じないことを示している。提案手法の一つとしてAdamにおける2次のモーメントの平方根を保存する方法や、算術演算の順序を再配置する手法を含み、メモリと計算量を削減しつつ、ハイパーパramータのチューニングなしに、連続制御環境においてフル精度と同等の報酬を達成できる。

ABSTRACT

Low-precision training has become a popular approach to reduce compute requirements, memory footprint, and energy consumption in supervised learning. In contrast, this promising approach has not yet enjoyed similarly widespread adoption within the reinforcement learning (RL) community, partly because RL agents can be notoriously hard to train even in full precision. In this paper we consider continuous control with the state-of-the-art SAC agent and demonstrate that a naïve adaptation of low-precision methods from supervised learning fails. We propose a set of six modifications, all straightforward to implement, that leaves the underlying agent and its hyperparameters unchanged but improves the numerical stability dramatically. The resulting modified SAC agent has lower memory and compute requirements while matching full-precision rewards, demonstrating that low-precision training can substantially accelerate state-of-the-art RL without parameter tuning.

研究の動機と目的

• 最先端のSACアルゴリズムを用いて、連続制御タスクにおける低精度強化学習を可能にすること。
• 教師あり学習の低精度技術をRLに直接適用する際に生じる数値不安定性を解消すること。
• フル精度パフォーマンスを維持しつつ、メモリと計算量を削減する、簡単で即座に適用可能な修正手法の開発。
• 再訓練やハイパーパramータ調整なしに、低精度RLが実現可能であることを実証すること。
• 将来の低精度RL分野における研究を促進するため、再現可能で明確なコードベースと手法を提供すること。

提案手法

• Adamオプティマイザにおいて、2次のモーメントの平方根を保存することで動的範囲を低減し、低精度での数値安定性を向上させる。
• 算術演算の順序を再配置することで、キャンセルエラーを最小限に抑え、勾配更新における精度を向上させる。
• 報酬スケーリングとNaN/無限大のハンドリング戦略をRL向けに適応するが、これら単体では不十分である。
• 重みと活性化にfp16を用いたミックスド・プレシジョントレーニングを実施するが、提案手法と組み合わせて初めて有効である。
• qtorchを用いてbfloat16や8ビット固定小数点など、さまざまな低精度フォーマットのシミュレーションを実施し、耐障害性を検証する。
• 鍵となるコンponents（例：ターゲットネットワーク）の統計量をフル精度で維持するように変更したトレーニングループを導入する。

実験結果

リサーチクエスチョン

• RQ1 教師あり学習で用いられる標準的な低精度トレーニング技術を、強化学習におけるSACに直接適用できるか？
• RQ2 SACのfp16トレーニングで数値安定性を達成するために、具体的にどのような修正が必要か？
• RQ3 ハイパーパramータチューニングなしに、低精度SACがフル精度パフォーマンスに一致できるか？
• RQ4 提案手法を用いて半精度でSACをトレーニングした場合、メモリと計算時間にどの程度の削減が見込めるか？
• RQ5 提案手法は、さまざまなハイパーパramータ設定や環境においても頑健か？

主な発見

• 教師あり学習の低精度技術（例：報酬スケーリング、ミックスド・プレシジョン）を素朴に適用した場合、SACトレーニングで競争力のある性能は得られない。
• 提案された6つの修正により、状態ベースおよびピクセルベースの連続制御環境において、fp16でのSACトレーニングがフル精度性能に一致することが確認された。
• 提案手法を用いた半精度トレーニングでは、約50%のメモリ使用量削減と、計算時間の顕著な高速化が達成された。
• ランダムなハイパーパramータ設定に対しても安定性を示しており、チューニング済みの設定にとどまらない安定性が裏付けられた。
• アブレーションスタディにより、各提案手法が数値安定性に寄与していることが確認され、特に2次のモーメントの平方根の保存が顕著な効果を示した。
• コードと実装が公開されており、再現性と将来的な低精度RL分野における研究を支援する。