# Tracking Object Positions in Reinforcement Learning: A Metric for Keypoint Detection (extended version)

arXiv: https://arxiv.org/abs/2312.00592 · HTML: https://mlsanity.com/paper/2312.00592/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 39 references that are in the map.
- First posted on arXiv 1 Dec 2023.

## Builds on (most cited first, 30 of 30)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Deep Spatial Autoencoders for Visuomotor Learning](https://mlsanity.com/paper/1509.06113/) (2015)
- [Unsupervised Learning of Object Landmarks through Conditional Image Generation](https://mlsanity.com/paper/1806.07823/) (2018)
- [3D Neural Scene Representations for Visuomotor Control](https://mlsanity.com/paper/2107.04004/) (2021)
- [Unsupervised Learning of Object Keypoints for Perception and Control](https://mlsanity.com/paper/1906.11883/) (2019)
- [Reinforcement Learning with Action-Free Pre-Training from Videos](https://mlsanity.com/paper/2203.13880/) (2022)
- [Offline Reinforcement Learning from Images with Latent Space Models](https://mlsanity.com/paper/2012.11547/) (2020)
- [Unsupervised Learning of Object Structure and Dynamics from Videos](https://mlsanity.com/paper/1906.07889/) (2019)
- [ALIKE: Accurate and Lightweight Keypoint Detection and Descriptor Extraction](https://mlsanity.com/paper/2112.02906/) (2021)
- [Unsupervised Discovery of Object Landmarks as Structural Representations](https://mlsanity.com/paper/1804.04412/) (2018)
- [KETO: Learning Keypoint Representations for Tool Manipulation](https://mlsanity.com/paper/1910.11977/) (2019)
- [Unsupervised Part-Based Disentangling of Object Shape and Appearance](https://mlsanity.com/paper/1903.06946/) (2019)
- [panda-gym: Open-source goal-conditioned environments for robotic learning](https://mlsanity.com/paper/2106.13687/) (2021)
- [Learning Actionable Representations from Visual Observations](https://mlsanity.com/paper/1808.00928/) (2018)
- [Deep Predictive Policy Training using Reinforcement Learning](https://mlsanity.com/paper/1703.00727/) (2017)
- [Self-Supervised Keypoint Discovery in Behavioral Videos](https://mlsanity.com/paper/2112.05121/) (2021)
- [Learning Representations by Maximizing Mutual Information in Variational Autoencoders](https://mlsanity.com/paper/1912.13361/) (2019)
- [KOVIS: Keypoint-based Visual Servoing with Zero-Shot Sim-to-Real Transfer for Robotics Manipulation](https://mlsanity.com/paper/2007.13960/) (2020)
- [Unsupervised Object Keypoint Learning using Local Spatial Predictability](https://mlsanity.com/paper/2011.12930/) (2020)
- [BKinD-3D: Self-Supervised 3D Keypoint Discovery from Multi-View Videos](https://mlsanity.com/paper/2212.07401/) (2022)
- [End-to-end Reinforcement Learning of Robotic Manipulation with Robust Keypoints Representation](https://mlsanity.com/paper/2202.06027/) (2022)
- [Learning Fast and Precise Pixel-to-Torque Control](https://mlsanity.com/paper/2208.02315/) (2022)
- [Learning of feature points without additional supervision improves reinforcement learning from images](https://mlsanity.com/paper/2106.07995/) (2021)

## Abstract (arXiv)

Reinforcement learning (RL) for robot control typically requires a detailed representation of the environment state, including information about task-relevant objects not directly measurable. Keypoint detectors, such as spatial autoencoders (SAEs), are a common approach to extracting a low-dimensional representation from high-dimensional image data. SAEs aim at spatial features such as object positions, which are often useful representations in robotic RL. However, whether an SAE is actually able to track objects in the scene and thus yields a spatial state representation well suited for RL tasks has rarely been examined due to a lack of established metrics. In this paper, we propose to assess the performance of an SAE instance by measuring how well keypoints track ground truth objects in images. We present a computationally lightweight metric and use it to evaluate common baseline SAE architectures on image data from a simulated robot task. We find that common SAEs differ substantially in their spatial extraction capability. Furthermore, we validate that SAEs that perform well in our metric achieve superior performance when used in downstream RL. Thus, our metric is an effective and lightweight indicator of RL performance before executing expensive RL training. Building on these insights, we identify three key modifications of SAE architectures to improve tracking performance.