TREND: Task-Oriented World Models for Visual Robotic Manipulation

Abstract

World models, the key component of model-based reinforcement learning (MBRL), enable sample-efficient learning by modeling the environment. However, for challenging visual robotic manipulation tasks with continuous action spaces, standard image reconstruction and uniform sampling schemes are suboptimal and may even impede learning. To address this, we propose TREND, a task-oriented world model that emphasizes robotic task-critical features and leverages motion-informative samples to accelerate environment modeling, thereby improving sample efficiency and policy convergence. Our method consists of two components: (1) Inter-Frame Difference Capture (IFDC): This mechanism replaces traditional image reconstruction and performs global and local difference capture to emphasize robotic task-critical regions, orienting the world model to reconstruct task-relevant features and facilitating task completion. (2) Motion-Density-guided Experience Replay (MDER): Instead of uniform sampling, this scheme employs kernel density estimation of end-effector positions to prioritize motion-informative trajectory samples, thereby enabling the world model to rapidly learn the dynamics of the robotic environment. Our experiments show that TREND significantly outperforms six other existing methods in terms of sample efficiency and success rate on ten simulated and five real-world visual robotic manipulation tasks.