Reward-free World Models for Online Imitation Learning
This work addresses the problem of enabling agents to learn from expert demonstrations in complex environments for robotics and control applications, representing an incremental improvement over prior online IL methods.
The paper tackles the challenge of online imitation learning in complex tasks with high-dimensional inputs and dynamics by proposing a reward-free world model approach that learns latent dynamics without reconstruction and uses inverse soft-Q learning for stable optimization. The method achieves expert-level performance across benchmarks like DMControl, MyoSuite, and ManiSkill2, outperforming existing approaches.
Imitation learning (IL) enables agents to acquire skills directly from expert demonstrations, providing a compelling alternative to reinforcement learning. However, prior online IL approaches struggle with complex tasks characterized by high-dimensional inputs and complex dynamics. In this work, we propose a novel approach to online imitation learning that leverages reward-free world models. Our method learns environmental dynamics entirely in latent spaces without reconstruction, enabling efficient and accurate modeling. We adopt the inverse soft-Q learning objective, reformulating the optimization process in the Q-policy space to mitigate the instability associated with traditional optimization in the reward-policy space. By employing a learned latent dynamics model and planning for control, our approach consistently achieves stable, expert-level performance in tasks with high-dimensional observation or action spaces and intricate dynamics. We evaluate our method on a diverse set of benchmarks, including DMControl, MyoSuite, and ManiSkill2, demonstrating superior empirical performance compared to existing approaches.