PaperlayerView exact revision on arXiv

Source-linked AI summary

UniMPA: A Unified Memory-Prediction-Action Model via Action-Grounded Transition Modeling

Wei Li, Rui Shao, Jie He, Lingsen Zhang, Ziwei Liu, Liqiang Nie

arXiv:2609.11875v1cs.RO

TL;DR

Observation-to-action learning remains limited by a transition realizability gap involving ambiguous phases, physically unrealizable predictions, and scene-mismatched historical actions. UniMPA couples future-supervised transition modeling with bidirectional visual-action memory and prototype-biased action refinement. It reports state-of-the-art performance across evaluated benchmarks while using 25–50% of π0.5’s training epochs.

  • Problem

    Observation-to-action policies must infer intended visual transitions and realize them with executable actions despite transition ambiguity, prediction–execution mismatch, and experience–realization mismatch.

  • Method

    UniMPA uses Persistent-Selective Future Prediction, bidirectional Visual-Action and Action-Visual Memory Banks, and Prototype-Biased Flow through a shared action-grounded transition interface.

  • Results

    UniMPA achieves state-of-the-art performance across evaluated benchmarks while requiring only 25–50% of the training epochs used by the π0.5 baseline.

  • Takeaways & Limitations

    UniMPA organizes prediction, memory, and action generation around a shared physical-transition interface rather than treating them as independent capabilities.

  • Takeaways & Limitations

    The transition-critical trigger remains binary; future work proposes uncertainty-aware allocation of prediction resolution, token count, decoding depth, or temporal horizon.

Abstract

from arXiv · show

Recent advances in Vision-Language-Action (VLA) models have improved robotic manipulation, yet observation-to-action learning remains limited by a fundamental transition realizability gap, manifested in three tightly coupled problems: (i) Transition ambiguity. Visually similar current observations may correspond to different manipulation phases and imply different subsequent transitions. (ii) Prediction--execution mismatch. A visually plausible predicted future observation does not necessarily correspond to a physically realizable transition. (iii) Experience--realization mismatch. A historically executable action pattern may not necessarily realize the intended transition in the current scene and therefore requires context-aware adaptation. Accordingly, we propose UniMPA, a Unified Memory-Prediction-Action model that addresses these problems through a shared action-grounded transition interface. (i) UniMPA introduces Persistent-Selective Future Prediction to resolve transition ambiguity by modeling the intended future state evolution. A persistent latent stream continuously tracks task-level progress, while a transition-critical pixel stream selectively resolves fine-grained interaction changes through memory-grounded prediction. (ii) To assess the physical executability of the anticipated transition, the predicted transition queries a temporal Visual-Action Memory Bank. The bank retrieves historically realized visual-action experience, grounding future prediction in executable evidence. (iii) To adapt executable experience to the current scene, an Action-Visual Memory Bank retrieves visually grounded action prototypes from historical action evolution. Prototype-Biased Flow then shifts the flow source toward a historically supported action manifold for context-aware refinement.

I. INTRODUCTION

UniMPA frames observation-to-action learning as a transition realizability problem: policies must infer intended state evolution, ground it in executable experience, and adapt actions to the current scene. It unifies future prediction, bidirectional memory, and action refinement through an action-grounded transition interface.

  • I. INTRODUCTION: The transition realizability gap comprises ambiguity between visually similar manipulation phases, prediction–execution mismatch, and experience–realization mismatch under changed scene geometry.These limitations concern inferring the intended transition, realizing it physically, and adapting historical action patterns to current poses or contacts.
  • I. INTRODUCTION: UniMPA closes the missing interface between intended state evolution and executable action generation by coupling prediction, memory, and action refinement.Its anticipate–ground–refine process uses prediction to form retrieval queries, memory to ground them in visual–action evolution, and retrieved experience as a contextual action prior.
  • I. INTRODUCTION: Persistent-Selective Future Prediction shapes World-Expert transition tokens with persistent latent progress and selectively triggered pixel supervision for interaction-critical changes.The decoded future supervises the tokens, while the Action Expert consumes pre-head tokens rather than the decoded future observation.
  • I. INTRODUCTION: Bidirectional Visual-Action and Action-Visual Memory Banks ground predicted transitions in executable history and construct scene-adapted action prototypes through Prototype-Biased Flow.The design avoids replaying historical actions as-is by shifting the flow source toward a supported action manifold before refinement.
  • I. INTRODUCTION: 25–50% of the training epochs used by π0.5 sufficed for UniMPA to achieve state-of-the-art performance across the evaluated benchmarks.The reported result comes from extensive simulation and real-world experiments assessing robustness and generalization.

