Source-linked AI summary

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

Hassan Ismkhan, Hamid Bouchahcia

arXiv:2606.15514v1cs.ROcs.LG

TL;DR

Missing sensors challenge imitation learning systems that assume intact modalities. RL4IL uses reinforcement-learning-guided retrieval and soft imputation, substantially outperforming baselines under camera dropout across three LIBERO suites.

  • Problem

    Imitation learning commonly assumes intact sensor modalities, yet camera failures or occlusions make missing-modality robustness necessary for deployment.

  • Method

    RL4IL uses PPO-guided retrieval over frozen demonstration libraries, soft cross-attention fusion, and per-modality retrieval with soft imputation for missing inputs.

  • Results

    Up to 0.733 success rate under complete camera dropout versus 0.295 for DisDP, with RL4IL substantially outperforming all baselines across three LIBERO suites.

  • Takeaways & Limitations

    The results support retrieval-based imitation learning with learned ranking and soft imputation as a practical strategy when sensor availability cannot be guaranteed.

  • Takeaways & Limitations

    The paper’s scope leaves online settings, richer temporal context, and larger, more diverse task distributions for future work.

Abstract

from arXiv · show

Robotic systems perceive the world through multiple input modalities -- including visual camera streams and natural language instructions -- and must select appropriate actions based on these signals. However, assuming the permanent availability of all input devices is unrealistic, as sensors may fail, become occluded, or drop out entirely during deployment. Robust handling of such missing-modality scenarios is therefore essential for real-world robot operation. This paper introduces RL4IL, a reinforcement learning guided method for imitation learning that selects the most suitable action for a given observation by identifying the most relevant expert demonstrations from a training library. A reinforcement learning policy, trained via Proximal Policy Optimisation over Breadth-First Search candidate sets, ranks candidate demonstrations and a soft cross-attention fusion head aggregates their action signals to produce the final prediction. When a modality is missing at inference time, a dedicated per-modality RL retrieval policy identifies donor demonstrations from the training library, and a soft imputation head reconstructs the missing embedding via cross-attention over the top-ranked donors -- without requiring any retraining of the system. Experiments on three LIBERO benchmark suites demonstrate that RL4IL substantially outperforms state-of-the-art imitation learning methods under sensor dropout conditions, while requiring no policy network training. The code can be found at https://github.com/h-ismkhan/Reinforcement-Learning-via-kNN-for-Robotic-Learning-with-Missing-Camera

1 Introduction

The introduction frames missing-modality robustness as essential for real-world robotic manipulation, while existing approaches remain poorly suited to heterogeneous robot inputs and zero-shot sensor dropout. It presents RL-guided demonstration retrieval and soft cross-attention fusion as the paper’s core contributions.

  • Motivation: Real-world robotic manipulation must integrate heterogeneous streams, including multiple camera viewpoints and natural-language instructions, despite unreliable sensor availability.The introduction characterizes missing-modality robustness as a fundamental deployment requirement.
  • Limitations of prior work: Existing missing-modality methods are limited by domain-specific modalities, incomplete coverage of vision-and-language robot learning, and the need for retraining or fine-tuning.These constraints make zero-shot sensor dropout at inference time infeasible.
  • Contributions: RL-guided demonstration retrieval uses a PPO-trained policy over BFS-augmented candidate sets to rank neighbours by usefulness for the current imitation-learning query.This replaces heuristic distance-weighted aggregation with a learned selection strategy.
  • Contributions: Soft cross-attention fusion avoids committing to a single retrieved demonstration by combining information from retrieved candidates.The supplied contribution passage introduces this mechanism but is truncated before further details.

2 Related Work

