Resumo Técnico: Busca em Árvore de Monte Carlo Epistêmica

Declaração do Problema

A família de algoritmos AlphaZero/MuZero (A/MZ) alcançou sucesso significativo ao integrar a Busca em Árvore de Monte Carlo (MCTS) com modelos aprendidos de valor e dinâmicas do ambiente. No entanto, existe uma limitação crítica: enquanto os modelos aprendidos introduzem incerteza epistêmica (incerteza decorrente da cobertura limitada dos dados de treinamento), a MCTS padrão não leva em conta a propagação dessa incerteza durante o processo de busca. Consequentemente, A/MZ não consegue aproveitar eficazmente a MCTS para exploração profunda em ambientes com recompensas esparsas. A exploração profunda exige que um agente se direcione para transições novas, independentemente de sua distância do estado atual, uma capacidade essencial para tarefas como design de algoritmos ou programação, onde as recompensas são esparsas e o espaço de estados é vasto. Sem levar em conta a incerteza epistêmica, a busca pode convergir para políticas subótimas baseadas em previsões de modelos imprecisas, falhando em explorar regiões necessárias do espaço de estados.

Metodologia: MCTS Epistêmica (EMCTS)

Os autores propõem a MCTS Epistêmica (EMCTS), um framework motivado teoricamente que integra a incerteza epistêmica no processo de MCTS para facilitar a exploração profunda. A metodologia envolve três componentes principais:

1. Formulação da Busca com Incerteza

Os autores modelam o modelo de ambiente aprendido $\hat{M}$ como uma variável aleatória. Eles derivam um Limite Superior de Confiança (UCB) para a função de valor ótima $Q^*$ com base na variância das previsões de valor dentro do modelo aprendido.

• Base Teórica: O Teorema 1 estabelece que, para um modelo aprendido $\hat{M}$, o valor ótimo verdadeiro $Q^*(s, a)$ é limitado pelo valor esperado máximo no modelo mais um termo proporcional ao desvio padrão desse valor, escalado por um parâmetro de confiança $\delta$.
• Política de Busca: A política de seleção padrão PUCT (Upper Confidence Bound do Preditor) é modificada para P/UCT Epistêmico (EP/UCT). O critério de seleção torna-se:
  $$a = \arg\max_{a} \left( q_{\hat{M}}(s, a) + \beta \sqrt{V[q_{\hat{M}}(s, a)]} + \text{Termo de Exploração} \right)$$
  Aqui, $q_{\hat{M}}$ representa o valor estimado, e $\sqrt{V[q_{\hat{M}}]}$ representa a incerteza epistêmica. O hiperparâmetro $\beta$ controla o trade-off entre exploração e exploração.

2. Propagação da Incerteza Epistêmica

Uma contribuição central é o mecanismo para propagar a incerteza através da árvore de busca, e não apenas o valor.

• Incerteza em Backups: A incerteza de uma etapa de backup $\nu$ é calculada somando as variâncias da recompensa imediata e da incerteza do valor futuro descontado.
• Incerteza do Valor do Nó: Como A/MZ usa o mesmo modelo durante todo o planejamento, os retornos de backup são correlacionados. Para evitar assumir independência, os autores propõem um limite superior para a variância do valor do nó $V[q_{\hat{M}}(s, a)]$ usando a soma dos desvios padrão dos retornos de backup individuais:
  $$V[q_{\hat{M}}(s, a)] \leq \left( \frac{1}{N(s, a)} \sum_{i=1}^{N(s, a)} \sqrt{V[\nu_i(s, a)]} \right)^2$$
• Estimadores: O método utiliza estimadores de incerteza existentes para recompensas (por exemplo, Distilação de Rede Aleatória (RND) ou contagem baseada em hash) e valores (por exemplo, Equação de Bellman de Incerteza (UBE)). Para transições não observadas, a variância é definida como a variância máxima possível para uma variável aleatória limitada.

3. Lidando com Modelos de Transição Aprendidos