D. Memory-Augmented Robot Manipulation

UniMPA uses current multi-view observations, language, proprioception, and action history to construct an action-grounded transition interface. Its interacting streams and bidirectional memory connect future transition representations with executable action generation.

  • D. Memory-Augmented Robot Manipulation: Historical trajectories are organized into complementary memory banks whose transition-aware retrieval links predicted visual transitions to executable action evolution.Episode and timestep information support temporally coherent retrieval, while action history provides visually grounded action priors.
  • D. Memory-Augmented Robot Manipulation: UniMPA takes multi-view observations, language instructions, proprioception, and historical actions as context for transition-aware action generation.The context is represented as Xt = {Vt, ℓ, qt, Ht}, with Ht containing the historical action window.
  • D. Memory-Augmented Robot Manipulation: UniMPA learns a latent conditional transition operator rather than representing the intended change as an observation difference.The future observation supervises training, while World-Expert features provide the representation used for action generation.
  • D. Memory-Augmented Robot Manipulation: The model couples Persistent-Selective Future Prediction with Visual-Action and Action-Visual Memory Banks in an anticipate–ground–refine pipeline.Future prediction supplies transition queries, visual-action memory retrieves executable evidence, and action-visual memory supplies visually grounded action prototypes.
  • D. Memory-Augmented Robot Manipulation: Three interacting Transformer streams—the VLM backbone, World Expert, and Action Expert—implement the action-grounded transition interface through cross-stream updates.A block mask controls information exchange, and the World-Expert stream transforms vision-language context into world keys and values.

B. Transition-Aligned Bidirectional Memory Pretraining

UniMPA pretrains bidirectional memory from historical trajectories by aligning visual state evolution with the actions that induce it. Stateful temporal encoding and cross-stream grounding preserve episode-level correspondence for transition-conditioned retrieval.

  • B. Transition-Aligned Bidirectional Memory Pretraining: UniMPA explicitly associates visual state evolution with the action patterns that induce it, creating a transition-level interface between prediction and action generation.Unlike memories storing observations or implicit history independently, the bidirectional entries retain aligned visual and action keys and values.
  • B. Transition-Aligned Bidirectional Memory Pretraining: Visual-Action Memory retrieves action-aware transition experience, while Action-Visual Memory retrieves visually grounded action prototypes from historical action evolution.Episode and timestep indices are retained to support temporally coherent retrieval.
  • B. Transition-Aligned Bidirectional Memory Pretraining: Stateful Vision-Mamba and Action-Mamba modules encode episode-level temporal correspondence in the visual and action memory values.The visual temporal state is maintained across the episode, enabling memory values to represent evolving interaction context.
  • B. Transition-Aligned Bidirectional Memory Pretraining: Cross-stream conditioning injects action dynamics into visual memory and multi-view visual dynamics into action memory.The resulting representations preserve action-conditioned visual context and visually grounded action context before memory-value construction.
  • B. Transition-Aligned Bidirectional Memory Pretraining: Future reconstruction trains visual values to preserve future images and action values to preserve executable future motion patterns.The reconstruction objective uses a future offset ∆ and decodes the two memory modalities into their corresponding future targets.

4) Retrieval-Simulation Training:

