Rank2Reward: Learning Shaped Reward Functions from Passive Video

TL;DR

Rank2Reward shows that well-shaped reinforcement learning rewards can be learned from passive video demonstrations by framing progress as a temporal frame ranking. It learns a frame-ranking utility to produce a progression signal p_RF(s) and couples it with a discriminator to form a reward \\hat{r}(s) that guides policy optimization while penalizing deviations from the expert data distribution. The method, which integrates offline ranking with online adversarial weighting, achieves strong performance on simulated and real robotic tasks and scales to web-scale Ego4D video data. While effective, it acknowledges limitations such as embodiment gap and single-task rewards, pointing to future work in invariant representations and multi-task generalization.

Abstract

Teaching robots novel skills with demonstrations via human-in-the-loop data collection techniques like kinesthetic teaching or teleoperation puts a heavy burden on human supervisors. In contrast to this paradigm, it is often significantly easier to provide raw, action-free visual data of tasks being performed. Moreover, this data can even be mined from video datasets or the web. Ideally, this data can serve to guide robot learning for new tasks in novel environments, informing both "what" to do and "how" to do it. A powerful way to encode both the "what" and the "how" is to infer a well-shaped reward function for reinforcement learning. The challenge is determining how to ground visual demonstration inputs into a well-shaped and informative reward function. We propose a technique Rank2Reward for learning behaviors from videos of tasks being performed without access to any low-level states and actions. We do so by leveraging the videos to learn a reward function that measures incremental "progress" through a task by learning how to temporally rank the video frames in a demonstration. By inferring an appropriate ranking, the reward function is able to guide reinforcement learning by indicating when task progress is being made. This ranking function can be integrated into an adversarial imitation learning scheme resulting in an algorithm that can learn behaviors without exploiting the learned reward function. We demonstrate the effectiveness of Rank2Reward at learning behaviors from raw video on a number of tabletop manipulation tasks in both simulations and on a real-world robotic arm. We also demonstrate how Rank2Reward can be easily extended to be applicable to web-scale video datasets.

Figures (7)

• Figure 1: Depiction of the problem setting in Rank2Reward - inferring well-shaped and calibrated reward functions from video demonstrations that enable effective policy optimization.
• Figure 2: A schematic depiction of reward inference using Rank2Reward . Given video demonstrations from a human supervisor, Rank2Reward learns a reward function by combining two distinct elements --- (1) a ranking function that temporally orders frames, providing a monotonically increasing reward signal $p_{RF}(s)$. Secondly, (2) a classifier $D_\phi$ between expert and on-policy data, so that expert data is weighted higher than on-policy data. When combined multiplicateively, they yield a well-shaped reward function for RL that pushes down on-policy data, and pushes up expert data.
• Figure 3: Simulation environments for evaluation in the Meta-world yu2020meta benchmark: (1) reach, (2) push, (3) hammer, (4) drawer open, (5) door open, (6) door close, (7) button press, (8) assembly
• Figure 4: Real-world environments including standard tasks (reaching and pushing), tasks where exploration is non-trivial (pushing with obstacles and drawer opening), and tasks where state estimation is non-trivial (sweeping and drawing). The blue arrows indicate the directions to go.
• Figure 5: Visualization of policy learning experiments in simulation with Rank2Reward . Our method - Rank2Reward (purple) distinctly outperforms other methods on hammer, drawer open, button press, and assembly, while performing similarly to the best baseline with reach, and door close and worse than the best baseline in push and door open. The plots show the episodic return from 10 evaluation episodes averaged over 3 seeds plotted over the course of training the DrQ-v2 agent for 1.5 million steps, with higher being better.