Nostra: Enabling Robust Robot Imitation via
Multimodal Latent Imagination

Abstract

Similar to humans, robots benefit from multiple sensing modalities when performing complex manipulation tasks. Current behavior cloning (BC) policies typically fuse learned observation embeddings from multimodal inputs before decoding them into actions. This approach suffers from two key limitations: 1) it requires all modalities to be present and in-distribution at test time, otherwise corrupting the latent state and leading to fragile execution; and 2) naive fusion across all inputs hinders learning from large-scale heterogeneous datasets, where only a subset of modalities may be informative at different phases of a task. We introduce Nostra, a multimodal state-space model that learns a modular per-modality latent representation, enabling flexible action prediction with or without specific inputs. BC-Nostra improves robustness to unseen noise by using KL divergence between inferred and imagined multimodal latents as a noise measure, and by employing latent imagination to predict action trajectories over arbitrary horizons. On a suite of MuJoCo-based tasks, BC-Nostra fits expert demonstrations up to six input modalities (multi-view RGB, depth, and proprioception), achieving over 20% higher performance under noisy evaluation. Furthermore, Nostra adaptively down-weights non-informative inputs, facilitating effective co-training on large heterogeneous robotics datasets with 𝒪(10k) demonstrations spanning diverse tasks and visual conditions. Finally, we demonstrate real-world deployment, where BC-Nostra achieves up to a 40% performance gain under camera occlusions on multiple manipulation tasks.

Key Ideas

Method

BC-Nostra maximizes an ELBO over a joint model of actions and observations. For every modality m it learns a closed-loop posterior q(z^m_t | z^1:M_t−1, o^m_t) and an open-loop prior p(z^m_t | z^1:M_t−1), sharing a GRU between the two. Per-modality latent samples are fused by a multi-head attention Modality Transformer before action decoding. Because an explicit prior is learned, the policy can roll out even when some or all observations are missing — unlike BC-LSTM or Diffusion Policy, which only model p(z_t | o_t).

Results

Robustness to unseen visual noise

Four MimicGen tasks × three noise types (black occlusion mask, camera pose jitter, unseen table texture), applied during critical task phases. Success rates averaged over 50 rollouts.

MLI lifts average success under noise from 27.0 → 60.0, beating Diffusion Policy by up to 40% on Stack D1 under masking, and BC-RSSM + LI by ~17% on average — evidence that a separable latent space, not merely latent imagination, is what filters noisy inputs.

AdaMLI handles noise in any modality

Eight MimicLabs tasks, with noise injected into RGB, then additionally depth (RGBD), then additionally zeroed proprioception (All) to emulate sensor failure.

AdaMLI beats Diffusion Policy by more than 20% on average and up to 60% on long-horizon tasks, with fewer than 3% false-positive switches when no noise is present.

Which modality is the policy using?

Real-robot deployment

Under wrist-camera occlusion Diffusion Policy fails outright on two of three tasks — the arm makes jerky or frozen motions and error accumulates. BC-Nostra detects the occlusion via AdaMLI and switches that modality to imagination, holding success near its unoccluded rate. BC-RSSM, which shares the imagination machinery but uses a joint latent, trains unstably and never succeeds at even the easy lift task — isolating modularity as the critical ingredient.

Learning from heterogeneous data

• Pre-training with a non-informative view. After pre-training on data where the robot self-occludes the agent view, fine-tuning on 10 demos of Stack D1 reaches 44% vs. 14% for BC-RSSM; on Coffee D2 with 50 demos, 46% vs. 8%.
• Co-training at scale. Adding 𝒪(10k) retrieved demonstrations, Diffusion Policy drops 20–22% on the long-horizon open drawer & place bowl task, while BC-Nostra gains 14% and beats BC-RSSM by 8%.
• Ignoring what doesn't help. When a view is uninformative, the information stored in its latent falls by up to 32% (measured in nats), with the slack picked up by the remaining modalities.

Experimental Setup

• Tasks: 12 MuJoCo tasks — 4 MimicGen (1000 demos each) + 8 MimicLabs (200 demos each); 4 real-robot tasks (30 demos each).
• Modalities (up to 6): agent-view RGB, wrist-view RGB, agent-view depth, EEF position, EEF orientation, gripper width.
• Baselines: BC-LSTM, Diffusion Policy, Diffusion Forcing, BC-RSSM (joint-latent ablation), plus modality-dropout variants.
• Training: 300k gradient steps on a single NVIDIA A40; ResNet18 encoder on 84×84 images with crop augmentation; latent size 80 per modality; sequence length 16.

BibTeX

@article{saxena2026nostra,
  title   = {{NOSTRA}: Enabling Robust Robot Imitation via Multimodal Latent Imagination},
  author  = {Saxena, Vaibhav and Luo, Yunhao and Koga, Yotto and Xu, Danfei},
  journal = {IEEE Robotics and Automation Letters},
  year    = {2026}
}