最初のデシジョン トランスフォーマーをトレーニングする

最初のデシジョン トランスフォーマーのトレーニング

以前の投稿で、transformersライブラリでのDecision Transformersのローンチを発表しました。この新しい技術は、Transformerを意思決定モデルとして使用するというもので、ますます人気が高まっています。

今日は、ゼロからオフラインのDecision Transformerモデルをトレーニングして、ハーフチータを走らせる方法を学びます。このトレーニングは、Google Colab上で直接行います。こちらで見つけることができます👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb

*ジムのHalfCheetah環境でオフラインRLを使用して学習された「専門家」Decision Transformersモデルです。

ワクワクしませんか？では、始めましょう！

  • Decision Transformersとは何ですか？
  • Decision Transformersのトレーニング
    • データセットの読み込みとカスタムデータコレータの構築
    • 🤗 transformers Trainerを使用したDecision Transformerモデルのトレーニング
  • 結論
  • 次は何ですか？
  • 参考文献

Decision Transformersとは何ですか？

Decision Transformerモデルは、チェンL.らによる「Decision Transformer: Reinforcement Learning via Sequence Modeling」で紹介されました。これは、強化学習を条件付きシーケンスモデリングの問題として抽象化したものです。

主なアイデアは、RLの手法を使用してポリシーをトレーニングする代わりに、返却を最大化するためにどの行動を取るかを教えてくれる価値関数をフィットするなどの方法を使用せず、希望する返却、過去の状態、行動を与えると、希望する返却を達成するために将来の行動を生成するシーケンスモデリングアルゴリズム（Transformer）を使用することです。これは、希望する返却、過去の状態、行動に条件づけられた自己回帰モデルであり、希望する返却を達成するための将来の行動を生成します。

これは、従来のRLアルゴリズムを置き換えるための強化学習パラダイムの完全な転換です。つまり、Decision Transformersでは、返却を最大化するのではなく、希望する返却を達成する一連の将来の行動を生成します。

プロセスは次のように進みます。

  1. 直近のKタイムステップをDecision Transformerに3つの入力として供給します：
    • Return-to-go
    • State
    • Action
  2. 状態がベクトルの場合は、線形層でトークンを埋め込みます。フレームの場合はCNNエンコーダを使用します。
  3. 入力はGPT-2モデルで処理されます。これにより、自己回帰モデリングによって将来の行動が予測されます。

Decision Transformerアーキテクチャ。状態、行動、返却は、モダリティ固有の線形埋め込みに供給され、位置エピソードタイムステップエンコーディングが追加されます。トークンはGPTアーキテクチャに供給され、因果自己注意マスクを使用して行動を自己回帰的に予測します。図は[1]から引用されています。

さまざまなタイプのDecision Transformersがありますが、今日はオフラインのDecision Transformerをトレーニングします。つまり、他のエージェントや人間のデモから収集されたデータのみを使用します。エージェントは環境と対話しません。オフラインとオンラインの強化学習の違いについて詳しく知りたい場合は、この記事を参照してください。

オフラインDecision Transformersの理論を理解したので、実際にトレーニングする方法を見てみましょう。

Decision Transformersのトレーニング

前の投稿では、transformersのDecision Transformerモデルを使用して、🤗 hubから事前学習済みの重みをロードする方法を示しました。

この部分では、🤗 Trainerとカスタムデータコレータを使用して、🤗 hubでホストされているオフラインRLデータセットを使用して、ゼロからDecision Transformerモデルをトレーニングします。このチュートリアルのコードは、このColabノートブックで見つけることができます。

私たちはオフラインRLを実行して、mujoco halfcheetah環境で次の動作を学習します。

*ジムのHalfCheetah環境でオフラインRLを使用して学習された「専門家」Decision Transformersモデルです。

データセットの読み込みとカスタムデータコレータの構築

私たちは、hub上でいくつかのオフラインRLデータセットをホストしています。今日はhalfcheetahの「専門家」データセットを使用してトレーニングを行います。こちらのリンクからホストされています。

最初に、🤗データセットパッケージからload_dataset関数をインポートし、データセットをマシンにダウンロードする必要があります。

from datasets import load_dataset
dataset = load_dataset("edbeeching/decision_transformer_gym_replay", "halfcheetah-expert-v2")

ハブ上のほとんどのデータセットはそのまま使用できますが、場合によってはデータセットの追加処理や変更を行いたいことがあります。この場合、著者の実装に合わせる必要があります。具体的には以下の処理を行います:

  • 各特徴量を平均を引いて標準偏差で割って正規化します。
  • 各軌跡の割引報酬を事前計算します。
  • 報酬と割引報酬を1000倍にスケーリングします。
  • エキスパートエージェントの軌跡の長さを考慮したデータセットのサンプリング分布を拡張します。

これらのデータセットの前処理を行うために、カスタムな🤗データコレータを使用します。

それでは、オフライン強化学習のためのカスタムデータコレータを作成しましょう。

