Embodied CoT Distillation From LLM To Off-the-shelf Agents

TL;DR

The paper tackles embodied task planning on capacity-limited devices by distilling LLM reasoning into small LM policies. It introduces DeDer, a two-tier policy framework (reasoning-policy and planning-policy) guided by an embodied knowledge graph and a contrastively prompted attention, with rationale data generated through MDP-featured in-context learning and LLM self-critique. On the ALFRED benchmark, DeDer outperforms state-of-the-art LLM-based planners and distillation baselines, achieving substantial gains in task success rate and goal-conditioned success while running on off-the-shelf LMs; the improvements are evidenced by an average SR and GC uplift of around 21.6 and 12.3 percentage points, respectively, over a strong baseline. This work demonstrates a practical route to real-time, robust embodied reasoning on modest hardware and offers a scalable blueprint for deploying LLM-like planning through compact models in real-world robotics and embodied AI settings.

Abstract

We address the challenge of utilizing large language models (LLMs) for complex embodied tasks, in the environment where decision-making systems operate timely on capacity-limited, off-the-shelf devices. We present DeDer, a framework for decomposing and distilling the embodied reasoning capabilities from LLMs to efficient, small language model (sLM)-based policies. In DeDer, the decision-making process of LLM-based strategies is restructured into a hierarchy with a reasoning-policy and planning-policy. The reasoning-policy is distilled from the data that is generated through the embodied in-context learning and self-verification of an LLM, so it can produce effective rationales. The planning-policy, guided by the rationales, can render optimized plans efficiently. In turn, DeDer allows for adopting sLMs for both policies, deployed on off-the-shelf devices. Furthermore, to enhance the quality of intermediate rationales, specific to embodied tasks, we devise the embodied knowledge graph, and to generate multiple rationales timely through a single inference, we also use the contrastively prompted attention model. Our experiments with the ALFRED benchmark demonstrate that DeDer surpasses leading language planning and distillation approaches, indicating the applicability and efficiency of sLM-based embodied policies derived through DeDer.

Figures (8)

• Figure 1: $\textsc{DeDer}$ framework with three phases: (i) In rationale dataset construction phase, the MDP-featured in-context learning and self-critic function are employed to extract rationales from the LLM; (ii) In policy distillation phase, the sLM-based policy consisting of reasoning-policy and planning-policy is trained using the extracted rationale data; (iii) In zero-shot deployment, the distilled sLM-based policy is evaluated in unseen environments.
• Figure 2: MDP-featured in-context learning in $\textsc{DeDer}$ for rationale extraction from the LLM: the examples of inputs, queries (in red), and rationales (in blue) for the desired plan are presented, wherein MDP-aligned ones are specifically emphasized.
• Figure 3: Distillation procedures in $\textsc{DeDer}$: During the rationale dataset construction phase, the LLM is iteratively prompted with queries $q_i$ and rationales $r_i$ to refine in-context examples $\mathcal{C}_k$ through retrieval augmentation. The LLM also serves as a critic, evaluating the validity of the extracted rationales $\mathcal{R}$. During the policy distillation phase, the embodied KG containing environment information as well as expert experiences is used as input $g$ to the sLM-based reasoning-policy with the prompted casual attention, which is trained through behavior-based contrastive learning. The structure of reasoning-policy is specifically designed to produce multiple rationales $\mathcal{R}$ in a single-step CoT process through the integration of the prompted attention $\Psi$ and the encoder-decoder architecture. The reasoning-policy is distilled from the embodied KG, which is continually updated from the dataset. Subsequently, the planning-policy $\Phi_\text{p}$ is trained to produce a timely action plan $a$, by immediately using the rationales $\mathcal{R}$ at each step.
• Figure 4: Ablation on each rationale
• Figure A.5: Task examples set within different indoor scenes. The observation includes a variety of objects with which the agent can interact and alter states to complete the given task.