RxBrain: Embodied Cognition Foundation Model with Joint Language-Visual Reasoning and Imagination

Haotian Liang, Mingkang Chen, Yufei Huang, Yuchun Guo, Xiaomeng Zhu, Xiangli Shi, Kaixuan Wang, Yunxuan Mao, Weijie Zhou, Ling Chen, Shirong Zeng, Yueyu Long, Yuchen Si, Yajuan Zhu, Xingyu Zhou, Minghui Wang, Wanjia He, Xin Yang, Lingzhu Xiang, Zhiqing Liu, Bohan Ma, Xiran Huang, Tianshuo Yang, Zhiheng Liu, Xuantang Xiong, Zisheng Lu, Ping Luo, Yao Mu, Han Hu, Zhengyou Zhang

RxBrain integrates textual reasoning and visual imagination into a unified planning sequence for embodied agents.

How can we unify language-visual reasoning and future-state imagination into a single foundation model for embodied agents?

Embodied agents struggle to bridge the gap between high-level textual instructions and the physical state changes required to execute them, often treating reasoning and visual prediction as disconnected processes. RxBrain solves this by interleaving language and visual imagination within a single planning sequence: language defines the task logic and constraints, while visual generation predicts the resulting physical state at each step. On the RxBrain-Bench evaluation, this coupled approach enables the model to produce coherent, multi-step plans that outperform models treating reasoning and generation as separate capabilities.

Paper Primer

RxBrain uses a modality-aware Mixture-of-Transformers (MoT) architecture to unify language and vision. It treats visual imagination as a flow-matching task in the VAE latent space, allowing the model to "imagine" the physical outcome of a planned action before proceeding to the next step.

The core mechanism is an interleaved generation loop: the model generates a reasoning text segment, then synthesizes the corresponding visual state, and re-encodes that state to inform the next step of the plan. This is like a navigator who sketches the next landmark on a map before deciding which turn to take next, ensuring the plan remains grounded in physical reality.

RxBrain achieves superior performance in joint language-visual planning compared to models that separate reasoning and generation.

Evaluations on the RxBrain-Bench-JointPlan track, which measures the model's ability to maintain consistency between textual subgoals and imagined visual states over long horizons. The model demonstrates successful real-robot performance without requiring large-scale action-data pretraining, validating the effectiveness of its joint planning representation.

Why is it insufficient to just use a standard Vision-Language Model (VLM) for embodied planning?

VLMs typically express their cognitive process in language, leaving the desired physical outcome only implicitly specified. They lack the explicit visual goal imagination required to verify if a planned action will actually result in the intended physical state.

What is the purpose of the RxBrain-Bench benchmark?

It evaluates whether a model can jointly use language and visual imagination to construct coherent plans, rather than measuring understanding or generation abilities in isolation.

RxBrain shifts the paradigm of embodied AI from "reasoning about actions" to "reasoning through coupled language-visual sequences," providing a blueprint for agents that plan by imagining the physical consequences of their decisions.

Introduction to RxBrain

Embodied agents need unified language‑visual reasoning to plan and act.

Embodied cognition demands that an agent connect high‑level task reasoning with the concrete physical states it must achieve. Existing systems either excel at textual reasoning or at visual prediction, but neither alone can specify a complete, executable plan.

The model provides a single, unified representation that simultaneously encodes abstract task logic and concrete world‑state changes.

Vision‑language models (VLMs) excel at interpreting observations and producing textual decisions, yet they leave the desired physical outcome implicit. Conversely, world‑model generators can forecast future frames but lack the causal, constraint‑aware reasoning that language provides. This separation prevents agents from forming executable intentions that are both logically sound and visually grounded.

RxBrain addresses this gap by embedding language reasoning and visual imagination within a single planning sequence. The architecture builds on a unified multimodal Mixture‑of‑Transformers (MoT) that processes language, images, and video tokens together, allowing the model to produce paired textual actions and corresponding visual subgoals at every step.

The core need is joint language‑visual reasoning and imagination so embodied agents can produce coherent, executable plans.

The RxBrain Architecture

