Subtask Aware End-to-End Learning for Visual Room Rearrangement

The goal of intelligent embodied agents is to learn how to explore within the environment, interact with objects, and understand the environment in order to achieve task objectives. There are two main approaches to training such agents: one is to train an action policy that performs the task goal through end-to-end learning, and the other is to construct a policy by implementing the necessary abilities according to the task goal in a modular manner. For complex and long-horizon tasks, such as visual room rearrangement, a modular approach that infers task sequence by identifying the causality of actions through prior knowledge shows higher performance. Based on this insight, we propose an Online Subtask Prediction Network (OSPNet) that determines the subtask to be performed at each moment based on the environment information and past subtask inference history to train an embodied agent for long-horizon tasks through an end-to-end manner, and also propose a Subtask Aware Policy Network (SAPNet) as the action policy that decides actions based on the reasoning of the OSPNet. We implement an embodied agent that performs visual room rearrangement using the proposed SAPNet and train it through imitation learning, demonstrating similar or better performance with much fewer training steps than previous works.