Blog de Amazon Web Services (AWS)

Formación de un agente de aprendizaje de refuerzo con Unity y Amazon SageMaker RL

Por Yohei Nakayama, Henry Wang y Yijie Zhuang

 

Unity es uno de los motores de juego más populares que ha sido adoptado no sólo para el desarrollo de videojuegos sino también por industrias como el cine y la automoción. Unity ofrece herramientas para crear entornos virtuales simulados con física, paisajes y personajes personalizables. Unity Machine Learning Agents Toolkit (ML-Agents) es un proyecto de código abierto que permite a los desarrolladores entrenar agentes de aprendizaje de refuerzo (RL) contra los entornos creados en Unity.

El aprendizaje de refuerzo es un área de aprendizaje automático (ML) que enseña a un agente de software cómo tomar acciones en un entorno con el fin de maximizar un objetivo a largo plazo. Para obtener más información, consulte Amazon SageMaker RL — Managed Reinforcement Learning with Amazon SageMaker. ML-Agents se está convirtiendo en una herramienta cada vez más popular entre muchas compañías de juegos para casos de uso como el diseño de dificultad del nivel de juego, la corrección de errores y la detección de trampas. Actualmente, ML-Agents se utiliza para capacitar a los agentes localmente, y no se puede escalar para utilizar de manera eficiente más recursos informáticos. Usted tiene que entrenar a los agentes de RL en un motor Unity local durante un largo período de tiempo antes de obtener el modelo entrenado. El proceso consume mucho tiempo y no es escalable para procesar grandes cantidades de datos.

En esta publicación, demostramos una solución mediante la integración de la interfaz ML-Agents Unity con Amazon SageMaker RL, lo que le permite capacitar a los agentes de RL en Amazon SageMaker de una manera totalmente gestionada y escalable.

 

Descripción general de la solución

SageMaker es un servicio totalmente gestionado que permite un rápido desarrollo de modelos. Proporciona muchas características integradas para ayudarle con la capacitación, el ajuste, la depuración y la implementación de modelos. SageMaker RL se construye sobre SageMaker, agregando bibliotecas RL preconstruidas y facilitando la integración con diferentes entornos de simulación. Puede utilizar frameworks de Deep learning integrados como TensorFlow y PyTorch con varios algoritmos RL integrados de la biblioteca RLLib para entrenar políticas de RL. Las infraestructuras para la formación y la inferencia son totalmente gestionadas por SageMaker, por lo que puede centrarse en la formulación de RL. SageMaker RL también proporciona un conjunto de  Jupyter notebooks, lo que demuestra diversas aplicaciones de RL de dominio en robótica, investigación de operaciones, finanzas y más.

El siguiente diagrama ilustra nuestra arquitectura de soluciones.

 

 

En este post, analizamos los detalles de la formación de un agente de RL en SageMaker interactuando con el entorno Unity de muestra. Para acceder al Jupyter notebook completo de esta publicación, consulte el ejemplo del Jupyter notebook de SageMaker en GitHub.

 

Configuración de sus entornos

Para empezar, importamos las bibliotecas de Python necesarias y configuramos entornos para permisos y configuraciones. El código siguiente contiene los pasos para configurar un depósito de Amazon Simple Storage Service (Amazon S3), definir el prefijo del trabajo de formación, especificar la ubicación del trabajo de formación y crear un rol de AWS Identity and Access Management (IAM):

 

import sagemaker
import boto3
 
# set up the linkage and authentication to the S3 bucket
sage_session = sagemaker.session.Session()
s3_bucket = sage_session.default_bucket()  
s3_output_path = 's3://{}/'.format(s3_bucket)
print("S3 bucket path: {}".format(s3_output_path))

# create a descriptive job name
job_name_prefix = 'rl-unity-ray'

# configure where training happens – local or SageMaker instance
local_mode = False

if local_mode:
    instance_type = 'local'
else:
    # If on SageMaker, pick the instance type
    instance_type = "ml.c5.2xlarge"

# create an IAM role
try:
    role = sagemaker.get_execution_role()
except:
    role = get_execution_role()

print("Using IAM role arn: {}".format(role))

 

Creación de un contenedor Docker

SageMaker utiliza contenedores Docker para ejecutar scripts, entrenar algoritmos e implementar modelos. Un contenedor Docker es un paquete independiente de software que administra todo el código y dependencias, e incluye todo lo necesario para ejecutar una aplicación. Empezamos construyendo sobre una imagen preconstruida de SageMaker Docker que contiene dependencias para Ray, luego instalamos los paquetes principales requeridos:

  • gym-unity — Unity proporciona un envoltorio para envolver el entorno Unity en una interfaz de  gym, una biblioteca de código abierto que le da acceso a un conjunto de entornos RL clásicos
  • mlagents-envs — Paquete que proporciona una API de Python para permitir la interacción directa con el motor de juego Unity