Embora a derivação teórica assuma um modelo de transição conhecido, os autores abordam os desafios das dinâmicas de transição aprendidas (como no MuZero). Eles propõem uma aproximação "maximamente otimista" onde, ao encontrar a primeira transição incerta em uma trajetória, todas as previsões subsequentes nessa trajetória são assumidas como tendo incerteza máxima. Isso garante que o UCB permaneça um limite superior válido para fins de exploração.

Principais Contribuições

1. MCTS Epistêmica (EMCTS): Um algoritmo inovador que estende a MCTS para estimar e propagar incerteza epistêmica a partir de modelos aprendidos de valor e/ou recompensa, permitindo que o processo de busca busque ativamente regiões incertas.
2. Framework Teórico: Uma derivação de políticas de busca baseadas em UCB (EP/UCT) que são fundamentadas teoricamente na variância de modelos aprendidos, fornecendo um mecanismo formal para exploração profunda.
3. Implementação: Uma implementação paralelizada em JAX da EMCTS acoplada a um agente AlphaZero, aplicada ao ambiente de linguagem Assembly subleq e ao benchmark Deep Sea.

Resultados Experimentais

Os autores avaliam a EMCTS em dois domínios desafiadores com recompensas esparsas:

1. Tarefa de Programação Subleq

• Tarefa: Escrever código na linguagem assembly subleq para resolver funções específicas (Negar Positivos e Função Identidade). Isso envolve buscar um espaço de estados de aproximadamente $16^{10}$ estados.
• Resultados: A EMCTS acoplada ao AlphaZero (E-AZ) superou significativamente o AlphaZero de base. O E-AZ resolveu a tarefa mais difícil "Função Identidade" com muito menos amostras do que a base. O método demonstrou que o uso de um estimador de incerteza apropriado (por exemplo, hash IO vs. hash de estado completo) melhorou ainda mais a eficiência de amostragem.

2. Benchmark Deep Sea

• Tarefa: Um ambiente de mundo em grade onde o agente deve encontrar uma trajetória ótima única com recompensas esparsas. A probabilidade de encontrar a solução via exploração aleatória decai exponencialmente com o tamanho da grade.
• Resultados:
  • Exploração Profunda: Agentes A/MZ de base falharam em resolver variações do Deep Sea (tanto recompensas determinísticas quanto estocásticas) dentro de orçamentos de treinamento razoáveis. Em contraste, agentes EMCTS (E-AZ e E-MZ) resolveram essas tarefas, demonstrando escalabilidade subexponencial da complexidade de amostragem com o tamanho do ambiente.
  • Benefício da Busca: A EMCTS superou significativamente uma ablação (A/MZ+UBE) que usava incerteza para seleção de ações, mas não usava busca para estimar essa incerteza. Isso confirma que a própria busca melhora a qualidade da estimativa de incerteza, levando a uma exploração mais eficiente.
  • Robustez: O método permaneceu eficaz mesmo ao usar dinâmicas de transição aprendidas do MuZero (abstração equivalente a valor) e na presença de recompensas estocásticas.

Significado e Alegações

O artigo afirma que a EMCTS aborda uma lacuna fundamental no aprendizado por reforço baseado em modelos: a incapacidade da MCTS padrão de utilizar a incerteza epistêmica para exploração. Ao integrar a propagação de incerteza na árvore de busca, o método permite que agentes A/MZ:

• Alcançem eficiência de amostragem significativamente maior em ambientes com recompensas esparsas.
• Resolvam benchmarks de exploração difícil (como Deep Sea) que são praticamente insolúveis por A/MZ de base.
• Potencialmente melhorem a confiabilidade em RL offline e geração de alvos off-policy, fornecendo melhores estimativas de incerteza para previsões de valor.

Os autores posicionam a EMCTS como um aprimoramento prático e teoricamente motivado para a família A/MZ, tornando esses algoritmos melhor equipados para aplicações do mundo real envolvendo design de algoritmos e recompensas esparsas, onde a exploração profunda é crítica.