RxBrain’s architecture unifies language and vision via a modality‑aware Mixture‑of‑Tokens design.

Building a single model that can both interpret visual observations and synthesize future images forces a design tension: visual tokens must serve a perception role and, later, a generation role, yet the model cannot afford two completely separate vision backbones.

MoT routes each incoming token to the transformer branch that matches its modality, then lets the branches exchange information through a shared self‑attention layer — like a postal hub that first directs letters to department‑specific sorting bins and then runs a global scan so every department can see what the others received.

How does MoT differ from a classic Mixture‑of‑Experts (MoE) layer?

MoE selects a subset of expert feed‑forward networks per token based on a gating score, but all tokens still share the same attention matrix. MoT, by contrast, first partitions tokens into modality‑specific transformer branches, each with its own attention projection, and only then performs a single global attention that mixes the branches. The routing is hard‑wired by modality rather than learned per‑token gating.

Modality labels are assigned: t₁, t₂ → “text”; vᵤ₁, vᵤ₂ → “vision‑understanding”; v𝓰₁, v𝓰₂ → “vision‑generation”.

The routing module sends t₁, t₂ to the Language MoT branch, vᵤ₁, vᵤ₂ to the Vision MoT “understanding” branch, and v𝓰₁, v𝓰₂ to the Vision MoT “generation” expert.

Each branch computes its own self‑attention over its tokens (e.g., the language branch attends only to t₁, t₂).

A global self‑attention layer then mixes all six tokens, allowing, for example, t₁ to attend to vᵤ₁.

Finally, each token passes through its branch‑specific feed‑forward expert, producing modality‑tailored representations.

Routing preserves modality‑specific computation while still granting every token a view of the full multimodal context, which is essential for joint reasoning and generation.

RxBrain applies two complementary attention masks: causal masks for autoregressive text tokens and full (bidirectional) masks within each visual frame, so text respects temporal order while each image can freely exchange information among its own pixels.

**Figure 2.** **Architecture overview of RxBrain.** The model adopts a modality-aware MoT architecture to jointly process visual input tokens, text tokens, and generated vision tokens. Visual tokens are handled by a shared vision backbone with full attention, where understanding and generation tokens share attention projections but are routed to separate FFN experts. Text tokens are processed by the Language Transformer with global causal attention for autoregressive generation. Generated vision tokens are decoded through a VAE decoder. Different token colors indicate visual inputs, text inputs/outputs, and generated visual outputs.

**Figure 3.** Hybrid attention patterns in RxBrain. RxBrain adopts hybrid attention masks for unified multimodal understanding and generation. For text generation (a), input images are encoded with bidirectional visual attention, while text tokens follow causal attention for autoregressive decoding. For multi frames generation (b), input images and text prompts follow the same pattern, while generated frames use bidirectional attention within each frame and causal attention across frames. For interleaved generation (c), text is generated autoregressively before image generation; generated images are then re-encoded as visual input for next steps.

Data Construction: Temporal Annotation

We slice raw embodied videos into short, annotated steps that pair text and visual state changes.

Standard text‑to‑image datasets teach open‑domain synthesis but never show how a robot’s hand moves or objects shift over time. Without such temporal supervision, an embodied model cannot learn to predict intermediate world states or plan multi‑step actions.

It cuts a long, untrimmed robot video into bite‑size planning steps, each anchored by a start frame and an end frame that visibly capture a state change.

$\Phi$ receives $G$ and proposes a coarse step $c_1$ with start index $p_1=1$ (frame $f_1$) and end index $q_1=2$ (frame $f_2$), naming it “pick‑up”.

Refinement samples frames $f_1$–$f_2$ and selects $I^{\text{start}}_1=f_1$ (object still on table) and $I^{\text{end}}_1=f_2$ (gripper holding object).

$\Phi$ proposes a second coarse step $c_2$ with $p_2=3$, $q_2=4$, labeled “place‑down”.

Refinement finds $I^{\text{start}}_2=f_4$ (object in hand) and $I^{\text{end}}_2=f_5$ (object on target surface).

