Resumo Técnico: Uma Framework de Regularização KL para Aprendizado de Planejamento com Priors Adaptativos

1. Declaração do Problema

No Aprendizado por Reforço Baseado em Modelo (MBRL), a exploração eficaz permanece um desafio crítico, particularmente em tarefas de controle contínuo de alta dimensão onde a eficiência de amostragem é primordial. Abordagens recentes, como TD-MPC e BMPC, utilizam políticas aprendidas como distribuições de proposta para o planejamento de Integral de Caminho Modelado por Previsão (MPPI). No entanto, uma limitação fundamental surge do desajuste de política: a política de amostragem ($\pi_{\theta_s}$) é frequentemente atualizada independentemente da distribuição de ação do planejador ($\pi_P$).

Este desacoplamento cria um desajuste de distribuição onde a função de valor de ação bootstrap é treinada em trajetórias geradas pelo planejador MPPI (que explora uma distribuição específica), enquanto a política de amostragem é atualizada para maximizar valores de ação locais, muitas vezes divergindo da distribuição do planejador. Este desalinhamento leva a estimativas de valor pouco confiáveis, particularmente para horizontes curtos, e degrada o desempenho de longo prazo. Tentativas anteriores de alinhar essas distribuições, como a minimização da divergência KL reversa (BMPC), sofrem de alta variância porque dependem de amostras do planejador "desatualizadas" armazenadas no buffer de replay, que representam uma mistura de políticas de planejamento desatualizadas em vez de uma distribuição coerente atual.

2. Metodologia: Framework PO-MPC

Os autores propõem Otimização de Política–Controle Preditivo Modelado (PO-MPC), uma framework unificada que trata a distribuição de ação do planejador como um prior na otimização de política. O PO-MPC enquadra o passo de aprendizado da política de amostragem como uma instância de Aprendizado por Reforço Regularizado por KL.

Formulação Central

O objetivo é aprender uma política de amostragem $\pi_{\theta_s}$ que maximize o retorno esperado enquanto permanece próxima de um prior adaptativo $\pi_p$ (derivado do planejador MPPI). A função objetivo é:

$$J(\pi_{\theta_s}) = \mathbb{E} \left[ \sum_{t=0}^{T-1} \gamma^t \left( r(z_t, a_t) - \lambda \text{KL}[\pi_{\theta_s}(\cdot | z_t) \| \pi_p(\cdot | z_t)] \right) \right]$$

Onde:

• $z_t$ é a representação do estado latente.
• $\lambda$ é um hiperparâmetro que controla a força da regularização KL.
• $\pi_p$ é o prior da política.

Componentes Chave

1. RL Regularizado por KL: A framework unifica métodos existentes como casos especiais desta formulação:

   • $\lambda = 0$: Recupera o TD-MPC2, onde a política é atualizada exclusivamente via maximização de valor de ação e regularização de entropia.
   • $\lambda \to \infty$: Recupera o BMPC (com KL reversa), onde a política é constrangida a imitar a distribuição do planejador.
   • $\lambda$ Intermediário: O artigo explora valores intermediários (por exemplo, $\lambda \in [0.1, 9]$) para equilibrar a maximização do retorno e o alinhamento de distribuição.

2. Prior Intermediário Adaptativo ($\pi_{\theta_p}$):
   Para abordar a alta variância causada pelo uso direto de amostras desatualizadas do planejador do buffer de replay, o PO-MPC introduz um prior intermediário aprendido. Em vez de clonar a distribuição bruta de amostras do planejador armazenadas (que pode ser uma mistura gaussiana variante no tempo), a framework treina uma rede neural $\pi_{\theta_p}$ para aproximar a distribuição atual induzida pelo MPPI $\pi_P$.

   • Este prior é destilado da saída do planejador, protegendo as atualizações da política de amostragem da variância de amostras desatualizadas.
   • O prior pode ser treinado usando KL Reversa (minimizando $\text{KL}[\pi_{\theta_p} \| \pi_P]$) ou KL Direta (minimizando $\text{KL}[\pi_P \| \pi_{\theta_p}]$).
     • KL Reversa: Viésa a política para corresponder a modos específicos do planejador (convergência mais rápida, potencialmente menos exploração).
     • KL Direta: Viésa a política para cobrir o suporte de todas as distribuições amostradas (exploração aprimorada, potencialmente convergência mais lenta).

