🤖 machine learning

Reinforcement Learning-Guided Retrieval with Soft Fusion for Robust Multimodal Imitation Learning under Missing Modalities

Original authors: Hassan Ismkhan, Hamid Bouchahcia

Published 2026-06-16

📖 4 min read☕ Coffee break read

Original authors: Hassan Ismkhan, Hamid Bouchahcia

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 a meal. Usually, you show it a video of a chef chopping vegetables (visual data) and you give it a spoken recipe (language data). The robot learns by watching and listening, then trying to copy the chef's moves.

But what happens if the camera breaks halfway through, or the microphone stops working? In the real world, sensors fail, get blocked by objects, or just stop sending signals. Most current robot learning methods panic when this happens; they assume the camera and microphone will always work perfectly. If one fails, the robot often freezes or makes a mistake.

This paper introduces a new method called RL4IL that acts like a super-smart librarian for robots. It allows the robot to keep working perfectly even when a camera or sensor goes dark, without needing to be retrained or taught new lessons.

Here is how it works, using simple analogies:

1. The "Smart Librarian" (Reinforcement Learning)

Usually, when a robot needs to figure out what to do, it looks at its library of past demonstrations (videos of experts doing tasks) and picks the one that looks most similar to what it is seeing right now. It's like a student trying to find the right page in a textbook by glancing at the pictures.

The problem is, if the picture is missing (because the camera broke), the student might pick the wrong page.

RL4IL replaces this simple "glance" with a Smart Librarian. This librarian is a Reinforcement Learning agent (a type of AI that learns by trial and error). Instead of just picking the "closest" match, the librarian learns to rank the top candidates in the library based on how useful they are for the current situation. It's like a librarian who doesn't just find the book with the most similar cover, but finds the book that actually contains the right instructions for the specific problem you have, even if the cover is damaged.

2. The "Group Chat" (Soft Fusion)

Old methods often pick just one best match from the library and say, "Okay, we are copying this specific video." If that one video is slightly noisy or imperfect, the robot fails.

RL4IL uses a Soft Fusion approach. Imagine instead of listening to just one expert, the robot asks the top 5 or 10 most relevant experts for advice. It doesn't just pick one; it listens to all of them and blends their advice together, giving more weight to the experts who seem most confident. This is like a "group chat" where the robot averages out the noise. If one expert is slightly off, the others correct them, resulting in a much smoother and more reliable action.

3. The "Magic Fill-In" (Missing Modality Imputation)

This is the paper's biggest trick. What if the camera is completely dead? The robot has no visual data at all.

Instead of giving up, RL4IL has a Magic Fill-In mechanism.

The Detective: A specialized AI detective looks at the other things the robot does have (like the language instructions or the hand camera) and says, "Based on these clues, which expert in the library had a similar situation?"
The Reconstruction: Once it finds those experts, it doesn't just copy their video. It uses a "cross-attention" mechanism to look at the missing parts of those experts' videos and mathematically reconstructs what the missing camera would have seen.
The Result: It creates a fake, but highly accurate, version of the missing camera feed. The robot then uses this reconstructed feed to find the right action, just as if the camera had never broken.

Why is this a big deal?

No Retraining: Most methods require you to retrain the robot from scratch every time you want it to handle a broken sensor. RL4IL is "zero-shot." You train it once, and if a camera breaks tomorrow, it handles it immediately without any new lessons.
Better than the Rest: The authors tested this on three different sets of robot tasks (moving objects, following goals, and spatial reasoning). When they simulated cameras failing, their method (RL4IL) succeeded about 70% to 73% of the time. The next best method only succeeded about 29% of the time.
Robustness: It works even when the robot is in a chaotic environment where sensors might get blocked or fail entirely.

Summary

Think of RL4IL as a robot that doesn't just memorize a script. It has a Smart Librarian to find the best help, a Group Chat to blend advice from multiple experts, and a Magic Fill-In tool to guess what a broken camera would have seen. This allows the robot to keep working smoothly in the messy, unpredictable real world where sensors often fail.

Technical Summary: RL4IL for Robust Multimodal Imitation Learning

Problem Statement

Robotic manipulation systems rely on heterogeneous sensory streams, typically combining visual camera feeds (e.g., agent-view and in-hand cameras) with natural language instructions. While recent imitation learning (IL) methods have achieved strong performance on structured benchmarks, they operate under the unrealistic assumption that all sensor modalities remain permanently available. In real-world deployment, sensors may fail, become occluded, or degrade, leading to "missing modality" scenarios. Existing robustness methods often require retraining or fine-tuning to handle new dropout patterns, and many retrieval-based approaches fail when the query observation is incomplete. Furthermore, current multimodal learning techniques for missing data are often confined to static tasks (e.g., classification) or domain-specific modalities (e.g., MRI), making them unsuitable for the sequential action outputs required in robotics.

Methodology: RL4IL

