Source-linked AI summary
Unsupervised Learning of Object Keypoints for Perception and Control
Tejas Kulkarni, Ankush Gupta, Catalin Ionescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, Volodymyr Mnih
TL;DR
Object representations for control and reinforcement learning require concise, reusable geometric information beyond task-specific visual features. Transporter learns keypoints from raw videos by transporting image features through a keypoint bottleneck, producing accurate long-term tracking and supporting data-efficient learning and deep exploration without extrinsic rewards.
Problem
Task-specific visual representations are difficult to repurpose, while existing keypoint methods struggle with varying object numbers, sizes, and motions in reinforcement-learning domains.
Method
Transporter discovers spatially, temporally, and geometrically aligned object keypoints from videos using feature transport through a differentiable keypoint bottleneck.
Results
Transporter tracks keypoints accurately over long horizons, improves data-efficient reinforcement learning with keypoint inputs, and enables exploration through keypoint control without rewards.
Takeaways & Limitations
Keypoint coordinates and features provide a flexible representation for control, while keypoint-based skills reduce exploration search complexity and can reach states unavailable to random primitive actions.
Takeaways & Limitations
Applying the approach to all Atari games requires training Transporter inside the reinforcement-learning loop because random-policy pretraining is insufficient when new objects or screens appear.
Abstract
from arXiv · showhide
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.
1 Introduction
Transporter learns object-centric keypoints from video by transporting features between frames, addressing tracking failures under varying object properties. The learned representations support accurate tracking, data-efficient reinforcement learning, and exploration through keypoint control.
- Existing representations are often task-specific, while object-centric representations aim to capture reusable objects, relations, geometry, and cardinality.
- Existing keypoint methods struggle to track objects accurately when their number, size, and motion vary in reinforcement-learning domains.
- Transporter discovers spatially, temporally, and geometrically aligned keypoints from videos by transporting image features between source and target frames.The architecture uses a keypoint bottleneck and feature transport during training.
- Transporter learns robust object keypoints across reinforcement-learning environments with varying numbers, sizes, and motions of objects.
- Control applications: Using learned keypoints as policy inputs yields better performance than state-of-the-art model-free and model-based methods with up to 100k environment interactions.
- Control applications: Controlling keypoint coordinates provides an efficient exploration action space, enabling random agents to play several Atari games without rewards or task-dependent learning.
2 Related Work
Transporter extends unsupervised keypoint discovery by adding feature transport to representations learned from video. This addresses limitations of single-frame reconstruction and improves temporal tracking, including in hard-exploration settings.
- Transporter adds feature transport to a differentiable keypoint bottleneck, imposing stronger spatial alignment than prior baselines.
- Single-frame autoencoders do not learn temporal geometry, while some control-oriented approaches can learn non-spatial latent embeddings.
- The learned keypoints remain temporally consistent across hundreds of environment steps and can identify the most controllable keypoint without environment rewards.
3 Method
Transporter discovers and tracks object keypoints from unlabeled video by explicitly transporting image features between frames. These keypoints then provide a concise representation for data-efficient control and keypoint-directed exploration.
- Unsupervised keypoint discovery: Transporter extracts K two-dimensional keypoints from frame pairs and reconstructs the target frame without manual location labels.CNN feature maps and keypoint coordinates are extracted from source and target frames, then used for reconstruction.
- Feature transport: Feature transport replaces source features at target keypoint positions with target features and zeros source positions before refinement.The refine network inpaints missing source features and cleans the image around target positions.
- Feature transport: Unlike stacked-heatmap reconstruction, explicit spatial transport strengthens image-location correlations and supports more robust long-term tracking.The architecture adds feature transport as an inductive bias relative to Jakab et al.'s approach.
- Keypoints for control: Learned keypoints form a concise visual state abstraction for control, while their coordinates and associated features can drive a policy.The control setup uses thermometer-encoded keypoint coordinates together with image features as inputs to a recurrent Q-function.
- Keypoint-based exploration: Keypoint-based exploration defines K × 4 intrinsic rewards for moving each keypoint up, down, left, or right and trains corresponding Q-functions.The coordinate changes define the rewards, and the Q-functions are trained with n-step Q(λ).
- Keypoint-based exploration: Random exploration in the keypoint action space reduces search relative to raw actions, with evaluation selecting the most controllable keypoint.The controllability policy favors keypoints whose actions produce greater prospective change in all spatial directions.
4 Experiments
The experiments evaluate Transporter’s long-term keypoint tracking, data-efficient Atari control, and exploration through controllable keypoint options. Transporter outperforms landmark baselines across environments, improves Atari learning with limited interactions, and reaches rare states through reward-free exploration.
- Data-efficient reinforcement learning: The agent uses keypoint coordinates and corresponding image features as inputs to a recurrent neural-fitted Q-learning controller.Transporter is trained offline from random-policy data for this control architecture.
- Long-term tracking: Keypoint trajectories are evaluated with precision and recall over lengths from 1 to 200 frames, using thresholded Euclidean matching.One-to-one trajectory assignments are computed with minimum-cost linear-sum assignment.
- Long-term tracking: Transporter consistently tracks salient object keypoints over long horizons and outperforms baseline methods across environments, except slightly on Ms. Pacman.The exception still achieves similar long-trajectory performance.
- Data-efficient reinforcement learning: KeyQN trained for 100,000 interactions, or 400,000 frames, outperforms SimPLe and Rainbow on four of five Atari games.The comparison uses the Atari mean-score evaluation with three random seeds.
- Scope boundary: Pretraining Transporter on random-policy data is insufficient when new objects or screens appear, requiring training inside the reinforcement-learning loop.This limits straightforward application across all Atari games.
- Efficient exploration: Keypoint options consistently outperform random actions for exploration and can reach rare, rewarding Atari states without extrinsic rewards.The evaluation compares episodic returns and percentile returns over a 4-billion-frame budget.
5 Conclusion
The conclusion presents stable, reward-independent object keypoints as reusable representations for efficient control and exploration. It identifies moving backgrounds and richer environments as important boundaries for future work.
- Conclusion: Stable object keypoints can be learned across thousands of environment steps without task-specific reward functions.The conclusion frames this as evidence for reusable visual abstractions in control.
- Conclusion: Object keypoints could provide flexible, reusable representations for efficient control and reinforcement learning.The conclusion connects stable keypoints with repurposable downstream use.
- Limitations and future work: The model does not currently handle moving backgrounds, while richer datasets and environments remain future research targets.The authors suggest integrating explicit camera or ego-motion reasoning to transport features globally.
Appendix A Implementation Details
The implementation uses convolutional feature and keypoint networks with fixed core hyperparameters, environment-specific keypoint counts, and normalized-coordinate evaluation thresholds. The code is publicly available.
- Network architecture: The feature extractor uses six Conv-BatchNorm-ReLU layers, while KeyNet adds a 1×1 regressor producing K keypoint feature maps.RefineNet transposes the feature extractor and uses bilinear upsampling to restore resolution.
- Evaluation: The keypoint evaluation threshold ϵ equals each environment’s average ground-truth entity extent in normalized [-1,1] coordinates.The threshold applies to both x and y dimensions.
- Reproducibility: Transporter code is available in the DeepMind Research GitHub repository.The implementation link is provided directly by the paper.
Appendix B Diverse Data Generation
Transporter training uses diverse observation pairs sampled from random-policy trajectories, with buffer replacement based on nearest-neighbor novelty. Lower-resolution GPU copies make diversity selection efficient.
- Motivation: The data-generation procedure seeks diverse observation pairs because random-policy trajectories may contain many similar frame pairs.Diversity is used to improve the training dataset rather than storing all pairs unconditionally.
- Pair construction: Trajectories contain up to 100 steps, and one observation is sampled from each half to form temporally separated frame pairs.With action repeat 4, these trajectories represent up to 400 environment frames.
- Buffer diversity: Existing pairs are replaced when newly generated pairs have greater nearest-neighbor distance, encouraging novelty in the buffer.Nearest-neighbor distances are computed using lower-resolution grayscale copies on the GPU while high-resolution RGB pairs remain on CPU RAM.
Appendix D Pixel Transport versus Feature Transport
The experiment compares transporting learned features with transporting pixels in Transporter. Feature transport produces substantially better reconstructions because pixel transport creates ambiguous intermediate representations.
- Transporting learned features significantly outperforms transporting pixels in the reconstruction experiment.
- Pixel transport requires the refinement network to resolve difficult and ambiguous computations when predicting the target frame.
- Feature representations encode higher-level information, allowing the decoder to learn a more abstract function for prediction.
Appendix E Temporal consistency of keypoints
The temporal-consistency evaluation visualizes inferred keypoints across episodes in Atari and Manipulator environments. Frames are sampled ten time steps apart, with colors kept consistent for each keypoint throughout an episode.
- Figures visualize inferred keypoints for Pong, Frostbite, Ms. Pac-Man, and stack_4 episodes.
- Selected frames are separated by 10 time steps to assess keypoint behavior over temporal intervals.
- Colors remain time consistent, so each color denotes the same keypoint throughout an episode.
- Videos provide inferred keypoints from the three methods for entire episodes.
Appendix F Reconstructions
The reconstruction figures compare Transporter with two existing keypoint-discovery methods across Atari and Manipulator domains. They show inputs and reconstructions while accounting for Transporter’s paired-frame input requirement.
- Reconstructions are visualized for Pong, Frostbite, Ms. Pac-Man, and Manipulator stack_4 domains.
- Figure rows correspond to Transporter, Jakab and Gupta et al., and Zhang et al.
- The first two columns show model inputs, while the third column shows the reconstruction.
- Transporter receives an image and a future_image pair, whereas the other two models receive only future_image.