Dependiendo del estado de la máquina, el proceso de construcción de Docker puede tardar hasta 10 minutos. Para todas las imágenes preconstruidas de SageMaker RL Docker, consulte el repositorio de GitHub.

 

Ejemplo de entorno Unity

En esta publicación, utilizamos un ejemplo sencillo de entorno de Unity llamado Basic. En la siguiente visualización, el agente que estamos controlando es el cuadro azul que se mueve hacia la izquierda o hacia la derecha. Por cada paso que toma, le cuesta al agente algo de energía, incurriendo en pequeñas recompensas negativas (-0.01). Las bolas verdes son objetivos con ubicaciones fijas. El agente se inicializa aleatoriamente entre las bolas verdes, y recoge recompensas cuando colisiona con las bolas verdes. La bola verde grande ofrece una recompensa de +1, y la pequeña bola verde ofrece una recompensa de +0.1. El objetivo de esta tarea es entrenar al agente para moverse hacia la pelota que ofrece las recompensas más acumuladas.

 

 

Modelo de capacitación, evaluación e implementación

En esta sección, le guiamos a través de los pasos para entrenar, evaluar e implementar modelos.

Escribir un guión de entrenamiento

Antes de iniciar el trabajo de formación de SageMaker RL, necesitamos especificar las configuraciones del proceso de formación. Por lo general, se logra en un solo guión fuera del cuaderno. El script de entrenamiento define la entrada (el entorno Unity) y el algoritmo para la formación RL. El siguiente código muestra cómo se ve el script:

 

import json
import os

import gym
import ray
from ray.tune import run_experiments
from ray.tune.registry import register_env

from sagemaker_rl.ray_launcher import SageMakerRayLauncher
from mlagents_envs.environment import UnityEnvironment
from mlagents_envs.exception import UnityWorkerInUseException
from mlagents_envs.registry import default_registry
from gym_unity.envs import UnityToGymWrapper

class UnityEnvWrapper(gym.Env):
    def __init__(self, env_config):
        self.worker_index = env_config.worker_index
        if 'SM_CHANNEL_TRAIN' in os.environ:
            env_name = os.environ['SM_CHANNEL_TRAIN'] +'/'+ env_config['env_name']
            os.chmod(env_name, 0o755)
            print("Changed environment binary into executable mode.")
            # Try connecting to the Unity3D game instance.
            while True:
                try:
                    unity_env = UnityEnvironment(
                                    env_name, 
                                    no_graphics=True, 
                                    worker_id=self.worker_index, 
                                    additional_args=['-logFile', 'unity.log'])
                except UnityWorkerInUseException:
                    self.worker_index += 1
                else:
                    break
        else:
            env_name = env_config['env_name']
            while True:
                try:
                    unity_env = default_registry[env_name].make(
                        no_graphics=True,
                        worker_id=self.worker_index,
                        additional_args=['-logFile', 'unity.log'])
                except UnityWorkerInUseException:
                    self.worker_index += 1
                else:
                    break
            
        self.env = UnityToGymWrapper(unity_env) 
        self.action_space = self.env.action_space
        self.observation_space = self.env.observation_space

    def reset(self):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)

class MyLauncher(SageMakerRayLauncher):

    def register_env_creator(self):
        register_env("unity_env", lambda config: UnityEnvWrapper(config))

    def get_experiment_config(self):
        return {
          "training": {
            "run": "PPO",
            "stop": {
              "timesteps_total": 10000,
            },
            "config": {
              "env": "unity_env",
              "gamma": 0.995,
              "kl_coeff": 1.0,
              "num_sgd_iter": 20,
              "lr": 0.0001,
              "sgd_minibatch_size": 100,
              "train_batch_size": 500,
              "monitor": True,  # Record videos.
              "model": {
                "free_log_std": True
              },
              "env_config":{
                "env_name": "Basic"
              },
              "num_workers": (self.num_cpus-1),
              "ignore_worker_failures": True,
            }
          }
        }

if __name__ == "__main__":
    MyLauncher().train_main()

 