The paper introduces RL4IL, a framework that addresses missing modalities through a retrieval-based approach guided by Reinforcement Learning (RL) and soft fusion, requiring no retraining of the core policy network at inference time.

1. Core Architecture

The system operates on a frozen library of expert demonstrations $\mathcal{D}_{tr}$ . Each demonstration consists of multimodal observations (visual frames, language) and a corresponding action sequence.

Frozen Encoders: Pretrained encoders map raw sensory inputs to fixed-dimensional embeddings. These weights are never updated.
Modality-Fair Normalization: To handle heterogeneous embedding dimensions, the system applies modality-specific z-score normalization and defines a distance metric that weights each modality's contribution equally.

2. RL-Guided Retrieval

Instead of using heuristic distance weighting to aggregate neighbors, RL4IL employs a learned selection strategy:

BFS-Augmented Candidate Sets: For a query observation, the system first identifies a $k$ -nearest neighbor (k-NN) set. This set is expanded via a Breadth-First Search (BFS) on a training k-NN graph to create a richer, more diverse candidate pool $B(s)$ .
PPO Policy: A policy trained via Proximal Policy Optimization (PPO) ranks candidates within the BFS-augmented set. The policy learns to select the most relevant demonstration (the "oracle") based on the query's state.
Reward Signal: The reward is defined relative to an oracle (the candidate with the same label and minimum graph distance), encouraging the policy to match the optimal selection.

3. Soft Cross-Attention Fusion

Rather than committing to a single "best" neighbor (hard argmax), which can be noisy, RL4IL utilizes a soft fusion head:

The top- $K'$ candidates ranked by the RL policy are aggregated using a cross-attention mechanism.
This head computes attention weights over the candidate embeddings and their associated action sequences (labels), producing a weighted average prediction.
This approach is shown to be more robust to retrieval noise than selecting a single demonstration.

4. Zero-Shot Missing Modality Imputation

When a sensor fails at inference time (e.g., a camera is occluded), the system reconstructs the missing embedding without retraining:

Per-Modality RL Retrieval: A dedicated RL policy, trained specifically for each modality, identifies "donor" demonstrations from the library where the missing modality is present.
Soft Imputation Head: Similar to the prediction fusion, a soft imputation head uses cross-attention over the top-ranked donors to reconstruct the missing embedding vector.
Pipeline Continuity: Once the missing embedding is reconstructed, the full observation is passed to the standard prediction pipeline. The system requires no modification to the downstream policy network.

Key Contributions

RL-Guided Demonstration Retrieval: The first application of reinforcement learning to select expert demonstrations for imitation learning, replacing heuristic aggregation with a learned ranking strategy over BFS-augmented candidate sets.
Soft Cross-Attention Fusion: A lightweight mechanism that aggregates signals from multiple top-ranked candidates, outperforming hard selection methods, particularly under noisy retrieval conditions.
Zero-Shot Robustness: The ability to handle sensor dropout at inference time without retraining the policy network. This is achieved via dedicated per-modality RL retrieval policies and soft imputation heads that reconstruct missing embeddings.
State-of-the-Art Performance: Demonstrated superior performance on the LIBERO benchmarks under complete camera dropout compared to existing baselines.

Experimental Results

The authors evaluated RL4IL on three LIBERO benchmark suites: LIBERO-Spatial, LIBERO-Object, and LIBERO-Goal. The experiments simulated complete camera dropout (either agent-view or in-hand camera absent).

Performance: RL4IL substantially outperformed all baselines, including Behavior Cloning (BC), Disentangled BC (DisBC), BESO-ACT, and Disentangled Diffusion Policy (DisDP).
- On LIBERO-Goal with the agent-view camera missing, RL4IL achieved a success rate of 0.700, whereas the strongest baseline (DisDP) achieved only 0.004.
- On LIBERO-Object with the in-hand camera missing, RL4IL achieved 0.677, compared to 0.226 for DisDP.
Ablation Studies:
- Soft Fusion: Consistently matched or outperformed hard argmax selection across all suites.
- Modality-Fair Normalization: Ensured consistent performance across suites by preventing any single modality with a large embedding norm from dominating the distance metric.
- Training Epochs: The system maintained competitive performance even after a single training epoch, surpassing all baselines reported in the literature.
- Hyperparameters: Performance remained stable across a reasonable range of top- $K$ fusion candidates ( $K' \in \{4, 8, 12, 16, 32\}$ ).

Significance and Claims

The paper claims that RL4IL offers a practical and effective strategy for deploying robotic systems in environments where sensor availability cannot be guaranteed. By decoupling the retrieval and imputation mechanisms from the policy network training, the method achieves robustness without the computational cost of retraining. The authors emphasize that their approach addresses a fundamental requirement for real-world deployment—handling sensor dropout—while maintaining high performance on complex manipulation tasks. The work suggests that retrieval-based imitation learning, when augmented with learned ranking and soft imputation, is a viable path toward robust autonomous systems.