🐾

Genesis環境でUnitree Go2の4足歩行強化学習をフルスクラッチ実装してみた

に公開3件
PyTorch
強化学習
genesis
unitreego2
tech

はじめに

皆さんこんにちは、一関高専 村上研究室所属 小山田です。

この記事では、私の卒業研究の一環として取り組んでいる、Genesisを用いたUnitree Go2の4足歩行強化学習について紹介します。

また、この強化学習の検証を行うにあたって、株式会社ハイレゾ様からGPUサーバーのレンタルサービス「GPUSOROBAN」の1ヶ月無料トライアルをご提供頂きました。これについても後ほど紹介させて頂きます。

https://soroban.highreso.jp/aispacon

対象読者

  • 強化学習を用いたロボットの歩行制御に興味がある方
  • Genesisを用いたロボットシミュレーションに興味がある方

※なおこの記事では、PPOについての基礎的な知識があることを前提に話を進めていきます。
※簡単な高校数学も一部必要になるかもしれません。

概要

この記事では、Genesis環境でUnitree Go2の4足歩行強化学習をフルスクラッチ実装した内容について紹介します。

学習フェーズでは、Actor-Criticから始まりPPOについて学習を行いました。
実装フェーズでは、強化学習モデルと学習アルゴリズム(PPO)をPytorchのみを用いて実装しました。
検証フェーズでは、ハイレゾ様よりご提供頂いたH200、学校に設置されているA100、研究室のPCに搭載されているRTX4060 Tiを用いて実験を行いました。

参考リンク集

今回の実装にあたって、ロボットの連続値制御に用いられる強化学習など、不足している知識がたくさんあったので以下のリンクを参考に学習を進めました。
特に、強化学習アルゴリズムの実装に関しては、Exampleで用いられているrsl-rlの実装を読み込みました。

学習

このセクションでは、4足歩行実装前の前提知識を学習するために行ったリサーチや実験について、順に紹介していきます。

Actor-Criticの理解

この時点では、AlphaZeroの実装経験はあったため、強化学習の基礎概念の理解や離散値で行動を選択するような環境での強化学習の実装に関する知識はあったものの、連続値で行動を選択する必要がある4足歩行ロボットの制御に関しては未経験でした。

そのため、まずは連続値制御での強化学習手法について調査を行い、Actor-Criticの理解を深めることから始めました。

要点

  • 価値関数と方策の2つのネットワークを用いる
  • 価値関数は状態の価値を推定し方策の更新に利用される
  • 方策関数は行動を選択する確率分布を出力する (例: ガウス分布などがよく使われる)

PPOの理解

PPO（Proximal Policy Optimization）はActor-Criticの派生手法であり、方策の更新を行う際に古い方策と新しい方策の間の距離をクリッピングして制限することで、学習の安定性を向上させる手法です。
ここでは、実際にPPOの論文を読み込みました😎

要点

  • 古い方策と新しい方策の間の距離をクリッピングして制限する
  • GAE (Generalized Advantage Estimation) を用いることでAdvantageの分散と誤差のトレードオフを調整できる

また、PPOを実装するにあたって、GymnasiumのLunarLanderContinuous-v2環境を用いて実装と動作確認を行いました。LunarLanderは2次元の連続値制御環境であるため、4足歩行制御(12次元の連続値制御)の前段階として適していると考え選択しました。

https://github.com/Oya-Tomo/lunar-lander-ppo-pt

実装

このセクションでは、Genesis環境でUnitree Go2の4足歩行強化学習をフルスクラッチ実装した内容について紹介します。

特に、この実装を行う上でrsl-rlをかなり参考にしました。rsl-rlは様々な機能を汎用的に実装しているため、コードが複雑になっており理解が少し難しいです。
私の実装は可読性重視で、PPOとRNNの実装を簡略化して行っていますが、データの流れは殆ど同じなのでrsl-rlを解読する際はこの記事の実装セクションが少しは参考になるかもしれません。

コード

実装したコードは、以下のリポジトリで公開しています。

https://github.com/Oya-Tomo/genesis-go2-rl-pt/tree/rnn-ppo

学習ループ

ロボットの強化学習に多用されるIsaacSimやMujoco MJX、今回用いるGenesisなどは並列で複数の環境を動かすことができます。
また、環境設計によってはそれぞれのエピソードで終了位置が異なることがあり、これに対応するために学習ループもシングルエピソードと比較して特徴的な実装が必要になります。

