# Distill Knowledge in Multi-task Reinforcement Learning with Optimal-Transport Regularization

arXiv: https://arxiv.org/abs/2309.15603 · HTML: https://mlsanity.com/paper/2309.15603/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 34 references that are in the map.
- First posted on arXiv 27 Sep 2023.

## Builds on (most cited first, 25 of 25)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [Safe and Efficient Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1606.02647/) (2016)
- [Distral: Robust Multitask Reinforcement Learning](https://mlsanity.com/paper/1707.04175/) (2017)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Taming the Noise in Reinforcement Learning via Soft Updates](https://mlsanity.com/paper/1512.08562/) (2015)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [COG: Connecting New Skills to Past Experience with Offline Reinforcement Learning](https://mlsanity.com/paper/2010.14500/) (2020)
- [Hierarchical and Interpretable Skill Acquisition in Multi-task Reinforcement Learning](https://mlsanity.com/paper/1712.07294/) (2017)
- [In Defense of the Unitary Scalarization for Deep Multi-Task Learning](https://mlsanity.com/paper/2201.04122/) (2022)
- [Conservative Data Sharing for Multi-Task Offline Reinforcement Learning](https://mlsanity.com/paper/2109.08128/) (2021)
- [Information asymmetry in KL-regularized RL](https://mlsanity.com/paper/1905.01240/) (2019)
- [Q(λ) with Off-Policy Corrections](https://mlsanity.com/paper/1602.04951/) (2016)
- [Imitation Learning with Sinkhorn Distances](https://mlsanity.com/paper/2008.09167/) (2020)
- [Wasserstein Unsupervised Reinforcement Learning](https://mlsanity.com/paper/2110.07940/) (2021)
- [Importance Weighted Transfer of Samples in Reinforcement Learning](https://mlsanity.com/paper/1805.10886/) (2018)

## Abstract (arXiv)

In multi-task reinforcement learning, it is possible to improve the data efficiency of training agents by transferring knowledge from other different but related tasks. Because the experiences from different tasks are usually biased toward the specific task goals. Traditional methods rely on Kullback-Leibler regularization to stabilize the transfer of knowledge from one task to the others. In this work, we explore the direction of replacing the Kullback-Leibler divergence with a novel Optimal transport-based regularization. By using the Sinkhorn mapping, we can approximate the Optimal transport distance between the state distribution of tasks. The distance is then used as an amortized reward to regularize the amount of sharing information. We experiment our frameworks on several grid-based navigation multi-goal to validate the effectiveness of the approach. The results show that our added Optimal transport-based rewards are able to speed up the learning process of agents and outperforms several baselines on multi-task learning.