💻 computer science

Key-Gram: Extensible World Knowledge for Embodied Manipulation

Key-Gram is a conditional-memory framework that decouples extensible linguistic world knowledge from visual reasoning in embodied manipulation by storing task-specific priors in an external memory for efficient retrieval and injection, thereby significantly improving compositional grounding, transfer, and real-world performance across diverse benchmarks.

Original authors: Jingjing Fan, Siyuan Li, Botao Ren, Zhidong Deng

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jingjing Fan, Siyuan Li, Botao Ren, Zhidong Deng

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to cook. You give it a complex instruction: "Put the yellow and white mug in the microwave and close the door."

Current robot brains (AI models) try to do two very different jobs at the exact same time, which causes a traffic jam in their thinking process:

The "What" Job: Remembering facts about the world (e.g., "Microwaves have doors," "Mugs are fragile," "Yellow and white mugs exist").
The "How" Job: Watching the video feed and figuring out the physical movements (e.g., "The arm is too far left, move it right," "Grab the handle").

In most existing robots, these two jobs are mashed together into one giant brain. The "facts" and the "video" fight for attention, like a noisy crowd trying to hear a single speaker. If you want the robot to learn a new fact (like "don't put metal in the microwave"), you often have to retrain the whole brain, which is slow and risky because it might forget how to move its arms.

The Solution: Key-Gram

The paper introduces Key-Gram, a new way to organize the robot's brain. Think of it as giving the robot a smart, external notebook that it can flip through instantly.

Here is how it works, using a simple analogy:

1. The "Index Card" System (Decomposition)

Instead of reading the whole sentence "Put the yellow and white mug in the microwave," the robot breaks it down into small, specific index cards called Key-Grams.

Card A: "Put mug in microwave"
Card B: "Close microwave door"
Card C: "Yellow and white mug"

2. The "Instant Lookup" (Hashing)

The robot doesn't read the whole library to find the answer. It uses a magic hash code (like a barcode scanner) to instantly jump to the exact page in its notebook where the answer is stored.

It looks up Card A and finds the rule: "Microwaves have doors."
It looks up Card C and finds the rule: "This is a specific type of mug."

This lookup is instant ( $O(1)$ ), meaning it takes the same amount of time whether the notebook has 10 pages or 10 million pages.

3. The "Smart Assistant" (Fusion)

Once the robot pulls these facts from the notebook, it doesn't just dump them on the screen. It uses a smart filter (a gate) to decide when and where to use them.

When the robot is looking at the microwave door, the filter says, "Hey, look at the 'Close door' fact!"
When the robot is looking at the mug, the filter says, "Hey, look at the 'Yellow and white mug' fact!"

The main robot brain (the "Vision Backbone") is now free to focus entirely on the hard part: watching the video and moving the arms. It doesn't have to waste energy trying to remember facts; it just asks the notebook when it needs them.

Why This is a Big Deal

The paper claims this approach solves three major problems:

No More "Forgetting": Because the facts are in a separate notebook, you can add new pages (new knowledge) without rewriting the robot's brain. If you want to teach the robot about a toaster, you just add a new page to the notebook. The robot's ability to move its arms stays exactly the same.
Better at Complex Tasks: In tests, robots using Key-Gram were much better at long, multi-step tasks (like sorting bread, then putting it in a box, then moving the box). They didn't get confused or forget the middle steps because the "facts" were always right there in the notebook.
Handling New Combinations: When the researchers mixed and matched objects the robot hadn't seen together before (e.g., picking up a pen and a lighter, which it had never done as a pair), Key-Gram robots succeeded much more often. They could combine the "pick pen" fact with the "pick lighter" fact instantly.

The Results

The researchers tested this on simulated robots and a real dual-arm robot.

In Simulation: The robots improved their success rate by about 30% on difficult tasks.
In the Real World: On a real robot arm, the improvement was even more noticeable for complex assembly tasks, jumping from a 52% success rate to 66% (a 27% relative gain).

The Bottom Line

Key-Gram is like giving a robot a separate, expandable encyclopedia for facts, while letting its main brain focus purely on doing the physical work. This separation makes the robot smarter, more adaptable to new instructions, and less likely to forget what it already knows.

Technical Summary: Key-Gram

Problem Statement

Embodied manipulation requires models to simultaneously preserve reusable world knowledge (objects, skills, relations, constraints) and continuously adapt to dynamic visual states under action. Current Vision-Language-Action (VLA) models and World-Action Models (WAMs) typically couple linguistic knowledge with visual computation within a shared backbone or conditioning pathway. This design leads to two primary issues:

Modality Competition: In dense fusion approaches, a small number of instruction tokens compete with a large, continuously changing set of visual tokens within the same attention stream, potentially diluting instruction-specific information.
Fragile Knowledge Extension: In conditional generative models, language often acts merely as a prompt-like steering signal, while compositional world knowledge remains implicitly baked into the generative backbone. Consequently, adapting to new physical regimes or acquiring new knowledge often requires gradient updates that can overwrite previously acquired world knowledge, making continual adaptation and modular extension inherently fragile.

Methodology: Key-Gram

