Addestra il tuo primo Decision Transformer

Addestra il Decision Transformer iniziale.

In un post precedente, abbiamo annunciato il lancio di Decision Transformers nella libreria Transformers. Questa nuova tecnica di utilizzo di un Transformer come modello di decisione sta diventando sempre più popolare.

Quindi oggi, imparerai ad addestrare il tuo primo modello Offline Decision Transformer da zero per far correre un mezzo-ghepardo. Lo addestreremo direttamente su un Google Colab che puoi trovare qui 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb

*Un modello “esperto” di Decision Transformers, appreso utilizzando l’apprendimento rinforzato offline nell’ambiente HalfCheetah di Gym.*

Suona emozionante? Cominciamo!

  • Cosa sono i Decision Transformers?
  • Addestramento di Decision Transformers
    • Caricamento del set di dati e costruzione del Data Collator personalizzato
    • Addestramento del modello Decision Transformer con il Trainer di 🤗 transformers
  • Conclusione
  • Cosa succede dopo?
  • Riferimenti

Cosa sono i Decision Transformers?

Il modello Decision Transformer è stato introdotto da “Decision Transformer: Reinforcement Learning via Sequence Modeling” di Chen L. et al. Astrae l’apprendimento rinforzato come un problema di modellazione di sequenze condizionali.

L’idea principale è che invece di addestrare una politica utilizzando metodi di apprendimento rinforzato, come adattare una funzione di valore che ci indicherà quale azione intraprendere per massimizzare il ritorno (ricompensa cumulativa), utilizziamo un algoritmo di modellazione di sequenze (Transformer) che, dato il ritorno desiderato, gli stati passati e le azioni, genererà azioni future per raggiungere questo ritorno desiderato. È un modello autoregressivo condizionato al ritorno desiderato, agli stati passati e alle azioni per generare azioni future che raggiungano il ritorno desiderato.

Questo rappresenta un cambiamento completo nel paradigma dell’apprendimento rinforzato, poiché utilizziamo la modellazione generativa delle traiettorie (modellazione della distribuzione congiunta della sequenza di stati, azioni e ricompense) per sostituire gli algoritmi tradizionali di apprendimento rinforzato. Ciò significa che nei Decision Transformers non massimizziamo il ritorno, ma generiamo piuttosto una serie di azioni future che raggiungono il ritorno desiderato.

Il processo avviene in questo modo:

  1. Alimentiamo gli ultimi K passaggi temporali nel Decision Transformer con tre input:
    • Ritorno da percorrere
    • Stato
    • Azione
  2. I token vengono incorporati con un livello lineare se lo stato è un vettore o un codificatore CNN se si tratta di frame.
  3. Gli input vengono elaborati da un modello GPT-2, che predice azioni future tramite modellazione autoregressiva.

Architettura Decision Transformer. Stati, azioni e ritorni vengono inseriti in incorporamenti lineari specifici della modalità, e viene aggiunta una codifica temporale episodica posizionale. I token vengono inseriti in un’architettura GPT che predice azioni in modo autoregressivo utilizzando una maschera di attenzione causale. Figura da [1].

Ci sono diversi tipi di Decision Transformers, ma oggi addestreremo un Decision Transformer offline, il che significa che utilizziamo solo dati raccolti da altri agenti o dimostrazioni umane. L’agente non interagisce con l’ambiente. Se vuoi saperne di più sulla differenza tra apprendimento rinforzato offline e online, consulta questo articolo.

Ora che comprendiamo la teoria dietro gli Offline Decision Transformers, vediamo come addestrarne uno nella pratica.

Addestramento di Decision Transformers

Nel post precedente, abbiamo dimostrato come utilizzare un modello transformers Decision Transformer e caricare pesi preaddestrati dall’hub 🤗.

In questa parte utilizzeremo il Trainer di 🤗 e un Data Collator personalizzato per addestrare un modello Decision Transformer da zero, utilizzando un set di dati Offline RL ospitato sull’hub 🤗. Puoi trovare il codice per questo tutorial in questo notebook Colab.

Eseguiremo RL offline per imparare il comportamento seguente nell’ambiente mujoco halfcheetah.

*Un modello “esperto” di Decision Transformers, appreso utilizzando RL offline nell’ambiente HalfCheetah di Gym.*

Caricamento del set di dati e costruzione del Data Collator personalizzato

Ospitiamo diversi set di dati Offline RL sull’hub. Oggi addestreremo con il set di dati “esperto” halfcheetah, ospitato qui sull’hub.

Prima di tutto dobbiamo importare la funzione load_dataset dal pacchetto 🤗 datasets e scaricare il dataset sul nostro computer.

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

Anche se la maggior parte dei dataset nel repository sono pronti all’uso, a volte desideriamo eseguire ulteriori elaborazioni o modifiche al dataset. In questo caso, desideriamo adattare l’implementazione dell’autore, ossia:

  • Normalizzare ogni caratteristica sottraendo la media e dividendo per la deviazione standard.
  • Calcolare in anticipo i ritorni scontati per ogni traiettoria.
  • Scala le ricompense e i ritorni di un fattore di 1000.
  • Estendere la distribuzione di campionamento del dataset in modo da tenere conto della lunghezza delle traiettorie dell’agente esperto.