3. Loop do Algoritmo:
   O treinamento segue um ciclo Planejar $\to$ Inferir $\to$ Regularizar:

   • Planejar: O MPPI gera uma distribuição de trajetória $\pi_P$ usando a política de amostragem atual e a função de valor.
   • Inferir: O prior adaptativo $\pi_{\theta_p}$ é atualizado para aproximar $\pi_P$.
   • Regularizar: A política de amostragem $\pi_{\theta_s}$ é atualizada usando o objetivo regularizado por KL, equilibrando a maximização do valor-Q e a penalidade KL em direção a $\pi_{\theta_p}$.

3. Contribuições Principais

• Unificação: O artigo unifica métodos de RL baseados em MPPI diversos (TD-MPC2, BMPC) sob uma única framework teórica (PO-MPC), esclarecendo suas escolhas de design e relações.
• Priors Adaptativos: Introduz um prior intermediário aprendido para substituir a clonagem direta de amostras do buffer de replay, demonstrando teórica e empiricamente que isso reduz a variância do gradiente e estabiliza o treinamento.
• Objetivos Flexíveis: A framework permite a exploração de variantes previamente não estudadas, especificamente ajustando a força do KL $\lambda$ e escolhendo entre perdas KL Direta/Reversa para o prior.
• Desempenho de Estado da Arte: O método alcança eficiência de amostragem superior e desempenho final em tarefas de controle contínuo de alta dimensão em comparação com baselines existentes.

4. Resultados Experimentais

Os autores avaliaram o PO-MPC em 21 tarefas desafiadoras de controle contínuo do DeepMind Control Suite (Humanoid e Dog) e HumanoidBench.

• Ganhos de Desempenho: O PO-MPC com valores intermediários de $\lambda$ superou consistentemente tanto o TD-MPC2 ($\lambda=0$) quanto o BMPC ($\lambda \to \infty$).
  • Em tarefas de alta dimensão (por exemplo, HumanoidBench), valores intermediários de $\lambda$ (particularmente $\lambda=1$ ou $\lambda=9$) produziram melhorias significativas no retorno final e na eficiência de amostragem.
  • Em alguns ambientes onde as baselines falharam em aprender (por exemplo, HumanoidBench "Balance Simple", "Slide", "Stair"), o PO-MPC convergiu com sucesso.
• Eficácia do Prior Intermediário: O uso do prior intermediário aprendido ($\pi_{\theta_p}$) igualou ou superou o desempenho do uso de amostras brutas do planejador do buffer de replay, confirmando a redução na variância de atualização.
• Sensibilidade da Função de Perda: A escolha da perda KL para treinar o prior importou. A KL Direta foi benéfica em tarefas pesadas de exploração (por exemplo, "Stair"), enquanto a KL Reversa acelerou a convergência em tarefas críticas de precisão (por exemplo, "Balance Simple").
• Ablação: Os resultados mostraram que o desempenho degrada se $\lambda$ for muito baixo (retorna desajuste de política) ou muito alto (colapso prematuro de modos), destacando a importância de ajustar este trade-off.

5. Significado e Alegações

O artigo alega que o PO-MPC representa uma unificação principiante de abordagens baseadas em MPPI. Ao alinhar explicitamente a política de amostragem com o planejador via um objetivo regularizado por KL e um prior adaptativo, a framework:

1. Fecha o ciclo entre planejamento e aprendizado de política, garantindo que a função de valor seja avaliada em distribuições relevantes para o planejador atual.
2. Melhora a estabilidade ao mitigar a variância introduzida por amostras de planejamento desatualizadas no buffer de replay.
3. Avança o estado da arte em RL baseado em modelo para controle contínuo, oferecendo uma framework flexível que pode ser ajustada para requisitos específicos de tarefa (exploração vs. exploração).

Os autores observam que, embora a framework seja poderosa, ela assume distribuições gaussianas para políticas e não explora totalmente todas as trajetórias simuladas para aprendizado de valor. Sugere-se trabalho futuro para estender distribuições de políticas além de Gaussianas e ajustar automaticamente o hiperparâmetro $\lambda$.