Retrieval-simulation training optimizes memory keys and values jointly so retrieval is evaluated by future prediction rather than embedding similarity alone. Episode-constrained retrieval and transition-critical triggering focus supervision on temporally compatible, interaction-relevant evidence.

  • 4) Retrieval-Simulation Training:: The retrieval objective uses temporally aligned visual-action pairs as positives and other valid batch entries as negatives.Cross-modal similarity is temperature-scaled, and the two directional objectives are averaged across camera views.
  • 4) Retrieval-Simulation Training:: Lret evaluates whether retrieved values predict the correct future, distinguishing UniMPA from metric-learning objectives based only on embedding similarity.Reconstruction, retrieval simulation, and cross-modal pairing are jointly optimized to learn predictive and aligned memory representations.
  • 4) Retrieval-Simulation Training:: Cross-modal pairing aligns visual and action values associated with the same physical transition in both retrieval directions.Visual-to-action and action-to-visual objectives provide complementary supervision across camera views.
  • 4) Retrieval-Simulation Training:: The Trigger Gate activates pixel-level future prediction when latent semantic or action-derived motion signals indicate a critical transition.Selected timestamps cluster around reaching, contact, pose adjustment, and gripper switching rather than redundant frames.
  • 4) Retrieval-Simulation Training:: Coarse-to-fine retrieval first selects a relevant episode, then performs soft retrieval only within that episode.Episode-constrained aggregation prevents averaging visually similar but temporally incompatible states across unrelated demonstrations.

C. Persistent-Selective Future Prediction

Persistent-Selective Future Prediction supervises transition tokens with continuous latent prediction and selectively triggered pixel prediction, using transition-aware memory retrieval to focus calibration on executable, interaction-critical changes.

  • Persistent-Selective Future Prediction: Persistent latent prediction remains active across the trajectory, while pixel prediction is selectively activated at transition-critical moments rather than serving as the action-generation representation.The explicit prediction heads can be removed at inference while the future-supervised transition representation remains active.
  • Persistent latent supervision: The World Expert produces pre-head transition tokens that are projected into V-JEPA 2 latent space and aligned with future latent targets during training.This objective shapes the transition tokens into a future-supervised latent representation without making them decoded future states.
  • Selective pixel supervision: A training-only Trigger Gate combines predicted semantic change with action-derived motion events to identify contact, gripper, orientation, and release transitions.Latent and action cues provide complementary evidence for capturing critical transitions beyond visual change or action magnitude alone.
  • Memory-grounded calibration: The predicted transition is projected into visual-memory key space as an endpoint-oriented query, retrieving action-aware transition values for memory-grounded pixel prediction.The retrieved visual values serve as decoder memory, and pixel loss is imposed only when the Trigger Gate activates.
  • Memory-grounded calibration: Pixel prediction calibrates transition tokens around critical visual transitions, while the Action Expert consumes the transition representation rather than decoded future images.The Action-Manifold Prototype Refinement figure describes the related prototype-guided action refinement mechanism.

D. Action-Manifold Prototype Refinement

Action-Manifold Prototype Refinement retrieves visually grounded action experience, shifts the flow source toward a historically supported motion region, and lets the Action Expert adapt it to current context and transition representations.

  • Prototype-Biased Flow: The Action Expert refines prior-biased samples through flow matching using vision-language context, future-supervised transition representations, and prototype-conditioned action context.The Stage-2 objective combines action denoising with persistent latent and trigger-gated pixel prediction.
  • Action-manifold prototype refinement: The Action-Visual Memory Bank retrieves a visually grounded action value from recent action evolution, encoding both an executable motion pattern and its aligned visual-transition semantics.The query represents current translation, rotation, and gripper evolution.
  • Action-manifold prototype refinement: The retrieved action value is fused with recent action history to form an action-manifold prior rather than being executed directly.This prior conditions subsequent flow generation.
  • Prototype-Biased Flow: Prototype bias translates the zero-mean Gaussian source toward a historically supported action region while preserving covariance and stochasticity.The prior strength is controlled by λp.
  • Inference: At inference, future prediction heads and their visual-retrieval path are disabled, while the World Expert and Action-Visual Memory Bank support iterative action generation.The source begins as x1 = ϵ + λpPa,t and is iteratively denoised into the final action chunk.

IV. EXPERIMENTS

