Entrena tu primer Decision Transformer

Entrena tu Decision Transformer inicial.

En una publicación anterior, anunciamos el lanzamiento de Decision Transformers en la biblioteca transformers. Esta nueva técnica de utilizar un Transformer como modelo de toma de decisiones está ganando cada vez más popularidad.

Por lo tanto, hoy aprenderás a entrenar tu primer modelo de Decision Transformer Offline desde cero para hacer correr a un semi-guepardo. Lo entrenaremos directamente en un Google Colab que puedes encontrar aquí 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb

*Un modelo “experto” de Decision Transformers, aprendido utilizando RL fuera de línea en el entorno HalfCheetah de Gym.*

¿Suena emocionante? ¡Empecemos!

• ¿Qué hay de nuevo en los Difusores? 🎨
• Cómo entrenar un Modelo de Lenguaje con Megatron-LM
• SetFit Aprendizaje Eficiente de Pocos Ejemplos Sin Indicaciones

• ¿Qué son los Decision Transformers?
• Entrenamiento de Decision Transformers
  • Cargando el conjunto de datos y construyendo el Collator de Datos Personalizado
  • Entrenando el modelo de Decision Transformer con un entrenador de 🤗 transformers
• Conclusión
• ¿Qué sigue?
• Referencias

¿Qué son los Decision Transformers?

El modelo Decision Transformer fue introducido por “Decision Transformer: Aprendizaje por Reforzamiento a través de Modelado de Secuencias” por Chen L. et al. Abstrae el Aprendizaje por Reforzamiento como un problema de modelado de secuencias condicionales.

La idea principal es que en lugar de entrenar una política utilizando métodos de RL, como ajustar una función de valor que nos dirá qué acción tomar para maximizar el retorno (recompensa acumulada), utilizamos un algoritmo de modelado de secuencias (Transformer) que, dada la recompensa deseada, los estados pasados y las acciones, generará acciones futuras para lograr este retorno deseado. Es un modelo autoregresivo condicionado al retorno deseado, los estados pasados y las acciones para generar acciones futuras que logren el retorno deseado.

Este es un cambio completo en el paradigma del Aprendizaje por Reforzamiento, ya que utilizamos el modelado generativo de trayectorias (modelando la distribución conjunta de la secuencia de estados, acciones y recompensas) para reemplazar los algoritmos convencionales de RL. Esto significa que en los Decision Transformers, no maximizamos el retorno, sino que generamos una serie de acciones futuras que logren el retorno deseado.

El proceso es el siguiente:

1. Alimentamos los últimos K pasos de tiempo en el Decision Transformer con tres entradas:
   • Retorno-a-ir
   • Estado
   • Acción
2. Los tokens se incrustan ya sea con una capa lineal si el estado es un vector o un codificador de CNN si son fotogramas.
3. Las entradas son procesadas por un modelo GPT-2, que predice acciones futuras mediante modelado autoregresivo.

Arquitectura del Decision Transformer. Los estados, acciones y retornos se alimentan en incrustaciones lineales específicas de la modalidad, y se agrega una codificación episódica posicional para cada paso de tiempo. Los tokens se alimentan en una arquitectura GPT que predice acciones autoregresivamente utilizando una máscara de atención causal. Figura de [1].

Existen diferentes tipos de Decision Transformers, pero hoy vamos a entrenar un Decision Transformer fuera de línea, lo que significa que solo utilizamos datos recopilados de otros agentes o demostraciones humanas. El agente no interactúa con el entorno. Si quieres saber más sobre la diferencia entre el aprendizaje por refuerzo en línea y fuera de línea, revisa este artículo.

Ahora que comprendemos la teoría detrás de los Decision Transformers Offline, veamos cómo vamos a entrenar uno en la práctica.

Entrenamiento de Decision Transformers

En la publicación anterior, demostramos cómo utilizar un modelo Decision Transformer de transformers y cargar pesos preentrenados desde el 🤗 hub.

En esta parte utilizaremos el 🤗 Trainer y un Collator de Datos personalizado para entrenar un modelo Decision Transformer desde cero, utilizando un conjunto de datos de RL fuera de línea alojado en el 🤗 hub. Puedes encontrar el código para este tutorial en este cuaderno Colab.

Realizaremos RL fuera de línea para aprender el siguiente comportamiento en el entorno del semi-guepardo de mujoco.

*Un modelo “experto” de Decision Transformers, aprendido utilizando RL fuera de línea en el entorno HalfCheetah de Gym.*

Cargando el conjunto de datos y construyendo el Collator de Datos Personalizado

Hospedamos varios conjuntos de datos de RL fuera de línea en el hub. Hoy entrenaremos con el conjunto de datos “experto” de medio guepardo, alojado aquí en el hub.

Primero necesitamos importar la función load_dataset del paquete 🤗 datasets y descargar el conjunto de datos en nuestra máquina.

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

Aunque la mayoría de los conjuntos de datos en el hub están listos para usar, a veces deseamos realizar algún procesamiento adicional o modificación del conjunto de datos. En este caso, deseamos coincidir con la implementación del autor, es decir, necesitamos:

• Normalizar cada característica restando la media y dividiendo por la desviación estándar.
• Calcular los retornos descontados para cada trayectoria.
• Escalar las recompensas y retornos por un factor de 1000.
• Aumentar la distribución de muestreo del conjunto de datos para tener en cuenta la longitud de las trayectorias del agente experto.

