[Paper Review] Unsupervised Learning of Object Keypoints for Perception and Control

Tejas Kulkarni, Ankush Gupta+5 more|arXiv (Cornell University)|Jun 19, 2019

Domain Adaptation and Few-Shot Learning42 references71 citations

TL;DR

Transporter learns unsupervised object keypoints from videos and uses them for data-efficient RL and exploration, by transporting features between frames through a keypoint bottleneck.

ABSTRACT

The study of object representations in computer vision has primarily focused on developing representations that are useful for image classification, object detection, or semantic segmentation as downstream tasks. In this work we aim to learn object representations that are useful for control and reinforcement learning (RL). To this end, we introduce Transporter, a neural network architecture for discovering concise geometric object representations in terms of keypoints or image-space coordinates. Our method learns from raw video frames in a fully unsupervised manner, by transporting learnt image features between video frames using a keypoint bottleneck. The discovered keypoints track objects and object parts across long time-horizons more accurately than recent similar methods. Furthermore, consistent long-term tracking enables two notable results in control domains -- (1) using the keypoint co-ordinates and corresponding image features as inputs enables highly sample-efficient reinforcement learning; (2) learning to explore by controlling keypoint locations drastically reduces the search space, enabling deep exploration (leading to states unreachable through random action exploration) without any extrinsic rewards.

Motivation & Objective

• Motivate learning object-centric representations that are useful for control and reinforcement learning instead of only classification or detection.
• Propose a fully unsupervised architecture (Transporter) to discover concise geometric keypoints that track objects over long horizons.
• Demonstrate that keypoints improve data efficiency in RL when used as state input.
• Show that controlling keypoints enables task-agnostic exploration and acquisition of skills for efficient exploration.

Proposed method

• Predict spatial keypoint coordinates Psi(x) for frame pairs using a CNN-based feature extractor Phi and a differentiable keypoint network (KeyNet).
• Compute Gaussian heatmaps around keypoints and transport features from a source frame to a target frame by replacing source heatmap locations with target features; zero-out source keypoint locations and refine the transport with a small-refinement network.
• Train end-to-end with pixelwise L2 reconstruction of the target frame to enforce geometric consistency of keypoints over time.
• Pre-train Transporter on data collected with a random policy, then fix transporter weights while learning control policies from rewards (data-efficient RL using KeyQN).
• Define intrinsic keypoint-based actions (K x 4 directional moves per keypoint) to learn exploration-friendly options via Q-functions, selecting the most controllable keypoint for exploration.

Experimental results

Research questions

• RQ1 Can unsupervised learning of object keypoints yield stable, long-term tracking across diverse RL environments?
• RQ2 Does using keypoint coordinates and localized features as inputs improve data efficiency for reinforcement learning compared to RGB observations?
• RQ3 Can learning to control keypoints provide a reward-free, task-agnostic mechanism to enhance exploration in hard-exploration environments?

Key findings

• Transporter learns state-of-the-art object keypoints that robustly track objects across long time horizons in Atari ALE and Manipulator domains.
• Using keypoint coordinates and localized features as inputs yields more data-efficient RL than strong baselines on several Atari games with around 100k interactions.
• Learning to control keypoint locations enables efficient exploration, significantly reducing the action-space search; random exploration in keypoint space can reach behaviors unattainable by random raw actions.
• The most controllable keypoint often corresponds to the avatar in games and remains consistently trackable across state variations.
• Keypoint-based exploration demonstrates substantial performance gains in hard exploration tasks such as Montezuma’s Revenge without extrinsic rewards.
• The approach yields temporally consistent object representations suitable for re-use across tasks without task-specific rewards.