@dataclass
class DecisionTransformerGymDataCollator:
    return_tensors: str = "pt"
    max_len: int = 20 #訓練に使用するエピソードのサブセット
    state_dim: int = 17  # 状態空間のサイズ
    act_dim: int = 6  # アクション空間のサイズ
    max_ep_len: int = 1000 # データセット内の最大エピソード長
    scale: float = 1000.0  # 報酬/割引報酬の正規化
    state_mean: np.array = None  # 状態の平均値を格納するための変数
    state_std: np.array = None  # 状態の標準偏差を格納するための変数
    p_sample: np.array = None  # 軌跡の長さを考慮した分布
    n_traj: int = 0 # データセット内の軌跡の数を格納するための変数

    def __init__(self, dataset) -> None:
        self.act_dim = len(dataset[0]["actions"][0])
        self.state_dim = len(dataset[0]["observations"][0])
        self.dataset = dataset
        # 状態の正規化のためにデータセットの統計量を計算する
        states = []
        traj_lens = []
        for obs in dataset["observations"]:
            states.extend(obs)
            traj_lens.append(len(obs))
        self.n_traj = len(traj_lens)
        states = np.vstack(states)
        self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6
        
        traj_lens = np.array(traj_lens)
        self.p_sample = traj_lens / sum(traj_lens)

    def _discount_cumsum(self, x, gamma):
        discount_cumsum = np.zeros_like(x)
        discount_cumsum[-1] = x[-1]
        for t in reversed(range(x.shape[0] - 1)):
            discount_cumsum[t] = x[t] + gamma * discount_cumsum[t + 1]
        return discount_cumsum

    def __call__(self, features):
        batch_size = len(features)
        # 非一様な分布からサンプリングするためのハック
        batch_inds = np.random.choice(
            np.arange(self.n_traj),
            size=batch_size,
            replace=True,
            p=self.p_sample,  # タイムステップに応じて再重み付け
        )
        # データセットの特徴量のバッチ
        s, a, r, d, rtg, timesteps, mask = [], [], [], [], [], [], []
        
        for ind in batch_inds:
            # for feature in features:
            feature = self.dataset[int(ind)]
            si = random.randint(0, len(feature["rewards"]) - 1)

            # データセットからシーケンスを取得
            s.append(np.array(feature["observations"][si : si + self.max_len]).reshape(1, -1, self.state_dim))
            a.append(np.array(feature["actions"][si : si + self.max_len]).reshape(1, -1, self.act_dim))
            r.append(np.array(feature["rewards"][si : si + self.max_len]).reshape(1, -1, 1))

            d.append(np.array(feature["dones"][si : si + self.max_len]).reshape(1, -1))
            timesteps.append(np.arange(si, si + s[-1].shape[1]).reshape(1, -1))
            timesteps[-1][timesteps[-1] >= self.max_ep_len] = self.max_ep_len - 1  # パディングのカットオフ
            rtg.append(
                self._discount_cumsum(np.array(feature["rewards"][si:]), gamma=1.0)[
                    : s[-1].shape[1]   # TODO ここで削除された+1を確認する
                ].reshape(1, -1, 1)
            )
            if rtg[-1].shape[1] < s[-1].shape[1]:
                print("if true")
                rtg[-1] = np.concatenate([rtg[-1], np.zeros((1, 1, 1))], axis=1)

            # パディングと状態と報酬の正規化
            tlen = s[-1].shape[1]
            s[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, self.state_dim)), s[-1]], axis=1)
            s[-1] = (s[-1] - self.state_mean) / self.state_std
            a[-1] = np.concatenate(
                [np.ones((1, self.max_len - tlen, self.act_dim)) * -10.0, a[-1]],
                axis=1,
            )
            r[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, 1)), r[-1]], axis=1)
            d[-1] = np.concatenate([np.ones((1, self.max_len - tlen)) * 2, d[-1]], axis=1)
            rtg[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, 1)), rtg[-1]], axis=1) / self.scale
            timesteps[-1] = np.concatenate([np.zeros((1, self.max_len - tlen)), timesteps[-1]], axis=1)
            mask.append(np.concatenate([np.zeros((1, self.max_len - tlen)), np.ones((1, tlen))], axis=1))

        s = torch.from_numpy(np.concatenate(s, axis=0)).float()
        a = torch.from_numpy(np.concatenate(a, axis=0)).float()
        r = torch.from_numpy(np.concatenate(r, axis=0)).float()
        d = torch.from_numpy(np.concatenate(d, axis=0))
        rtg = torch.from_numpy(np.concatenate(rtg, axis=0)).float()
        timesteps = torch.from_numpy(np.concatenate(timesteps, axis=0)).long()
        mask = torch.from_numpy(np.concatenate(mask, axis=0)).float()

        return {
            "states": s,
            "actions": a,
            "rewards": r,
            "returns_to_go": rtg,
            "timesteps": timesteps,
            "attention_mask": mask,
        }

それはたくさんのコードでしたが、TLDRは、データセットを取り、必要な前処理を行い、バッチ単位で状態、アクション、報酬、リターン、タイムステップ、およびマスクを返すクラスを定義したことです。これらのバッチは、🤗 transformers Trainerで直接使用してDecision Transformerモデルをトレーニングするために使用することができます。

🤗 transformers Trainerを使用してDecision Transformerモデルをトレーニングする。

