[Paper Review] Keypoints into the Future: Self-Supervised Correspondence in Model-Based Reinforcement Learning

Lucas Manuelli, Yunzhu Li+2 more|arXiv (Cornell University)|Sep 10, 2020

Reinforcement Learning in Robotics27 references34 citations

TL;DR

The paper introduces self-supervised visual correspondence as the latent state for a model-based RL system, enabling closed-loop control via MPC and demonstrating strong sim-to-real transfer and category-level generalization for vision-based manipulation.

ABSTRACT

Predictive models have been at the core of many robotic systems, from quadrotors to walking robots. However, it has been challenging to develop and apply such models to practical robotic manipulation due to high-dimensional sensory observations such as images. Previous approaches to learning models in the context of robotic manipulation have either learned whole image dynamics or used autoencoders to learn dynamics in a low-dimensional latent state. In this work, we introduce model-based prediction with self-supervised visual correspondence learning, and show that not only is this indeed possible, but demonstrate that these types of predictive models show compelling performance improvements over alternative methods for vision-based RL with autoencoder-type vision training. Through simulation experiments, we demonstrate that our models provide better generalization precision, particularly in 3D scenes, scenes involving occlusion, and in category-generalization. Additionally, we validate that our method effectively transfers to the real world through hardware experiments. Videos and supplementary materials available at https://sites.google.com/view/keypointsintothefuture

Motivation & Objective

• Motivate and enable visual-based model learning for robotic manipulation with high-dimensional observations.
• Propose a latent state formed from tracked keypoints via self-supervised dense visual descriptors.
• Learn dynamics in the latent keypoint space and perform online planning with MPC for closed-loop control.
• Demonstrate improved generalization and real-world transfer over autoencoder- and transport-based baselines.
• Show robustness to occlusion and category-level variation in both simulated and hardware experiments.

Proposed method

• Learn a latent state z from dense visual descriptors via tracked keypoints extracted by a self-supervised visual-correspondence model.
• Train a forward dynamics model z_{t+1}=f(z_t, a_t) to minimize multi-step prediction error over horizon H (L_dynamics).
• Construct z from y (keypoint locations) using variants: Descriptor Set (DS), Spatial Descriptor Set (SDS), Weighted Descriptor Set (WDS), and their combinations (WSDS).
• Use a learnable mapping phi(y; α) to form z, with convex weights over keypoints to handle unreliable points (WDS).
• Keep the dense correspondence network fixed during dynamics learning.
• Perform online planning with MPC (MPPI-based) to maximize reward over a horizon (H) using the learned dynamics, with a goal latent state z*.
• Evaluate both in simulation (top-down, angled, occlusions, mugs/category) and on hardware with a Kuka manipulator using 10 minutes of interaction data for dynamics learning and a single demonstration for goal specification.

Experimental results

Research questions

• RQ1 Can self-supervised dense visual correspondences serve as an informative and transferable latent state for model-based RL in vision-based manipulation?
• RQ2 How do different descriptor-based latent state constructions (DS, SDS, WDS, WSDS) affect dynamics learning and closed-loop control under occlusion and category variation?
• RQ3 Does the approach outperform autoencoder and transport-based baselines in sim and real hardware, especially under partial occlusion and category-level generalization?
• RQ4 Can the method achieve real-world transfer with limited interaction data and a single demonstration for goal specification?

Key findings

• Dense visual-correspondence based latent states enable more accurate and sample-efficient model-based RL than baselines across simulated tasks and hardware.
• Under occlusion, SDS/WDS/WSDS outperform DS due to robustness to occluded keypoints, improving dynamics accuracy and MPC performance.
• On category-level mugs tasks, methods with more keypoints (DS/WDS) perform better than sparse variants due to capturing shape variation; SDS/WDS retain robustness.
• Compared to Transporter and Autoencoder baselines, SDS/WDS achieve better pose/angle deviations across tasks, with largest gains in occlusions and mugs categories.
• Hardware experiments show the method stabilizes long-horizon trajectories with a 2-second MPC horizon and 5 Hz command rate, using ~10 minutes of interaction data and a single demonstration, across multiple trajectories.