🤖

視覚的連続制御の習得:データ拡張強化学習の改良版という論文のコード実行環境構築

に公開

今、コード付きの強化学習の論文を読んでます
(京都大学人工知能研究会:KaiRAというサークルの社会人メンバーとして論文読み会で発表のためですが)。 
今回はUbuntu22.04での環境設定方法を行って行きます。
論文の要約やコード実行方法の詳細やその解釈は発表後にもし私が理論や設計はある程度正しいことを書けるだけの私の理解が進めばですが詳細は別記事で書けたら書くことにします。
https://arxiv.org/pdf/2107.09645v1

実装の注意事項:
結論から言うと難易度が低いものはGOOGLE COLABの無料版のT4(16GB VRAM)でも動くと思います。ただし難易度高のHUMANOID系はVRAM以上にかなりのRAMが必要なようです。20GB VRAM 256 RAMの環境でバッチサイズを論文で使われたデフォルト256から128に落として、
Quadruped Walk で6時間, Humanoid Run で3日ほど学習に時間がかかりました。
COLABだとより高いRAMと、バックでJOBを動かし続けられるプランにはいってないと再現は厳しいかもしれません。

内容はともかくコードは結構古い。(正確には古くはないかもしれませんがあまりにも早く変わるすぎる!整合性維持が大変です! 特にNVIDIA のドライバー!CUDA!,cudnn!, pytorch! & numpy!)
論文のコードは以下ですが、そのまま動かしても、結構エラーでて大変なので、忘備録としてまとめました。CONDAとPIPの混在という本来あまり推奨されない環境です。(この記事の動かし方もいつまでもつか疑問です。AWSの‘EC2では標準でUbuntu20.04の選はなくなりました、以下はPYTHON3.8というすでにサポート切れのバージョンで構築されているようです。)
https://github.com/facebookresearch/drqv2
まずここに書いてますがdrq v2とは何ぞやということで

DrQ-v2は、ピクセル入力に対するモデルフリー・オフポリシー連続制御手法DrQを改良したアルゴリズム

  • SACからDDPGへベース学習器を変更
  • nステップリターンでTD誤差を推定
  • 探索ノイズを減衰スケジュール化
  • 実装を3.5倍高速化
  • ハイパーパラメータを最適化

環境設定

  • まずMuJoCoの導入
    MuJoCo(Multi-Joint dynamics with Contact)は、ロボティクスやバイオメカニクス、グラフィックス、アニメーションなど、高速かつ高精度な物理シミュレーションを必要とする分野の研究開発を支援する無料のオープンソース物理エンジンです。

MuJoCoのライセンスキーをとります
https://www.roboti.us/license.html

バイナリをダウンロードして解凍して展開します。(mujoco200)
https://www.roboti.us/download.html

システムのどこかに配置してライセンスキーをそのフォルダーにいれます。(私は/optにいれました)

# pwd
/opt
# tree -L 1 mujoco200
mujoco200
├── bin
├── doc
├── include
├── mjkey.txt
├── model
└── sample

その後環境変数を設定します。
私はユーザーでなくシステム全体に反映させるため、/etc/environment にいれました、一応念のため/etc/profileの最後にも追加してます。

# cat /etc/environment
PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/usr/local/cuda/bin"
LD_LIBRARY_PATH="/usr/local/cuda/lib64:/opt/mujoco200/bin:$LD_LIBRARY_PATH"
MUJOCO_PY_MUJOCO_PATH="/opt/mujoco200"
MUJOCO_PY_MJKEY_PATH="/opt/mujoco200/mjkey.txt"
# cat /etc/profile
~~~
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/opt/mujoco200/bin:$LD_LIBRARY_PATH
export MUJOCO_PY_MUJOCO_PATH=/opt/mujoco200
export MUJOCO_PY_MJKEY_PATH=/opt/mujoco200/mjkey.txt
  • パッケージ導入とPYTHON仮想環境
    その後aptでパッケージを導入します
# apt install libosmesa6-dev libgl1-mesa-glx libglfw3
# apt install ffmpeg