Prior imitation-learning and retrieval-based policies perform strongly when all modalities are available, but existing approaches remain vulnerable to camera or complete sensor dropout. Related robustness and missing-modality methods typically require policy retraining or are limited to specific modalities and domains.

  • Imitation learning for manipulation: Behaviour Cloning, diffusion policies, action-chunking transformers, and Vision-Language-Action models achieve strong results on LIBERO but assume every training modality is available at deployment.Camera position is identified as the largest single source of generalisation failure.
  • Robustness to sensor dropout: Masked-modality reinforcement learning and multi-view disentanglement improve camera-failure robustness, while DisDP reduces reliance on individual cameras but requires full policy retraining and collapses under complete camera dropout.DisDP uses shared and private representations for sensor inputs.
  • Retrieval-based policies: Retrieval-augmented imitation-learning methods condition policies on retrieved demonstrations but assume complete observations at query time.
  • Multimodal Learning: Domain-specific missing-modality methods for MRI sequences and incomplete physiological signals require full retraining and do not generalise beyond their target modalities.M3AE and ShaSpec address missing MRI sequences, while CIMSleepNet addresses incomplete physiological signals.

3 Method

RL4IL freezes modality encoders, retrieves and ranks demonstrations with PPO over BFS-expanded candidate sets, and uses soft attention for prediction and missing-modality imputation. The same pipeline supports supervised prediction and multimodal imitation learning from recorded demonstrations.

  • Imitation learning: In imitation learning, each library item is a multimodal expert demonstration pairing visual and language observations with its recorded action sequence.The same retrieval-and-fusion pipeline applies when demonstrations replace labelled input–output pairs.
  • Representation and retrieval: Frozen pretrained encoders produce modality embeddings, while modality-specific normalization enables a distance weighting modalities equally despite differing embedding dimensions.Encoder weights remain fixed; learning occurs only in retrieval and imputation modules.
  • Candidate construction: For each query, the method forms a k-nearest-neighbor seed set and expands training candidates through depth-limited graph search, retaining the nearest same-label oracle.Candidate sets are shuffled during training, and points without a same-label candidate are excluded from RL training.
  • RL retrieval policy: An attention-style policy scores every BFS candidate, selects via softmax, and is trained with PPO using reward 0 for matching the oracle and −1 otherwise.Only candidate sets with at least two members provide a meaningful choice.
  • Soft fusion: A fusion head attends over the RL policy’s top-K′ candidates, combines attended labels with query context, and outputs predictions for regression or classification.Classification uses one-hot labels, cross-entropy loss, and the same attention and context computation as regression.
  • Missing-modality imputation: For a missing modality, a dedicated RL policy retrieves an eligible donor using present modalities, then the pipeline imputes the missing embedding before label prediction.Eligible donors contain the missing modality, and the imputation oracle minimizes squared ℓ2 distance to the ground-truth embedding.

4 Experiments

RL4IL is evaluated on three LIBERO manipulation suites under agent-view and in-hand camera dropout, using task success rate across repeated rollouts. Ablations show benefits from soft fusion, modality-fair normalisation, longer training for LIBERO-Goal, and stable performance across moderate top-K values.

  • Benchmark setup: RL4IL is evaluated on LIBERO-Spatial, LIBERO-Object, and LIBERO-Goal, each containing ten robotic manipulation tasks with distinct reasoning requirements.LIBERO-Spatial focuses on tabletop spatial reasoning, LIBERO-Object on object-centric manipulation, and LIBERO-Goal on longer-horizon goal-conditioned instructions.
  • Evaluation protocol: Experiments mask either the agent-view camera or in-hand camera, reconstruct the missing embedding before retrieval, and measure success within 260 steps.Success rate is averaged over 3 random seeds with 25 rollouts per task per seed.
  • Effect of soft fusion: Soft cross-attention fusion consistently matches or outperforms hard-argmax selection across all suites and camera-dropout conditions.Aggregating top-ranked candidates improves reliability when individual demonstrations are imperfect and retrieval noise is present.
  • Effect of modality-fair normalisation: Modality-fair z-score normalisation and dimension scaling prevent large-norm modalities from dominating retrieval distances, despite occasional higher unnormalised scores.Those isolated gains can occur when a single modality happens to suffice for a task, rather than indicating genuine improvement.
  • Effect of training epochs: Performance remains relatively stable across training epochs on LIBERO-Spatial and LIBERO-Object, while LIBERO-Goal improves with longer training.The result indicates that additional epochs provide meaningful gains particularly for goal-conditioned tasks.
  • Effect of top-K fusion candidates: Performance is generally stable across moderate values of K′ across all three LIBERO suites, indicating limited sensitivity to this hyperparameter.The tested fusion-candidate set is K′_imp ∈ {4, 8, 12, 16, 32}.