El script de entrenamiento tiene dos componentes:

  • UnityEnvWrapper — El entorno Unity se almacena como un archivo binario. Para cargar el entorno, necesitamos usar la API Python de Unity ML-Agents. UnityEnvironment toma el nombre del entorno y devuelve un objeto de entorno interactivo. Luego envolveremos el objeto con UnityToGymWrapper y devolvemos un objeto que es entrenable usando Ray-RLLIB y SageMaker RL.
  • MyLauncher — Esta clase hereda la clase base SageMakerrayLauncher para que las aplicaciones SageMaker RL utilicen RAY-RLLIB. Dentro de la clase, registramos el entorno para ser reconocido por Ray y especificamos las configuraciones que queremos durante el entrenamiento. Los hiperparámetros de ejemplo incluyen el nombre del entorno, el factor de descuento en las recompensas acumuladas, la tasa de aprendizaje del modelo y el número de iteraciones para ejecutar el modelo. Para obtener una lista completa de los hiperparámetros utilizados comúnmente, consulte Parámetros comunes.

Entrenamiento del modelo

Después de configurar la configuración y la personalización del modelo, estamos listos para iniciar el trabajo de formación de SageMaker RL. Consulte el siguiente código:

 

metric_definitions = RLEstimator.default_metric_definitions(RLToolkit.RAY)
    
estimator = RLEstimator(entry_point="train-unity.py",
                        source_dir='src',
                        dependencies=["common/sagemaker_rl"],
                        image_name=custom_image_name,
                        role=role,
                        train_instance_type=instance_type,
                        train_instance_count=1,
                        output_path=s3_output_path,
                        base_job_name=job_name_prefix,
                        metric_definitions=metric_definitions,
                        hyperparameters={
				# customize Ray parameters here
                        }
                    )

estimator.fit(wait=local_mode)
job_name = estimator.latest_training_job.job_name
print("Training job: %s" % job_name)

 

  • entry_point — La ruta al script de entrenamiento que escribimos que especifica el proceso de entrenamiento
  • source_dir — La ruta al directorio con otras dependencias de código fuente de entrenamiento aparte del archivo de punto de entrada
  • dependencies : una lista de rutas a directorios con bibliotecas adicionales que se exportarán al contenedor

Además, se indica el nombre de la imagen del contenedor, la información de la instancia de entrenamiento, la ruta de salida y las métricas seleccionadas. También podemos personalizar cualquier parámetro relacionado con rayos usando el argumento hyperparameters. Lanzamos el trabajo de formación de SageMaker RL llamando a estimator.fite iniciamos el proceso de formación del modelo basado en las especificaciones del guión de entrenamiento.

En un nivel alto, el trabajo de formación inicia una red neuronal y actualiza gradualmente la red hacia la dirección en la que el agente recopila una recompensa mayor. A través de múltiples pruebas, el agente finalmente aprende cómo navegar a la ubicación altamente gratificante de manera eficiente. SageMaker RL gestiona todo el proceso y le permite ver el estado del trabajo de formación en la página Trabajos de formación de la consola de SageMaker.

También es posible supervisar el rendimiento del modelo examinando los registros de formación registrados en Amazon CloudWatch. Debido a la simplicidad de la tarea, el modelo completa el entrenamiento (10.000 movimientos de agentes) con aproximadamente 800 episodios (número de veces que el agente alcanza una bola objetivo) en menos de 1 minuto. La siguiente trama muestra que la recompensa media recogida converge alrededor de 0.9. La recompensa máxima que el agente puede obtener de este entorno es 1, y cada paso cuesta 0,01, por lo que una recompensa media alrededor de 0,9 parece ser el resultado de una política óptima, ¡lo que indica que nuestro proceso de entrenamiento es exitoso!

 

 

Evaluación del modelo

Cuando el entrenamiento del modelo está completo, podemos cargar el modelo entrenado para evaluar su rendimiento. Similar a la configuración en el guión de entrenamiento, envolvemos el entorno Unity con un gym wrapper. A continuación, creamos un agente cargando el modelo entrenado.

Para evaluar el modelo, ejecutamos el agente entrenado varias veces contra el entorno con un agente fijo y inicializaciones objetivo, y sumamos las recompensas acumuladas que el agente recopila en cada paso para cada episodio.

De los cinco episodios, la recompensa promedio de episodios es de 0,92 con la recompensa máxima de 0,93 y la recompensa mínima de 0,89, lo que sugiere que el modelo entrenado tiene un buen rendimiento.

 

Implementación del modelo

Podemos implementar la política de RL entrenada con solo unas pocas líneas de código utilizando la API de implementación del modelo de SageMaker. Puede pasar una entrada y obtener las acciones óptimas basadas en la política. La forma de entrada debe coincidir con la forma de entrada de observación del entorno.

Para el entorno Basic, implementamos el modelo y pasamos una entrada al predictor:

 

