# DisCo RL: Distribution-Conditioned Reinforcement Learning for General-Purpose Policies

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 55 references that are in the map.
- First posted on arXiv 23 Apr 2021.

## Builds on (most cited first, 30 of 45)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Conditional Image Generation with PixelCNN Decoders](https://mlsanity.com/paper/1606.05328/) (2016)
- [Behavioral Cloning from Observation](https://mlsanity.com/paper/1805.01954/) (2018)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Data-Efficient Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1805.08296/) (2018)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [FeUdal Networks for Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1703.01161/) (2017)
- [Sim-to-Real: Learning Agile Locomotion For Quadruped Robots](https://mlsanity.com/paper/1804.10332/) (2018)
- [Multi-Goal Reinforcement Learning: Challenging Robotics Environments and Request for Research](https://mlsanity.com/paper/1802.09464/) (2018)
- [Automatic Goal Generation for Reinforcement Learning Agents](https://mlsanity.com/paper/1705.06366/) (2017)
- [Visual Reinforcement Learning with Imagined Goals](https://mlsanity.com/paper/1807.04742/) (2018)
- [Learning Agile Robotic Locomotion Skills by Imitating Animals](https://mlsanity.com/paper/2004.00784/) (2020)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [Go-Explore: a New Approach for Hard-Exploration Problems](https://mlsanity.com/paper/1901.10995/) (2019)
- [Efficient Exploration via State Marginal Matching](https://mlsanity.com/paper/1906.05274/) (2019)
- [Generative Adversarial Imitation from Observation](https://mlsanity.com/paper/1807.06158/) (2018)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)
- [Skew-Fit: State-Covering Self-Supervised Reinforcement Learning](https://mlsanity.com/paper/1903.03698/) (2019)
- [Imitation from Observation: Learning to Imitate Behaviors from Raw Video via Context Translation](https://mlsanity.com/paper/1707.03374/) (2017)
- [End-to-End Robotic Reinforcement Learning without Reward Engineering](https://mlsanity.com/paper/1904.07854/) (2019)
- [Zero-Shot Visual Imitation](https://mlsanity.com/paper/1804.08606/) (2018)
- [Universal Successor Features Approximators](https://mlsanity.com/paper/1812.07626/) (2018)
- [Unsupervised Control Through Non-Parametric Discriminative Rewards](https://mlsanity.com/paper/1811.11359/) (2018)
- [Imitating Latent Policies from Observation](https://mlsanity.com/paper/1805.07914/) (2018)

## Built on by (most cited first, 9 of 9)

- [Learning Language-Conditioned Robot Behavior from Offline Data and Crowd-Sourced Annotation](https://mlsanity.com/paper/2109.01115/) (2021)
- [Generalized Decision Transformer for Offline Hindsight Information Matching](https://mlsanity.com/paper/2111.10364/) (2021)
- [How Far I'll Go: Offline Goal-Conditioned Reinforcement Learning via f-Advantage Regression](https://mlsanity.com/paper/2206.03023/) (2022)
- [Learning Multimodal Behaviors from Scratch with Diffusion Policy Gradient](https://mlsanity.com/paper/2406.00681/) (2024)
- [AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents](https://mlsanity.com/paper/2310.09971/) (2023)
- [Adversarial Intrinsic Motivation for Reinforcement Learning](https://mlsanity.com/paper/2105.13345/) (2021)
- [The Challenges of Exploration for Offline Reinforcement Learning](https://mlsanity.com/paper/2201.11861/) (2022)
- [DISCOVER: Automated Curricula for Sparse-Reward Reinforcement Learning](https://mlsanity.com/paper/2505.19850/) (2025)
- [Deriving time-averaged active inference from control principles](https://mlsanity.com/paper/2208.10601/) (2022)

## Abstract (arXiv)

Can we use reinforcement learning to learn general-purpose policies that can perform a wide range of different tasks, resulting in flexible and reusable skills? Contextual policies provide this capability in principle, but the representation of the context determines the degree of generalization and expressivity. Categorical contexts preclude generalization to entirely new tasks. Goal-conditioned policies may enable some generalization, but cannot capture all tasks that might be desired. In this paper, we propose goal distributions as a general and broadly applicable task representation suitable for contextual policies. Goal distributions are general in the sense that they can represent any state-based reward function when equipped with an appropriate distribution class, while the particular choice of distribution class allows us to trade off expressivity and learnability. We develop an off-policy algorithm called distribution-conditioned reinforcement learning (DisCo RL) to efficiently learn these policies. We evaluate DisCo RL on a variety of robot manipulation tasks and find that it significantly outperforms prior methods on tasks that require generalization to new goal distributions.