5 Conclusion

RL4IL is a reinforcement learning-guided retrieval framework designed to handle sensor dropout without policy retraining. It combines learned demonstration ranking with soft fusion and modality-specific imputation, achieving strong performance across LIBERO benchmarks and motivating extensions to online and larger-scale settings.

  • Core framework: RL4IL handles sensor dropout without policy retraining by using a PPO-trained policy over BFS-augmented candidate sets to select expert demonstrations from a frozen library.This is presented as the first application of reinforcement learning to nearest-neighbour selection in imitation learning.
  • Soft fusion and imputation: Soft cross-attention fusion aggregates signals from top-ranked candidates and consistently outperforms hard argmax selection across evaluated conditions.For missing modalities, dedicated per-modality retrieval policies select donor demonstrations and a soft imputation head reconstructs missing embeddings via cross-attention.
  • Benchmark results: 0.733 success rate is achieved by RL4IL under complete camera dropout, compared with 0.295 for DisDP across three LIBERO benchmark suites.The evaluated suites are LIBERO-Spatial, LIBERO-Object, and LIBERO-Goal.
  • Ablations: Ablations show that soft fusion, modality-fair normalisation, and the RL retrieval policy each contribute meaningfully to overall performance.The conclusion also reports competitive results under camera-dropout conditions.
  • Future work: Future work may extend RL4IL to online settings, incorporate richer temporal context into retrieval queries, and scale the demonstration library to larger, more diverse task distributions.These directions follow the reported practicality of retrieval-based imitation learning when sensor availability cannot be guaranteed.

A Modality-Specific Normalisation and Modality-Fair Distance

The method normalizes each modality using training-set z-score statistics to prevent scale and variance differences from distorting retrieval distances. It then computes modality-fair errors by averaging per-dimension and per-modality contributions, while handling near-constant dimensions safely and preventing test-data leakage.

  • Normalization: Modality-specific z-score normalization places heterogeneous embeddings on a common scale, preventing large raw magnitudes from dominating distances.Means and standard deviations are computed element-wise over all N training samples for each modality.
  • Normalization: Dimensions with standard deviation below ϵ > 0 are clamped to 1, preventing division by zero and minimizing near-constant dimensions’ influence.The clamping is applied before normalized embeddings are computed.
  • Inference: At inference, test embeddings use the training statistics µ_m and σ_m, preventing test information from entering normalization.No test-set statistics are used for this step.
  • Modality-fair distance: The distance averages each modality’s squared ℓ2 error per dimension and then averages modality contributions, preventing larger embeddings or modality counts from dominating comparisons.Dividing by d_m yields mean squared error per dimension, while the outer 1/M factor averages across present modalities.

B State and Candidate Features … Encoders and Embeddings

RL4IL represents each BFS candidate set with normalized state and candidate features, estimates test-time set sizes from neighboring training sets, and uses frozen CLIP embeddings for all modalities. Hyperparameters are shared across the three LIBERO suites unless specified otherwise.

  • B State and Candidate Features: The state representation includes the normalized concatenated embedding fs, label variance Var(yB(s)), and BFS set size |B(s)|.
  • B State and Candidate Features: Candidate features include maximum similarity, maximum depth, and a distance-based rank within the BFS candidate set.
  • B State and Candidate Features: Shortest-path depth, 0-indexed ascending-distance rank, and mean label are defined for the BFS features, whose normalized values lie in [0, 1].
  • C Test-Time Set Size Estimation: At test time, the appropriate BFS set size is estimated from the training sizes of initial UNN or k-NN neighbours, weighted by inverse distance.
  • C Test-Time Set Size Estimation: The BFS expands until exactly the estimated number of nodes is collected, matching test-time set-size distributions to those encountered during training.
  • D Hyperparameters: Hyperparameters are shared across LIBERO-Spatial, LIBERO-Object, and LIBERO-Goal unless stated otherwise.
  • Encoders and Embeddings: Agent-view camera, in-hand camera, and natural-language instruction use frozen CLIP ViT-B/32 encoders, producing 512-dimensional embeddings per modality.
  • Encoders and Embeddings: Visual demonstration embeddings average CLIP image features from 8 uniformly sampled frames, and no encoder weights are updated.

