HuggingFace ブログ : 最初の Decision Transformer の訓練

by

HuggingFace ブログ : 最初の Decision Transformer の訓練 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 11/17/2022

* 本ページは、HuggingFace Blog の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

クラスキャット 人工知能 研究開発支援サービス

◆ クラスキャット は人工知能・テレワークに関する各種サービスを提供しています。お気軽にご相談ください :

◆ 人工知能とビジネスをテーマに WEB セミナーを定期的に開催しています。スケジュール。
  • お住まいの地域に関係なく Web ブラウザからご参加頂けます。事前登録 が必要ですのでご注意ください。

◆ お問合せ : 本件に関するお問い合わせ先は下記までお願いいたします。

  • 株式会社クラスキャット セールス・マーケティング本部 セールス・インフォメーション
  • sales-info@classcat.com  ;  Web: www.classcat.com  ;   ClassCatJP

 

 

HuggingFace ブログ : 最初の Decision Transformer の訓練

前の記事 で、transformers ライブラリの Decision Transformers の launch をアナウンスしました。Transformer を意思決定モデルとして使用する この新しいテクニックは段々とポピュラーになっています。

そして本日は、half-チーターを走らせる最初のオフライン Decision Transformer モデルをゼロから訓練することを学習します。それを Google Colab で直接的に訓練します、それはここで見つけられます 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb

*An “expert” Decision Transformers model, learned using offline RL in the Gym HalfCheetah environment.*

 

Decision Transformer とは何か？

Decision Transformer モデルは “Decision Transformer: Reinforcement Learning via Sequence Modeling” by Chen L. et al. により紹介されました。それは強化学習を 条件付きシークエンス・モデリング問題 として抽象化しています。

主要なアイデアは、価値関数を適合させるような、RL 手法を使用して (どのアクションを取ればリターン (累積報酬) を最大化するかを教えてくれる) ポリシーを訓練する代わりに、望ましいリターン, 過去の状態とアクションが与えられたとき、この望ましいリターンを獲得する未来のアクションを生成する シークエンスモデリング・アルゴリズム (Transformer) を使用する ということです。それは、望ましいリターン, 過去の状態とアクションで条件付けられた自己回帰モデルで、その望ましいリターンを獲得する未来のアクションを生成します。

これは強化学習パラダイムにおける完全なシフトです、何故ならば従来の RL アルゴリズムを置き換えるために 生成的軌道 (= generative trajectory) モデリング (状態, アクションと報酬のシークエンスの同時分布のモデリング) を使用するからです。それは、Decision Transformer では、リターンを最大化するのではなく、望ましいリターンを獲得する未来のアクションの系列を生成することを意味します。

そのプロセスはこのように進みます :

  1. 最後の K 個の時間ステップ を Decision Transformer に 3 つの入力とともに供給します :
    • Return-to-go
    • 状態
    • アクション

  2. トークンは 状態がベクトルであれば線形層で 埋め込まれ、それがフレームであれば CNN エンコーダで 埋め込まれます。

  3. 入力は GPT-2 モデルにより処理されます、これは自己回帰モデリングを通して未来のアクションを予測します。

Decision Transformer アーキテクチャ。状態, アクションとリターンがモダリティ固有の線形埋め込みに供給され、位置的エピソードの (= positional episodic) 時間ステップ・エンコーディングが追加されます。トークンは GPT アーキテクチャに供給されます、これは causal 自己アテンションマスクを使用してアクションを自己回帰的に予測します。図は [1] から引用。

様々なタイプの Decision Transformers がありますが、本日は、オフライン Decision Transformer を訓練していきます、つまり他のエージェントや人間の実演から集めたデータだけを使用します。エージェントは環境と相互作用しません。オフラインとオンライン強化学習の間の違いについて更に知りたい場合は、この記事 を確認してください。

オフライン Decision Transformers の背後のカテゴリー理論を理解したので、実際にそれを訓練していく方法を見ましょう。

 

Decision Transformers の訓練

前の記事で、transformers Decision Transformer モデルを使用して 🤗 ハブから事前訓練済み重みをロードする方法を実演しました。

このパートでは、🤗 ハブにホストされているオフライン RL データセットを使用し、Decision Transformer モデルをゼロから訓練するために 🤗 Trainer とカスタム Data Collator を使用します。
このチュートリアルのコードは この colab ノートブック で見つけられます。