UniMPA is evaluated across complementary simulation benchmarks and real-world manipulation suites designed to test standard performance, distribution shifts, semantic reasoning, and physical interaction across embodiments.

  • Simulation benchmarks: The evaluation covers four simulation benchmarks spanning basic manipulation, controlled perturbations, randomized bimanual manipulation, and language-conditioned reasoning.These settings include LIBERO, LIBERO-Plus, RoboTwin 2.0, and VLABench.
  • Simulation benchmarks: LIBERO measures four suites emphasizing spatial relationships, object-centric transfer, goal-conditioned behavior, and long-horizon composition.The benchmark contains Spatial, Object, Goal, and Long suites.
  • Simulation benchmarks: LIBERO-Plus evaluates zero-shot robustness across camera, robot, language, illumination, background, observation-noise, and object-layout perturbations.The evaluation reports original performance, perturbation-averaged success, and performance drop.
  • Simulation benchmarks: RoboTwin 2.0 tests out-of-distribution generalization by training on Clean demonstrations and evaluating randomized bimanual tasks in the Hard setting.Its tasks contain diverse object configurations and bimanual manipulation demands.
  • Real-world evaluation: Real-world evaluation spans seven manipulation suites and uses final task completion plus progressive checkpoints to measure execution progress.Experiments use 25 independent trials per task and report TSR and CSR.

1) LIBERO:

UniMPA achieves strong results across simulation and real-world evaluations, with gains under distribution shift, randomized manipulation, semantic reasoning, and bimanual execution; ablations support its complementary prediction, memory, and temporal-retrieval components.

  • 1) LIBERO:: 98.6% average LIBERO success surpasses X-VLA, while using only 25% of the training epochs required by π0 and π0.5.The result covers the Spatial, Object, Goal, and Long suites.
  • 2) LIBERO-Plus:: 85.3% average success on LIBERO-Plus ranks first, exceeding the strongest competing average by 5.6 points and π0.5 by 11.7 points, with the smallest 13.3-point degradation.The benchmark covers seven perturbation dimensions in zero-shot evaluation.
  • 3) RoboTwin 2.0:: 58.2% average success on randomized RoboTwin 2.0 exceeds HALO and π0.5 by 25.7 and 18.5 points, with UniMPA ranking first on 8/11 tasks.All policies are trained on Clean demonstrations and evaluated in the Randomized Hard setting.
  • 4) VLABench:: 44.0% average VLABench success is highest, exceeding π0.5 by 4.3 points and achieving best or tied-best performance on four of six tasks.The tasks test semantically diverse manipulation and language-conditioned understanding.
  • Real-world evaluation: On GALAXEA R1 Lite, UniMPA reaches 77.7% TSR and 86.3% CSR, while on AgileX Cobot Magic it reaches 74.9% TSR and 86.4% CSR.These results span seven task suites and distinct bimanual embodiments.
  • Ablations: Ablations show complementary benefits from persistent-selective prediction, bidirectional memory, and temporal retrieval, with removing temporal modeling reducing LIBERO and real-world TSR/CSR by 2.7 and 9.2/7.0 points.Memory removal causes 4.1-point LIBERO and 13.8/10.7-point real-world drops, while removing both prediction branches also reduces performance.

3) Future-Oriented Memory Pretraining:

UniMPA trains future-oriented visual and action memories to preserve task evolution, executable motion, and cross-modal correspondence. The resulting memories support future prediction, temporal retrieval, and coarse action prototypes that are adapted to the current scene.

  • Future-Oriented Memory Pretraining: Current-state reconstruction instead of future-oriented training degrades LIBERO performance by 2.5 points and real-world TSR/CSR by 12.0/9.7 points, showing that memory training must preserve transition semantics.Removing visual future reconstruction further decreases performance by 2.2 and 10.3/7.9 points, while removing action future reconstruction also reduces performance.
  • Future-Oriented Memory Pretraining: Cross-modal pairing further reduces performance by 1.5 points on LIBERO and 5.8/4.1 points in real-world TSR/CSR when removed, supporting compatible vision-action transition representations.
  • Memory Bank Analysis: Memory values form smooth ordered trajectories across visual and action t-SNE spaces, while aligned visual-action pairs show cross-modal retrieval correspondence.
  • Action-Visual Memory Bank: Action-Visual Memory prototypes capture major translation, rotation, gripper-state, and stage-transition trends in one forward pass, providing a coarse executable manifold for later adaptation.
  • Future Prediction: Pixel-level future prediction generates interaction-relevant cues for robot motion, object displacement, and contact-related scene evolution in a single forward pass.
  • World-Transition Modeling: World-expert attention concentrates on manipulation-relevant regions, including the gripper, manipulated object, and surrounding contact area, across trajectory timesteps.