Neighbourhood and Graph Construction

The method seeds BFS from approximate nearest-neighbour candidates, expands over a fixed k-NN graph under depth and visitation limits, and uses a validation split for training-set evaluation. These settings define the neighbourhood and graph construction procedure.

  • Neighbourhood initialization: The initial neighbour set uses kapprox = 20 approximate nearest neighbours to form the candidate pool and kseed = 5 seed neighbours for BFS.The five seed neighbours are passed directly to the BFS.
  • Graph and BFS limits: The k-NN graph uses kgraph = 5 edges per node, while BFS is bounded by D = 6 hops or 200 visited nodes per source.The first limit reached terminates expansion.
  • Data split: The training set uses a validation fraction of 0.15, with BFS statistics computed on training embeddings.The passage specifies the split fraction and the embeddings used for BFS statistics.

Prediction RL Policy (PPO)

The prediction RL policy uses PPO with a clipped surrogate objective and trains for 30 epochs. Optimization uses specified entropy regularization, Adam settings, gradient clipping, and 128-sample rollout minibatches.

  • Prediction RL Policy (PPO): The RL policy is trained with Proximal Policy Optimisation for 30 epochs using a clipped surrogate objective with clipping threshold ε = 0.2.Training follows PPO with the stated clipping threshold.
  • Prediction RL Policy (PPO): The entropy bonus coefficient is set to αent = 10−4, and Adam uses a learning rate of 3 × 10−4.Both values specify the policy optimization configuration.
  • Prediction RL Policy (PPO): Gradient norm clipping is set to 1.0, while rollout minibatches contain 128 samples.These settings constrain optimization and define rollout minibatch size.

Soft Prediction Fusion Head · Imputation RL Policy (PPO) · Soft Imputation Head

The method uses soft cross-attention heads for prediction fusion and missing-camera imputation, with separately trained PPO retrieval policies and fixed training configurations. Language is always available, while the imputation policy remains frozen during imputation-head training.

  • Soft Prediction Fusion Head: The prediction fusion head attends over the top-K′ = 32 RL-ranked candidate demonstrations.It uses shared projection dimension df = 128 across H = 4 attention heads.
  • Soft Prediction Fusion Head: The prediction fusion head is trained for 30 epochs using Adam at learning rate 3 × 10−4, batch size 32, and gradient clipping 1.0.These settings specify the optimization procedure for aggregating candidate action signals.
  • Imputation RL Policy (PPO): A dedicated imputation PPO policy is trained independently for each camera modality over 50 PPO epochs with minibatches of 64 samples.The language modality is always present and is never imputed.
  • Imputation RL Policy (PPO): The imputation policy uses entropy coefficient αent = 10−4, learning rate 3 × 10−4, and gradient norm clipping at 1.0.Its clipping threshold and optimiser match those of the prediction RL policy.
  • Soft Imputation Head: The soft imputation head aggregates K′The supplied passage begins this statement and continues in the following passage.
  • Soft Imputation Head: imp = 32 RL-ranked donor embeddings via cross-attention, using dimension dimp = 64 across Himp = 2 attention heads.It is trained for 30 epochs with supervised MSE, Adam, learning rate 3 × 10−4, batch size 32, and gradient clipping 1.0; the imputation policy is frozen.

Rollout Evaluation Protocol

Task success is evaluated over 3 random seeds with 25 rollouts per task per seed, using a 260-step maximum per rollout; success rate is the seed-averaged fraction completing within that budget.

  • Rollout Evaluation Protocol: Evaluation uses 3 random seeds and 25 rollouts per task per seed, totaling 75 rollouts per task.The protocol follows DisDP.
  • Rollout Evaluation Protocol: Each rollout permits a maximum of 260 environment steps, and success rate averages across seeds the fraction completed within this budget.Success is reported according to the evaluation protocol of DisDP.
Loading 2606.15514v1…