UniWALM: Unified World Action Language Model for Embodied Robot Control

What it claims

UniWALM, from Brown University, Amazon and collaborators (arXiv, October 2026), adds language generation to a video-action diffusion model. Its language and action experts are trained from scratch on a video-only backbone, with no pretrained language or reasoning model. The baseline it builds on is FastWAM, which the paper describes as a pretrained Wan2.2 video DiT coupled to an action DiT. Against that same-backbone baseline, UniWALM gains +14.8% on average across three MolmoSpaces compositional pick-and-place variants. On RMBench's memory-dependent M(n) subset it gains +21.2 points (13.3% to 34.5%).

How it works

The model is a mixture of transformers with three DiTs (video, action, language) that share features through unified self-attention in every layer. The action and language DiTs copy the video DiT's depth of 30 layers at a smaller hidden dimension and start from random weights. Inputs are the current RGB views (2 in MolmoSpaces, 3 in RMBench and real-world), proprioception, the T5-encoded task prompt, an optional visual memory and the previous textual memory.

• Language head: discrete language diffusion. Ground-truth text is randomly token-replaced according to the timestep, and a cross-entropy loss trains the head to recover it. It emits "subtask || textual memory" in the same denoising pass as the action chunk.
• Loss: L1 flow-matching loss on video velocity, plus the same on action velocity, plus α times the language loss. The authors found L1 more stable than L2 for actions. Subsets of the experts can be trained depending on which data is available.
• Task-completion signal β: a 1-D value attached to each action that fires at subtask completion. Textual memory only advances when β fires, so intermediate reasoning does not overwrite the progress record.
• Visual memory: one or more keyframes with RoPE-shifted tokens. All experiments use a single anchor frame from episode start.
• FMC: during training, action trajectories are perturbed and the model is supervised on a corrective velocity back toward the clean path. It adds no parameters and no sampling steps.

Video generation is disabled at inference. Only the language and action DiTs produce outputs.

flowchart LR
  O[Observation + proprio] --> M[Shared MoT]
  P[Prompt + textual memory] --> M
  F[Anchor frame] --> M
  M --> A[Action chunk + β]
  M --> L[Subtask || memory]
  A -->|β fires| U{Update memory?}
  L --> U
  U -->|yes| P
  U -->|no: keep old| P

Evidence

• RMBench M(n): 34.5% for UniWALM against 13.3% for FastWAM. Mem-0 scores 28.5%, but it uses per-task privileged keyframes, while UniWALM uses one anchor frame.
• Real-world bimanual Find-the-Item: generating textual memory lifts success from 11.1% (FastWAM, same backbone, no language) to 61.1%. The state-of-the-art VLAs tested stay at or below chance, which is 33.3%.
• MolmoSpaces: UniWALM has the highest success on all three compositional pick-and-place variants.

Limits

The fetched text does not give a total parameter count, training compute, a code or weights release, or a licence. Only a project page is linked (rahulsajnani.github.io/projects/uniwalm/). The authors state that textual memory alone struggles with scene state that is hard to put into words, such as precise locations or occluded layouts. The visual memory exists to cover that gap. Training the language head needs subtask and memory strings as ground truth, and the page excerpt does not say how these labels are produced. The text cuts off in Section 3.3, so the FMC ablation and the per-variant tables are not shown here.