stable-baselines3を使ってみる (環境構築)

参考リンク

• リポジトリ (DLR-RM/stable-baselines3)
• ドキュメント

環境作成

pip install stable-baselines3[extra]

[extra]なしで最低限のモジュールだけインストールすることも可能。

TensorFlow2.xには対応していないとのこと。
これまではTensorFlow2.4に合わせたCUDAとcuDNNの設定をしていたが、TensorFlow1.xに合わせたカスタマイズも必要になる（1.14の場合はここが参考になりそう）。

(間が空いてしまったけど、前回作業した分を思い出して書いてみる)

公式の対応表 をベースに設定。
tensorflow_gpu-1.14.0の対応バージョンをインストールした。

• cuDNN 7.4
• CUDA 10.0

↓

設定できたと思ったが、プログラムを実行してみるとCPUで動いているっぽいログが出た。プログラム中ではCUDAを指定してはいるはず。。。

model = DQN("CnnPolicy", env, device="cuda")

そもそもstable-baselines3はPyTorchをバックエンドにしているため、PyTorchのバージョンに応じた設定が必要。

Stable-Baselines3 requires python 3.7+ and PyTorch >= 1.1

ということで、いったん新しく環境を作ることにする（これまでは、keras-rl2を使っていた環境をそのまま拡張しようとしていた）。公式ドキュメントから、対応のバージョンを確認できる。

• CUDA 11.7
• cuDNN 8.8.1

をインストールしてみる。

仮想環境にもPyTorchを入れる。

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu117

torchaudioは不要なのでインストールしなかった。

パソコンを再起動後、GPUが認識できているかをPythonのインタラクティブモードで確認する。

>>> import torch
>>> print("Is CUDA available?", torch.cuda.is_available())
Is CUDA available? True
>>> print("Number of available GPUs:", torch.cuda.device_count())
Number of available GPUs: 1

大丈夫そう。

↓

DQNのプログラムを実行。GPUを使って計算してくれるようになった！