Transformer-based Imagination with Slot Attention

World models have been proposed for improving the learning efficiency of deep reinforcement learning (RL) agents.However, it remains challenging for world models to effectively replicate environments that are high-dimensional, non-stationary, and comprising multiple objects and their interactions. We propose Transformer-based Imagination with Slot Attention (TISA), an RL agent that integrates a Transformer-based object-centric world model, policy function, and value function. The world model in TISA uses a Transformer-based architecture to handle each object's state, actions, and rewards (or costs) separately, effectively managing high-dimensional observations and preventing the combinatorial explosion of dynamics. Also, the Transformer-based policy and value functions can make decisions by considering the dynamics of individual objects and their interactions. In Safety-Gym benchmark, TISA outperforms a previous Transformer-based world model method.