Policy from a Single Demo: Robot Policy Learning with

Temporal Optimal Transport Reward

Yuwei Fu¹ ² * Haichao Zhang² Di Wu¹ Wei Xu² Benoit Boulet¹

¹McGill University ²Horizon Robotics

*Work done during an internship at Horizon Robotics

NeurIPS 2024

Abstract

Reward specification is one of the most tricky problems in Reinforcement Learning, which usually requires tedious hand engineering in practice. One promising approach to tackle this challenge is to adopt existing expert video demonstrations for policy learning. Some recent work investigates how to learn robot policies from only a single/few expert video demonstrations. For example, reward labeling via Optimal Transport (OT) has been shown to be an effective strategy to generate a proxy reward by measuring the alignment between the robot trajectory and the expert demonstrations. However, previous work mostly overlooks that the OT reward is invariant to temporal order information, which could bring extra noise to the reward signal. To address this issue, in this paper, we introduce the Temporal Optimal Transport (TemporalOT) reward to incorporate temporal order information for learning a more accurate OT-based proxy reward. Extensive experiments on the Meta-world benchmark tasks validate the efficacy of the proposed method

An illustration of the pipeline of applying OT-based reward in RL. In this toy example, we rollout two agent for five steps of transitions. Both agents start from the initial state and take same actions a0 and a1 at the first two states. Then the two agents take different actions a a 2 and a b 2 to generate different trajectories τa = (o0, a0, o1, a1, o2, aa 2 , oa 3 , aa 3 , oa 4 , aa 4 , oa 5 ) and τb = (o0, a0, o1, a1, o2, ab 2 , ob 3 , ab 3 , ob 4 , ab 4 , ob 5 ) The OT rewards for (o0, a0) and (o1, a1) in τ a and τ b are different even though the state-action pairs are exactly the same.

Illustration of the TemporalOT method

Instead of using a pair-wise cosine similarity as the transport cost, we use a group-wise cosine similarity to learn a more accurate cost matrix (left).
We use a temporal mask to enforce the OT reward to focus on a narrow scope to avoid potential distractions from observations outside of the mask window (right).

Performance Comparison

TemporalOT outperforms all baselines on average, which proves the effectiveness of the proposed method.

Related Publications and Resources

Robot Policy Learning with Temporal Optimal Transport Reward

Yuwei Fu, Haichao Zhang, Di Wu, Wei Xu, Benoit Boulet

Advances in Neural Information Processing Systems (NeurIPS), 2024

[PDF] [Code]

@inproceedings{furl, title={Robot Policy Learning with Temporal Optimal Transport Reward}, author={Yuwei Fu and Haichao Zhang and Di Wu and Wei Xu and Benoit Boulet}, booktitle={Advances in Neural Information Processing Systems}, year={2024}}