O blog da AWS

Treinando um agente de aprendizado por reforço com Unity e Amazon SageMaker RL

Por Yohei Nakayama, Henry Wang e Yijie Zhuang

 

 

Unity é um dos motores de jogos mais populares que foi adotado não só para o desenvolvimento de videogames, mas também por indústrias como a cinematográfica e automotiva. Unity oferece ferramentas para criar ambientes virtuais simulados com física, paisagens e personagens customizáveis. O Unity Machine Learning Agents Toolkit (ML-Agents) é um projeto de código aberto que permite aos desenvolvedores treinar agentes de aprendizado por reforço (Reinforcement Learning – RL) nos ambientes criados no Unity.

O aprendizado por reforço é uma área do aprendizado de máquina (ML) que ensina um agente de software a tomar ações em um ambiente para maximizar um objetivo de longo prazo. Para obter mais informações, consulte Amazon SageMaker RL — Aprendizado por reforço gerenciado com o Amazon SageMaker. ML-Agents está se tornando uma ferramenta cada vez mais popular entre muitas empresas de jogos para casos de uso como design de dificuldade de nível de jogo, correção de bugs e detecção de fraude. Atualmente, o ML-Agents é usado para treinar agentes localmente e não consegue escalar para usar eficientemente mais recursos computacionais. Você tem que treinar agentes de RL com a Unity engine localmente por um longo período de tempo antes de obter o modelo treinado. O processo é demorado e não escalável para processar grandes quantidades de dados.

Neste post, demonstramos uma solução integrando a interface Unity do ML-Agents com o Amazon SageMaker RL, permitindo treinar agentes de RL no Amazon SageMaker de forma totalmente gerenciada e escalável.

 

Visão geral da solução

O SageMaker é um serviço totalmente gerenciado que permite o rápido desenvolvimento de modelos. Ele fornece muitos recursos integrados para ajudá-lo com treinamento, ajuste, depuração e implantação de modelos. O SageMaker RL se baseia no SageMaker, adicionando bibliotecas de RL pré-configuradas e facilitando a integração com diferentes ambientes de simulação. Você pode usar frameworks integrados de aprendizado profundo, como TensorFlow e PyTorch, com vários algoritmos RL integrados da biblioteca RLLib para treinar políticas de RL. A infraestrutura para treinamento e inferência é totalmente gerenciada pelo SageMaker, para que você possa se concentrar na formulação do problemas de RL. O SageMaker RL também fornece um conjunto de notebooks Jupyter, demonstrando variedades de aplicações de RL em domínios como robótica, pesquisa operacional, finanças e muito mais.

O diagrama a seguir ilustra nossa arquitetura da solução.

 

 

Neste post, analisamos as especificidades do treinamento de um agente de RL no SageMaker interagindo com o ambiente Unity de exemplo. Para acessar o notebook completo deste post, consulte o exemplo de notebook do SageMaker no GitHub.

 

Configurando seus ambientes

Para começar, importamos as bibliotecas Python necessárias e configuramos ambientes e permissões. O código a seguir contém as etapas para configurar um bucket do Amazon Simple Storage Service (Amazon S3), definir o prefixo do trabalho de treinamento, especificar o local do trabalho de treinamento e criar uma role do 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))

 

Criando um contêiner Docker

O SageMaker usa contêineres do Docker para executar scripts, treinar algoritmos e implantar modelos. Um contêiner Docker é um pacote independente de software que gerencia todo o código e dependências, e inclui tudo o que é necessário para executar um aplicativo. Começamos criando em cima de uma imagem Docker pré-criada do SageMaker que contém dependências para o Ray e, em seguida, instalamos os pacotes principais necessários:

  • gym-unity — Unity fornece um wrapper para envolver o ambiente Unity em uma interface do gym, uma biblioteca de código aberto que oferece acesso a um conjunto de ambientes RL clássicos
  • mlagents-envs — Pacote que fornece uma API Python para permitir interação direta com o motor de jogo Unity

Dependendo da máquina, o processo de build do Docker pode levar até 10 minutos. Para todas as imagens do SageMaker RL Docker pré-criadas, consulte o repositório do GitHub.

 

Exemplo de ambiente Unity

