Source-linked AI summary
Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, Sergey Levine
TL;DR
Meta-RL seeks rapid adaptation with limited experience, but existing methods rely heavily on inefficient on-policy data and do not adequately represent task uncertainty. PEARL disentangles probabilistic task inference from control and integrates latent context variables with off-policy RL. It achieves 20-100X higher meta-training sample efficiency and improved asymptotic performance on several benchmarks.
Problem
Meta-RL must reduce both the extensive training data used across tasks and the experience needed to adapt to new tasks, while existing methods often rely on on-policy data.
Method
PEARL infers probabilistic latent task contexts online, conditions actor and critic policies on them, and separates encoder context sampling from off-policy RL batches.
Results
20-100X improvement in meta-training sample efficiency and substantial asymptotic-performance increases were demonstrated over prior state-of-the-art methods on six continuous-control meta-learning environments.
Takeaways & Limitations
Probabilistic contexts support posterior sampling for structured, temporally extended exploration and rapid adaptation to new tasks.
Takeaways & Limitations
The method assumes tasks are drawn from a task distribution of MDPs with bounded rewards and unknown transition and reward functions.
Abstract
from arXiv · showhide
Deep reinforcement learning algorithms require large amounts of experience to learn an individual task. While in principle meta-reinforcement learning (meta-RL) algorithms enable agents to learn new skills from small amounts of experience, several major challenges preclude their practicality. Current methods rely heavily on on-policy experience, limiting their sample efficiency. The also lack mechanisms to reason about task uncertainty when adapting to new tasks, limiting their effectiveness in sparse reward problems. In this paper, we address these challenges by developing an off-policy meta-RL algorithm that disentangles task inference and control. In our approach, we perform online probabilistic filtering of latent task variables to infer how to solve a new task from small amounts of experience. This probabilistic interpretation enables posterior sampling for structured and efficient exploration. We demonstrate how to integrate these task variables with off-policy RL algorithms to achieve both meta-training and adaptation efficiency. Our method outperforms prior algorithms in sample efficiency by 20-100X as well as in asymptotic performance on several meta-RL benchmarks.
1. Introduction
Meta-RL can adapt quickly to new tasks, but prevailing methods require extensive on-policy data and lack efficient uncertainty-aware exploration. PEARL addresses both challenges through probabilistic task inference, off-policy learning, and posterior sampling.
- Conventional RL often requires millions of environment interactions per task, making large behavior repertoires prohibitive.
- Meta-learning exploits shared task structure to adapt quickly from small amounts of experience, but meta-training still requires massive data across tasks.
- Most current meta-RL methods rely on on-policy data during meta-training and adaptation, limiting sample efficiency.
- PEARL integrates online probabilistic context inference with off-policy RL, separating task inference from action selection.
- 20-100X improvement in meta-training sample efficiency and substantial asymptotic-performance gains were demonstrated across six continuous-control meta-learning environments.
2. Related Work
Prior meta-RL methods use recurrent, gradient-based, or contextual representations, while PEARL combines probabilistic task contexts with off-policy learning. The approach also connects to posterior sampling and structured inference in partially observed tasks.
- Context-based meta-RL methods aggregate experience into latent representations that condition the policy for adaptation.
- PEARL represents task contexts with probabilistic latent variables, enabling reasoning over task uncertainty beyond deterministic contextual embeddings.
- Gradient-based meta-RL methods commonly focus on on-policy meta-learning, whereas PEARL focuses on learning from off-policy data.
- Posterior sampling provides temporally extended exploration by acting according to a sampled task or MDP hypothesis.
- Test-time adaptation can be viewed as a POMDP problem in which the task is unobserved, motivating belief estimation and posterior sampling.
3. Problem Statement
The problem is to learn policies that adapt efficiently across a distribution of tasks while minimizing both prior-task and new-task experience. Tasks are modeled as MDPs, with adaptation based on transition history.
- Sample efficiency concerns both meta-training samples from previous tasks and adaptation experience required in a new task.
- The task distribution contains MDPs with states, actions, transition functions, and bounded reward functions that may vary across tasks.
- Transition and reward functions are unknown but can be sampled through environment interaction.
- The policy adapts by conditioning on the history of past transitions, called context c.
4. Probabilistic Latent Context
PEARL infers a probabilistic latent context from experience and conditions the policy on it, using permutation-invariant encoding and posterior sampling for exploration. Its off-policy training procedure separates context-data sampling from actor-critic updates.
- Modeling and Learning Latent Contexts: The policy πθ(a|s, z) is conditioned on a latent probabilistic context Z inferred from recent task experience.
- Modeling and Learning Latent Contexts: The inference network qφ(z|c) estimates the posterior over Z using amortized variational inference.
- Modeling and Learning Latent Contexts: An information bottleneck constrains Z to retain task-relevant context information and mitigate overfitting to training tasks.
- Modeling and Learning Latent Contexts: The encoder uses a permutation-invariant representation because task-relevant transition collections do not depend on observation order.
- Posterior Sampling and Exploration via Latent Contexts: Posterior sampling holds a sampled latent context fixed during an episode, producing temporally extended exploration and updating beliefs with collected experience.
- Off-Policy Meta-Reinforcement Learning: The policy trains on off-policy replay data while the encoder uses separately sampled context batches to reduce train-test distribution mismatch.
5. Off-Policy Meta-Reinforcement Learning
The method enables off-policy meta-RL by separating context inference from policy learning, using probabilistic latent contexts with SAC. Context data and replay-buffer data are sampled separately to support efficient training while limiting distribution mismatch.
- Off-policy learning: The approach separates encoder training from policy training, allowing the actor and critic to use off-policy replay data while the encoder receives separately sampled context batches.The context sampler uses recently collected data, whereas actor and critic batches are drawn from the replay buffer.
- Off-policy learning: The actor and critic condition on a sampled latent context z, treating it as part of the state in an off-policy reinforcement-learning loop.Uncertainty in q(z|c) supplies stochasticity for exploration.
- Meta-training: During meta-training, each task collects trajectories using contexts sampled from qφ(z|c), then samples context and reinforcement-learning batches for optimization.The algorithm initializes per-task replay buffers and repeatedly gathers data before drawing training batches.
- Meta-training: Sampling encoder contexts from recently collected data rather than the entire replay buffer retains on-policy performance while improving efficiency.The recent-data buffer is recollected every 1000 meta-training optimization steps, while actor and critic batches use the entire replay buffer.
- Implementation: The method builds on SAC, an off-policy maximum-entropy actor-critic algorithm with a probabilistic interpretation suited to latent contexts.The actor and critic are optimized jointly with the inference network using reparameterization gradients through sampled contexts.
- Implementation: The inference network is trained through the critic’s Bellman update, and the reported experiments found state-action value recovery superior to alternative encoder objectives.The alternatives included maximizing actor returns or reconstructing states and rewards.
6. Experiments
PEARL is evaluated on continuous-control benchmarks, sparse navigation, and ablations of its probabilistic context and data-sampling design. Across these experiments, it combines strong sample efficiency and performance with rapid exploration and adaptation.
- Benchmark evaluation: PEARL is evaluated against prior meta-RL methods on six continuous-control environments involving reward-function or dynamics variation.The experiments assess benchmark performance, posterior-sampling exploration, and algorithmic design choices through ablations.
- Posterior sampling for exploration: PEARL began adapting to unseen sparse-navigation goals after an average of only 5 trajectories and outperformed MAESN.The task provides reward only within the goal radius, with distinct training and test goals.
- Benchmark evaluation: 20-100x fewer meta-training samples yielded improved final asymptotic performance by 50-100% in five of six domains.The comparison uses prior meta-RL approaches as baselines.
- Ablations: Replacing PEARL’s encoder with an RNN produced comparable performance but slower optimization, whereas trajectory-sampled RL batches caused a steep performance drop.The ablation compares unordered-transition and trajectory sampling.
- Ablations: Sampling context fully off-policy significantly hurt performance, while reusing the RL batch helped in this setting, underscoring the importance of data-sampling choices.PEARL instead uses recent context transitions distinct from the RL mini-batch.
- Ablations: A deterministic latent context was unable to solve sparse-reward navigation because time-invariant policy stochasticity hindered temporally extended exploration.The probabilistic variant supports a distribution over tasks and posterior-sampling exploration.
7. Conclusion
PEARL decouples task inference from control by conditioning the policy on a probabilistic latent context. This supports off-policy meta-training and posterior-sampling exploration, while achieving stronger results with less experience across continuous-control meta-RL domains.
- Contribution: PEARL infers a latent context variable on which the policy is conditioned, decoupling task inference from solving the task.This design minimizes mismatch between training and test context distributions.
- Contribution: Probabilistic context enables posterior sampling and temporally extended exploration during test-time adaptation.The policy can use off-policy meta-training while the context remains probabilistic.
- Conclusion: PEARL achieves superior results to prior meta-RL algorithms while requiring far less experience across diverse continuous-control domains.
A. Experimental Details
The experimental details cover six MuJoCo meta-learning domains spanning task variation in rewards, goals, directions, velocities, and system dynamics. The benchmark families include both small discrete task sets and larger train/test task collections.
- Benchmark domains: The six MuJoCo domains cover locomotion tasks that vary reward functions, target goals, or randomized system dynamics.The environments include Half-Cheetah, Ant, Humanoid, and Walker variants.
- Task splits: Half-Cheetah-Dir and Ant-Fwd-Back each contain two tasks, while Half-Cheetah-Vel, Humanoid-Dir-2D, and Ant-Goal-2D use 100 training and 30 test tasks.Walker-2D-Params uses 40 training and 10 test tasks with randomized dynamics parameters.
- Task splits: Walker-2D-Params requires forward movement under randomized system dynamics parameters.