LLM-Based Offline Learning for Embodied Agents via Consistency-Guided Reward Ensemble

TL;DR

This work presents a consistency-guided reward ensemble framework (CoREN), designed for tackling difficulties in grounding LLM-generated estimates to the target environment domain and enabling effective offline learning of embodied agents in different environment domains.

Abstract

Employing large language models (LLMs) to enable embodied agents has become popular, yet it presents several limitations in practice. In this work, rather than using LLMs directly as agents, we explore their use as tools for embodied agent learning. Specifically, to train separate agents via offline reinforcement learning (RL), an LLM is used to provide dense reward feedback on individual actions in training datasets. In doing so, we present a consistency-guided reward ensemble framework (CoREN), designed for tackling difficulties in grounding LLM-generated estimates to the target environment domain. The framework employs an adaptive ensemble of spatio-temporally consistent rewards to derive domain-grounded rewards in the training datasets, thus enabling effective offline learning of embodied agents in different environment domains. Experiments with the VirtualHome benchmark demonstrate that CoREN significantly outperforms other offline RL agents, and it also achieves comparable performance to state-of-the-art LLM-based agents with 8B parameters, despite CoREN having only 117M parameters for the agent policy network and using LLMs only for training.

Figures (2)

• Figure 1: $\textsc{CoREn}$, a framework for LLM-based reward estimation and offline learning. In (i), an LLM estimates rewards based on spatio-temporal (i.e., contextual, structural, and temporal) consistencies; In (ii), these rewards are integrated into a single domain-grounded reward via an ensemble. Using the reward-augmented dataset, offline RL can be conducted effectively to achieve embodied agents with resource efficiency and low latency.
• Figure 2: Two-staged reward estimation in $\textsc{CoREn}$. In (i), spatio-temporally consistent rewards, constrained by contextual, structural, and temporal consistencies, are calculated. (a) Contextual consistency is achieved through majority voting across the responses from different prompts $\mathcal{P}_n$, resulting in contextually consistent rewards $r^C$. (b) Structural consistency is achieved by presenting MDP-specific queries to the LLM. If the LLM incorrectly answers these queries (indicated by a red 'X'), the rewards estimated from these particular prompts are removed from majority voting. The successfully verified rewards contribute to structurally consistent rewards $r^S$. (c) Temporal consistency involves collecting high-value actions $H_n(\tau)$ and subjecting them to backward verification through LLM queries. Actions that fail this verification are excluded from the candidates for majority voting. Otherwise, they contribute to temporally consistent rewards $r^T$. In (ii), a trajectory $(i, \tau)$ with success flag $f_s(i, \tau)$ is sampled from the given offline dataset $\mathcal{D}$. The spatio-temporally consistent rewards $(r^C, r^S, r^T)$ in (i) are combined using weights $(w^C, w^S, w^T)$, which are generated by the reward orchestrator $\Psi_\theta$. This combined result renders a unified stepwise, more domain-grounded reward $\hat{r}$. The orchestrator $\Psi_\theta$ is trained to align the trajectory's return of accumulating stepwise rewards $\hat{r}$ with the sparse reward $f_s(i, \tau)$ annotated on the trajectory.