An Investigation into Pre-Training Object-Centric Representations for Reinforcement Learning

TL;DR

This paper addresses whether unsupervised object-centric representation (OCR) pre-training improves image-based reinforcement learning (RL) beyond indirect metrics like segmentation quality. It introduces a simple, objective OCR pretraining benchmark with 2D Spriteworld tasks and a 3D CausalWorld task, and systematically compares multiple OCR models (e.g., IODINE, Slot-Attention, SLATE, MAE-based encoders) against baselines using PPO. Key contributions include a scalable benchmark for validating OCR pre-training in RL, an empirical evaluation across object-centric tasks and generalization scenarios, and insights into when OCR helps (notably relational reasoning) along with the importance of transformer-based pooling/decoding in SLATE. The findings clarify the conditions under which OCR pre-training yields benefits, guiding future use of OCR in RL and suggesting directions for more complex visual environments and auxiliary training signals.

Abstract

Unsupervised object-centric representation (OCR) learning has recently drawn attention as a new paradigm of visual representation. This is because of its potential of being an effective pre-training technique for various downstream tasks in terms of sample efficiency, systematic generalization, and reasoning. Although image-based reinforcement learning (RL) is one of the most important and thus frequently mentioned such downstream tasks, the benefit in RL has surprisingly not been investigated systematically thus far. Instead, most of the evaluations have focused on rather indirect metrics such as segmentation quality and object property prediction accuracy. In this paper, we investigate the effectiveness of OCR pre-training for image-based reinforcement learning via empirical experiments. For systematic evaluation, we introduce a simple object-centric visual RL benchmark and conduct experiments to answer questions such as ``Does OCR pre-training improve performance on object-centric tasks?'' and ``Can OCR pre-training help with out-of-distribution generalization?''. Our results provide empirical evidence for valuable insights into the effectiveness of OCR pre-training for RL and the potential limitations of its use in certain scenarios. Additionally, this study also examines the critical aspects of incorporating OCR pre-training in RL, including performance in a visually complex environment and the appropriate pooling layer to aggregate the object representations.

Figures (16)

• Figure 1: The model architectures for representation types.
• Figure 2: The performance comparison of unsupervised object-centric representation (OCR) pre-training against other representation types in object-centric tasks. The results indicate that OCR pre-training demonstrates a significant performance gap compared to other representations and slightly worse performance than ground truth states in the comparison tasks where relational reasoning is a crucial aspect. However OCR pre-training performance is similar to or worse than baselines for other tasks.
• Figure 3: Samples from the dataset and five tasks in our benchmark; Object Goal / Object Interaction / Object Comparison / Property Comparison / Object Reaching tasks. In the 2D tasks ((b) - (e)), the red ball is always the agent. See main text for details about each task. In the robotics task (f), the goal is to use the green robotic finger to touch the blue object before touching any of the distractor objects.
• Figure 4: Success Rates for Object Goal, Object Interaction, Object Comparison, and Property Comparison Tasks. The specific representation types and training regimes used for each model are outlined in Table \ref{['tab:rep_regime_summ']}.
• Figure 5: The comparison of success rate against the number of interaction steps with the environments. Note that SLATE is compared with baselines for the Object Interaction task, where averaged performance of OCR pre-training is hard to be compared because other OCR methods failed to solve.