# Task Aware Dreamer for Task Generalization in Reinforcement Learning

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 41 of its 51 references that are in the map.
- First posted on arXiv 9 Mar 2023.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [On the Properties of Neural Machine Translation: Encoder-Decoder Approaches](https://mlsanity.com/paper/1409.1259/) (2014)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Learning to Drive in a Day](https://mlsanity.com/paper/1807.00412/) (2018)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Multi-Task Reinforcement Learning with Soft Modularization](https://mlsanity.com/paper/2003.13661/) (2020)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [Reinforcement Learning with Action-Free Pre-Training from Videos](https://mlsanity.com/paper/2203.13880/) (2022)
- [Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1910.05396/) (2019)
- [Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck](https://mlsanity.com/paper/1910.12911/) (2019)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [Observational Overfitting in Reinforcement Learning](https://mlsanity.com/paper/1912.02975/) (2019)

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

- [PEAC: Unsupervised Pre-training for Cross-Embodiment Reinforcement Learning](https://mlsanity.com/paper/2405.14073/) (2024)
- [ManiBox: Enhancing Embodied Spatial Generalization via Scalable Simulation Data Generations](https://mlsanity.com/paper/2411.01850/) (2024)
- [PH-Dreamer: A Physics-Driven World Model via Port-Hamiltonian Generative Dynamics](https://mlsanity.com/paper/2605.18303/) (2026)
- [ValuePilot: A Two-Phase Framework for Value-Driven Decision-Making](https://mlsanity.com/paper/2503.04569/) (2025)
- [ValuePilot: A Two-Phase Framework for Value-Driven Decision-Making](https://mlsanity.com/paper/2512.13716/) (2025)

## Abstract (arXiv)

A long-standing goal of reinforcement learning is to acquire agents that can learn on training tasks and generalize well on unseen tasks that may share a similar dynamic but with different reward functions. The ability to generalize across tasks is important as it determines an agent's adaptability to real-world scenarios where reward mechanisms might vary. In this work, we first show that training a general world model can utilize similar structures in these tasks and help train more generalizable agents. Extending world models into the task generalization setting, we introduce a novel method named Task Aware Dreamer (TAD), which integrates reward-informed features to identify consistent latent characteristics across tasks. Within TAD, we compute the variational lower bound of sample data log-likelihood, which introduces a new term designed to differentiate tasks using their states, as the optimization objective of our reward-informed world models. To demonstrate the advantages of the reward-informed policy in TAD, we introduce a new metric called Task Distribution Relevance (TDR) which quantitatively measures the relevance of different tasks. For tasks exhibiting a high TDR, i.e., the tasks differ significantly, we illustrate that Markovian policies struggle to distinguish them, thus it is necessary to utilize reward-informed policies in TAD. Extensive experiments in both image-based and state-based tasks show that TAD can significantly improve the performance of handling different tasks simultaneously, especially for those with high TDR, and display a strong generalization ability to unseen tasks.