Appendix of UniMPA

The appendix details UniMPA’s joint-space trigger gate, shared cross-expert attention, cached inference pathway, and implementation configuration. For bimanual control, semantic and arm-wise motor evidence jointly determine when pixel prediction is activated.

  • Joint-Space Trigger Gate: The gate combines predicted latent semantic change with left- and right-arm 6-DoF joint changes and gripper-state switches to identify motor transitions.
  • Joint-Space Trigger Gate: The bimanual joint-space trigger gate activates pixel-level future prediction when semantic change, arm motion, or gripper switching indicates a critical transition.It replaces Cartesian indicators with left- and right-arm joint-position variation while retaining latent-scene and gripper-state criteria.
  • Cross-Expert Attention: UniMPA uses shared layer-wise self-attention across backbone, World Expert, and Action Expert streams, with hierarchical visibility allowing progressively broader information access.Backbone tokens attend only to the backbone, World tokens also access the backbone, and Action tokens access all three streams.
  • Inference Caching: During inference, the backbone and World Expert run once per control step and cache layer-wise key/value tensors, while only the Action Expert is recomputed during denoising.This preserves the anticipate–ground–refine pathway while avoiding repeated backbone and World-Expert computation.
  • Implementation Configuration: The model uses a π0.5 backbone, an 18-layer World Expert with 16 transition queries, frozen V-JEPA 2 representations, and a lightweight pixel-prediction decoder.

3) Model Configuration:

The training configuration uses two stages with AdamW optimization and standardized memory-pretraining procedures across benchmarks. Stage 1 pretrains memory components with truncated temporal backpropagation and retrieval simulation.

  • Training Configuration: Stage 1 pretrains the memory components for 10 epochs across all benchmarks using truncated backpropagation through time over 8 steps.Retrieval simulation excludes the query entry itself from its candidate set.
  • Training Configuration: Both training stages use AdamW with β1 = 0.9 and β2 = 0.95, with fixed reconstruction, retrieval, and cross-modal pairing loss weights in Stage 1.
  • Configuration Scope: Table XV and Table XVI specify the model architecture, optimization, data-processing, benchmark, and computational-budget settings used for reproducible experiments.

B. Real-World Task Organization

The real-world evaluation organizes manipulation tasks by dominant physical interaction and orthogonal capability challenges. It uses standardized randomized protocols, ordered checkpoints, and complementary success metrics across two bimanual platforms.

  • Evaluation Organization: Real-world evaluation spans seven interaction suites across AgileX Cobot Magic and GALAXEA R1 Lite, covering rearrangement, articulated interaction, assembly, deformable manipulation, coordination, dynamics, and long-horizon recovery.GALAXEA evaluates 21 tasks, while AgileX evaluates seven representative tasks.
  • Evaluation Protocol: All methods share randomized initial states, language instructions, camera configurations, observation modalities, and control interfaces, with each task generally evaluated over N = 25 independent trials.
  • Checkpoint-Based Evaluation: Tasks are represented by ordered physical checkpoints, where later checkpoints count as successful only after all preceding checkpoints are completed.
  • Evaluation Metrics: TSR measures final-checkpoint success, whereas CSR averages cumulative checkpoint success to distinguish early failures from failures near task completion.
  • Task Properties: Orthogonal tags identify transition ambiguity, critical transitions, long-horizon progress, visual-to-action and action-to-visual retrieval, bimanual coordination, dynamics, and recovery.These properties can co-occur within a task rather than defining mutually exclusive categories.
  • Task Descriptions: Natural-language commands specify each task goal, while ordered checkpoints define its observable physical progression and are shared across compared methods.

