إسكا: ترسيخ الوكلاء المتجسدين باستخدام الرسوم البيانية للمشهد وJAX لتحقيق أداء في الوقت الفعلي

Embodied AI, where AI agents interact with the physical world, faces significant challenges in accurately perceiving and interpreting visual information. Multi-Modal Language Models (MLLMs), despite their advancements, often struggle with the fine-grained links between visual features and textual semantics. Researchers at the University of Pennsylvania, in collaboration with Google, have developed ESCA (Embodied and Scene-Graph Contextualized Agent) to address this perception bottleneck. By leveraging scene graphs and the JAX framework, they’ve achieved significant improvements in accuracy and performance, demonstrating the power of structured data and efficient infrastructure for real-world AI applications.

The Perception Problem in Embodied AI

Existing MLLMs often fail to accurately capture the relationship between visual inputs and textual instructions. This leads to errors in navigation and interaction tasks. A key finding from their analysis of navigation tasks revealed that a significant proportion of agent failures are rooted in perception errors.

• High Error Rate: 69% of agent failures in EmbodiedBench navigation tasks stem from perception errors.
• Types of Errors: These errors include hallucination, wrong recognition, and spatial understanding issues.
• Fine-grained Understanding Required: Models struggle to identify precise spatial relationships between objects needed to complete tasks.

ESCA: Illuminating the Agent’s Environment

ESCA addresses the perception problem by using scene graphs to provide structured contextual information to MLLMs. This approach is likened to an anglerfish using its bioluminescent lure to reveal its surroundings.

• Scene Graphs: ESCA generates a structured map of objects, attributes, and relationships within the agent’s environment.
• Selective Grounding: Instead of overwhelming the model with all available information, ESCA selectively identifies and uses the most relevant objects and relations to the current task.
• Probabilistic Reasoning: ESCA employs probabilistic reasoning to create prompts enriched with the contextual details the agent needs to act effectively.

LASER and Scallop: The Engine Behind ESCA

The core of ESCA is LASER, a CLIP-based foundation model trained on a large dataset of video-caption pairs. LASER works with Scallop, a neurosymbolic programming language with JAX backends, to connect scene graphs to logical rules.

• LASER: A CLIP-based foundation model trained on video-caption pairs.
• Scallop: A neurosymbolic programming language that supports JAX backends.
• Automated Graph Generation: This pipeline enables the training of perception models that produce detailed graphs without extensive manual annotation.

JAX: Performance and Modularity

The transition to JAX proved crucial in achieving real-time performance. JAX’s functional design and Just-In-Time (JIT) compilation capabilities were instrumental in optimizing the pipeline.

• Statelessness: JAX’s functional architecture enabled modularity and facilitated effective JIT compilation, reducing GPU overhead.
• Complex Control Flow: JAX provides primitives like jax.lax.cond for managing control flow within probabilistic graphs efficiently.
• Debugging and Transparency: JAX’s debugging tools, like jax.debug.print and jax.disable_jit, allow for easy inspection and debugging of compiled functions. The open-source nature of JAX also provides deep insights into its workings.
• Seamless Integration with Flax: The NNX library integrates seamlessly, allowing easy structuring of models and management of parameters.

JAX Performance Gains

The switch from PyTorch to JAX resulted in a substantial performance boost, critical for real-time embodied agent applications.

• 25% Speedup: JAX reduced the average time per frame on an NVIDIA H100 GPU from 18.15 ms (PyTorch) to 14.55 ms.
• Increased FPS: Correspondingly, Frames Per Second (FPS) increased from 55.15 to 68.82.
• Real-Time Viability: This speedup makes the ESCA framework more viable for real-time applications.

Conclusion

ESCA demonstrates the importance of combining structured data (scene graphs) with efficient infrastructure (JAX) to overcome the perception bottleneck in Embodied AI. The research highlights how JAX’s speed, transparency, and modularity facilitate the development of real-time, reliable reasoning agents. The success of ESCA paves the way for more advanced embodied AI systems capable of navigating and interacting with the physical world more effectively.

Source: Google Open Source Blog