Per eseguire questa elaborazione del dataset, useremo un Data Collator personalizzato di 🤗.

Ora iniziamo con il Data Collator personalizzato per il Reinforcement Learning offline.

@dataclass
class DecisionTransformerGymDataCollator:
    return_tensors: str = "pt"
    max_len: int = 20 #sottoinsieme dell'episodio che usiamo per l'addestramento
    state_dim: int = 17  # dimensione dello spazio di stato
    act_dim: int = 6  # dimensione dello spazio di azione
    max_ep_len: int = 1000 # lunghezza massima dell'episodio nel dataset
    scale: float = 1000.0  # normalizzazione delle ricompense/ritorni
    state_mean: np.array = None  # per memorizzare le medie degli stati
    state_std: np.array = None  # per memorizzare le deviazioni standard degli stati
    p_sample: np.array = None  # una distribuzione per tenere conto delle lunghezze delle traiettorie
    n_traj: int = 0 # per memorizzare il numero di traiettorie nel 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
        # calcola le statistiche del dataset per la normalizzazione degli stati
        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)
        # questa è una piccola trucco per poter campionare da una distribuzione non uniforme
        batch_inds = np.random.choice(
            np.arange(self.n_traj),
            size=batch_size,
            replace=True,
            p=self.p_sample,  # ri-bilancia in modo che campioniamo in base ai passaggi temporali
        )
        # un batch di caratteristiche del dataset
        s, a, r, d, rtg, timesteps, mask = [], [], [], [], [], [], []
        
        for ind in batch_inds:
            # per ogni caratteristica in features:
            feature = self.dataset[int(ind)]
            si = random.randint(0, len(feature["rewards"]) - 1)

            # ottieni sequenze dal 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  # taglio di padding
            rtg.append(
                self._discount_cumsum(np.array(feature["rewards"][si:]), gamma=1.0)[
                    : s[-1].shape[1]   # TODO controllare il +1 rimosso qui
                ].reshape(1, -1, 1)
            )
            if rtg[-1].shape[1] < s[-1].shape[1]:
                print("se vero")
                rtg[-1] = np.concatenate([rtg[-1], np.zeros((1, 1, 1))], axis=1)

            # padding e normalizzazione dello stato e della ricompensa
            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,
        }

Questo è stato molto codice, in breve abbiamo definito una classe che prende il nostro dataset, esegue la preelaborazione richiesta e ci restituirà batch di stati, azioni, ricompense, ritorni, istanti di tempo e maschere. Questi batch possono essere utilizzati direttamente per addestrare un modello Decision Transformer con un Trainer dei transformers 🤗.

Addestramento del modello Decision Transformer con un Trainer dei transformers 🤗.

Per addestrare il modello con la classe Trainer dei 🤗, prima dobbiamo assicurarci che il dizionario restituito contenga una perdita, in questo caso la norma L-2 delle previsioni delle azioni del modello e degli obiettivi. Per farlo, creiamo una classe TrainableDT, che eredita dal modello Decision Transformer.

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

    def forward(self, **kwargs):
        output = super().forward(**kwargs)
        # aggiungi la perdita 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)

La classe Trainer dei transformers richiede un certo numero di argomenti, definiti nella classe TrainingArguments. Utilizziamo gli stessi iperparametri dell’implementazione originale degli autori, ma addestriamo per meno iterazioni. Ci vogliono circa 40 minuti per addestrare in un notebook Colab, quindi prenditi un caffè o leggi il post del blog 🤗 Annotated Diffusion mentre aspetti. Gli autori addestrano per circa 3 ore, quindi i risultati che otteniamo qui non saranno così buoni come i loro.

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()

Ora che abbiamo spiegato la teoria dietro Decision Transformer, il Trainer e come addestrarlo. Sei pronto per addestrare il tuo primo modello Decision Transformer offline da zero per far correre un mezzo ghepardo 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb Il Colab include visualizzazioni del modello addestrato, così come come salvare il tuo modello sul 🤗 hub.

Conclusioni

Questo post ha dimostrato come addestrare il Decision Transformer su un dataset RL offline, ospitato su 🤗 datasets . Abbiamo utilizzato un Trainer dei transformers 🤗 e un data collator personalizzato.

Oltre ai Decision Transformer, vogliamo supportare più casi d’uso e strumenti dalla community dell’apprendimento automatico per rinforzo . Pertanto, sarebbe bello sentire il tuo feedback sul modello Decision Transformer e, più in generale, su qualsiasi cosa possiamo costruire insieme a te che potrebbe essere utile per RL. Non esitare a contattarci .

Cosa c’è di nuovo?

Nelle prossime settimane e mesi, abbiamo intenzione di supportare altri strumenti dell’ecosistema :

  • Ampliando il nostro repository di modelli Decision Transformer con modelli addestrati o raffinati in un ambiente online [2]
  • Integrazione della versione 2.0 di sample-factory

Il modo migliore per rimanere in contatto è unirsi al nostro server Discord per scambiare informazioni con noi e con la community.

Riferimenti

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

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