オフライン強化学習のためのDecision Transformersのトレーニング

Decision Transformers が RL をシーケンスモデリングへシフトする

Decision Transformers は、強化学習 (RL) を従来のポリシー最適化タスクではなく、条件付きシーケンスモデリング問題として扱います。累積報酬を最大化するために価値関数をフィットさせる代わりに、Decision Transformer はシーケンスモデリングアルゴリズム（Transformer）を用いて、望ましいリターン、過去の状態、過去の行動に基づいて将来の行動を生成します。

このアプローチは、生成的軌道モデリング（状態、行動、報酬の結合分布をモデリング）を活用することで、従来の RL アルゴリズムに取って代わるパラダイムシフトを表しています。モデルは従来の意味でリターンを最大化するのではなく、特定の目標リターンを達成すると予測される行動を生成します。

モデルアーキテクチャとプロセス

Decision Transformer のアーキテクチャは以下の手順に従います:

  1. Input Sequence: モデルは最後の $K$ タイムステップを受け取り、各ステップは Return-to-go、State、Action の3つの要素で構成されます。
  2. Embedding: トークンはベクトル状態の場合は線形層、画像フレームの場合は CNN エンコーダを用いて埋め込みます。
  3. Prediction: GPT-2 ベースのモデルがこれらの入力を処理し、自己回帰的モデリングにより将来の行動を予測します。

オフライン Decision Transformers の実装

オフライン Decision Transformers は、環境との直接的な相互作用なしに、他のエージェントや人間のデモンストレーションから収集されたデータセット上でトレーニングされます。

データセットの準備とカスタムデータコレータ

Decision Transformer をトレーニングするには、生の RL データをシーケンスモデリングに適した形式に前処理する必要があります。Hugging Face Hub の halfcheetah-expert-v2 データセットを使用する場合、以下の前処理ステップが必要です:

  • Normalization: 各特徴量は平均を引き、標準偏差で割って正規化します。
  • Return Computation: 各軌道について割引リターンを事前に計算します。
  • Scaling: 報酬とリターンは 1000 倍にスケーリングします。
  • Sampling Distribution: エキスパートエージェントの軌道長さを考慮してサンプリング分布を拡張します。

これらのステップはカスタム Data Collator によって実装され、状態、行動、報酬、Return-to-go、タイムステップ、アテンションマスクを含むバッチを返します。

Hugging Face Trainer を用いたトレーニング

トレーニングは Hugging Face の Trainer クラスを使用して実行されます。標準の DecisionTransformerModel は損失値を自動的に返さないため、ラッパークラス (TrainableDT) を用いて、モデルの行動予測とデータセットのターゲット行動間の L-2 ノルム（平均二乗誤差）を計算します。

主要なトレーニングハイパーパラメータ:

  • Epochs: 120
  • Batch Size: 64
  • Learning Rate: $1e-4$
  • Weight Decay: $1e-4$
  • Warmup Ratio: 0.1
  • Optimizer: AdamW
  • Max Gradient Norm: 0.25

Transformers における RL の今後の方向性

Hugging Face は、以下の計画されたイニシアチブを通じて Deep Reinforcement Learning コミュニティへのサポートを拡大することを目指しています:

  • Online Training: オンライン設定でトレーニングまたはファインチューニングされた Decision Transformer モデルをリポジトリに追加します。
  • Tool Integration: sample-factory バージョン 2.0 をエコシステムに統合します。

Sources