Resumo Técnico: Planejamento Analítico sob Incerteza com Fechamento de Momentos

Declaração do Problema

O aprendizado por reforço baseado em modelos (MBRL) eficaz em ambientes estocásticos requer um planejamento que leve em conta a incerteza preditiva. As abordagens tradicionais enfrentam um compromisso (trade-off):

1. Amostragem Estocástica: O MBRL profundo moderno frequentemente depende da amostragem de trajetórias imaginadas (Monte Carlo) para propagar a incerteza. Embora escalável, isso introduz uma variância significativa no alvo, tornando a estimativa de valor sensível ao orçamento de amostragem e potencialmente instável.
2. Aproximações Determinísticas: Métodos que colapsam modelos estocásticos para sua média ou aprendem dinâmicas determinísticas descartam a covariância preditiva. Isso leva a um viés sistemático quando as funções de valor são não lineares.
3. Propagação Analítica: Embora a propagação analítica de distribuições completas de estado seja um princípio sólido, historicamente ela exigiu estruturas de política ou recompensa restritivas (ex: configurações lineares-quadráticas) para permanecer tratável.

A questão central abordada é se o planejamento consciente da distribuição é possível sem essas restrições, especificamente através do cálculo analítico da expectativa de Bellman sob dinâmicas estocásticas aprendidas.

Metodologia: Planejamento Analítico Compatível com Momentos (MoCA)

Os autores propõem um framework que desacopla a tratabilidade da seleção de ação gulosa da tratabilidade da expectativa preditiva. O método baseia-se em dois componentes estruturais fundamentais:

1. Parametrização Quadrática de Valor-Ação (Desacoplamento da Maximização)

