Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

1. Problema e Motivação

Os agentes baseados em Grandes Modelos de Linguagem (LLMs) têm demonstrado capacidades notáveis de raciocínio e interação, mas frequentemente sofrem de "fundamentação rasa" (shallow grounding). Isso significa que, embora percebam o ambiente, eles carecem de uma compreensão causal profunda de como suas ações atuais remodelarão o estado futuro.

• Limitações Atuais: Métodos existentes geralmente realizam rolagens (rollouts) de um único passo ou com horizonte fixo. Estratégias rígidas falham em capturar dependências de longo prazo em tarefas complexas (como tarefas domésticas ou científicas) e podem incorrer em custos computacionais desnecessários para ações triviais.
• Consequência: Sem a capacidade de projetar o futuro, os agentes são propensos a falhas catastróficas, descobrindo erros ou conflitos de estado apenas após executá-los irreversivelmente no ambiente real.

2. Metodologia: Imagine-then-Plan (ITP)

O artigo propõe o ITP, um quadro unificado que permite aos agentes planejar tarefas através de uma "imaginação" de lookahead (olhar para frente) adaptativa, utilizando modelos de mundo aprendidos.

Conceitos Fundamentais

• POIMDP (Processo de Decisão de Markov Parcialmente Observável e Imaginável): O trabalho estende o POMDP tradicional para um POIMDP. A política do agente não é otimizada apenas com base no estado observável atual, mas também em trajetórias futuras "imaginadas".
• Modelo de Mundo: Um LLM é fine-tuned para atuar como um modelo de mundo ($M_\phi$), prevendo o próximo estado textual ($s_{t+1}$) dado o estado atual e a ação ($s_t, a_t$).

Mecanismo de Lookahead Adaptativo

A inovação central é a capacidade de ajustar dinamicamente o horizonte de imaginação ($K_t$) com base no progresso da tarefa e no objetivo final, em vez de usar um horizonte fixo.

• Troca (Trade-off): O sistema equilibra a necessidade de prever consequências de longo prazo (para evitar conflitos) contra o custo computacional e o acúmulo de erros em previsões muito longas.

Variantes do ITP

O framework é implementado em duas versões:

1. ITP-I (In-Imagination Learning - Treinamento Livre):

   • Funcionamento: Ocorre no tempo de inferência. Tanto a política quanto o modelo de mundo permanecem congelados.
   • Processo: Em cada passo, o agente seleciona adaptivamente um horizonte $K_t$, gera uma trajetória imaginada com o modelo de mundo e reflete sobre essa trajetória para identificar conflitos ou progresso antes de executar a ação real.
   • Objetivo: Melhorar agentes LLM existentes sem necessidade de treinamento adicional.

2. ITP-R (Reinforced Training - Treinamento Reforçado):

   • Funcionamento: Otimiza conjuntamente a política de ação e a seleção do horizonte de lookahead.
   • Fases de Treinamento:
     1. Rótulos Pseudo: Gera rótulos para o horizonte ideal de lookahead usando o modelo de mundo e a política inicial.
     2. Aquecimento (Warm-up): Fine-tuning supervisionado para aprender a agir sob trajetórias condicionadas e prever o horizonte necessário.
     3. Otimização Online: Usa Aprendizado por Reforço (A2C - Advantage Actor-Critic) para maximizar a recompensa da tarefa enquanto penaliza o custo computacional do lookahead e a redundância de raciocínio.

3. Contribuições Principais

1. Formalização POIMDP: Introduz um novo formalismo matemático que integra futuros imaginados e interações históricas no processo de decisão do agente.
2. Mecanismo Adaptativo: Propõe um mecanismo de lookahead que escala dinamicamente o horizonte de imaginação, superando as limitações de abordagens de horizonte fixo.
3. Validação Empírica: Demonstra que tanto a variante sem treinamento (ITP-I) quanto a treinada (ITP-R) superam significativamente as linhas de base competitivas em múltiplos benchmarks.

4. Resultados Experimentais

Os experimentos foram conduzidos em quatro benchmarks representativos: ALFWorld (tarefas domésticas), ScienceWorld (experimentos científicos), WebShop (navegação web) e StableToolBench (uso de ferramentas).

• Desempenho Superior:
  • ITP-I: Superou métodos de prompting fortes (como ReAct e RAP) sem treinamento adicional, melhorando as taxas de sucesso em cenários zero-shot.
  • ITP-R: Superou consistentemente métodos baseados em treinamento (como SFT, WKM e IWM). Por exemplo, no ALFWorld com o modelo Qwen3-8B, alcançou 88,57% de taxa de sucesso, contra 82,86% do melhor baseline treinado.
• Eficiência e Custo:
  • O lookahead adaptativo alcançou taxas de sucesso mais altas com um custo computacional (orçamento de tokens) significativamente menor do que o lookahead de horizonte fixo ou aleatório.
  • Estudos de ablação mostraram que remover o treinamento reforçado (RT) degrada drasticamente o desempenho, confirmando que o aprendizado de quando e quanto imaginar é crucial.
• Robustez do Modelo de Mundo: O ITP mostrou ser agnóstico ao modelo, funcionando bem com diferentes backbones (Qwen, Llama, DeepSeek), embora o treinamento específico do modelo de mundo seja benéfico para ambientes interativos específicos.

5. Significado e Impacto

O trabalho representa um avanço principial na direção de agentes autônomos mais deliberativos.

• Superação da Fundamentação Rasa: Ao permitir que o agente "rehearse" (ensaie) ações em um "sandbox mental" antes de agir, o ITP fecha o ciclo entre planejamento e estimativa de consequências.
• Escalabilidade: A capacidade de ajustar o esforço de raciocínio (número de passos de imaginação) conforme a complexidade da tarefa oferece um caminho eficiente para lidar com problemas de longo horizonte sem explodir os custos computacionais.
• Futuro: O framework sugere que a integração de modelos de mundo com mecanismos de controle adaptativo é essencial para resolver tarefas complexas no mundo real, embora o artigo reconheça desafios futuros em ambientes multimodais e robóticos.

Em resumo, o Imagine-then-Plan oferece uma estrutura robusta para transformar agentes reativos em agentes proativos, capazes de planejar com base em previsões futuras adaptativas, resultando em maior sucesso em tarefas complexas e maior eficiência computacional.