cs.ROOct 20, 2026

UniWALM: Unified World Action Language Model for Embodied Robot Control

Authors: Rahul Sajnani, Diego Romeres, Mingyuan Zhang, Kai Gao, Carl Qi, Baichuan Huang, Srinath Sridhar, Dylan Randle, 

Organizations: Brown University · Amazon · Northeastern University · UT Austin · 

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Motivated by the human ability to integrate experience across language, perception, and action, we introduce UniWALM, a unified world-action-language model that learns video, action, and language in a single shared representation, where each modality is both a conditioning signal and a supervision target. We ask: can jointly generating language enrich this shared representation, improving manipulation and supporting memory-dependent behavior? Unlike prior approaches that attach a reasoning or vision-language model to a video-action backbone and run them as distinct stages, UniWALM unifies all three modalities inside a single mixture-of-transformers backbone whose per-modality diffusion heads co-denoise synchronously, decoding subtask and textual memory in the same pass as action with no separate planning stage; the generated language is re-consumed on the next step, closing a language-as-memory loop without an external planner or memory bank. To keep the action stream reliable under this richer objective, we further introduce Flow Matching with Correction (FMC), a training objective that teaches recovery from perturbed action trajectories at no inference-time cost. UniWALM consistently improves over same-backbone baselines (e.g., FastWAM) on compositional pick-and-place (MolmoSpaces) and across the memory-dependent subset of RMBench, where it reaches 34.5% without any pretrained language or reasoning model and exceeds Mem-0's keyframe-equipped setup (28.5%) using only a single anchor frame. On a real-world Find-the-Item task, generating textual memory lifts success from 11.1% to 61.1%, while state-of-the-art VLAs stay at or below chance (<=33.3%).

Figures & tables