Neste post, usamos um exemplo simples ambiente Unity chamado Basic. Na visualização a seguir, o agente que estamos controlando é a caixa azul que se move para a esquerda ou para a direita. Cada passo que é tomado custa ao agente alguma energia, incorrendo em pequenas recompensas negativas (-0,01). Bolas verdes são os alvos com locais fixos. O agente é inicializado aleatoriamente entre as bolas verdes, e recolhe recompensas quando colide com as bolas verdes. A bola verde grande oferece uma recompensa de +1, e a pequena bola verde oferece uma recompensa de +0,1. O objetivo desta tarefa é treinar o agente para se mover em direção à bola que oferece as recompensas maiores.

 

 

Treinamento, avaliação e implantação de modelos

Nesta seção, vamos ajudar você a navegar pelas etapas de treino, avaliação e implantação de modelos.

Escrevendo um script de treinamento

Antes de iniciar o trabalho de treinamento no SageMaker RL, precisamos especificar as configurações do processo de treinamento. Isso é geralmente é feito em um script fora do notebook. O script de treinamento define a entrada (o ambiente Unity) e o algoritmo para treinamento RL. O código a seguir mostra como é o 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()

 

O script de treinamento tem dois componentes:

  • UnityEnvWrapper — O ambiente Unity é armazenado como um arquivo binário. Para carregar o ambiente, precisamos usar a API Python do Unity ML-Agents. UnityEnvironment recebe como entrada o nome do ambiente e retorna um objeto de ambiente interativo. Em seguida, criamos um wrapper ao objeto com UnityToGymWrapper e retornamos um objeto que é treinável usando Ray-RLLib e SageMaker RL.
  • MyLauncher — Essa classe herda a classe base SageMakerRayLauncher para aplicativos SageMaker RL que usam o Ray-RLLib. Dentro da classe, registramos o ambiente para ser reconhecido pelo Ray e especificamos as configurações que queremos durante o treinamento. Os hiperparâmetros de exemplo incluem o nome do ambiente, o fator de desconto em recompensas cumulativas, a taxa de aprendizado do modelo e o número de iterações para executar o modelo. Para obter uma lista completa de hiperparâmetros comumente usados, consulte Parâmetros comuns.

Treinando o modelo

Depois de criar a configuração e customizar o modelo, estamos prontos para iniciar o trabalho de treinamento do SageMaker RL. Veja o seguinte 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)

 

  • Dentro do código, especificamos alguns parâmetros:
    • entry_point — O caminho para o script de treinamento que escrevemos que especifica o processo de treinamento
    • source_dir — O caminho para o diretório com outras dependências do código-fonte de treinamento, além do arquivo de entry point
    • dependencies — Uma lista de caminhos para diretórios com bibliotecas adicionais a serem exportadas para o contêiner

    Além disso, nós configuramos o nome da imagem do contêiner, as informações da instância de treinamento, o caminho de saída e as métricas selecionadas. Também temos permissão para personalizar quaisquer parâmetros relacionados ao Ray usando o argumento hyperparameters. Lançamos o trabalho de treinamento do SageMaker RL chamando estimator.fit e iniciamos o processo de treinamento do modelo com base nas especificações do script de treinamento.

    Em um nível alto, o trabalho de treinamento inicia uma rede neural e atualiza a rede gradualmente para a direção em que o agente coleta recompensas mais altas. Após várias interações, o agente eventualmente aprende como navegar para o local de altas recompensas de forma eficiente. O SageMaker RL lida com todo o processo e permite visualizar o status do trabalho de treinamento na página Trabalhos de treinamento na console do SageMaker.

    Também é possível monitorar o desempenho do modelo examinando os logs de treinamento no Amazon CloudWatch. Devido à simplicidade da tarefa, o modelo completa o treinamento (10.000 movimentos do agente) com aproximadamente 800 episódios (número de vezes que o agente atinge uma bola alvo) em menos de 1 minuto. O gráfico a seguir mostra que a recompensa média coletada converge em torno de 0,9. A recompensa máxima que o agente pode obter neste ambiente é 1, e cada passo custa 0,01, portanto, uma recompensa média em torno de 0,9 parece ser o resultado de uma política ideal, indicando que nosso processo de treinamento é bem sucedido!

 

Avaliando o modelo

Quando o treinamento do modelo estiver completo, podemos carregar o modelo treinado para avaliar seu desempenho. Semelhante à configuração no script de treinamento, criamos um wrapper do Gym sobre o ambiente Unity. Em seguida, criamos um agente carregando o modelo treinado.

Para avaliar o modelo, executamos o agente treinado várias vezes no ambiente com um agente fixo e objetivo inicializado, e adicionamos as recompensas cumulativas que o agente coleta em cada etapa para cada episódio.

