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

🕒 Published (v1): 2026-07-15 00:00 UTC · Source: HuggingFace · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

RxBrain (Hy-Embodied-RxBrain) is a 6.2B-parameter embodied cognition foundation model that produces a single interleaved planning sequence where language tokens specify task decomposition and constraints while generated image tokens depict the resulting physical world states. Unlike VLMs that reason purely in text and world models that generate video without causal structure, RxBrain couples both within one autoregressive pass. An automatic pipeline converts 50,177 hours of embodied video into joint text-visual supervision to train this capability.

Problem

Existing VLMs express embodied plans in language alone, leaving goal world states only implicitly specified, while generative world models predict visual futures but lack explicit causal reasoning and constraint-aware decision-making. No prior model organizes textual reasoning and visual goal imagination as complementary components within the same planning representation; even recent unified frameworks (e.g., Cosmos 3) treat understanding and generation as separate capabilities rather than a coupled cognitive process.

Method

Architecture. RxBrain uses a modality-aware Mixture-of-Transformers (MoT) with two specialized branches: a Language Transformer (~1.5B parameters, causal attention) and a Vision Transformer (~1.5B for understanding). For generation, a dedicated visual-generation FFN expert with widened intermediate dimension (6,144 → 12,288) contributes ~2.4B parameters and shares QKV projections with the understanding branch, routing via modality-conditioned gating. HY-ViT 2.0 extracts both semantic and VAE-aligned fine-grained features from visual inputs; visual generation decodes in the FLUX VAE latent space via flow-matching.

Hybrid attention. Visual tokens receive intra-image bidirectional attention (replacing causal outputs at visual positions) while the overall multimodal sequence remains causal, preserving temporal order across steps.

Three generation modes. - Text generation: standard causal next-token prediction over the multimodal prefix. - Multi-frame generation: 4 future frames predicted in parallel via flow-matching in VAE latent space, with frame-wise bidirectional attention and causal cross-frame masking. - Interleaved generation: planning trajectory \(r_1 \langle\text{Image}\rangle z_{1,1:N}\langle/\text{Image}\rangle \; r_2 \langle\text{Image}\rangle z_{2,1:N}\langle/\text{Image}\rangle \cdots\) where \(r_i\) is autoregressive text and \(z_{i,1:N}\) is flow-matched image latent conditioned on full context; generated frames are re-encoded by ViT and appended before the next text step.

Data pipeline. Raw embodied videos (real-robot, UMI, simulation, egocentric human) are processed in three stages: (1) Temporal Segment Annotation — an MLLM \(\Phi\) proposes planning steps with start/end frame anchors capturing visual state transitions; (2) Quality Verification — filtering at ~75.18% pass rate yields 21,506,919 segments from ~28.61M candidates; (3) Segment Structuring — segments organized into four supervision tasks at varying granularity (continuous prediction, planning, high-level planning, final-state imagining). Total corpus: 50,177.2 hours across 46 source splits.

Action extension. RxBrain is extended to continuous robot action generation by appending an action head; deployed on real hardware without large-scale action-data pretraining.

Key Contributions

• Joint language-visual planning formulation in a single autoregressive sequence where language and image tokens play complementary, not redundant, roles.
• Unified MoT architecture (6.2B) supporting text, image, and video understanding/generation within one model via shared visual attention with separate FFN experts.
• Automatic video-to-supervision pipeline producing multi-granularity joint text-visual planning data from 50K+ hours of embodied video.
• RxBrain-Bench: a new benchmark evaluating whether models jointly use language and visual imagination for task understanding, world-state prediction, and coherent plan construction, rather than measuring understanding and generation separately.
• Real-robot action generation without large-scale robot action pretraining, as a downstream extension of the foundation model.

Results

Results section (§6) is not included in the provided text excerpt. The abstract and introduction state: - RxBrain produces joint planning outputs on RxBrain-Bench where language specifies actions/constraints and images depict intermediate/final world states. - "Promising real-robot performance" for continuous action generation without large-scale action-data pretraining. - Model maintains competitive embodied understanding and generation performance on "a broad set of embodied and multimodal benchmarks" (specific numbers not available in the provided excerpt).

Limitations

• No quantitative benchmark numbers are available in the provided text; claims of "promising" performance are not numerically substantiated.
• Data pipeline relies heavily on an MLLM \(\Phi\) for annotation quality, which may introduce systematic errors or biases in segmentation.
• Training requires ~50K hours of diverse embodied video, which is expensive to collect and process.
• Authors describe results as "an initial step" — the architecture is not yet demonstrated to scale or generalize robustly.
• The 75.18% pass rate implies ~25% of candidate segments are discarded; quality of the retained supervision is pipeline-dependent.
• Action generation extension evaluated only on real-robot; no ablation of which planning representation components contribute to downstream task success.

Relevance to Vision-Language Models

RxBrain directly extends the VLM paradigm by replacing text-only planning outputs with interleaved text-image sequences, demonstrating that visual generation (not just visual understanding) can be embedded into the VLM reasoning chain as a first-class output. This challenges the standard VLM assumption that language is the sole medium for expressing goals and intermediate plans, and connects to a growing line of work on unified understanding-generation models (e.g., BAGEL, \(\pi_{0.7}\)). For VLM researchers, the MoT architecture — particularly the sharing of QKV projections between visual understanding and generation branches with separate FFN experts — offers a concrete design pattern for adding generative visual imagination to existing VLM backbones without full architectural replacement. The RxBrain-Bench benchmark also provides a new evaluation surface beyond standard VQA and scene-understanding tasks, targeting the complementarity of modalities rather than isolated capabilities.