Figure 1: (a) UniWALM is a unified World Action Language Model that extends a single video-model backbone to act and reason by co-generating video, action, and language through one shared diffusion process — no separate language model fused onto a world-action model. This supports (b) single and (c) bimanual dexterous manipulation with (d) multi-stage manipulation and (e) memory-dependent reasoning . Prior work (f) composes a VLM with a VLA and a video model in pipeline or parallel; UniWALM (g) tunes a single unified representation end-to-end. Each manipulation example displays the input task instruction at the bottom and the generated video, action, and subtask at the top. In (e) the predicted subtask and textual memory (shown as ‘ subtask || memory ’) are re-consumed as the next-step prompt.
Figure 2: Overview of UniWALM . UniWALM is a unified mixture-of-transformers (MoT) world-action-language model that co-denoises video, language, and action in a single shared representation. Given the current observation, task instruction, proprioception, an optional visual memory F\mathcal{F} (one or more frames; in our experiments a single anchor at episode start), and the previously generated textual memory , UniWALM processes inputs through a (a) video , (b) language , and (c) action DiT that share features via unified self-attention. The language stream emits a subtask together with a textual memory string, which is re-consumed as the prompt on the next control step, closing a language-as-memory loop across time. Video generation is disabled at inference, so the language and action DiTs run without producing future frames.
Figure 3: Flow Matching with Correction (FMC) . A clean action x0x_{0} (green path) and its perturbed counterpart x0+ϵ1x_{0}+\epsilon_{1} (red path) are noised to the same time tt . The model is trained to transport the perturbed sample back onto the clean trajectory within Δt\Delta t (yellow arrow), while the clean branch follows standard flow matching. This requires no extra conditioning or model capacity and leaves inference unchanged.
Figure 4: Mixture of Transformers Layer . Each modality is processed by its own transformer that performs adaptive layer norm followed by Q,K,V projection. The (a) unified self-attention enables feature sharing between different modalities. Each modality transformer then cross attends to the (b) task instruction and proprioception for conditional generation of video, language, and action in a unified architecture.
Figure 5: Qualitative Results of UniWALM . Our work can perform various tasks including: (a) pick and place tasks, (b) find the item memory task, (c) simple high-dexterous task, and (d) complex dexterous bi-manual task involving difficult items to grasp . For each demonstration, the top row displays the subtasks and memory predicted by the model, and the task instruction is shown below the images. Note, in (b) we make the task difficult by shifting items mid way during roll out for the policy to update its memory. In (d) UniWALM reasons about both the hands and retries to remove objects with varying geometries.
Press Button (target: 4 left, 1 middle, then confirm)
next subtask: left button. memory: visible 4 1 || left button 0/4 middle button 0/1 || left button.next subtask: left button. memory: visible 4 1 || left button 1/4 middle button 0/1 || left button.next subtask: left button. memory: visible 4 1 || left button 2/4 middle button 0/1 || left button.next subtask: left button. memory: visible 4 1 || left button 3/4 middle button 0/1 || left button. next: middle button.next subtask: middle button. memory: visible 4 1 || left button 4/4 middle button 0/1 || middle button. next: confirm.
Battery Try (find working (battery, slot, orientation) pairs)
next subtask: dual positive. memory: left none right none || dual positive. next: left negative.next subtask: dual positive. memory: left none right none || dual positive. next: left negative.next subtask: left negative. memory: left pos right pos || left negative. next: done.next subtask: right negative. memory: left neg right pos || right negative. next: left.next subtask: right negative. memory: left neg right pos || right negative. next: done.
Blocks Ranking Try (find target permutation by trial and press to confirm)
next subtask: press. memory: press 0 swap 0 checking || press. next: swap middle right.next subtask: swap middle right. memory: press 1 swap 0 swapping || swap middle right. next: press.next subtask: press. memory: press 1 swap 1 checking || press. next: swap left right.next subtask: swap left right. memory: press 2 swap 1 swapping || swap left right. next: press.next subtask: press. memory: press 2 swap 1 checking || press. next: done.
Cover Blocks (cover blocks then uncover in red – green – blue order)
Figure 6: RMBench M(n)M(n) qualitative rollouts. Each row is a successful UniWALM rollout on one M(n)M(n) task; we display a subsampled sequence of inference steps for clarity. Below each panel we show the predicted next subtask and the generated textual memory that is re-consumed as the prompt on the next control step; within the memory string, “ || ” separates memory components (task state, progress counters, and the next subtask). Reading left-to-right shows how the language-memory loop accumulates task-relevant state (button counts, battery orientations, tried permutations, cover progress) that memoryless baselines lose across steps.
Press Button, target: 3 left, 4 middle, then confirm.
next subtask: left button. memory: visible 3 4 || left 0/3 mid 0/4 || left button.next subtask: left button. memory: visible 3 4 || left 2/3 mid 0/4 || left button. next: middle.next subtask: middle button. memory: visible 3 4 || left 3/3 mid 0/4 || middle button.next subtask: middle button. memory: visible 3 4 || left 3/3 mid 1/4 || middle button.next subtask: middle button. memory: visible 3 4 || left 3/3 mid 3/4 || middle button. next: confirm.next subtask: confirm. memory: visible 3 4 || left 3/3 mid 4/4 || confirm. next: done.
Figure 7: Language-as-memory loop on Press Button. At each step UniWALM emits a next subtask and textual memory that is re-consumed as the next-step prompt; press counters transition cleanly (left 0/3→3/30/3\to 3/3 ; middle 1/4→4/41/4\to 4/4 ) until confirm. Intermediate frames skipped for clarity; see Fig. 6 for all four M(n)M(n) tasks.
Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Empty Container (real, bi-manual) rollouts. Each row is a successful UniWALM rollout. The subtask under each frame is generated by the language DiT in tandem with the action (co-denoised in the same forward pass).
Figure 9: Real-world Find the Item rollouts. Row 1: black box ; Row 2: white tissue box ; Row 3: green parmesan package (ends at the cover-complete moment). Items are hidden mid-episode under cardboard covers; the subtask and textual memory under each frame are generated by the language DiT in tandem with the action , and the memory string is re-consumed as the next-step prompt to preserve the item–region (L/M/R) mapping across the occlusion.
Figure 10: 2D robot-trajectory benchmark used for FMC ablation. Each shape class contains multiple noisy demonstrations sharing a canonical motion pattern (blue curves show ground-truth trajectories). Shapes span repeating closed motion (Square, Figure-8), unbounded expansion (Spiral), and multi-modal starts (Line’s radial variants). The policy is trained to generate trajectories matching the full distribution of each class.
Figure 11: FMC vs vanilla flow matching on 2D robot trajectories (overlay view). Blue curves are ground-truth demonstrations; red are model-generated trajectories, overlaid for direct comparison. Top row: vanilla flow matching frequently drifts off the demonstration manifold, truncates the shape, or collapses mode coverage. Bottom row: FMC tracks the ground-truth distribution more faithfully across all three shape classes, consistent with the quantitative Chamfer improvements in Tab. 6 .
Failure case: Press Button (target: 5 left, 2 middle, 1 confirm; actual: 5- 3 -1)
next subtask: left button. memory: visible 5 2 || left button 0/5 middle button 0/2 || left button.next subtask: middle button. memory: visible 5 2 || left button 5/5 middle button 0/2 || middle button.next subtask: middle button. memory: visible 5 2 || left button 5/5 middle button 1/2 || middle button. next: confirm.next subtask: confirm. memory: visible 5 2 || left button 5/5 middle button 2/2 || confirm. next: done.next subtask: confirm. memory: visible 5 2 || left button 5/5 middle button 2/2 || confirm. next: done.
Appendix
Figure 12: Timing-lag failure on Press Button. Target counts are (left=5,middle=2,confirm=1)(\text{left}=5,\text{middle}=2,\text{confirm}=1) . The memory readout tracks the intended count correctly ( left 5/5, middle 2/2, confirm, next: done ), but the action executor over-shoots by one press on middle, yielding 3 physical middle presses instead of 2. The language-memory loop is correct; the task-completion channel fires one step late, so the action expert continues the current subtask for one extra press.
Figure 13: MolmoSpaces Pick & Place rollouts (sim). Each row is a successful UniWALM rollout across diverse scenes. The subtask under each frame is generated by the language DiT in tandem with the action (co-denoised in the same forward pass), not produced autoregressively before action.
Figure 14: MolmoSpaces Pick & Place rollouts (sim). Each row is a successful UniWALM rollout across diverse scenes. The subtask under each frame is generated by the language DiT in tandem with the action (co-denoised in the same forward pass), not produced autoregressively before action.
More in this topic

Vision-Language-Action Models

Vision-language-action models turn visual observations and natural-language instructions into actions for robots and other embodied agents.

See all 1,171 papers

Explore similar work

CardsList

  1. UniWAM: Unified World-Action Model

    Oct 1, 2026Wenxuan Song, Jiayi Chen, Jingbo Wang +13Vision-Language-Action ModelsWorld Action Models

  2. World-Language-Action Model for Unified World Modeling, Language Reasoning, and Action Synthesis

    Jun 4, 2026Yi Yang, Zhihong Liu, Siqi Kou +9World Models for RoboticsRobot Foundation Models

    SJTU · SII · HUST +4

  3. Qwen-VLA: Unifying Vision-Language-Action Modeling across Tasks, Environments, and Robot Embodiments

    May 28, 2026Qiuyue Wang, Mingsheng Li, Jian Guan +37Language-Conditioned Robot ManipulationRobot Foundation Models

    Qwen Team