The final segment list is $S_1=\langle f_1,f_2,\text{pick-up},\dots\rangle$, $S_2=\langle f_4,f_5,\text{place-down},\dots\rangle$.

Even with only five frames, the pipeline isolates the exact moments of grasp and release, producing clean start/end anchors that downstream models can learn from.

How does Temporal Segment Annotation differ from generic video segmentation?

Generic segmentation merely partitions a video into homogeneous clips, often ignoring the semantics of an action. Temporal Segment Annotation explicitly requires a visible state change, attaches a textual action label and description, and refines boundaries so that each segment is a self‑contained planning step rather than an arbitrary visual chunk.

**Figure 4.** Data construction overview and statistics. This figure summarizes the verified post-training corpus by scale, source composition, and scene coverage. Panel A reports total video hours across four data categories: Real-Robot Data, UMI Data, Simulation Data, and Egocentric Human Data. Panel B shows the source composition, with source categories in the inner ring and individual datasets in the outer ring. Panel C ranks the fourteen top-level scene categories by segment count.

**Figure 5.** Data construction pipeline. Raw embodied task videos are converted into joint text-visual planning samples through three stages: Temporal Segment Annotation, Quality Verification, and Segment Structuring. The final samples are organized into L0–L3 tasks, covering intra-step world state prediction, step-level planning, subgoal planning, and final-state imagination.

Data Construction: Quality Verification

Verified segments are turned into hierarchical training samples for multi‑scale planning.

Training data harvested from raw videos contain many noisy text‑visual pairs; without a filter the model would learn spurious correlations.

A single MLLM call decides whether a candidate planning segment is clean enough to supervise the model.

Segment 1 (score 5, similarity 0.98) passes all checks and is kept.

Segment 2 (score 3) fails the quality threshold and is rejected.

Segment 3 (score 4, similarity 0.996 > 0.995) is flagged as near‑duplicate and removed.

Segment 4 (score 2) is rejected; its description is also corrected but the segment is still discarded because the quality score is too low.

The verifier removes almost half of low‑quality or redundant segments while preserving the majority of useful data, which explains the 75 % retention figure.

How does this verifier differ from a standard binary classifier that only accepts or rejects a segment?

Beyond a binary decision, the verifier returns a fine‑grained quality score, flags specific failure reasons (semantic, grounding, consistency), and can automatically rewrite an inaccurate textual description, providing interpretability and data‑preserving corrections.

Verified segments are then packed into hierarchical samples that let the model learn at four temporal scales, from a single step to the whole task.

Verified segments are organized into a four‑level hierarchy (L0–L3) so the model can practice both fine‑grained state changes and long‑range goal imagination.

Why not train directly on raw video frames without introducing the L0–L3 hierarchy?

Training on raw frames forces the model to infer both fine‑grained dynamics and high‑level goals simultaneously, which overwhelms the learning signal. The hierarchy isolates each granularity, giving the model focused objectives and enabling joint world‑state prediction and subgoal planning.

**Figure 6.** **Hierarchical data construction pipeline.** Verified planning segments are organized into four levels of supervision. L0 learns the visual state change inside a single planning step, while L1 links consecutive steps for step-level joint planning. L2 groups neighboring steps into subgoals, and L3 connects the initial observation with the final task state.

Training Objectives

Training aligns language, vision, and embodied generation through a single multimodal loss.

Embodied agents often stumble because their language and visual modules are trained in isolation, leading to mismatched representations at inference time. RxBrain resolves this by training all modalities together under a single objective, ensuring that the backbone sees a coherent multimodal signal throughout learning.

The loss ties language, image, and interleaved generation together, so the shared MoT backbone learns a representation that is simultaneously useful for text prediction, visual synthesis, and embodied reasoning.

How does this unified objective differ from training separate text‑only and image‑only models?

Separate models keep their parameters disjoint, so each learns a modality‑specific representation. The unified objective forces a single MoT backbone to serve all modalities, aligning their latent spaces and allowing the model to reason jointly over language and vision during both training and inference.

Compute the weighted text contribution: $0.6 \times 1.2 = 0.72$.