The authors propose Key-Gram, a conditional-memory framework that functionally decouples language-derived world knowledge from visual-state reasoning. The core philosophy is that visual computation should focus on scene dynamics (state evolution and action), while language serves as a compact, reusable index over abstract task priors.

1. Sparse Key-Gram Retrieval

Instead of processing raw tokens, the system decomposes a language instruction $I$ into a small set of task-specific key-grams (e.g., "put mug in microwave," "close door").

Extraction: A lightweight parser extracts fixed-length, knowledge-bearing units from the instruction.
Hashed Lookup: Each key-gram is converted into a fixed-length integer key and mapped to static memory embeddings via deterministic hashed lookup. This involves multiplicative-XOR hashing to compute memory indices across multiple hash heads ( $H$ ) and layers ( $\ell$ ).
Extensibility: The memory is organized as a logical table composed of multiple physical sub-tables. This allows the system to scale with instruction length (by adding key-gram slots) and task complexity (by adding parallel sub-tables) without modifying the backbone.

2. Context-Adaptive Fusion

Retrieved embeddings are injected into selected Transformer layers of the visual backbone through a residual module:

Gating: A token-wise gate is computed using a scaled dot-product between the visual hidden states ( $H^{(\ell)}$ ) and the projected memory keys ( $K_m^{(\ell)}$ ). This gate ( $A^{(\ell)}$ ) determines the relevance of the retrieved linguistic memory to each visual token, ensuring priors are injected only where useful.
Convolutional Refinement: The gated memory features are passed through a lightweight long-span convolution to mix heterogeneous instruction-level knowledge components before injection.
Residual Injection: The refined features are added to the original hidden states ( $\tilde{H}^{(\ell)} = H^{(\ell)} + \Delta H^{(\ell)}$ ) before standard attention computation.

3. Prediction Heads

The memory-guided backbone supports two downstream pathways:

Vision Reasoning: Final-layer hidden states are adapted via cross-attention to predict compact future visual latents (visual foresight).
Action Inference: A lightweight action expert uses the backbone's KV cache as conditional context to decode action trajectories via flow-matching denoising.

4. System Efficiency

The memory table is designed for O(1) lookup relative to total memory size. During inference, indices are determined solely by the instruction and fixed hash functions, allowing large knowledge tables to reside in host (CPU) memory, while the backbone remains focused on visual reasoning.

Key Contributions

Functionally Decoupled Framework: Key-Gram separates instruction-side world-knowledge retrieval from vision-side physical reasoning, reducing modality competition in embodied control.
Extensible External Memory: It stores reusable linguistic priors in a structured memory accessed via deterministic hashed lookup. This enables modular knowledge expansion (adding new objects, skills, or relations) without rewriting the backbone weights.
Consistent Performance Gains: The approach demonstrates improved compositional grounding, transfer, and adaptation across diverse linguistic formulations and environments.

Experimental Results

The authors evaluated Key-Gram on RoboTwin2.0, LIBERO, LIBERO-Plus, and real-world dual-arm manipulation tasks using $\pi0$ and $\pi0.5$ backbones.

RoboTwin2.0: Key-Gram improved average success rates by 29.5% for $\pi0$ and 9.9% for $\pi0.5$ . Gains were particularly pronounced in "Hard" settings and with longer horizons.
LIBERO-Plus Transfer: In zero-shot transfer from LIBERO-trained models to LIBERO-Plus (without target-domain fine-tuning), Key-Gram achieved a 35.8% relative gain for $\pi0$ and 4.5% for $\pi0.5$ .
Real-World Long-Horizon Tasks: On a Piper dual-arm platform, Key-Gram improved average performance by 15.4% ( $\pi0$ ) and 8.1% ( $\pi0.5$ ). The largest gains were observed in assembly tasks, suggesting the memory helps maintain object relations and sub-goal structures.
Expansion Tasks: In unseen compositional pairings (e.g., recombining objects and targets), Key-Gram substantially improved performance (e.g., +34.6% and +41.7% for $\pi0$ on specific transfer tasks), whereas base models often failed on these recombination scenarios.
Ablation Studies: Layer placement analysis revealed that injecting memory at early layers (e.g., Layer 1) is most critical, with diminishing returns for deeper injections. This supports the hypothesis that external priors reduce the instruction-related burden on the backbone early in the processing pipeline.

Significance and Claims

The paper claims that separating language-side world-knowledge retrieval from vision-side physical reasoning provides a more effective inductive bias for embodied intelligence than dense token fusion or coarse generative conditioning.

Key-Gram addresses the tension between static world knowledge and online physical reasoning by treating language as an extensible index rather than a competing token stream. By externalizing knowledge into a modular memory table, the framework allows for:

Modular Knowledge Growth: New knowledge can be appended as new entries without overwriting existing weights.
Protection from Gradient Interference: Existing memory is protected during backbone adaptation, making the architecture naturally compatible with open-world deployment and continual accumulation.
Efficient Scaling: The O(1) lookup pattern and separation of memory from computation allow for efficient handling of large knowledge bases without increasing the computational cost of the visual backbone.

The authors conclude that this lightweight and extensible direction offers a practical mechanism for improving compositional grounding and transfer in embodied manipulation.