[Paper Review] Meta reinforcement learning as task inference

Jan Humplik, Alexandre Galashov+4 more|arXiv (Cornell University)|May 15, 2019

Reinforcement Learning in Robotics56 references63 citations

TL;DR

This paper frames meta-reinforcement learning as inferring the unknown task within a POMDP, and trains a two-stream agent where a belief module estimates the task posterior while the policy acts on state and belief. It shows that supervised belief learning with privileged task information speeds learning and improves performance, including in long-horizon continuous control with sparse rewards.

ABSTRACT

Humans achieve efficient learning by relying on prior knowledge about the structure of naturally occurring tasks. There is considerable interest in designing reinforcement learning (RL) algorithms with similar properties. This includes proposals to learn the learning algorithm itself, an idea also known as meta learning. One formal interpretation of this idea is as a partially observable multi-task RL problem in which task information is hidden from the agent. Such unknown task problems can be reduced to Markov decision processes (MDPs) by augmenting an agent's observations with an estimate of the belief about the task based on past experience. However estimating the belief state is intractable in most partially-observed MDPs. We propose a method that separately learns the policy and the task belief by taking advantage of various kinds of privileged information. Our approach can be very effective at solving standard meta-RL environments, as well as a complex continuous control environment with sparse rewards and requiring long-term memory.

Motivation & Objective

• Motivate learning to learn by treating the task as a hidden variable in a multi-task RL setting.
• Propose a two-stream agent that separately learns task belief and control policy.
• Leverage privileged supervision to train a belief network that estimates the task posterior.
• Show off-policy learning efficiency and benefit of information bottleneck regularization.

Proposed method

• Formulate meta-RL as a POMDP with states (x, w) where w is the unobserved task; the task posterior b_t(w) is updated from trajectory data.
• Introduce a belief network that outputs an approximate belief astcall b_t(h_t| au_{0:t}) using privileged task information h_t (e.g., task description, expert actions, task embeddings).
• Train the belief network with auxiliary supervised losses to mimic the true task information, minimizing a log loss or KL divergence to the posterior.
• Use a two-stream architecture where the policy (and optionally a critic) conditions on (x_t, ast b_t) rather than the full history, with gradients decoupled from the belief network.
• Explore off-policy SVG(0) (with entropy regularization) and compare to PPO, incorporating information bottleneck (IB) regularization to improve generalization in replay data.
• Compare belief-network agents to baseline LSTM and auxiliary-head variants, and evaluate on multiple meta-RL environments including sparse rewards.

Experimental results

Research questions

• RQ1 Can privileged information about the training task be leveraged to learn a useful belief representation that speeds meta-RL adaptation?
• RQ2 Does separating task inference from control improve sample efficiency, especially in off-policy learning?
• RQ3 How does information bottleneck regularization affect generalization and learning speed in belief-based meta-RL?
• RQ4 Can the framework scale to complex continuous control tasks with long-term memory and sparse rewards?

Key findings

• Supervised belief learning with task descriptions speeds up learning and improves final performance across several meta-RL environments.
• Off-policy SVG(0) with a belief stream is more sample-efficient than on-policy PPO in simple tasks, and IB regularization reduces the generalization gap.
• Belief-network agents outperform auxiliary-head and baseline LSTM agents in harder environments, including tasks with sparse rewards and long-term memory demands.
• Supervision with task descriptions generally yields better posterior estimation than predicting training-task IDs or unstructured targets.
• In the Numpad task, the belief-based approach achieves faster adaptation than Thompson-sampling-like methods (e.g., PEARL) by guiding broader exploration across episodes.