あと、今回はPYTHON3.8ということもあり、あまり私は普段使わないですがANACONDAを導入しました。以下のサイトからダウンロードしてインストールください。(Ubuntuの場合はSHをダウンロードしてそれを実行するだけです)

# wget https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-x86_64.sh
# chmod +x Anaconda3-2024.10-1-Linux-x86_64.sh
# ./Anaconda3-2024.10-1-Linux-x86_64.sh
# source ~/anaconda3/bin/activate

そしてGITからソースコードをダウンロードします.

# cd /opt
# git clone https://github.com/facebookresearch/drqv2.git
# cd drqv2

以下はAnacondaのbaseをactivateした状態で以下のコマンドを実行ください

(base)# conda env create -f conda_env.yml
(base)# conda activate drqv2

ここで環境を作るときに一部パッケージのインストールでエラーがでます。ここでpipで導入します

(drqv2)# pip install dm_env dm_control
(drqv2)# pip install hydra-core
(drqv2)# pip install hydra-submitit-launcher
(drqv2)# pip install omegaconf
(drqv2)# pip install termcolor
(drqv2)# pip install tensorboard
(drqv2)# pip install opencv-python
(drqv2)# pip install imageio-ffmpeg
(drqv2)# pip install --upgrade "numpy<2"

さてgitのREADMEではこれで準備OKのはずで、以下が動くはずなのですが... 実際に実行するとエラーになります。

(drqv2)# python train.py task=quadruped_walk

ソースコード一部改修

  • drqv2.py
    obs = の部分をコメントして、float32に変えてます
    return で返り値を変更してます
    def act(self, obs, step, eval_mode):
        #obs = torch.as_tensor(obs, device=self.device)
        obs = torch.tensor(obs, dtype=torch.float32, device=self.device)
        obs = self.encoder(obs.unsqueeze(0))
        stddev = utils.schedule(self.stddev_schedule, step)
        dist = self.actor(obs, stddev)
        if eval_mode:
            action = dist.mean
        else:
            action = dist.sample(clip=None)
            if step < self.num_expl_steps:
                action.uniform_(-1.0, 1.0)
        #return action.cpu().numpy()[0]
        return action.cpu().numpy()
  • replay_buffer.py
    __init__でself._specs = data_specsを追加
class ReplayBufferStorage:
    def __init__(self, data_specs, replay_dir):
        self._data_specs = data_specs
        self._specs = data_specs ## この行を追加
        self._replay_dir = replay_dir

以下のaddを改修

    def add(self, time_step):
        for spec in self._data_specs:
            value = time_step[spec.name]
            if np.isscalar(value):
                value = np.full(spec.shape, value, spec.dtype)
            assert spec.shape == value.shape and spec.dtype == value.dtype
            self._current_episode[spec.name].append(value)
        if time_step.last():
            episode = dict()
            for spec in self._data_specs:
                value = self._current_episode[spec.name]
                episode[spec.name] = np.array(value, spec.dtype)
            self._current_episode = defaultdict(list)
            self._store_episode(episode)

から、以下に

    def add(self, time_step):
        # バッチ次元 (先頭に長さ1の軸) を除去して shape を spec に合わせる
        def prep(v, spec):
            arr = np.array(v, dtype=spec.dtype)
            # 先頭軸が余分なら squeeze
            if arr.ndim > len(spec.shape) and arr.shape[0] == 1:
                arr = np.squeeze(arr, axis=0)  # → (12,) などに
            return arr

        values = (
            prep(time_step.observation, self._specs[0]),
            prep(time_step.action,      self._specs[1]),
            prep([time_step.reward],    self._specs[2]),
            prep([time_step.discount],  self._specs[3]),
        )
        for spec, value in zip(self._specs, values):
            assert spec.shape == value.shape and spec.dtype == value.dtype
            self._current_episode[spec.name].append(value)
        # エピソード終端ならファイルにストアしてバッファをクリア
        if time_step.last():
            # 各キーに対し numpy array にまとめる
            episode = {
                spec.name: np.array(self._current_episode[spec.name], dtype=spec.dtype)
                for spec in self._specs
            }
            # ディスクへ保存
            self._store_episode(episode)
            # バッファをリセット
            self._current_episode = defaultdict(list)		
  • video.py
    def save(self, file_name)で保存の形式を完全にffmpegを利用する形に書き換える
    #def save(self, file_name):
    #    if self.enabled:
    #        path = self.save_dir / file_name
    #        imageio.mimsave(str(path), self.frames, fps=self.fps)
    def save(self, file_name):
        if not self.enabled:
            return
        path = self.save_dir / file_name
        # FFMPEG プラグインを使って動画を書き出し
        writer = imageio.get_writer(
            str(path),
            format='FFMPEG',    # 動画向け FFmpeg バックエンドを使う
            mode='I',           # イメージを逐次追加
            fps=self.fps
        )
        for frame in self.frames:
            writer.append_data(frame)
        writer.close()

