3D-aware Disentangled Representation for Compositional Reinforcement Learning

ICLR 2026 Conference Submission16912 Authors

Keywords: object-centric learning, compositional generalization, goal-conditioned reinforcement learning

Abstract:

Vision-based reinforcement learning can benefit from object-centric scene representation, which factorizes the visual observation into individual objects and their attributes, such as color, shape, size, and position.
While such object-centric representations can extract components that generalize well for various multi-object manipulation tasks, they are prone to issues with occlusions and 3D ambiguity of object properties due to their reliance on single-view 2D image features.
Furthermore, the entanglement between object configurations and camera poses complicates the object-centric disentanglement in 3D, leading to poor 3D reasoning by the agent in vision-based reinforcement learning applications.
To address the lack of 3D awareness and the object-camera entanglement problem, we propose an enhanced 3D object-centric representation that utilizes multi-view 3D features and enforces more explicit 3D-aware disentanglement.
The enhancement is based on the integration of the recent success of multi-view Transformer and the shared concept memory among the object-centric representations.
The representation, therefore, can stably identify and track 3D trajectories of individual objects along with their semantic and physical properties, exhibiting excellent interpretability and controllability.
Then, our proposed block transformer policy effectively performs novel tasks by assembling desired properties adaptive to the new goal states.
We demonstrate that our 3D-aware block representation is scalable to compose diverse novel scenes and enjoys superior performance in out-of-distribution tasks with multi-object manipulations compared to existing methods.

Primary Area: reinforcement learning

Submission Number: 16912