# Trajectory-wise Multiple Choice Learning for Dynamics Generalization in Reinforcement Learning

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

## Key facts

- Cited by 35 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 26 Oct 2020.

## Built on by (most cited first, 30 of 35)

- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Augmented World Models Facilitate Zero-Shot Dynamics Generalization From a Single Offline Environment](https://mlsanity.com/paper/2104.05632/) (2021)
- [Trustworthy Reinforcement Learning Against Intrinsic Vulnerabilities: Robustness, Safety, and Generalizability](https://mlsanity.com/paper/2209.08025/) (2022)
- [Self-Paced Context Evaluation for Contextual Reinforcement Learning](https://mlsanity.com/paper/2106.05110/) (2021)
- [CARL: A Benchmark for Contextual and Adaptive Reinforcement Learning](https://mlsanity.com/paper/2110.02102/) (2021)
- [Dynamics Generalisation in Reinforcement Learning via Adaptive Context-Aware Policies](https://mlsanity.com/paper/2310.16686/) (2023)
- [A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2206.04551/) (2022)
- [The Evolution of Reinforcement Learning in Quantitative Finance: A Survey](https://mlsanity.com/paper/2408.10932/) (2024)
- [From Samples to Scenarios: A New Paradigm for Probabilistic Forecasting](https://mlsanity.com/paper/2509.19975/) (2025)
- [EUCLID: Towards Efficient Unsupervised Reinforcement Learning with Multi-choice Dynamics Model](https://mlsanity.com/paper/2210.00498/) (2022)
- [Self-Activating Neural Ensembles for Continual Reinforcement Learning](https://mlsanity.com/paper/2301.00141/) (2022)
- [Scenic4RL: Programmatic Modeling and Generation of Reinforcement Learning Environments](https://mlsanity.com/paper/2106.10365/) (2021)
- [Zero-Shot Reinforcement Learning Under Partial Observability](https://mlsanity.com/paper/2506.15446/) (2025)
- [Learning Action Translator for Meta Reinforcement Learning on Sparse-Reward Tasks](https://mlsanity.com/paper/2207.09071/) (2022)
- [Inferring Behavior-Specific Context Improves Zero-Shot Generalization in Reinforcement Learning](https://mlsanity.com/paper/2404.09521/) (2024)
- [Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization](https://mlsanity.com/paper/2508.20294/) (2025)
- [Decoupling Meta-Reinforcement Learning with Gaussian Task Contexts and Skills](https://mlsanity.com/paper/2312.06518/) (2023)
- [Winner-takes-all for Multivariate Probabilistic Time Series Forecasting](https://mlsanity.com/paper/2506.05515/) (2025)
- [Prototypical context-aware dynamics generalization for high-dimensional model-based reinforcement learning](https://mlsanity.com/paper/2211.12774/) (2022)
- [PrivilegedDreamer: Explicit Imagination of Privileged Information for Rapid Adaptation of Learned Policies](https://mlsanity.com/paper/2502.11377/) (2025)
- [Curiosity-Driven Imagination: Discovering Plan Operators and Learning Associated Policies for Open-World Adaptation](https://mlsanity.com/paper/2503.04931/) (2025)
- [TimePre: Bridging Accuracy, Efficiency, and Stability in Probabilistic Time-Series Forecasting](https://mlsanity.com/paper/2511.18539/) (2025)
- [Guided Cooperation in Hierarchical Reinforcement Learning via Model-based Rollout](https://mlsanity.com/paper/2309.13508/) (2023)
- [Residual Learning and Context Encoding for Adaptive Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2406.08238/) (2024)
- [VIPO: Value Function Inconsistency Penalized Offline Reinforcement Learning](https://mlsanity.com/paper/2504.11944/) (2025)
- [Pessimism Principle Can Be Effective: Towards a Framework for Zero-Shot Transfer Reinforcement Learning](https://mlsanity.com/paper/2505.18447/) (2025)
- [Multiple Choice Learning of Low-Rank Adapters for Language Modeling](https://mlsanity.com/paper/2507.10419/) (2025)
- [Open-Ended Diverse Solution Discovery with Regulated Behavior Patterns for Cross-Domain Adaptation](https://mlsanity.com/paper/2209.12029/) (2022)
- [Decomposed Mutual Information Optimization for Generalized Context in Meta-Reinforcement Learning](https://mlsanity.com/paper/2210.04209/) (2022)
- [Skill-aware Mutual Information Optimisation for Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2406.04815/) (2024)

## Abstract (arXiv)

Model-based reinforcement learning (RL) has shown great potential in various control tasks in terms of both sample-efficiency and final performance. However, learning a generalizable dynamics model robust to changes in dynamics remains a challenge since the target transition dynamics follow a multi-modal distribution. In this paper, we present a new model-based RL algorithm, coined trajectory-wise multiple choice learning, that learns a multi-headed dynamics model for dynamics generalization. The main idea is updating the most accurate prediction head to specialize each head in certain environments with similar dynamics, i.e., clustering environments. Moreover, we incorporate context learning, which encodes dynamics-specific information from past experiences into the context latent vector, enabling the model to perform online adaptation to unseen environments. Finally, to utilize the specialized prediction heads more effectively, we propose an adaptive planning method, which selects the most accurate prediction head over a recent experience. Our method exhibits superior zero-shot generalization performance across a variety of control tasks, compared to state-of-the-art RL methods. Source code and videos are available at https://sites.google.com/view/trajectory-mcl.