Compute the weighted vision contribution: $0.4 \times 0.8 = 0.32$.

Sum the two contributions to obtain the interleaved loss for this step: $0.72 + 0.32 = 1.04$.

The weighting lets the trainer prioritize the modality that is currently the bottleneck; adjusting $\lambda_t$ and $\lambda_v$ directly trades off language fluency against visual fidelity.

Training proceeds in two curriculum stages. Stage 1 jointly pre‑trains vision and language on massive image‑text pairs and instruction data, establishing a shared representation. Stage 2 fine‑tunes the model on embodied trajectories, interleaving text and visual states so the agent can plan and imagine future frames.

**Table 6.** Stage 1 data mixture. Two buckets mixed at integer ratio 6:4 ($\Sigma$=10); the per-step share is each bucket's ratio divided by 10.

Benchmarking Embodied Capabilities

RxBrain‑Bench evaluates interleaved text‑visual planning for embodied agents.

Embodied agents need a unified representation of task semantics and physical state; RxBrain addresses this by jointly training on language‑visual reasoning and future‑state visual imagination. RxBrain‑Bench is introduced to evaluate the full interleaved text‑visual generation loop.

RxBrain 6B attains an 82.4 score on the Text‑to‑image generation benchmark, the highest among all evaluated models.

Table 2 shows RxBrain 6B achieving 82.4, outperforming the next best model by 3.1 points.

Beyond raw generation, RxBrain‑Bench comprises three tracks—EVQA, WorldPred, and JointPlan—each probing a different facet of the interleaved loop. Tables 2 and 3 present detailed model‑wise comparisons, showing RxBrain’s competitive performance across embodied and spatial reasoning, robot affordance, and multi‑view understanding.

RxBrain‑Bench is a three‑track suite that jointly evaluates language‑based action planning, visual state imagination, and iterative closed‑loop reasoning in embodied scenarios.

**Table 2.** Performance comparison on various benchmarks across different models.

**Table 3.** **RxBrain-Bench-WorldPred: method comparison.** Held-out, short-horizon (4 frames). $S_{gen}$ is the weighted MLLM-judge score; the five columns after it are its judge criteria. Baselines: $Wan2.2-TI2V-5B$ = a general-purpose I2V diffusion model; $Cosmos3-Nano$ = an embodied omni world model (I2V). Both baselines and RxBrain are scored on the same held-out balanced subset with identical 4-frame alignment; $n$ differs due to judge content-filtering.

Empirical Evaluation

RxBrain outperforms baselines on world prediction (Sgen = 0.62) and joint planning (Splan = 0.68).

On RxBrain‑Bench‑WorldPred RxBrain attains $S_{\text{gen}} = 0.62$, a $+0.06$ absolute gain over the strongest baseline.

Baseline Cosmos3‑Nano scores $0.591$ and Wan2.2‑TI2V‑5B scores $0.429$ on the same held‑out set.

On RxBrain‑Bench‑JointPlan RxBrain reaches $S_{\text{plan}} = 0.68$, a $+0.159$ absolute improvement over the best baseline.

Cosmos3‑Nano (agent) scores $0.521$, BAGEL‑7B‑MoT $0.503$, Qwen‑Agent $0.431$ on the same held‑out rollout.

**Figure 7.** **RxBrain-Bench-WorldPred (UMI).** Four-frame future prediction for handheld manipulation. Rows compare the ground truth (GT), RxBrain (HY-Unified), the general-purpose I2V model Wan2.2-TI2V, and the embodied world model Cosmos3. RxBrain preserves the scene layout and reproduces the intended grasping behavior.

**Figure 8.** **RxBrain-Bench-JointPlan (BridgeV2)**, task “move to the right edge of the blue cloth, grasp it, pull it leftward, and release it.” rows are models, columns are planning steps, and each cell pairs the imagined goal frame with the model’s generated sub-step text. RxBrain reproduces the grasp–pull–release sequence and the cloth’s changing shape in step with the ground-truth plan (GT goal).

Extending to Action Generation