Para evitar a intratabilidade de integrar um operador de maximização ($\max_{a'} Q(s', a')$) sobre uma distribuição estocástica, os autores utilizam uma parametrização quadrática de valor-ação (inspirada em NAF e BNAF).

• Estrutura: A função Q é decomposta como:
  $$Q(s, a) = V(s) - \frac{1}{2}(a - \tilde{\mu}(s))^\top P(s) (a - \tilde{\mu}(s))$$
  onde $V(s)$ é o valor do estado, $\tilde{\mu}(s)$ é a ação gulosa não restrita, e $P(s)$ é uma matriz de curvatura definida positiva.
• Efeito: Esta estrutura permite que a seleção da ação gulosa seja resolvida de forma fechada ($\arg\max_a Q(s,a) = \tilde{\mu}(s)$) e que o valor máximo seja exatamente $V(s)$. Consequentemente, o backup de Bellman reduz-se de uma expectativa sobre uma função Q maximizada para uma simples expectativa sobre a função de valor de estado sozinha:
  $$y(s, a) = r(s, a) + \gamma \mathbb{E}_{s' \sim p(\cdot|s,a)}[V(s')]$$
  Isso remove a maximização interna da expectativa, restando apenas o requisito de calcular $\mathbb{E}[V(s')]$ analiticamente.

2. Princípio da Compatibilidade de Momentos

A tratabilidade da expectativa restante $\mathbb{E}[V(s')]$ depende da compatibilidade entre a distribuição de transição preditiva e a classe da função de valor.

• Modelo de Dinâmica: A dinâmica do ambiente é aprendida como uma mistura gaussiana heterocedástica (um conjunto de $K$ modelos). Para evitar a amostragem durante o planejamento, o conjunto é aproximado por uma única distribuição Gaussiana cuja média e covariância correspondem aos dois primeiros momentos da mistura (capturando tanto a incerteza aleatória quanto a epistêmica).
• Representação de Valor: A função de valor de estado $V(s)$ é parametrizada como uma mistura de funções de base radial (RBF) gaussianas.
• Fechamento Analítico: Como a distribuição preditiva é Gaussiana e a função de valor é uma mistura de Gaussianas, a expectativa $\mathbb{E}[V(s')]$ admite uma solução de forma fechada em termos dos momentos da distribuição (média e covariância). Isso permite que o alvo de Bellman seja computado exatamente sem amostragem de Monte Carlo.

Procedimento de Aprendizado

O algoritmo (MoCA) integra esses componentes em um ciclo de aprendizado off-policy:

1. Aprendizado de Dinâmica: Um conjunto de modelos de dinâmica Gaussianos é treinado via máxima verossimilhança em amostras de bootstrap do replay buffer.
2. Construção do Alvo: Para cada transição, as previsões do conjunto são agregadas via correspondência de momentos para formar uma Gaussiana preditiva $\mathcal{N}(\bar{\mu}, \bar{\Sigma})$. O alvo de Bellman é computado analiticamente usando a expectativa de forma fechada da função de valor RBF sob essa Gaussiana.
3. Atualização do Crítico: Críticos duplos (para mitigar o viés de superestimação) são atualizados usando estimativa dupla truncada (clipped double estimation) nos alvos analíticos.

Principais Contribuições

1. Backup de Bellman Analítico: O artigo demonstra que expectativas de Bellman de forma fechada são possíveis sob dinâmicas estocásticas aprendidas ao desacoplar a maximização (via parametrização Q quadrática) da expectativa.
2. Framework de Compatibilidade de Momentos: Introduz um princípio de design onde a representação de valor é escolhida especificamente para ser compatível com os momentos da distribuição preditiva, permitindo a propagação analítica da incerteza sem restringir a classe da política.
3. Planejamento Consciente da Incerteza: O método propaga tanto a média preditiva quanto a covariância deterministicamente, removendo a variância de Monte Carlo da construção do alvo, enquanto preserva os benefícios da modelagem estocástica.

Resultados Empíricos

Experimentos foram conduzidos em tarefas de controle contínuo (Cartpole Balance e Pendulum Swingup) sob ruído de observação Gaussiano controlado (tanto homocedástico quanto heterocedástico).

• Robustez ao Ruído: O MoCA superou consistentemente os baselines (SAC, BNAF e uma versão de equivalência de certeza do MoCA que ignora a covariância) conforme o ruído de observação aumentava. Ele manteve retornos assintóticos mais altos e dinâmicas de aprendizado mais estáveis.
• Analítico vs. Monte Carlo: Quando comparado a variantes de Monte Carlo da mesma arquitetura (MoCA-MC) com diferentes quantidades de amostras ($K=4, 16$), o método analítico (MoCA) alcançou convergência mais rápida e desempenho assintótico superior. Os métodos de Monte Carlo sofreram com alta variância, especialmente com amostras limitadas, levando a uma seleção de ações subótima.
• Calibração: A abordagem analítica produziu incerteza preditiva bem calibrada (a cobertura de $1\sigma$ empírica permaneceu próxima ao nível nominal de 68%) e desvio padrão epistêmico estável.

Significância e Alegações

O artigo afirma fornecer um framework principiado para o planejamento com modelos de distribuição aprendidos que evita a variância da amostragem e o viés da substituição pela média.

• Estabilidade: Ao remover o ruído de Monte Carlo da construção do alvo, o método fornece planejamento estável mesmo quando o ruído de observação aumenta.
• Generalidade: O princípio de compatibilidade de momentos não está preso a uma única forma funcional; o pareamento Gaussiano-RBF é apresentado como uma instanciação concreta que equilibra tratabilidade e expressividade.
• Praticidade: A abordagem oferece uma alternativa computacionalmente eficiente ao planejamento baseado em amostragem em controle contínuo, demonstrando que a propagação analítica de momentos pode servir como um substituto prático para trajetórias estocásticas (rollouts).

Os autores reconhecem limitações, observando que a instanciação RBF enfrenta a "maldição da dimensionalidade" em relação ao número de bases necessárias em espaços de estado de alta dimensão, e que a implementação atual depende de suposições de covariância diagonal para eficiência computacional. Sugere-se como trabalho futuro a exploração de representações de espaço latente para escalar o método.