Source-linked AI summary
Meta reinforcement learning as task inference
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A. Ortega, Yee Whye Teh, Nicolas Heess
TL;DR
Efficient adaptation to new tasks is difficult when agents have little prior knowledge, and meta-RL treats hidden task information as a task-inference problem. This paper trains a separate belief network using privileged task information available during meta-training. The authors report improved performance across meta-RL algorithms and environments, including difficult continuous control, and find that structured task descriptions are useful supervision.
Problem
Method
The approach separates policy learning from task-belief estimation and supervises a belief network with task information available during meta-training.
Source: §1 · p. 2, §2.1 · p. 4Results
The authors report that privileged task information improves on- and off-policy meta-RL performance, and that posterior prediction of structured task descriptions works better than predicting task IDs or other auxiliary targets.
Source: §4 · p. 9Takeaways & Limitations
The experiments suggest that the natural structure in task descriptions can facilitate learning task beliefs.
Source: §4 · p. 9For tasks that are themselves POMDPs, the policy may need additional information about the task-specific belief state beyond the current state and task belief.
Source: §2 · p. 3
Abstract
from arXiv · showhide
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.
1 Introduction
Meta-RL aims to make agents learn new tasks efficiently by using prior experience and task structure. This paper frames unknown-task learning as task inference and separates estimating task beliefs from acting, using task information available during meta-training to supervise the belief network.
1 Introduction
- Meta-RL seeks to use prior knowledge from a distribution of tasks to speed learning, whereas learning with little prior knowledge can be computationally intensive.The paper contrasts this aim with human learning that draws on prior knowledge from similar problems.Source: §1 · p. 1
- Unknown-task RL can be represented as a POMDP whose static, hidden state is the task specification.The agent observes task interactions but not the task specification itself.Source: §1 · p. 2
- A belief state summarizes experience as a posterior over the unknown task, supporting decisions under task uncertainty.Task-level uncertainty matters for balancing exploration to identify the task with exploitation to maximize rewards given current knowledge.Source: §1 · pp. 1–2
- The method separates task inference from acting by training a belief network with task specifications available during meta-training; those specifications are not needed at meta-test.The paper identifies efficient learning on difficult continuous-control tasks, including sparse-reward problems requiring memory over more than 100 steps, among its contributions.Source: §1 · p. 2
- The paper reports improved meta-RL performance from cheap task information, efficient training with recurrent policies using off-policy algorithms, and solutions to difficult continuous-control problems.It also reports discovering a Bayes-optimal search strategy and comparing it with strategies motivated by Thompson sampling.Source: §1 · p. 2
2 Meta reinforcement learning and task inference
The section formulates meta-RL as inference over a hidden, fixed task and motivates policies that act using both the current state and an estimated task belief. It then describes supervised belief learning from privileged task information, the agent architectures and training algorithms, and related approaches.
2 Meta reinforcement learning and task inference
- The optimal meta-learner can act using the current state and task belief, but computing the exact belief is generally intractable, motivating a learned approximate belief module.For tasks that are themselves POMDPs, the policy may also need additional history, represented with an LSTM.Source: §2 · p. 3
2.1 Learning the belief network
- A belief module can be trained with task specifications or other privileged task information available during meta-training, without requiring that information at test time.The paper considers task descriptions, expert actions, task IDs, and pretrained task embeddings as supervision targets.Source: §2.1 · p. 4
- Supervised belief learning uses an auxiliary log loss, with task information from each training episode serving as an exact posterior sample and enabling training from off-policy replay data.The belief network is trained alongside the policy using trajectories and task information stored in a replay buffer.Source: §2.1 · p. 4
2.2 Architectures and algorithms
- The belief-network agent augments the baseline with a separate module that predicts a posterior over auxiliary task information and supplies its penultimate-layer representation to the actor and critic.Actor and critic gradients are not backpropagated into the belief network, keeping belief prediction separate from the reinforcement-learning objectives.Source: §2.2 · pp. 4–5, Figure 1 · p. 5
- The experiments use off-policy SVG(0), sometimes with Retrace, and compare it with on-policy PPO; information bottleneck regularization is also explored across the architectures.The supervised belief loss is optimized with policy and value losses during training.Source: §2.2 · p. 5
2.3 Related works
- The approach connects probabilistic and recurrent meta-RL by using privileged information to learn task representations, while related work also studies auxiliary losses and task structure.The discussion contrasts a concurrent approach that uses task-ID privilege and Thompson-sampling-like exploration with the paper’s approach.Source: §2.3 · pp. 5–6
3 Experiments
The experiments evaluate privileged supervision, off-policy learning, and information bottleneck regularization across simple and complex meta-RL tasks. Supervised belief agents outperform unsupervised agents across four environments, and the approach scales to sparse-reward tasks requiring long-term memory.
3 Experiments
- The experiments target faster learning and stronger final performance through privileged supervision, efficient off-policy recurrent training, and solving sparse-reward continuous-control tasks requiring memory beyond 100 steps.The study trained on 100 tasks and evaluated on 1,000 holdout tasks, reporting accumulated rewards during the exploration-and-adaptation episodes.Source: §3 · p. 6
3.1 Off-policy and on-policy learning, and information bottleneck
- Off-policy SVG(0) was more sample-efficient than PPO on Bandit, while belief agents generally beat baselines; information bottleneck regularization reduced the replay-to-on-policy generalization gap and improved sample efficiency.PPO on Semicircle was an exception, where insufficient off-policy replay data limited efficient belief-network training; on-policy data sufficed for Bandit.Source: §3.1 · pp. 6–7, Figure 3 · p. 7
3.2 Further evaluation of the role of supervision
- Task-description-supervised belief networks outperformed agents without supervision in all four evaluated meta-RL environments.The evaluation included Cheetah velocity, Semicircle variants, and Noisy target; Figure 3 summarizes the comparison on training and validation tasks.Source: §3.2 · p. 7, Figure 3 · p. 7
3.3 Scaling up to a complex continuous control environment
- In Numpad, task-description supervision performed best, belief networks outperformed auxiliary heads, and information bottleneck regularization was essential for efficient solving.Expert-action and task-embedding targets were useful but more prone to overfitting; generalization to validation tasks remained incomplete despite roughly 600 training tasks.Source: §3.3 · pp. 7, 9, Figure 4 · p. 8
3.4 Behavioral analysis
- In sparse 2D navigation, the Belief agent adapted faster than PEARL by exploring unvisited semicircle regions across episodes instead of repeatedly sampling one target.Figure 5 also visualizes task-sequence beliefs: the likelihood of the true sequence rises rapidly after discovering a new tile, and marginals reflect sequence structure.Source: Figure 6 · p. 8, Figure 5 · p. 8, §3.4 · p. 9
4 Discussion
Privileged information about hidden tasks improves on- and off-policy meta-RL, and structured task descriptions provide better belief targets than task IDs or other auxiliary targets. The discussion also contrasts the agents’ Bayes-optimal adaptation with Thompson-sampling exploration and notes an open question about information bottleneck regularization and generalization.
4 Discussion
- Privileged task information improves both on-policy and off-policy meta-RL, while off-policy belief estimation from trajectories can be unbiased.The authors argue that information bottleneck regularization can mitigate overfitting from replay-buffer and current-policy trajectory-distribution mismatch, but leave its relationship to generalization for future work.Source: §4 · p. 9
- Structured task descriptions such as instructions or goal locations work better as belief targets than training task IDs or other auxiliary targets.The authors suggest that natural task structure facilitates learning; limited posterior parameterization may explain the result for some other targets.Source: §4 · p. 9
- The agents’ Bayes-optimal adaptation strategy differs from the suboptimal Thompson-sampling exploration strategies often studied in meta-RL.Source: §4 · p. 9
Appendix A Task posterior and belief-MDP derivation
The appendix derives a belief-MDP representation for meta-RL: the task posterior can be updated from trajectory transitions and rewards, independently of the policy given the trajectory. Together with the current state, the belief provides Markovian inputs for restricting the policy.
Appendix A Task posterior and belief-MDP derivation
- The task belief at the next step depends on the previous belief and the current transition evidence: states, action, and reward.Source: Appendix A · p. 13
- Given the trajectory, the task posterior is independent of the policy, since the policy is independent of the task.Source: Appendix A · p. 13
- Because the belief is a deterministic function of history, the current state and belief form Markovian inputs on which the policy can depend.Source: Appendix A · pp. 12–13
Appendix B Learning task embedding
The Numpad task embedding is learned from task-ID inputs through a stochastic information-bottleneck layer. The appendix describes information-bottleneck regularization for supervised mappings and applies stochastic encoding and KL regularization to policy and critic networks trained with reinforcement-learning losses.
Appendix B Learning task embedding
- For Numpad, task-ID embeddings from a jointly trained multitask policy, passed through a two-layer MLP and stochastic information-bottleneck layer, serve as belief-network targets.Source: Appendix B · p. 14
Appendix C Regularization via information bottleneck
- The information bottleneck introduces a stochastic latent encoding and regularizes the supervised mapping to reduce mutual information between encoding and input.The appendix describes an upper bound for estimating the generally intractable mutual information and uses a standard normal reference distribution.Source: Appendix C · p. 14
- The authors also apply stochastic-encoder and KL information-bottleneck regularization to policy and critic networks trained with reinforcement-learning losses [31].Source: Appendix C · p. 14
D.1 Single-threaded experiments (Section 3.1)
The appendix reports how learning curves were aggregated and repeated for the bandit and semicircle experiments.
D.1 Single-threaded experiments (Section 3.1)
- Bandit curves average 100 episodes and semicircle curves 20, with five-point smoothing; experiments were repeated 15 and 8 times, respectively.Error bars show standard deviations of the smoothed curves.Source: §D.1 · p. 14
D.2 Distributed setting experiments (Section 3.1 and 3.3)
The experiments use distributed trajectory collection and replay-buffer learning for SVG(0), with auxiliary belief and regularization losses. Section 3.1 instead uses a single-threaded update schedule with TD(0).
D.2 Distributed setting experiments (Section 3.1 and 3.3)
- Learning curves report mean episodic return across distributed agents at each learner iteration, smoothed over 50 iterations, with standard-deviation error bars across three runs.The horizontal axis is learner-processed transitions.Source: §D.2 · p. 14
Distributed SVG(0)
- Distributed workers send trajectories to a shared replay buffer, from which the learner samples batches and updates actor, critic, and belief networks.The buffer holds 1e6 trajectories; recurrent internal states are stored for initialization, and updated parameters are shared with workers.Source: Distributed SVG(0) · pp. 14–15
- The distributed learner augments SVG(0) losses with auxiliary belief, KL, and policy-entropy terms, and periodically updates target parameters.Agents with information bottleneck may sample stochastically during acting or use the encoder mean; the implementation tried both settings.Source: Distributed SVG(0) · p. 15
Single-threaded SVG(0)
- Section 3.1 uses a single-threaded implementation, collecting a trajectory of length unroll length before n_train_iters actor and critic updates and one belief-network update; TD(0), not Retrace, estimates critic-training targets.Source: Single-threaded SVG(0) · p. 16
Network architectures
The appendix specifies MLP, recurrent, and optional information-bottleneck components for the actor, critic, and belief networks. It also lists separate hyperparameter settings for distributed, single-process, and PPO experiments.
Network architectures
- Actor, critic, and belief networks encode inputs with MLPs and pass representations through LSTMs; an optional information bottleneck projects LSTM outputs to a diagonal Gaussian.The belief network predicts the auxiliary target distribution, while actor and critic heads produce policy parameters or a scalar value.Source: Network architectures · p. 16
Hyperparameters–distributed setting
- Experiments use default actor, critic, and belief-network learning rates of 0.0003 and batch size 512; settings are further tuned per environment via grid and informal search over learning rates, entropy bonus, IB parameters, and bottleneck size.Defaults also include 200 parallel actors, a 50-step unroll, and a belief bottleneck dimension of 20.Source: Hyperparameters–distributed setting · p. 16
Hyperparameters–single process
- Single-process experiments use batch size 100, 10-step unrolls, and a 16-dimensional belief bottleneck, with hyperparameters chosen through grid and informal search.The listed actor and critic learning rates are 0.0005 for MLPs and 0.00005 for LSTMs.Source: Hyperparameters–single process · p. 17
Appendix F Algorithmic details for PPO
- PPO experiments clip the policy loss but update value and belief networks without clipping; bandit belief features are omitted from the recurrent value network.Search ranges differ by bandit versus semicircle experiments, with their settings listed separately.Source: Appendix F · pp. 17–18
Appendix G Implementation details of our environments
The appendix describes environments in which task inference arises from hidden or changing task information during interaction. Unlike an explicit K-episode formulation, the experiments use a single episode, with any K-shot structure implicit in environment dynamics.
Appendix G Implementation details of our environments
- The environments use a single episode rather than an explicit K-episode adaptation protocol, with K-shot structure implicit in dynamics such as resetting after a goal.Source: Appendix G · p. 18
G.1 Point mass semicircle
- In point-mass semicircle, the agent moves from the origin toward a target on a radius-0.2 semicircle, earns +1 on reaching it, and is teleported back.The episode lasts 100 action steps, and the task is the target angle.Source: §G.1 · p. 18
G.2 Multi-armed bandit
- In the bandit environment, each pull yields a Bernoulli reward with an unknown arm-specific probability, and the task is the vector of those probabilities.The task distribution is uniform on the K-dimensional unit hypercube, and beliefs use a Beta distribution.Source: §G.2 · p. 18
G.3 Cheetah
- The half-cheetah task is to run at an unknown target speed sampled uniformly from 3.0 to 10.0, with beliefs parameterized by a Gaussian.Source: §G.3 · p. 19
G.4 Rolling ball/quadruped semi-circle
- The rolling-ball and quadruped semicircle environment replaces the point mass with a robot and uses 1000-step episodes and Gaussian belief distributions.The rolling ball has continuous torque and forward-acceleration controls; observations include position, orientation, and proprioceptive features.Source: §G.4 · p. 19
G.5 Noisy target
- In the noisy-target environment, the ball is not teleported to the origin after reaching the target, so task inference requires integrating noisy rewards rather than remembering target location.Beliefs are parameterized with a Gaussian distribution.Source: §G.5 · p. 19
G.6 Numpad
- Numpad tasks require visiting a prescribed sequence on a 3x3 grid; correct-order tiles light up, while an out-of-order visit resets progress.Rewards mark first completion of valid subsequences; tasks are contiguous neighboring-tile sequences of length at most four, among 704 tasks. The reward is history-dependent, and task cues may mask digits.Source: §G.6 · pp. 19–20, Figure 7 · p. 19
Appendix H Generalization across tasks
The appendix examines training-set-size generalization in bandit, quadruped semicircle, and Numpad settings. Its figures cover train/validation performance, generalization gaps, and alternative task-visibility regimes; the supplied passages do not state comparative outcomes.
Appendix H Generalization across tasks
- The appendix varies training-set size to investigate generalization across tasks and clarifies that figure labels differ from the main text.Appendix labels map Belief to Belief (noIB), Belief + IB to Belief, LSTM to Baseline (LSTM, noIB), and LSTM + IB to Baseline (LSTM).Source: Appendix H · p. 20
H.1 Multi-armed bandit
- The bandit figure compares training and validation performance after 500 iterations for 5- and 20-arm problems.The caption identifies the baseline LSTM and best-performing agent as the compared agents.Source: §H.1 · p. 20, Figure 8 · p. 20
H.2 Quadruped semi-circle
- For quadruped semicircle, the figures show how generalization gap depends on training-set size and provide full training/validation learning curves.The gap figure reports performance after 5e9 learner updates.Source: §H.2 · p. 20, Figure 9 · p. 21
H.3 Numpad
- Numpad evaluation distinguishes fully visible, partially visible, and fully hidden task descriptions, with fully hidden as the default regime.Figures cover training-set-size effects on the generalization gap and full training/validation learning curves; performance is reported after 4e9 learner updates for the gap figure.Source: §H.3 · pp. 20, 22, Figure 11 · p. 21, Figure 10 · p. 21