IQL-TD-MPC: Implicit Q-Learning for Hierarchical Model Predictive Control

• Published in: 2024 IEEE International Conference on Robotics and Automation (ICRA) pp. 9154 - 9160
• Main Authors: Chitnis, Rohan, Xu, Yingchen, Hashemi, Bobak, Lehnert, Lucas, Dogan, Urun, Zhu, Zheqing, Delalleau, Olivier
• Format: Conference Proceeding
• Language: English
• Published: IEEE 13.05.2024
• Subjects: Benchmark testing; Planning; Prediction algorithms; Q-learning; Robotics and automation; Task analysis; Temporal difference learning
• DOI: 10.1109/ICRA57147.2024.10611711

Model-based reinforcement learning (RL) has shown great promise due to its sample efficiency, but still struggles with long-horizon sparse-reward tasks, especially in offline settings where the agent learns from a fixed dataset. We hypothesize that model-based RL agents struggle in these environments due to a lack of long-term planning capabilities, and that planning in a temporally abstract model of the environment can alleviate this issue. In this paper, we make two key contributions: 1) we introduce an offline model-based RL algorithm, IQL-TD-MPC, that extends the state- of-the-art Temporal Difference Learning for Model Predictive Control (TD-MPC) with Implicit Q-Learning (IQL); and 2) we propose to use IQL-TD-MPC as a Manager in a hierarchical setting with any off-the-shelf offline RL algorithm as a Worker. More specifically, we pre-train a temporally abstract IQL-TD-MPC Manager to predict "intent embeddings", which roughly correspond to subgoals, via planning. We show that augmenting state representations with intent embeddings generated by an IQL-TD-MPC manager significantly improves off-the-shelf offline RL agents' performance on some of the most challenging D4RL benchmark tasks. For instance, the offline RL algorithms AWAC, TD3-BC, DT, and CQL all get zero or near-zero normalized evaluation scores on the medium and large antmaze tasks, while our modification gives an average score over 40.