これで実行しました。学習が開始されてます。

(drqv2) # python train.py
train.py:207: UserWarning:
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path='cfgs', config_name='config')
/root/anaconda3/envs/drqv2/lib/python3.8/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
  ret = run_job(
/opt/drqv2/train.py:207: UserWarning:
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path='cfgs', config_name='config')
workspace: /opt/drqv2/exp_local/2025.05.19/222623_
| eval  | F: 0 | S: 0 | E: 0 | L: 1000 | R: 104.0364 | T: 0:00:00
| train | F: 5000 | S: 2500 | E: 5 | L: 1000 | R: 7.7750 | BS: 2500 | FPS: 14.3452 | T: 0:01:09
| train | F: 6000 | S: 3000 | E: 6 | L: 1000 | R: 6.9955 | BS: 3000 | FPS: 51.2946 | T: 0:01:29
| train | F: 7000 | S: 3500 | E: 7 | L: 1000 | R: 95.7961 | BS: 3500 | FPS: 51.3039 | T: 0:01:48
| train | F: 8000 | S: 4000 | E: 8 | L: 1000 | R: 181.1696 | BS: 4000 | FPS: 51.2025 | T: 0:02:08
| train | F: 9000 | S: 4500 | E: 9 | L: 1000 | R: 21.8619 | BS: 4500 | FPS: 51.1218 | T: 0:02:27
# nvidia-smi
Mon May 19 22:54:08 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 575.51.03              Driver Version: 575.51.03      CUDA Version: 12.9     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  Tesla T4                       Off |   00000000:00:1E.0 Off |                    0 |
| N/A   47C    P0             76W /   70W |    1154MiB /  15360MiB |     81%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A            1946    C+G   python                                 1150MiB |
+-----------------------------------------------------------------------------------------+

**追記
しばらくは知らせて一晩たって確認したら以下のエラーでした!

RuntimeError: DataLoader worker (pid 2425) is killed by signal: Killed

TASKはquadruped_walkで
WORKERを4から1にして再度実行です!
*cfgs/config.yaml

defaults:
  - _self_
  - task@_global_: quadruped_walk
  - override hydra/launcher: submitit_local
#replay_buffer_num_workers: 4
replay_buffer_num_workers: 1

無事終了、様子をみると報酬平均(出力でいうRの値)が上がってますので一応学習は順調であると思われます。

(drqv2) # tensorboard --logdir exp_local --host 0.0.0.0

!
これは論文で発表されたものとほぼ同じような効率で学習されてます。

ちなみに最終のポイントで生成された動画が以下です。
https://www.youtube.com/watch?v=woqyRk06Sww

(追記)3日たってようやく、HUMANOID_RUNの学習も終了して無事発表にはデモ動画が用意できました。ここにも追加しておきます。
学習の様子は以下、なかなか報酬がでず最初はあきらめかけてました。

ちなみに論文で発表された結果は以下、

明らかに今回は立ち上がりに時間がかかり最終的な報酬も論文の結果より悪かったです。バッチサイズを256から128に落としたのが影響したかもしれません。
https://www.youtube.com/watch?v=qBS5Pi_Ca4w

Discussion