Rustで作る強化学習エージェント

Transcript

1. 0 Rustで作る強化学習エージェント 2024-02-14 第113回NearMe技術勉強会 Takuma KAKINOUE
2. 1 概要 • Rust製のOSS強化学習フレームワークを開発中 ◦ https://github.com/kakky-hacker/reinforcex • 深層学習にはpytorchのrust版であるtchを使⽤ • 今回の勉強会では、reinforcexのサンプルを動かす
3. 2 なぜRustなのか？ • メモリ効率性が良い ◦ Pythonだとメモリリークで学習が中断されてしまうことがよくある ◦ Rustだと所有権システムによって使われなくなったメモリが即解放される • 並列処理に強い ◦ 複数エージェントで並列学習するアルゴリズムにおいてPythonより有利に • Rust製の深層学習フレームワークの開発が活発になってきている ◦ https://github.com/huggingface/candle ◦ https://github.com/tracel-ai/burn
4. 3 強化学習アルゴリズム • REINFORCE（実装済） ◦ episodicなアルゴリズム ▪ step単位ではなくepisode単位で学習する ◦ ⽅策ベース ◦ 実装が簡単 • DQN（実装中） ◦ 価値ベース ◦ 実装はやや⼤変 ▪ Replay Buffer ▪ ε-Greedy • その他 ◦ DDPG, TD3, SAC, PPO, A3Cなど追って実装予定
5. 4 サンプルを動かす • rust環境が未構築の場合は以下を実⾏してください ◦ brew install rustup-init ◦ rustup-init ◦ シェルの再起動 • 実験⽤にOpenAI-Gymを使うのでPythonをinstallして以下を実⾏してください ◦ pip3 install gymnasium==0.26.3 • リポジトリ（https://github.com/kakky-hacker/reinforcex）をcloneして、 cargo runを実⾏すればサンプルが動きます
6. 5 インターフェース設計 • サンプルコードを元に解説 ◦ https://github.com/kakky-hacker/reinforcex/blob/master/src/examples/train_c artpole_with_reinforce.rs • 意識したこと ◦ Pythonの強化学習フレームワークであるChainerrlやPFRLと近い設計にした
7. 6 ⼯夫したポイント • NNの重みの初期化メソッドの実装 ◦ https://github.com/kakky-hacker/reinforcex/blob/master/src/misc/weight_ini tializer.rs ◦ ランダム初期化よりもXavierやHeを使うことで学習がちゃんと進むようになった ◦ XavierとHeの使い分けは、該当の層の直後に噛ませる活性化関数の種類による ▪ Xavierの場合、tanhなどと相性が良い ▪ Heの場合、reluなどと相性が良い
8. 7 ⼯夫したポイント • リプレイバッファにおいて、Reference Counterによる変数のライフタイム管理 ◦ https://github.com/kakky-hacker/reinforcex/blob/master/src/replay_buffer/replay_bu ffer.rs ◦ nステップ分の状態遷移を保持するとき、各状態に直前nステップの状態への参照を持たせ ることで、無駄にメモリを使うことなく保持できる ◦ Reference Counterを使わない場合、nステップ分の状態のコピーが逐⼀保存されるので、 無駄にメモリを使う上に計算コストも増えてしまう
9. 8 今後の展望 • アルゴリズムの拡充 ◦ DQN, PPO, SACなど • tchをcandle（or burn）に移⾏ • 並列学習周りの実装 • Decision Transformerの実装 スター付けてもらえると、やる気が出ます！ https://github.com/kakky-hacker/reinforcex
10. 9 Thank you