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

TL;DR

Real-world embodied reasoning requires reasoning over long temporal horizons, and current models struggle with extreme long-context dependencies. The authors introduce $\infty$-Thor to generate, train, and evaluate long-horizon embodied tasks and present the NiEH benchmark for multimodal, memory-intensive QA. Key contributions include a scalable trajectory-generation framework, a static NiEH QA task, a large paired trajectory dataset, and architectural strategies such as interleaved Goal–State–Action modeling, long-context extensions (YaRN, LongRoPE), and Context Parallelism. Experimental results show substantial gains from longer-context exposure during training and provide insights into model behavior and training strategies under long-horizon conditions. Overall, the work lays a reproducible foundation for robust long-horizon reasoning and planning in embodied AI, with open simulators and datasets to drive future progress.

Abstract

We introduce $\infty$-THOR, a new framework for long-horizon embodied tasks that advances long-context understanding in embodied AI. $\infty$-THOR provides: (1) a generation framework for synthesizing scalable, reproducible, and unlimited long-horizon trajectories; (2) a novel embodied QA task, Needle(s) in the Embodied Haystack, where multiple scattered clues across extended trajectories test agents' long-context reasoning ability; and (3) a long-horizon dataset and benchmark suite featuring complex tasks that span hundreds of environment steps, each paired with ground-truth action sequences. To enable this capability, we explore architectural adaptations, including interleaved Goal-State-Action modeling, context extension techniques, and Context Parallelism, to equip LLM-based agents for extreme long-context reasoning and interaction. Experimental results and analyses highlight the challenges posed by our benchmark and provide insights into training strategies and model behaviors under long-horizon conditions. Our work provides a foundation for the next generation of embodied AI systems capable of robust, long-term reasoning and planning.

Figures (7)

• Figure 1: Example of the trajectory and a long-horizon embodied task generated from $\infty$-Thor. The final goal ("Put the tomato on the counter top" at t=670) requires recalling both the tomato (seen at t=17) and the counter (seen at t=560) to solved the long-horizon task. Context size refers to the input token length when converting the trajectory into the LLM input space.
• Figure 2: Example of N(s)iEH task and Ground-truth steps.
• Figure 3: Agent–environment interaction through interleaved Goal-State-Action modeling.
• Figure 4: Results of Needle(s) in the Embodied Haystack. The white dashed line denotes the maximum input context length of the model. Qwen2.5-VL was pre-trained initially with an 8K token context window and incrementally scaled up to 32K tokens in subsequent stages bai2025qwen25vltechnicalreport. The gray area indicates contexts not applicable (N/A) due to the model's smaller image token size limiting sequences to under 128K tokens. Context Parallelism is applied to all experiments with the context size over 384K.
• Figure 5: Agent's reward across different experimental configurations for high-level planning tasks. We compare (a) context extension methods at fixed scaling (x4), (b) varying YaRN scaling factors, and (c) fine-tuning with different context lengths using Context Parallelism. (d) summarizes the most effective strategies, highlighting that exposure to longer contexts during training significantly improves performance. Non-planner models cannot generate valid actions after around 250 steps ($\approx$376K in context size). More configurations are in Figure \ref{['fig:agent_reward']} in Appendix.