C. Supplementary RoboTwin 2.0 Results

Supplementary RoboTwin 2.0 results broaden the randomized zero-shot evaluation across 44 tasks, while additional real-world analyses examine task progress, ablations, and joint-space execution. UniMPA’s reported gains span diverse interaction patterns, though difficult behaviors remain challenging.

  • C. Supplementary RoboTwin 2.0 Results: The 44-task supplementary evaluation complements the representative randomized evaluation by characterizing performance across broader task-level distributions.Each task uses 100 rollouts under the same randomized protocol, and task groups report average performance.
  • C. Supplementary RoboTwin 2.0 Results: 35.84% average success across 44 RoboTwin 2.0 tasks exceeds π0.5 by 8.27 points, with gains distributed across diverse interaction patterns.UniMPA also surpasses HALO, BagelVLA, and π0, while ranking, handover, and stacking remain difficult cases.
  • D. Supplementary Real-World Ablation Results: 74.9% TSR / 86.4% CSR for persistent-selective prediction exceeds latent-only, pixel-only, and dense-pixel variants in the real-world ablation.The full formulation preserves its advantage at final completion and intermediate checkpoints.
  • D. Supplementary Real-World Ablation Results: Bidirectional memory, coarse-to-fine retrieval, full memory pretraining, and Prototype-Biased Flow each outperform their corresponding ablations.Removing temporal modeling, reducing memory directions, or directly copying nearest-neighbor actions produces larger degradations, especially at later checkpoints.
  • D. Supplementary Real-World Ablation Results: Joint-space diagnostics show temporally coherent generated commands and trigger pixel prediction around joint variation, gripper switching, or latent visual change.These diagnostics connect executed bimanual actions with transition-critical supervision.
  • D. Supplementary Real-World Ablation Results: Real-world rollout visualizations expose semantic transitions, coordination, disturbances, recovery, and multi-stage dependencies behind the progressive checkpoints.The qualitative sequences complement quantitative metrics by making physical transitions visible across the evaluated suites.

2) Simulation Benchmark Rollout Visualization:

Simulation rollout visualizations show UniMPA maintaining coherent, task-directed behavior across distribution shifts, language-conditioned tasks, and difficult zero-shot execution. The supplementary discussion also identifies asynchronous memory, fixed temporal scales, binary triggering, and single-prototype guidance as current boundaries.

  • 2) Simulation Benchmark Rollout Visualization:: Qualitative simulation rollouts cover LIBERO shifts, VLABench language-conditioned manipulation, and RoboTwin 2.0 Hard zero-shot execution.Together, these visualizations provide evidence across distribution shifts, manipulation categories, and challenging execution settings.
  • 2) Simulation Benchmark Rollout Visualization:: Real-world trigger diagnostics activate pixel prediction when joint-position variation, gripper switching, or latent visual change indicates a critical transition.The trigger evidence concentrates around motion reconfiguration and interaction-sensitive phases.
  • 2) Simulation Benchmark Rollout Visualization:: RoboTwin 2.0 Hard rollouts show continued execution after intermediate deviations rather than restarting the task.Paired instances illustrate consistency under scene and initialization changes, while error-recovery examples show corrective behavior.
  • 2) Simulation Benchmark Rollout Visualization:: New memory entries are available only to subsequent rollouts because current-rollout experience cannot update decisions within the ongoing execution.Memory construction is asynchronous and does not require retraining the policy or memory encoders.
  • 2) Simulation Benchmark Rollout Visualization:: A fixed future offset and action horizon may inadequately represent phases ranging from slow free-space motion to fine-grained contact and recovery.The paper proposes adaptive multi-horizon prediction as a future extension.
  • 2) Simulation Benchmark Rollout Visualization:: Binary triggering and single-prototype guidance simplify computation but may not capture graded uncertainty or genuinely multimodal manipulation trajectories.Future work proposes uncertainty-aware computation and multiple retrieved action hypotheses.
Loading 2609.11875v1…