ATK: Automatic Task-driven Keypoint Selection for Robust Policy Learning
Automatically selects a minimal set of task-relevant keypoints for policy learning: the selection mask and a policy operating on the selected subset are jointly optimized end-to-end. The resulting keypoint-based policies remain robust to distractors, lighting changes and visual variation, and transfer across visually distinct scenarios.