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

TL;DR

RxBrain (Hy-Embodied-RxBrain-1.0) is a 6.2B-parameter unified multimodal foundation model for embodied planning that couples language reasoning with visual imagination in a single autoregressive sequence: a learned <Image> token decides when to emit a flow-matched imagined frame between text reasoning steps. It targets three embodied capabilities in one model — VQA/CoT over multi-frame observations, near-future world-state prediction, and joint subgoal planning where each step emits both a language action and the goal image it should reach. The architecture is a modality-routed Mixture-of-Transformers with shared vision attention but separate understanding/generation FFN experts, built on HY-Embodied-0.5 and a FLUX VAE. Training uses an automatic pipeline that decomposes 50,170+ hours of embodied video (real-robot, UMI, simulation, egocentric human) into 21.5M joint text-visual planning segments. Weights and code are released under Apache-2.0.

Key claims

RxBrain organizes textual reasoning and visual imagination into one interleaved planning sequence <assistant> r_1 <Image> z_{1,1:N} </Image> r_2 <Image> ... <eos>, where the <Image> token is a learned modality-transition marker deciding when to emit a flow-matched latent frame [§2.2, Fig. 3c].
The backbone is a modality-aware Mixture-of-Transformers with a language branch (~1.5B), a vision branch (~1.5B), and a widened visual-generation FFN expert (~2.4B) whose intermediate dim is expanded from 6,144 to 12,288; understanding and generation tokens share vision QKV projections but route to separate FFNs [§2.1].
Hybrid attention: causal attention across the multimodal sequence for autoregressive text, plus an additional intra-image bidirectional self-attention pass that replaces causal outputs at visual-token positions, giving each image global self-mixing while preserving cross-image causality [§2.2].
The training corpus spans four categories — Real-Robot (17,292 h), UMI (17,506 h), Simulation (1,317 h), and Egocentric Human (14,062 h) — for a total of 50,177.2 hours across 46 source splits, yielding 21,506,919 trainable segments at a 75.18% verification pass rate from ~28.61M candidates [§3, Fig. 4].
Data construction is a three-stage automatic pipeline (Temporal Segment Annotation → Quality Verification → Segment Structuring) that decomposes untrimmed videos into localized planning steps S_i = (I_start, I_end, action_name, description, metadata) and produces four post-training tasks: continuous prediction, planning, high-level planning, and final-state imagining [§3, §3.1].
Multi-frame world-state prediction generates 4 future VAE-latent frames in parallel under a flow-matching objective, with an internal attention mask preserving causal structure across frames while each frame is denoised in parallel [§2.2].
RxBrain-Bench is introduced to evaluate joint textual-visual embodied planning rather than measuring understanding or generation in isolation [§5].
The model extends to continuous robot action generation and shows real-robot performance without large-scale action-data pretraining [§7].
Reuses HY-ViT 2.0 as the single vision encoder (supervised in the VAE latent space during HY-Embodied-0.5 pretraining) and a frozen FLUX VAE (83.8M) for visual generation tokens [§2.1].

Method

RxBrain is a Mixture-of-Transformers that shares a global attention across modalities but routes tokens to modality-specific pathways. Visual understanding tokens (from HY-ViT 2.0 on input images) and visual generation tokens (VAE latents from a frozen FLUX VAE) share the same vision transformer and QKV projections; a modality-conditioned router dispatches them to separate FFN experts, with the generation-side FFN widened to 2× the standard intermediate dimension. Text tokens live in a parallel language-transformer branch with causal attention for autoregressive next-token prediction. On top of the standard causal attention, an additional intra-image bidirectional attention pass over each image’s visual tokens replaces the causal outputs at those positions, giving each image global mixing while keeping cross-image causality intact.

Three generation modes share this backbone (Fig. 3): (a) text generation (autoregressive over multimodal prefix), (b) parallel 4-frame world-state prediction in the VAE latent space under a flow-matching objective with masked cross-frame attention preserving temporal causality, and (c) interleaved generation — the core capability — where reasoning text and imagined frames alternate under a single left-to-right decoder, with the <Image> learned token triggering flow-matched latent synthesis conditioned on the preceding multimodal context. After each frame is synthesized, it is re-encoded through the ViT visual encoder and appended to the context before the next reasoning step. Training supervision comes from a scalable video-to-plan pipeline that turns embodied videos into (start-frame, end-frame, action-name, description) segments at multiple granularities, filtered by an MLLM-based quality gate.

Results

Total corpus: 50,177.2 hours across 46 source splits → 21,506,919 verified trainable segments from ~28.61M candidates at 75.18% pass rate [§3, Fig. 4].
Reports maintained embodied understanding and generation across “a broad set of embodied and multimodal benchmarks,” with RxBrain-Bench specifically demonstrating coupled textual-reasoning + world-state-prediction + subgoal-planning outputs (specific numbers in §6.2) [§6].
Continuous action-generation extension shows “encouraging real-robot performance without relying on large-scale action-data pretraining” [§7.2].
Model size: 6.2B total (text ~1.5B, visual-understanding ~1.5B, widened visual-generation expert ~2.4B, SigLIP tower ~0.45B, tied embeddings ~0.25B, flow-matching head ~7M) [§2.1].

Why it’s interesting

Sits at the intersection of three threads active on the wiki: a unified-multimodal architecture in the Emerging Properties in Unified Multimodal Pretraining and Context Unrolling in Omni Models lineage, a VLA that generates actions from vision + language like Xiaomi-Robotics-U0: Unified Embodied Synthesis with World Foundation Model and the From Foundation to Application: Improving VLA Models in Practice (LingBot-VLA 2.0) LingBot-VLA family, and an interleaved thinking-with-images formulation directly comparable to InterleaveThinker: Reinforcing Agentic Interleaved Generation. The specific bet — that embodied plans should be represented as (reasoning_text, goal_image) pairs per step rather than a language-only chain-of-thought or a language-only VLA-plus-world-model stack — is a sharper claim than most unified embodied models make; RxBrain-Bench being explicitly designed to measure that coupling is the useful contribution beyond another checkpoint drop. The MoT with shared vision-attention but separate understanding/generation FFN experts also matches the Mamoda2.5: Enhancing Unified Multimodal Model with DiT-MoE DiT-MoE direction without going full expert-per-modality.