Dos cinco episódios, a recompensa média do episódio é de 0,92 com a recompensa máxima de 0,93 e recompensa a mínima de 0,89, sugerindo que o modelo treinado realmente tem um bom desempenho.

Implantando o modelo

Podemos implantar a política de RL treinada com apenas algumas linhas de código usando a API de implantação do modelo no SageMaker. Você pode enviar uma entrada e obter as ações ideais com base na política. A forma da entrada precisa seguir o formato de entrada de observação do ambiente.

Para o ambiente Basic, implementamos o modelo e passamos uma entrada para o 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'])

 

O modelo prevê um movimento para a esquerda ou para a direita. A direção de movimento recomendada para o agente caixa azul sempre aponta para a bola verde maior.

 

Limpando o que foi criado

Quando terminar de executar o modelo, chame predictor.delete_endpoint() para excluir o endpoint de implantação do modelo e evitar cobranças futuras.

 

Customizando algoritmos, modelos e ambientes de treinamento

Além do caso de uso anterior, recomendamos que você explore as funcionalidades de customização que essa solução oferece suporte.

No exemplo de código anterior, especificamos a Proximal Policy Optimization (PPO) para ser o algoritmo de treinamento. A PPO é um algoritmo RL popular que tem performance comparável às abordagens estado da arte, mas é muito mais simples de implementar e fazer tuning. Dependendo do seu caso de uso, você pode escolher o algoritmo mais apropriado para treinamento de uma lista abrangente de algoritmos já implementados no RLLib ou construindo um algoritmo personalizado a partir do zero.

Por padrão, o RLLib aplica uma rede neural convolucional pré-definida ou rede neural totalmente conectada. No entanto, você pode criar um modelo personalizado para treinamento e teste. Seguindo os exemplos do RLLib, você pode registrar o modelo personalizado chamando ModelCatalog.register_custom_model e, em seguida, usar o modelo recém-registrado usando o argumento custom_model.

Em nosso exemplo de código, invocamos um ambiente Unity pré-definido chamado Basic, mas você pode experimentar outros ambientes Unity pré-criados. No momento de escrita deste post, nossa solução suporta apenas um ambiente de agente único. Quando novos ambientes são construídos, registre-os chamando register_env  e use o ambiente com o parâmetro env.

 

Conclusão

Neste post, examinamos como treinar um agente de RL para interagir com ambientes de jogos Unity usando o SageMaker RL. Usamos um exemplo de ambiente Unity pré-criado para a demonstração, mas encorajamos você a explorar usando ambientes Unity personalizados ou outros pré-criados.

O SageMaker RL oferece uma maneira escalável e eficiente de treinar agentes de RL a jogar em ambientes Unity. Para o notebook que contém o código completo, consulte Unity 3D Game com Amazon SageMaker RL.

Se você quiser ajuda para acelerar seu uso de ML em seus produtos e processos, entre em contato com o Amazon ML Solutions Lab.

 

Esse artigo foi tadruzido do Blog da AWS em Inglês

 

Sobre los Autores

Yohei Nakayama é arquiteto de Deep Learning no Amazon Machine Learning Solutions Lab, onde trabalha com clientes em diferentes verticais para acelerar seu uso de inteligência artificial e serviços de nuvem AWS para resolver seus desafios de negócios. Ele está interessado em aplicar tecnologias ML/AI para a indústria espacial.

 

 

 

Henry Wang é Cientista de Dados no Amazon Machine Learning Solutions Lab. Antes de ingressar na AWS, ele foi estudante de pós-graduação em Harvard em Ciência e Engenharia Computacional, onde trabalhou em pesquisa em saúde com aprendizado por reforço. Em seu tempo livre, ele gosta de jogar tênis e golfe, ler e assistir torneios de StarCraft II.

 

 

 

Yijie Zhuang é Engenheiro de Software do Amazon SageMaker. Ele fez seu mestrado em Engenharia da Computação na Duke. Seus interesses estão na construção de algoritmos escaláveis e sistemas de aprendizagem por reforço. Ele contribuiu no desenvolvimento de algoritmos pré-configurados no Amazon SageMaker e Amazon SageMaker RL.

 

 

 

Revisor

Marcelo Cunha é um arquiteto de soluções da AWS especialista em AI/ML. Em sua função, ele ajuda os clientes a criar soluções de ML para resolver os seus desafios de negócio utilizando a plataforma AWS. De uma família de médicos e profissionais da saúde, ele também tem interesse em aplicações de ML nesse setor.