cut_out_episode

実際の学習ループでは、環境のリセットは各環境で独立して行い、エピソードが終了した環境のみをリセットします。
図の中のReal Episodeは、実際のエピソードを示しています。
対してShort Episodeは1回の学習ループで切り取るエピソードを示しています。

実装のPython疑似コードを以下に示します。

for loop in range(num_loops):
    for step in range(num_steps): # 各環境で並列に Short Episode を収集
        # 行動の選択
        action = policy(state)
        # 環境のステップ
        next_state, reward, done = env.step(action)
        # バッファにデータを格納
        replay_buffer.store(state, action, reward, done, next_state)
        state = next_state

    replay_buffer.compute_advantages_and_targets(next_state) # GAE

    for batch in replay_buffer:
        policy.update(batch)
        value.update(batch)

    replay_buffer.clear()

# Data Shape, Type
# state: (num_envs, state_dim) float
# action: (num_envs, action_dim) float
# reward: (num_envs,) float
# done: (num_envs,) bool

実装部分 ↓↓↓
https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/train.py#L138-L372

rsl-rlの学習ループ実装 ↓↓↓
https://github.com/leggedrobotics/rsl_rl/blob/main/rsl_rl/runners/on_policy_runner.py#L61-L175

方策関数と価値関数の実装

PPOでは、方策関数と価値関数の2つのニューラルネットワークを用います。
また、ロボットの強化学習では、過去の状態や行動を考慮する必要のある複雑な環境を学習するためにRNN (Recurrent Neural Network) を使用することがあります。
この実装でも、障害物フィールド上での歩行を想定しているためGRU (Gated Recurrent Unit) を使用しました。

model_arch

実装部分 ↓↓↓
https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/model.py

ちなみに、rsl-rlではMLPモデルまたはRNNモデルを選択できるようになっています。

学習アルゴリズム

PPOの学習ループは、上の疑似コードにも記載していますが、以下の3段階に分かれています。

  • Episodeの収集
  • GAEとValue Targetの計算
  • PolicyとValueの更新

Generalized Advantage Estimation (GAE)とValue Targetの計算

Advantageの推定には、GAE (Generalized Advantage Estimation) を使用します。
様々な方法がありますが、PPOでは以下の式で計算されます。

\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)
\hat{A}_t = \delta_t + \gamma \lambda \hat{A} _{t+1}

Value Targetは以下の式で計算されます。

V^{\text{target}}_t = \hat{A}_t + V(s_t)

実装部分 ↓↓↓
https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/buffer.py#L68-L89

rsl-rlの実装 ↓↓↓
https://github.com/leggedrobotics/rsl_rl/blob/main/rsl_rl/storage/rollout_storage.py#L127-L149

Proximal Policy Optimization (PPO)

PPOは、 L^{\text{CLIP}}(\theta) を最大化することを目的とした強化学習アルゴリズムです。
これを符号反転したものをPolicyの損失関数として最小化します。

L^{\text{CLIP}}(\theta) = \hat{\mathbb{E}}_t \left[ \min \left( r_t(\theta) \hat{A}_t,\ \text{clip}\left(r_t(\theta),\ 1 - \epsilon,\ 1 + \epsilon\right) \hat{A}_t \right) \right]

クリッピング前のLearning Objectiveは以下の式で表されます。
これにクリッピングを入れることで、方策の急激な変化を防ぎ学習の安定性を向上させます。

L^{\text{CPI}}(\theta) = \hat{\mathbb{E}}_t \left[ r_t(\theta) \hat{A}_t \right]

方策の変化率を表す r_t(\theta) は以下の式で定義されます。

r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)}

Valueの損失関数は、以下の式で表されます。

L^{\text{Value}}(\theta) = \text{MSE} \left( V(s_t) - V^{\text{target}}_t \right)

実装部分 ↓↓↓
https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/train.py#L282-L296

rsl-rlの実装 ↓↓↓
https://github.com/leggedrobotics/rsl_rl/blob/main/rsl_rl/algorithms/ppo.py#L296-L302

報酬設計

報酬設計は強化学習を実装する上である種、秘伝のタレのようなものです。
工夫に工夫を重ね、比率を調整しながら制作するので、実装には非常に時間がかかります。

