Overview

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 serves as both a conditioning signal and a supervision target.

We ask: can jointly generating language enrich the shared representation of a video-action model, yielding a unified representation that improves manipulation and supports memory-dependent behavior?

Unlike prior approaches that attach a separate 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 with per-modality diffusion heads that denoise synchronously. Subtask and textual-memory tokens are decoded in the same denoising pass as action — no separate planning stage — and action / language generation at inference does not require synthesizing future video.

Architecture

UniWALM is built on a mixture-of-transformers backbone with three expert streams that share co-attention:

1. Video expert — Wan2.2 video diffusion over future frames.
2. Action expert — a Diffusion Transformer over action chunks.
3. Language expert — a discrete diffusion head over subtask / memory tokens.

All three streams co-denoise in parallel at every timestep. Action and language can be sampled at inference without producing a full future video.

Language-as-Memory Loop

At every control step, the policy generates a short natural-language memory describing its own progress — e.g., which subtask has been completed, which object was just picked up. On the next step, the policy re-reads its own generated memory as conditioning.

This closes a tight generate → remember → act loop inside a single network, without an external planner, VLM, or memory bank. The result is memory-dependent behavior that emerges as a natural consequence of the unified formulation.

Flow Matching with Correction (FMC)

To improve action reliability, we introduce Flow Matching with Correction (FMC) — a training objective that teaches recovery from perturbed action trajectories. The perturbation is applied only during training, so there is no additional inference cost.

On a 2D toy benchmark of robot end-effector trajectories (circle, star, spiral, line, arc, figure-8, square), FMC reduces unconditional Chamfer distance by 4.0% and wins on 5/7 shape classes (71.4% win rate) over the linear flow-matching baseline.

Qualitative Results

We evaluate UniWALM across real-world and simulated manipulation tasks: single-arm, dexterous bimanual, dexterous bimanual sequential, and reasoning-dependent (Find-the-Item).

Rollouts across tasks

RMBench — per-task rollouts

RMBench — rollout videos

Quantitative Results

MolmoSpaces (simulation)

UniWALM achieves the highest success rate on all three compositional pick-and-place variants, improving over the closest world-action model (FastWAM) sharing the same Wan2.2 backbone by +14.8% on average.

RMBench — Memory-Dependent M(N)

Without any pretrained language or reasoning model, UniWALM reaches 34.5% on RMBench's memory-dependent M(N) subset — improving over same-backbone FastWAM (13.3%) by +21.2 points and exceeding Mem-0's keyframe-equipped setup (28.5%) using only a single episode-start anchor frame rather than per-task privileged keyframes.

Real-World Find-the-Item

On a real-world Find-the-Item task designed to require memory, generating textual memory lifts success from 11.1% (FastWAM, which shares our video backbone but has no language generation) to 61.1%. State-of-the-art VLAs stay at or below chance (≤ 33.3%).

Together, these results answer our question affirmatively: a single model generating language, video, and action can improve control, and memory-dependent behavior emerges as a natural consequence of the unified formulation.

Citation

Acknowledgements & Contact

Work done during Rahul's internship at Amazon Robotics.

Questions? Reach out at rahul_sajnani@brown.edu.