WALL-E: World Alignment by Rule Learning Improves World Model-based LLM Agents

TL;DR

This study reveals that the gaps between the prior knowledge of LLMs and the specified environment's dynamics can be bridged by aligning an LLM with its deployed environment and such"world alignment" can be efficiently achieved by rule learning on LLMs.

Abstract

Can large language models (LLMs) directly serve as powerful world models for model-based agents? While the gaps between the prior knowledge of LLMs and the specified environment's dynamics do exist, our study reveals that the gaps can be bridged by aligning an LLM with its deployed environment and such "world alignment" can be efficiently achieved by rule learning on LLMs. Given the rich prior knowledge of LLMs, only a few additional rules suffice to align LLM predictions with the specified environment dynamics. To this end, we propose a neurosymbolic approach to learn these rules gradient-free through LLMs, by inducing, updating, and pruning rules based on comparisons of agent-explored trajectories and world model predictions. The resulting world model is composed of the LLM and the learned rules. Our embodied LLM agent "WALL-E" is built upon model-predictive control (MPC). By optimizing look-ahead actions based on the precise world model, MPC significantly improves exploration and learning efficiency. Compared to existing LLM agents, WALL-E's reasoning only requires a few principal rules rather than verbose buffered trajectories being included in the LLM input. On open-world challenges in Minecraft and ALFWorld, WALL-E achieves higher success rates than existing methods, with lower costs on replanning time and the number of tokens used for reasoning. In Minecraft, WALL-E exceeds baselines by 15-30% in success rate while costing 8-20 fewer replanning rounds and only 60-80% of tokens. In ALFWorld, its success rate surges to a new record high of 95% only after 6 iterations.

Figures (6)

• Figure 1: Illustration of WALL-E mining a diamond in Minecraft. Step 1-2: the agent makes a plan via MPC with the initial unaligned world model, resulting in a failed action for mining iron ore. Step 3: by comparing real trajectories with the world model predictions, WALL-E learns a critical rule that if the tool is not proper to the material being mined, the action will fail. Step 4-5: the learned rule helps the world model make accurate predictions for transitions that were predicted mistakenly in MPC. Step 6: the agent accordingly modifies its plan and replaces stone pickaxe with an iron pickaxe toward completing the task.
• Figure 2: Overview of WALL-E. The agent's action per step is controlled by MPC, where the agent model plans actions in a look-ahead window based on the LLM+rules world model's predictions.
• Figure 3: Rule Learning details. The rule learning module iteratively refines the rules by comparing the world model predicted trajectories with the agent's actual trajectories in the environment. The rule learning takes five steps: (1) comparing predicted and actual trajectories; (2) learning new rules from real trajectories; (3) refining learned rules; (4) translating natural language rules to code; and (5) rule set pruning via solving a maximum coverage problem. (2)-(4) are handled by LLMs, while (1) and (5) are executed by programs.
• Figure 4: Comparison of WALL-E and baselines on 134 testing tasks from the ALFWorld benchmark.
• Figure 5: Cover rate of LLM failed predictions across different actions over iteration times during training. The cover rate represents the probability that the LLM's failed predictions are addressed by the rules obtained during the rule learning process. The predictions and rules are categorized by action type: craft, mine, gather and fight. The learnt rules at each iteration are displayed in black under each node, labeled with their respective rule IDs.