この記事では、現時点での報酬の採点項目をざっと紹介します。

  • xy移動速度: 指定された速度で移動することに対する報酬
  • yaw移動速度: 指定された速度で回転することに対する報酬
  • z速度安定化: 上下運動を抑制するための報酬
  • 足の先端位置: なるべく肩の真下に足を着地させるための報酬
  • 足の伸び: 足を適切な長さに伸ばすための報酬
  • 足の設置時間: 足が地面に設置している時間を適切に調整するための報酬
  • 円滑性: 足の急激な動きを抑制するための報酬
  • 力の効率: なるべく少ない力で動作するための報酬
  • 体の安定性: 体が不安定に傾かないようにするための報酬
  • 体の高さ: 体が適切な高さを維持するための報酬
  • 衝突回避: 障害物に接触しないように移動することに対する報酬

また、報酬を加算する際に、ロボットの姿勢を考慮して各項目にマスクをかけています。
こうすることで、ロボットが転倒している場合などに不適切な報酬が加算されることを防ぎます。

実装部分 ↓↓↓
https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/env.py#L492-L525

検証

モデルのタスク評価

この検証では、障害物フィールドを歩行するタスクを設定しました。
しかし、現時点では障害物に当たるとロボットが進路を勝手に変更してしまうため、平面の走行のみしかできません。
現時点では、ロボットに速度コマンドしか与えていないため、タスク設計はまだ大いに改善の余地があります。

build_terrain = Falseにすると、障害物をOFFにして平面のみで学習できます。

https://github.com/Oya-Tomo/genesis-go2-rl-pt/blob/rnn-ppo/src/train.py#L75

GPUのパフォーマンス比較

以下の3つで実行した学習のパフォーマンスを比較しました。

  • NVIDIA GeForce RTX 4060 Ti
  • NVIDIA A100 80GB
  • NVIDIA H200 141GB

※悲しいことにGenesisは記事執筆時点でMulti-GPUに対応していないため、単一のGPUでの比較になります。

主要な条件設定

  • n_envs 並列環境数: 40960
    ※ ただし、RTX4060 Tiではメモリ不足で動作しないため、20480に設定
  • steps_per_loop 1ループあたりのステップ数 (Short Episodeの長さ): 25
  • epochs_per_loop 1ループあたりのエポック数: 5
  • スループット比較のためパラメタを統一
ハイパーパラメタの詳細
{
  "env_params": {
    "n_envs": 40960,
    "command": {
      "x_vel_range": [-1.5, 1.5],
      "y_vel_range": [-1, 1],
      "init_rpy_max": [150, 0, 180],
      "init_rpy_min": [-180, 0, -180],
      "init_xyz_max": [25, 25, 2.5],
      "init_xyz_min": [5, 5, 2.5],
      "yaw_vel_range": [-2, 2]
    },
    "frequency": 50,
    "camera_pos": [-10, -10, 10],
    "show_viewer": false,
    "build_terrain": true,
    "camera_lookat": [15, 15, 0.5],
    "n_rendered_envs": 1,
    "avg_rollout_time": 10,
    "max_rollout_time": 20,
    "min_rollout_time": 8,
    "resetable_step_period": 25,
    "simulate_action_latency": true
  },
  "train_params": {
    "seed": 0,
    "value_lr": 0.001,
    "desired_kl": 0.01,
    "hidden_dim": 512,
    "num_layers": 1,
    "save_period": 50,
    "clip_epsilon": 0.2,
    "entropy_coef": 0.001,
    "global_loops": 1000000,
    "weight_decay": 0.0001,
    "max_grad_norm": 1,
    "max_policy_lr": 0.001,
    "min_policy_lr": 0.00001,
    "steps_per_loop": 25,
    "epochs_per_loop": 5,
    "default_policy_lr": 0.001,
    "normalize_advantages": true
  },
  "buffer_params": {
    "gamma": 0.99,
    "batch_num": 32,
    "gae_lambda": 0.95
  }
}

結果の算出方法

Wandbに記録している学習ログを元に、平均値を算出しています。

  • ループ速度: 1ループあたりの平均時間 (秒/ループ) を500レコードの平均値から算出
  • GPU使用率: GPU Utilization(%)の最初の100レコードを除外し、100 ~ 500レコードまでの平均値を算出
  • VRAM使用量: GPU Memory Allocated(Bytes)の安定した部分の値