from sagemaker.tensorflow.model import TensorFlowModel

model = TensorFlowModel(model_data=estimator.model_data,
              framework_version='2.1.0',
              role=role)

predictor = model.deploy(initial_instance_count=1, 
                         instance_type=instance_type)

input = {"inputs": {'observations': np.ones(shape=(1, 20)).tolist(),
                    'prev_action': [0, 0],
                    'is_training': False,
                    'prev_reward': -1,
                    'seq_lens': -1
                   }
        }    

result = predictor.predict(input)
print(result['outputs']['actions'])

 

El modelo predice un indicador correspondiente al movimiento hacia la izquierda o hacia la derecha. La dirección de movimiento recomendada para el agente de caja azul siempre apunta hacia la bola verde más grande.

 

Limpieza

Cuando haya terminado de ejecutar el modelo, llame al predictor.delete_endpoint () para eliminar el extremo de implementación del modelo y evitar incurrir en cargos futuros.

 

Personalización de algoritmos, modelos y entornos de formación

Además del caso de uso anterior, le recomendamos que explore las capacidades de personalización que admite esta solución.

En el ejemplo de código anterior, especificamos Optimización de Políticas Proximales (PPO) como algoritmo de entrenamiento. PPO es un algoritmo RL popular que funciona comparablemente con los enfoques de vanguardia, pero es mucho más simple de implementar y ajustar. Dependiendo de su caso de uso, puede elegir el algoritmo más ajustado para el entrenamiento seleccionando de una lista de algoritmos integrales ya implementados en RLLib o construyendo un algoritmo personalizado desde cero.

De forma predeterminada, RLLib aplica una red neuronal convolucional predefinida o una red neuronal completamente conectada. Sin embargo, puede crear un modelo personalizado para entrenamiento y pruebas. Siguiendo los ejemplos de RLLib, puede registrar el modelo personalizado llamando a ModelCatalog.register_custom_modely, a continuación, hacer referencia al modelo recién registrado utilizando el argumento custom_model .

En nuestro ejemplo de código, invocamos un entorno Unity predefinido llamado Basic, pero puede experimentar con otros entornos Unity preconstruidos. Sin embargo, a partir de este momento, nuestra solución solo admite un entorno de un solo agente. Cuando se construyan nuevos entornos, regístrelo llamando a register_env y consulte el entorno con el parámetro env.

 

Conclusión

En este post, explicamos cómo entrenar a un agente de RL para que interactúe con entornos de juego Unity usando SageMaker RL. Utilizamos un ejemplo de entorno de Unity precompilado para la demostración, pero le animamos a explorar el uso de entornos Unity personalizados u otros preconstruidos.

SageMaker RL ofrece una forma escalable y eficiente de entrenar a los agentes de juego RL para jugar a entornos de juego con tecnología Unity. Para ver el cuaderno que contiene el código completo, consulte Unity 3D Game with Amazon SageMaker RL.

Si desea ayuda para acelerar el uso de ML en sus productos y procesos, póngase en contacto con Amazon ML Solutions Lab.

 

Este artículo fue traducido del Blog de AWS en Inglés

 


Sobre los Autores

Yohei Nakayama es arquitecto de aprendizaje profundo en Amazon Machine Learning Solutions Lab, donde trabaja con clientes en diferentes verticales para acelerar el uso de la inteligencia artificial y los servicios en la nube de AWS para resolver sus desafíos empresariales. Está interesado en aplicar tecnologías ML/AI a la industria espacial.

 

 

 

Henry Wang es científico de datos en Amazon Machine Learning Solutions Lab. Antes de incorporarse a AWS, era un estudiante graduado en Harvard en Ciencias Computacionales e Ingeniería, donde trabajó en investigación sanitaria con aprendizaje reforzado. En su tiempo libre, le gusta jugar tenis y golf, leer y ver torneos de StarCraft II.

 

 

 

Yijie Zhuang es ingeniero de software con Amazon SageMaker. Hizo su MS en Ingeniería Informática de Duke. Sus intereses están en construir algoritmos escalables y sistemas de aprendizaje de refuerzo. Ha contribuido a los algoritmos integrados de Amazon SageMaker y a Amazon SageMaker RL.

 

 

 

 

Revisor

Sergio Beltrán es arquitecto de soluciones especialista en AI/ML en AWS. En su rol apoya a los clientes a encontrar la mejor solución y arquitectura para sus necesidades y aprovechar los servicios de  AI/ML para generar innovación y mejorar la productividad. Antes de AWS, trabajó como data scientist y gerente de business development en la industria Telco.