Para realizar este preprocesamiento del conjunto de datos, utilizaremos un 🤗 Data Collator personalizado.

Ahora, comencemos con el Data Collator personalizado para el Aprendizaje por Reforzamiento sin Conexión.

@dataclass
class DecisionTransformerGymDataCollator:
    return_tensors: str = "pt"
    max_len: int = 20 #subconjuntos del episodio que usamos para el entrenamiento
    state_dim: int = 17  # tamaño del espacio de estado
    act_dim: int = 6  # tamaño del espacio de acción
    max_ep_len: int = 1000 # longitud máxima del episodio en el conjunto de datos
    scale: float = 1000.0  # normalización de recompensas/retornos
    state_mean: np.array = None  # para almacenar las medias de los estados
    state_std: np.array = None  # para almacenar las desviaciones estándar de los estados
    p_sample: np.array = None  # una distribución para tener en cuenta las longitudes de las trayectorias
    n_traj: int = 0 # para almacenar el número de trayectorias en el conjunto de datos

    def __init__(self, dataset) -> None:
        self.act_dim = len(dataset[0]["actions"][0])
        self.state_dim = len(dataset[0]["observations"][0])
        self.dataset = dataset
        # calcular estadísticas del conjunto de datos para la normalización de los estados
        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)
        # esto es un poco de trampa para poder muestrear una distribución no uniforme
        batch_inds = np.random.choice(
            np.arange(self.n_traj),
            size=batch_size,
            replace=True,
            p=self.p_sample,  # reajusta para muestrear según los pasos de tiempo
        )
        # un lote de características del conjunto de datos
        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)

            # obtener secuencias del conjunto de datos
            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  # corte de relleno
            rtg.append(
                self._discount_cumsum(np.array(feature["rewards"][si:]), gamma=1.0)[
                    : s[-1].shape[1]   # TODO verificar el +1 eliminado aquí
                ].reshape(1, -1, 1)
            )
            if rtg[-1].shape[1] < s[-1].shape[1]:
                print("si es verdadero")
                rtg[-1] = np.concatenate([rtg[-1], np.zeros((1, 1, 1))], axis=1)

            # relleno y normalización de estados + recompensas
            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,
        }

Eso fue mucho código, el TLDR es que definimos una clase que toma nuestro conjunto de datos, realiza el preprocesamiento requerido y nos devolverá lotes de estados, acciones, recompensas, rendimientos, timesteps y máscaras. Estos lotes se pueden utilizar directamente para entrenar un modelo Decision Transformer con un Entrenador de 🤗 transformers.

Entrenando el modelo Decision Transformer con un Entrenador de 🤗 transformers.

Para entrenar el modelo con la clase 🤗 Trainer, primero debemos asegurarnos de que el diccionario que devuelve contenga una pérdida, en este caso, la norma L-2 de las predicciones de acción del modelo y los objetivos. Logramos esto haciendo una clase TrainableDT, que hereda del modelo Decision Transformer.

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

    def forward(self, **kwargs):
        output = super().forward(**kwargs)
        # agregar la pérdida de 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 clase 🤗 Trainer de transformers requiere una serie de argumentos, definidos en la clase TrainingArguments. Utilizamos los mismos hiperparámetros que en la implementación original de los autores, pero entrenamos durante menos iteraciones. Esto lleva alrededor de 40 minutos de entrenamiento en un cuaderno de Colab, así que tómate un café o lee la publicación del blog 🤗 Annotated Diffusion mientras esperas. Los autores entrenan durante aproximadamente 3 horas, por lo que los resultados que obtengamos aquí no serán tan buenos como los suyos.

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

Ahora que hemos explicado la teoría detrás de Decision Transformer, el Entrenador y cómo entrenarlo. Estás listo para entrenar tu primer modelo Decision Transformer fuera de línea desde cero para hacer correr un medio guepardo 👉 https://github.com/huggingface/blog/blob/main/notebooks/101_train-decision-transformers.ipynb El Colab incluye visualizaciones del modelo entrenado, así como cómo guardar tu modelo en el 🤗 hub.

Conclusión

Esta publicación ha demostrado cómo entrenar el Decision Transformer en un conjunto de datos de RL fuera de línea, alojado en 🤗 datasets. Hemos utilizado un Entrenador de 🤗 transformers y un consolidador de datos personalizado.

Además de Decision Transformers, queremos admitir más casos de uso y herramientas de la comunidad de Deep Reinforcement Learning. Por lo tanto, sería genial recibir tus comentarios sobre el modelo Decision Transformer y, en general, cualquier cosa que podamos construir contigo y que sea útil para RL. No dudes en ponerte en contacto con nosotros.

¿Qué sigue?

En las próximas semanas y meses, planeamos admitir otras herramientas del ecosistema:

• Ampliar nuestro repositorio de modelos Decision Transformer con modelos entrenados o ajustados en un entorno en línea [2]
• Integrar la versión 2.0 de sample-factory

La mejor manera de mantenerse en contacto es unirse a nuestro servidor de Discord para interactuar con nosotros y con la comunidad.

Referencias

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

[2] Zheng, Qinqing y Zhang, Amy y Grover, Aditya “Transformador de Decisión en Línea” (preimpresión de arXiv, 2022)

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