結果

GPU ループ速度
(秒/ループ)		 GPU使用率 (%) VRAM使用量 (GB) 備考
NVIDIA GeForce RTX 4060 Ti 8.429883324147157 97.017500 11.757748224 環境数が半分
NVIDIA A100 80GB 10.209379255988182 96.202500 23.756865536
NVIDIA H200 141GB 7.743022614307726 78.955000 23.239458816

この結果を見てわかる通り、A100と比較してH200のサーバーで実行した場合の方が 24.16% 実行時間が短縮されました。
また、GPUの平均使用率もA100と比較してまだ余裕があります。
※ 1 - 7.743022614307726 / 10.209379255988182 = 0.24157753178 → 24.16%

この検証では3Dシミュレーションを実施したため、単純なGPUのスペックだけでなく、CPUやメモリ、ストレージの性能も大きく影響していると考えられます。
単純にGPUのスペックだけで比較できないことに注意してください。
だとしても、H200の性能は圧倒的であることは間違いないです。

特に強化学習では、実行時間が膨大になりがちなので、このように高速化を図れるのは非常にありがたいです。

GPUSOROBANの紹介

最後に、今回のご厚意への感謝を込めて、GPUSOROBANについて簡単に紹介させて頂きます。

https://soroban.highreso.jp/aispacon

今回、株式会社ハイレゾ様からGPUサーバーのレンタルサービス「GPUSOROBAN」の1ヶ月無償トライアルを提供していただきました。
提供して頂いたインスタンスのスペックは以下の通りです。

  • GPU: H200 x 8
  • vCPU: 112コア (論理コア数: 224)
  • メモリ: 2TB
  • ストレージ: 30TB

そして、Dockerなどはプリインストールのため、すぐに開発を始めることができました。
開発時は、vscode tunnelを利用して接続し、快適に開発を行うことができました。

特に、最近ではLLM開発などでGPUサーバーへの注目が集まっていますが、実はこの検証のようなロボットの強化学習分野でも大きな需要があります。
検証を行う際に、高速にプログラムを動作させられる環境を利用させて頂けたのは非常にありがたかったです。

おわりに

まず、今回の実験を行うにあたって、H200をご提供いただいた株式会社ハイレゾ様に深く感謝申し上げます。このような、高性能なGPUサーバーで計算を回すことができたのは非常に貴重な経験でした。本当にありがとうございました !!

現時点で障害物フィールドを歩行させることは叶っていませんが、検証で得た結果を元に改善を行っていきたいと考えています。

Discussion

ShimbeiShimbei

はじめまして、しがない高専卒です。
突然のご連絡失礼致します。私も現在趣味で4脚歩行ロボットをPPOで歩かせることを目標としています。その際に1番大きな壁としてシミュレーション環境(特にGPUスペック)が大きいと感じています。そこでお伺いしたいのですが、OyaTomoさんはGPUSOROBAN以外での環境(google colabや自作PCなど)でシミュレーションをされていたかお伺いしたいです。もしされていたらどのようなシミュレーション環境だったかお伺いしたいです。
突然のご連絡で恐縮ですがご教示いただけたら幸いです。

OyaTomoOyaTomo

まず初めに、記事を読んで頂きありがとうございます。コメント頂けてとても嬉しいです。

ご質問頂きました学習環境についてですが、普段はrtx4060tiを搭載したPC、rtx4090を搭載したPCの2台を使用しております。また、記事にも記載したA100は、先生との交渉の後にお借りすることが出来ました。

因みに、この様なコンシューマ向けGPUの環境でも動作させられる要因としては、Genesisが高速であるというのが大きいと思います。MujocoやIsaacLab、Pybulletなどは計算の安定性・レンダリングの綺麗さなどが高いと思いますが、Genesisはそこよりも並列性能・軽量性に重きを置いている印象があります。

是非、 @Shimbei さんの強化学習が完成されたら、ご共有いただけると幸いです！！

ShimbeiShimbei

ご返信いただきありがとうございます。GPUスペックが思っていたより軽く済みそうで安心しました。ご教示いただきありがとうございます。また今までpybulletやmujocばかり使っていてgenesis環境は使ったことがなかったので試してみます！

1