Extending RxBrain with a lightweight action branch that reuses pretrained generation knowledge.

Robotic agents typically need a large amount of task‑specific action data, yet RxBrain already encodes rich world‑state and planning priors from multimodal pretraining. The challenge is to expose those priors to an action‑generation head without discarding the pretrained knowledge.

The action branch reuses the pretrained generation feed‑forward expert via a gated addition, letting the model start from a function‑preserving baseline and gradually inject predictive features as training progresses.

Compute $\text{FFN}_a(LN_a(x_a)) = [1,2,3,4]$.

Compute $\text{FFN}_g(LN_g(x_a)) = [1,2,3,4]$.

Scale the generation output by $g$: $g\odot[1,2,3,4]=[0,0,0,0]$.

Fuse: $[1,2,3,4] + [0,0,0,0] = [1,2,3,4]$.

Residual addition: $x^{\text{out}}_{a}=x_a +$ fused $= [2,4,6,8]$.

The gate starts at zero, so the branch behaves exactly like a vanilla action expert; as training proceeds, non‑zero entries in $g$ let the generation expert inject planning signals channel‑wise.

Euler ODE inference for flow‑matching action generation.

**Figure 9.** Action model architecture. A modality-specialized Action MoT branch uses dedicated attention and feed-forward parameters while sharing global self-attention with language, vision, and state tokens. A zero-initialized channel-wise gate transfers features from the pretrained generation expert to the action expert, whose outputs are decoded into action chunks.

We evaluate the action policy on three household manipulation tasks—Set the Table, Fold and Store Glasses, and Pick Trash—using two robot platforms. Each task is run for 100 trials per method.

**Figure 10.** Real-world robot evaluation on three manipulation tasks. We compare our method with $\pi_0$ and $\pi_{0.5}$ on three manipulation tasks. Our method achieves the highest success rate on average.

Related Work

Survey of vision‑language, world, and unified models that shape embodied cognition.

Vision‑language models (VLMs) have progressed from contrastive alignment (e.g., CLIP) to generative hybrids (Flamingo, BLIP‑2, LLaVA) and large‑scale multimodal systems (Qwen‑VL, InternVL, GPT‑4o).

These models now support robust scene understanding, enabling embodied agents to ground perception in physical environments.

Specialized embodied VLMs such as PaLM‑E, RT‑2, RoboBrain, RynnBrain, MolmoAct, and Cosmos‑Reason1 extend this capability to planning, affordance perception, and action reasoning, yet they still output language‑level decisions rather than explicit future visual states.

Parallel to VLMs, generative world models predict environment dynamics, from pixel‑level video synthesis to latent predictive encoders like V‑JEPA and V‑JEPA 2.

World‑action models (VPP, Ctrl‑World, LingBot‑VA, DreamZero) tie future visual prediction to robot control, but they focus on policy learning and lack integrated language‑visual inspection.

Unified multimodal approaches fuse understanding and generation in a single transformer, exemplified by Transfusion, Chameleon, BAGEL, Janus, and Cosmos3.

Cosmos3 demonstrates a Mixture‑of‑Tokens (MoT) architecture that jointly handles visual perception, generation, simulation, and action modeling for embodied agents.

Recent unified action‑generation works (UP‑VLA, MM‑ACT, RynnVLA‑002, Motus, InternVLA‑A1, BagelVLA) integrate visual, linguistic, and motor modalities, often building on the Cosmos3 foundation.

Cosmos3‑Nano is a compact variant of the Cosmos3 foundation model that retains the full Mixture‑of‑Tokens routing while reducing parameter count for faster inference on edge devices.

Wan2.2‑TI2V‑5B couples a 5‑billion‑parameter language core with a vision encoder that produces tokenized image‑to‑image (TI2V) embeddings, enabling bidirectional visual‑language translation.

Contributors and Acknowledgements

Lists project contributors and details the diverse data sources comprising the training corpus.

Project sponsors are Zhengyou Zhang† and Han Hu; project leaders are Yufei Huang and Yuchun Guo.

