Knowing More, Acting Better: Hierarchical Representation for Embodied Decision-Making (2025)

Abstract: Modern embodied AI uses multimodal large language models (MLLMs) as policy models, predicting actions from final-layer hidden states. This widely adopted approach, however, assumes that monolithic last-layer representations suffice for decision-making—a structural simplification at odds with decades of cognitive science, which highlights the importance of distributed, hierarchical processing for perception and action. Addressing this foundational asymmetry, we introduce a hierarchical action probing method that explicitly aggregates representations from all layers, mirroring the brain’s multi-level organization. Experiments reveal that early layers facilitate spatial grounding, middle layers support contextual integration, and later layers enable abstract generalization—which shows MLLMs inherently encode distributed action-relevant structures. These layer-wise features are integrated by a lightweight probe for spatial reasoning and contextual understanding, without costly backbone fine-tuning. This hierarchical solution shows significant improvements over standard last-layer embodied models in physical simulators, achieving a 46.6% success rate and a 62.5% gain in spatial reasoning tasks. These findings challenge conventional assumptions in embodied AI, establishing hierarchical probing as a principled alternative grounded in both cognitive theory and empirical evidence.

## introduction in embodied ai, multimodal large language models (mllms) enable agents to perceive visuallanguage inputs and generate actionable outputs @xcite @xcite . their current application in embodied tasks predominantly relies on single-layer representations-specifically, the final hidden state-for action prediction, mirroring traditional language generative modeling. however, this reliance on monolithic, final-layer representations presents a structural asymmetry in embodied ai: this largely unchallenged convention contrasts sharply with decades of cognitive science emphasizing hierarchical, distributed perception @xcite . this simplification, while perhaps adequate for static language tasks, overlooks rich intermediate-layer features in embodied multimodal tasks. our work challenges the common assumption that final-layer outputs are sufficient for complex embodied policies. we hypothesize that mllm layers inherently specialize (e.g., early for spatial relations, middle for contextual dependencies, later for abstract semantics), necessitating stratified grounding. discarding this internal structure potentially creates a bottleneck in aligning with the cognitive complexity of embodied environments. indeed, even in pure language tasks, leveraging multiple layers can improve performance @xcite @xcite , yet many embodied models @xcite still predominantly use final-layer outputs. to explore this, we introduce hierarchical action probing (hap) as a conceptual pivot. to the best of our knowledge, this is the first attempt to employ probing techniques-originally for model interpretability-to uncover the phenomenon of how distributed hierarchical representations, which are fundamental to human cognition @xcite , also naturally emerge and can be harnessed in embodied policy models. in particular, hap recenters the internal layer-wise states of mllm-based policy models for grounding behavior, using the full representational capacity. rather than simply combining multiple layers as in prior nlp work, our novel contribution lies in improving embodied policy models through cognitive insights. we show that the hierarchical distributed representation is crucial in both human cognition and multimodal embodied tasks. as a direct consequence of this insight, our technical approach consists of (i) a visual encoder, (ii) the mllm backbone generating layer-wise states, and (iii) a multi-layer action probe hierarchically aggregating these states. this hierarchy allows policy models to capture fine-grained spatial relationships and high-level task semantics, mirroring the brain's distributed processing. the methodological improvements are valuable outcomes stemming from this cognitive perspective. we examine hap's effectiveness on language-guided rearrangement tasks in the physical simulators. experiments show hap achieves strong performance, with our 7b model matching llarp-7b's 42% success rate and our 13b model reaching 46.6%, surpassing llarp-13b's 46%. hap particularly excels in spatial reasoning (13% vs. 8% for llarp-7b). detailed ablation studies further reveal layer-wise specializations-intermediate layers for spatial understanding, later layers for object affordances-empirically supporting our hierarchical design principles and the benefits of probing these internal structures. ## related work recent research has explored using large pretrained mllms @xcite @xcite @xcite @xcite @xcite @xcite as policies for embodied tasks, often termed vision-languageaction (vla) models. approaches like rt-2-x @xcite and @xcite scale to billion-parameter policies, leveraging large datasets like open x-embodiment for training. other works, such as @xcite @xcite , focus on single-robot or simulated setups, while @xcite lack efficient fine-tuning for new configurations. parameter-efficient techniques, such as llarp @xcite and @xcite , address computational challenges by fine-tuning mllms, enabling generalization to novel tasks. in contrast to existing vla models that rely on monolithic last-layer representations, we show the role of hierarchical representations in embodied policy models-mirroring the human brain's hierarchical and distributed cognitive processes. layer n layer 1 … task instruction: "fetch a glass of milk and put it on the table" text tokenizer visual encoder action predictor visual observation 𝑂 ! visual observation 𝑂 " ## methodology the embodied task is formulated as a partially observable markov decision process (pomdp), defined by (s, o, a, p, r, ρ 0 , γ), where s is the state space, o the observation space (e.g., egocentric rgb images), a the action space, p transition dynamics, r the reward function, ρ 0 the initial state distribution, and γ the discount factor. the agent, given a natural language goal @xmath0, generate actions a @xmath1, g). the objective is to learn a goal-conditioned policy π(a t |o 1 , . . . , o t , g). while pomdps formally use belief states, conditioning on history o 1 , . . . , o t is a practical approach with sequence models like lms, which implicitly summarize history to approximate the belief state, addressing partial observability. this policy model can generalize to novel goal distributions g ′ . ## experiments we assess our approach on challenging languageguided rearrangement tasks within the habitat 2.0 simulator @xcite . for fair comparisons, we follow the usage of frozen visualcortex vit @xcite and llama-2 (touvron et al., 2023)/qwen2.5-vl (qwen team, 2025) backbones from @xcite , training only a lightweight hierarchical probe using ppo; baselines include: 1) zs-chatgpt, iteratively refining actions with environment feedback; 2) zs-llama, generating a single plan with llm @xcite ; 3) zs-flamingo, a multimodal zero-shot planner using idefics @xcite ; 4) llarp-scratch, a 7b-llm variant trained from scratch; 5) lstm-flan, combining flan-t5 @xcite with an lstm; and 6) lstm-llama, using llama with a perceiver resampler @xcite . further details on simulators, baselines, and setups are in app. a. ## conclusion inspired by cognitive science, this work challenges prior mllm's use in embodied ai. our hierarchical action probing, by engaging internal mllm structures beyond just final layers, improves pol-icy models-notably in spatial reasoning and contextual understanding. we also reveal mllms' inherent, human brain-like layer-wise specializations, affirming distributed hierarchical processing as pivotal for more capable embodied agents.