🤗 Trainerクラスでモデルをトレーニングするためには、まず返される辞書に損失が含まれていることを確認する必要があります。この場合、モデルのアクションの予測とターゲットのL-2ノルムです。これを実現するために、Decision Transformerモデルを継承したTrainableDTクラスを作成します。

class TrainableDT(DecisionTransformerModel):
    def __init__(self, config):
        super().__init__(config)

    def forward(self, **kwargs):
        output = super().forward(**kwargs)
        # DTの損失を追加
        action_preds = output[1]
        action_targets = kwargs["actions"]
        attention_mask = kwargs["attention_mask"]
        act_dim = action_preds.shape[2]
        action_preds = action_preds.reshape(-1, act_dim)[attention_mask.reshape(-1) > 0]
        action_targets = action_targets.reshape(-1, act_dim)[attention_mask.reshape(-1) > 0]
        
        loss = torch.mean((action_preds - action_targets) ** 2)

        return {"loss": loss}

    def original_forward(self, **kwargs):
        return super().forward(**kwargs)

transformers Trainerクラスには、TrainingArgumentsクラスで定義されるいくつかの引数が必要です。著者の元の実装と同じハイパーパラメータを使用しますが、より少ないイテレーションでトレーニングします。これはColabノートブックで約40分かかるため、待っている間にコーヒーを飲んだり、🤗 Annotated Diffusionブログ記事を読んだりしてください。著者は約3時間トレーニングしているため、ここで得られる結果は彼らのものほど良くなりません。

training_args = TrainingArguments(
    output_dir="output/",
    remove_unused_columns=False,
    num_train_epochs=120,
    per_device_train_batch_size=64,
    learning_rate=1e-4,
    weight_decay=1e-4,
    warmup_ratio=0.1,
    optim="adamw_torch",
    max_grad_norm=0.25,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    data_collator=collator,
)

trainer.train()

Decision Transformer、Trainer、およびそれをトレーニングする方法の理論を説明しました。 スクラッチからオフラインのDecision Transformerモデルをトレーニングしてハーフチータを走らせる準備ができました 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb Colabにはトレーニング済みモデルの可視化と、🤗ハブにモデルを保存する方法も含まれています。

結論

この投稿では、🤗 datasetsにホストされたオフラインRLデータセットでDecision Transformerをトレーニングする方法を示しました。 🤗 transformers Trainerとカスタムデータコレータを使用しました。

Decision Transformer以外にも、Deep Reinforcement Learningコミュニティのさまざまなユースケースやツールをサポートしたいと考えています。したがって、Decision Transformerモデルに関するフィードバックや、RLに役立つものを一緒に構築できる他の何かについてのフィードバックをお待ちしています。お気軽にお問い合わせください。

次は何ですか？

今後数週間から数か月以内に、エコシステムの他のツールをサポートする予定です：

  • オンライン設定でトレーニングまたはファインチューニングされたモデルを使用して、Decision Transformerモデルのリポジトリを拡充する[2]
  • sample-factoryバージョン2.0の統合

連絡を取り合うための最良の方法は、discordサーバーに参加することです。そこで私たちやコミュニティと交流することができます。

参考文献

[1] Chen, Lili, et al. “Decision transformer: Reinforcement learning via sequence modeling.” Advances in neural information processing systems 34 (2021).

[2] 鄭, 金清、張, エイミー、グローバー, アディティア、「オンライン決定トランスフォーマー」(arXiv プレプリント, 2022)

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Related articles

Discover more

AIテクノロジー

「LXTのテクノロジーバイスプレジデント、アムル・ヌール・エルディン - インタビューシリーズ」

アムル・ヌール・エルディンは、LXTのテクノロジー担当副社長ですアムルは、自動音声認識（ASR）の文脈での音声/音響処理と機...

人工知能

「Ntropyの共同創設者兼CEO、ナレ・ヴァルダニアンについて - インタビューシリーズ」

「Ntropyの共同創設者兼CEOであるナレ・ヴァルダニアンは、超人的な精度で100ミリ秒以下で金融取引を解析することを可能にす...

人工知能

ムーバブルインクのCEO兼共同創設者であるヴィヴェク・シャルマ氏についてのインタビュー・シリーズ

ビヴェクは2010年にムーバブルインクを共同設立し、急速な成長を遂げながら、600人以上の従業員を擁し、世界有数の革新的なブ...

機械学習

3つの質問：大規模言語モデルについて、Jacob Andreasに聞く

CSAILの科学者は、最新の機械学習モデルを通じた自然言語処理の研究と、言語が他の種類の人工知能をどのように高めるかの調査...

AIニュース

Q&A：ブラジルの政治、アマゾンの人権、AIについてのGabriela Sá Pessoaの見解

ブラジルの社会正義のジャーナリストは、MIT国際研究センターのフェローです

データサイエンス

「Seerの最高データオフィサーであるDr. Serafim Batzoglouによるインタビューシリーズ」

セラフィム・バツォグルはSeerのチーフデータオフィサーですSeerに加わる前は、セラフィムはInsitroのチーフデータオフィサー...

Want to read more? Go here