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

TL;DR

The paper addresses the problem of assessing whether spatial autoencoders (SAEs) produce reliable, spatially meaningful keypoints for reinforcement learning (RL) in robotics. It introduces a lightweight, trajectory-based metric that accounts for unknown 3D-to-2D offsets by fitting a time-invariant affine transform to align SAE keypoints with ground-truth object trajectories, and then uses per-object tracking capability (TC) to summarize performance across runs. Through systematic evaluation of base SAE architectures and three targeted modifications, the study shows substantial variation in tracking quality, with KeyNet-vel-std-bg achieving near-perfect tracking and RL performance approaching that of full-state or ground-truth representations. The results validate the metric as a predictive, low-cost indicator of RL success and offer practical guidance for SAE design in robotic RL, including publicly available code for reproducibility.

Abstract

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.

Figures (12)

• Figure 1: SAE extracts 2D positions from images via a spatial bottleneck. The SAE encoder is then integrated into an RL framework to obtain a state representation for immeasurable objects.
• Figure 2: We consider keypoints (red) to be equally informative about object positions as the ground truth CM (black). Motion of both points results in a varying offset in the image plane. We evaluate with transformed keypoints (blue) minimizing the offset.
• Figure 3: The PandaPush-v3 task with three object positions $x_k$ marked (left). Selected ground truth, keypoint, and transformed keypoint trajectories are shown in white, red, and blue for the end effector (middle) and cube (left).
• Figure 4: Basic-kp32 and KeyNet-vel-std-bg tracking errors $e_{n^*_k,k}$ for $K=3$ objects over epochs.
• Figure 5: Box plots of the tracking error $e_{n^*_k,k}$ for $K=3$ ground truth objects.