Beyond Needle(s) in the Embodied Haystack: Environment, Architecture, and Training Considerations for Long Context Reasoning

1. Il Problema

L'intelligenza artificiale incarnata (Embodied AI) richiede la capacità di prendere decisioni sequenziali su orizzonti temporali lunghi (long-horizon), dove il successo di un compito dipende dalla memorizzazione e dal ragionamento su eventi distanti nel tempo.
Le attuali Vision-Language Models (VLM) affrontano sfide critiche in questo ambito:

• Finestre di contesto limitate: La maggior parte dei modelli non può gestire input superiori a 1 milione di token, insufficienti per traiettorie di migliaia di passi.
• Mancanza di benchmark specifici: Esistono ambienti simulati, ma mancano dataset standardizzati e suite di benchmark per valutare la capacità di ragionamento su contesti lunghi e multimodali (visivo + linguistico).
• Difficoltà nel recupero delle informazioni: I modelli faticano a recuperare e collegare "aghi" (informazioni chiave) dispersi in un "fieno" (la lunga traiettoria) di dati visivi e azioni, specialmente quando le dipendenze temporali sono distanti centinaia di passi.

2. Metodologia: Il Framework ∞-THOR

Gli autori introducono ∞-THOR, un nuovo framework completo per la generazione, l'addestramento e la valutazione di compiti incarnati a lungo termine.

A. Ambiente e Generazione Dati

• Base: Costruito sopra il simulatore AI2-THOR.
• Generazione di Traiettorie: Utilizza un pianificatore basato su PDDL per concatenare sequenze di sott-compiti singoli in traiettorie continue e lunghe (centinaia di passi).
• Dipendenza Temporale: Per forzare il ragionamento a lungo termine, il compito finale (goal sintetico) richiede di collegare oggetti o stati osservati solo nella prima e nell'ultima parte della traiettoria (es. osservare un pomodoro a $t=16$ e un piano di lavoro a $t=560$, per poi collocare il pomodoro sul piano a $t=670$).
• Dataset: Viene rilasciato un dataset su larga scala con episodi di training (>400 passi) e test/dev (>600 passi), con lunghezze di contesto che raggiungono 1,3 milioni di token.

B. Nuovo Benchmark: Needle(s) in the Embodied Haystack (NiEH)

Viene proposta una nuova task di valutazione statica (QA) che va oltre il classico "Needle in a Haystack" testuale:

• NiEH (Single-evidence): Richiede di recuperare un singolo indizio visivo/linguistico disperso.
• NsiEH (Multi-evidence): Richiede di aggregare e ragionare su più indizi sparsi temporalmente (es. "Quante volte hai spostato l'oggetto X?" o "Dove era l'oggetto Y prima di essere messo nel contenitore Z?").
• Validazione: Le domande sono generate tramite template e validate da un ensemble di quattro VLM multimodali per garantire la rispondibilità.

C. Architetture e Strategie di Addestramento

Per gestire contesti estremi, gli autori esplorano diverse configurazioni:

1. Architetture di Policy:
   • Interleaved Goal-State-Action: Codifica l'intera storia di percezione-azione come un'unica sequenza multimodale intercalata ($g_0, s_0, a_0, ..., s_t, a_t$).
   • Memory-Augmented Goal-State: Sostituisce la storia completa con riassunti testuali o immagini recuperate (retrieval) basate sulla similarità degli embedding, riducendo il carico di contesto.
2. Estensione del Contesto:
   • Implementazione di tecniche come Linear Interpolation, Dynamic Scaling, YaRN e LongRoPE per adattare le posizioni dei token oltre i limiti pre-addestrati.
3. Parallelismo del Contesto (Context Parallelism):
   • Utilizzo di Ring Attention per distribuire il calcolo dell'attenzione su più dispositivi, permettendo l'addestramento e l'inferenza su sequenze ultra-lunghe con overhead di memoria ridotto.

3. Risultati Chiave

Gli esperimenti sono stati condotti su modelli come LLaVA-OneVision, Qwen2.5-VL, DeepSeek-VL e Gemini 2.0 Flash.

• Prestazioni su NiEH (Statico):
  • I modelli standard falliscono su contesti completi (>1M token).
  • YaRN si è dimostrato la tecnica di estensione più efficace, permettendo di gestire contesti fino a 384K token con buone prestazioni.
  • L'approccio Interleaved ha superato quello basato su Memory-Augmented (recupero), suggerendo che la storia continua è cruciale per il ragionamento incarnato.
  • C'è un calo significativo delle prestazioni nel passaggio da single-evidence a multi-evidence, evidenziando la difficoltà nel ragionamento temporale complesso.
• Valutazione Online (Task a lungo termine):
  • L'addestramento su contesti più lunghi (es. 128K vs 32K) porta a guadagni significativi nel tasso di successo (da 6.4% a 9.2%).
  • L'uso combinato di addestramento esteso e inferenza con YaRN (x4) ha raggiunto il 11.5% di successo totale.
  • I compiti di "Put" (posizionamento) rimangono difficili a causa dei fallimenti cumulativi nei passi precedenti.
• Transfer Sim-to-Real:
  • L'addestramento sul dataset sintetico NiEH ha migliorato le prestazioni su benchmark fotorealistici (VSI-Bench) del +11.2% per il task di ordinamento temporale degli oggetti, dimostrando che le capacità di ragionamento apprese in simulazione si trasferiscono al mondo reale.
• Controllo a Basso Livello:
  • L'integrazione con ManipulaTHOR per il controllo robotico ha mostrato prestazioni basse, limitate dalla discrepanza tra le viste della telecamera e le aspettative dei modelli VLA esistenti.

4. Contributi Principali

1. Framework ∞-THOR: Un sistema scalabile per generare traiettorie infinite e complesse per l'AI incarnata.
2. Benchmark NiEH: Una nuova task di QA multimodale che sfida gli agenti a ragionare su prove sparse e temporali in traiettorie lunghe.
3. Dataset e Ambiente: Rilascio di un dataset di traiettorie (fino a 1.3M token) e un ambiente interattivo per apprendimento offline e online.
4. Analisi Architetturale e di Addestramento: Confronto sistematico tra modelli Interleaved e Memory-Augmented, e validazione di tecniche come Context Parallelism e YaRN per contesti estremi.
5. Validazione Empirica: Dimostrazione del successo del transfer Sim-to-Real e dell'importanza cruciale dell'esposizione a contesti lunghi durante l'addestramento (non solo all'inferenza).

5. Significato e Impatto

Questo lavoro segna un passo fondamentale verso la prossima generazione di agenti incarnati capaci di ragionamento e pianificazione a lungo termine.

• Dimostra che le attuali VLM, sebbene potenti, non sono ancora pronte per contesti di milioni di token senza adattamenti specifici (architettonici e di training).
• Evidenzia che l'addestramento su dati a lungo termine è più critico della semplice estensione della finestra di contesto all'inferenza.
• Fornisce una base solida (dataset e benchmark) per la ricerca futura, spostando il focus dalla semplice comprensione visiva a una memoria spaziale e temporale robusta, essenziale per la robotica reale e l'interazione complessa con l'ambiente.

Il codice e i dataset sono disponibili pubblicamente, facilitando la riproducibilità e l'ulteriore sviluppo nel campo dell'AI incarnata.