mujoco halfcheetah 環境 で以下の動作を学習するためにオフライン RL を実行していきます。

*An “expert” Decision Transformers model, learned using offline RL in the Gym HalfCheetah environment.*

 

データセットのロードとカスタム Data Collator の構築

私たちは幾つかのオフライン RL データセットをハブ上にホストしています。本日は、ハブ上にここでホストされている、halfcheetah “expert” データセットで訓練していきます。

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

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

ハブの殆どのデータセットがそのまま使用できる準備ができている一方で、ときどきデーセットにある追加の処理や変更を行ないたいことがあります。この場合、作者の実装にマッチさせたいです、つまり以下が必要です :

  • 平均を引いて標準偏差で除算することにより各特徴を正規化します。

  • 各軌跡に対する割引リターンを事前計算します。

  • 報酬とリターンを 1000 の因子でスケールする。

  • expert エージェントの軌跡の長さを考慮するように、データセットのサンプリング分布を増強します。

このデータセットの前処理を実行するため、カスタムの 🤗 Data Collator を使用します。

それではオフライン強化学習のためのカスタム Data Collator 上で始めましょう。

@dataclass
class DecisionTransformerGymDataCollator:
    return_tensors: str = "pt"
    max_len: int = 20 #subsets of the episode we use for training
    state_dim: int = 17  # size of state space
    act_dim: int = 6  # size of action space
    max_ep_len: int = 1000 # max episode length in the dataset
    scale: float = 1000.0  # normalization of rewards/returns
    state_mean: np.array = None  # to store state means
    state_std: np.array = None  # to store state stds
    p_sample: np.array = None  # a distribution to take account trajectory lengths
    n_traj: int = 0 # to store the number of trajectories in the dataset

    def __init__(self, dataset) -> None:
        self.act_dim = len(dataset[0]["actions"][0])
        self.state_dim = len(dataset[0]["observations"][0])
        self.dataset = dataset
        # calculate dataset stats for normalization of states
        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)
        # this is a bit of a hack to be able to sample of a non-uniform distribution
        batch_inds = np.random.choice(
            np.arange(self.n_traj),
            size=batch_size,
            replace=True,
            p=self.p_sample,  # reweights so we sample according to timesteps
        )
        # a batch of dataset features
        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)

            # get sequences from dataset
            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  # padding cutoff
            rtg.append(
                self._discount_cumsum(np.array(feature["rewards"][si:]), gamma=1.0)[
                    : s[-1].shape[1]   # TODO check the +1 removed here
                ].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)

            # padding and state + reward normalization
            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 norm を含むことを保証する必要があります。Decision Transformer モデルから継承した TrainableDT クラスを作成することでこれを実現します。

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

    def forward(self, **kwargs):
        output = super().forward(**kwargs)
        # add the DT loss
        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 クラスで定義された幾つかの引数を必要としました。作者のオリジナル実装のものと同じハイパーパラメータを使用しますが、より少ない反復で訓練します。This takes around 40 minutes to train in a colab notebook, so grab a coffee or read the 🤗 Annotated Diffusion blogpost while you wait. 作者はおよそ 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 の背後の理論、そしてそれを訓練する方法を說明しました。貴方は half-cheetah を走らせるために最初のオフライン Decision Transformer モデルをゼロから訓練する準備ができました 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb
colab は訓練済みモデルの可視化とモデルを 🤗 ハブにセーブする方法を含みます。

 

まとめ

この記事は Decision Transformer を (🤗 datasets にホストされている) オフライン RL データセットで訓練する方法を実演しました。🤗 transformers Trainer とカスタム data collator を使用しました。

In addition to Decision Transformers, we want to support more use cases and tools from the Deep Reinforcement Learning community. Therefore, it would be great to hear your feedback on the Decision Transformer model, and more generally anything we can build with you that would be useful for RL. Feel free to reach out to us.

 

What’s next?

In the coming weeks and months, we plan on supporting other tools from the ecosystem:

  • Expanding our repository of Decision Transformer models with models trained or finetuned in an online setting [2]

  • Integrating sample-factory version 2.0

The best way to keep in touch is to join our discord server to exchange with us and with the community.

 

References

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

  • [2] Zheng, Qinqing and Zhang, Amy and Grover, Aditya “Online Decision Transformer” (arXiv preprint, 2022)

 

以上