Core contributors include Haotian Liang∗, Mingkang Chen∗, Xiaomeng Zhu, Xiangli Shi, Kaixuan Wang, Yunxuan Mao, Weijie Zhou, and Ling Chen.

Additional contributors are Shirong Zeng, Yueyu Long, Yuchen Si, Yajuan Zhu, Xingyu Zhou, Minghui Wang, Wanjia He, Xin Yang, Lingzhu Xiang, Zhiqing Liu, Bohan Ma, Xiran Huang, Xuantang Xiong, Zisheng Lu, Tianshuo Yang, and Zhiheng Liu.

We thank Ping Luo and Yao Mu for fruitful discussions.

The final corpus comprises 46 splits totaling 50,177 hours, distributed as 20 Real‑Robot, 1 UMI, 3 Simulation, and 22 Egocentric Human splits.

Real‑Robot data include RealOmni‑Open (9,005.9 h), JAKA (2,558.2 h), Xtrainer (1,515.9 h), RoboMIND (888.0 h), Galaxea Open‑World (536.0 h), MolmoAct2 RT‑1 (365.7 h), LET Base / LejuRobot (318.0 h), RoboCOIN (315.9 h), AgiBot World (278.7 h), DROID (255.6 h), MolmoAct2 SO100/101 (180.3 h), MolmoAct2 BC‑Z (148.9 h), BridgeData V2 (109.1 h), MolmoAct2 Bridge (105.1 h), MolmoAct2 DROID (50.8 h), HumanoidEveryday (31.8 h), CyberOrigin (30.3 h), RH20T robot (596.4 h), MolmoAct2 Tabletop (0.5 h), and MolmoAct2 Household (0.6 h).

UMI data, self‑collected with universal manipulation grippers, contribute 17,505.8 hours after verification.

Simulation sources comprise BEHAVIOR‑1K (1,102.4 h), InternData‑A1 (197.8 h), and SynData (17.2 h).

Egocentric Human data span many collections: Ego4D (3,670.0 h), AI Data curated (3,584.5 h), AI Data raw (1,918.2 h), Egocentric‑10K (1,300.0 h), EgoVerse (1,132.8 h), EgoDex (829.0 h), Xperience (800.0 h), Something‑Something V2 (240.0 h), Ego‑Exo4D (155.1 h), EPIC‑KITCHENS‑100 (91.2 h), EgoMe (43.9 h), CaptainCook4D (39.4 h), HD‑EPIC (38.1 h), HoloAssist (32.2 h), HOT3D (5.5 h), TACO (3.0 h), ARCTIC (2.3 h), HOI4D (0.9 h), HO‑Cap (0.6 h), TASTE‑Rob (2.6 h), RH20T human (98.4 h), and OakInk2 (74.6 h).

Appendix: Data Construction Details

Full notation, derivations, and algorithmic steps for the temporal segment annotation pipeline.

C.1 expands the three‑step Temporal Segment Annotation pipeline with full notation, derivations, and algorithmic detail.

An atomic planning step is a single manipulation whose effect is a clearly visible change of object state.

A segment pins one atomic planning step to two anchor frames that bound the visible state change.

The sampler places anchors only where visual content changes, avoiding wasteful uniform sampling of static intervals.

Bidirectional boundary refinement for one segment $\tilde{S}_i$

Quality Verification inspects each pre‑verified segment $S_i$ using the checklist in Table 5.

Appendix: Per-Scene Taxonomy

Per‑scene activity distribution across the pretraining corpus.

The sunburst grid reveals that Kitchen, Living Room, and Robot Manipulation together account for the bulk of trainable segments, while niche domains such as Science Experiment or Outdoor / Garden contribute only a few thousand examples each.

**Figure.** A series of sunburst charts illustrating the distribution of activities across nine different environments: Kitchen, Office/Desktop, Other/Uncategorized, Living Room, Bedroom, Bathroom, Robot Manipulation, Workshop/Garage, and Retail/Store. Each chart displays the primary environment in the center, with sub-categories radiating outward, labeled with their respective frequency counts.

Appendix: Training and Evaluation Config

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers