# Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation

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

## Key facts

- Cited by 35 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 85 references that are in the map.
- First posted on arXiv 6 Jun 2021.

## Builds on (most cited first, 29 of 29)

- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [Identifying and attacking the saddle point problem in high-dimensional non-convex optimization](https://mlsanity.com/paper/1406.2572/) (2014)
- [Model-Based Reinforcement Learning with Value-Targeted Regression](https://mlsanity.com/paper/2006.01107/) (2020)
- [MT-Opt: Continuous Multi-Task Robotic Reinforcement Learning at Scale](https://mlsanity.com/paper/2104.08212/) (2021)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Value Iteration Networks](https://mlsanity.com/paper/1602.02867/) (2016)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [A Game Theoretic Framework for Model Based Reinforcement Learning](https://mlsanity.com/paper/2004.07804/) (2020)
- [The Value Function Polytope in Reinforcement Learning](https://mlsanity.com/paper/1901.11524/) (2019)
- [The Differentiable Cross-Entropy Method](https://mlsanity.com/paper/1909.12830/) (2019)
- [Policy-Aware Model Learning for Policy Gradient Methods](https://mlsanity.com/paper/2003.00030/) (2020)
- [Gradient-Aware Model-based Policy Search](https://mlsanity.com/paper/1909.04115/) (2019)

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

- [A Survey of Contextual Optimization Methods for Decision Making under Uncertainty](https://mlsanity.com/paper/2306.10374/) (2023)
- [Bridging State and History Representations: Understanding Self-Predictive RL](https://mlsanity.com/paper/2401.08898/) (2024)
- [Mastering Memory Tasks with World Models](https://mlsanity.com/paper/2403.04253/) (2024)
- [Mismatched No More: Joint Model-Policy Optimization for Model-Based RL](https://mlsanity.com/paper/2110.02758/) (2021)
- [Value Gradient weighted Model-Based Reinforcement Learning](https://mlsanity.com/paper/2204.01464/) (2022)
- [Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective](https://mlsanity.com/paper/2209.08466/) (2022)
- [Beyond backpropagation: bilevel optimization through implicit differentiation and equilibrium propagation](https://mlsanity.com/paper/2205.03076/) (2022)
- [Synthesis of Model Predictive Control and Reinforcement Learning: Survey and Classification](https://mlsanity.com/paper/2502.02133/) (2025)
- [Scaling physics-informed hard constraints with mixture-of-experts](https://mlsanity.com/paper/2402.13412/) (2024)
- [Functional Bilevel Optimization for Machine Learning](https://mlsanity.com/paper/2403.20233/) (2024)
- [Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning](https://mlsanity.com/paper/2206.02072/) (2022)
- [When does Self-Prediction help? Understanding Auxiliary Tasks in Reinforcement Learning](https://mlsanity.com/paper/2406.17718/) (2024)
- [TaskMet: Task-Driven Metric Learning for Model Learning](https://mlsanity.com/paper/2312.05250/) (2023)
- [A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2310.06253/) (2023)
- [Lyapunov Exponents for Diversity in Differentiable Games](https://mlsanity.com/paper/2112.14570/) (2021)
- [Operator Splitting Value Iteration](https://mlsanity.com/paper/2211.13937/) (2022)
- [Model-Value Inconsistency as a Signal for Epistemic Uncertainty](https://mlsanity.com/paper/2112.04153/) (2021)
- [Scaling up and Stabilizing Differentiable Planning with Implicit Differentiation](https://mlsanity.com/paper/2210.13542/) (2022)
- [On Rate-Distortion Theory in Capacity-Limited Cognition & Reinforcement Learning](https://mlsanity.com/paper/2210.16877/) (2022)
- [DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning](https://mlsanity.com/paper/2402.05421/) (2024)
- [Bayesian Reinforcement Learning with Limited Cognitive Load](https://mlsanity.com/paper/2305.03263/) (2023)
- [Learning Theory for Kernel Bilevel Optimization](https://mlsanity.com/paper/2502.08457/) (2025)
- [Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2505.13709/) (2025)
- [Calibrated Value-Aware Model Learning with Probabilistic Environment Models](https://mlsanity.com/paper/2505.22772/) (2025)
- [Non-Stationary Functional Bilevel Optimization](https://mlsanity.com/paper/2601.15363/) (2026)
- [Decision-Focused Model-based Reinforcement Learning for Reward Transfer](https://mlsanity.com/paper/2304.03365/) (2023)
- [λ-models: Effective Decision-Aware Reinforcement Learning with Latent Models](https://mlsanity.com/paper/2306.17366/) (2023)
- [Learning Set Functions with Implicit Differentiation](https://mlsanity.com/paper/2412.11239/) (2024)
- [Bi-Level Reinforcement Learning Pathway for Sim-to-Real Optimality](https://mlsanity.com/paper/2510.17709/) (2025)
- [Beyond Distributions: Geometric Action Control for Continuous Reinforcement Learning](https://mlsanity.com/paper/2511.08234/) (2025)

## Abstract (arXiv)

The shortcomings of maximum likelihood estimation in the context of model-based reinforcement learning have been highlighted by an increasing number of papers. When the model class is misspecified or has a limited representational capacity, model parameters with high likelihood might not necessarily result in high performance of the agent on a downstream control task. To alleviate this problem, we propose an end-to-end approach for model learning which directly optimizes the expected returns using implicit differentiation. We treat a value function that satisfies the Bellman optimality operator induced by the model as an implicit function of model parameters and show how to differentiate the function. We provide theoretical and empirical evidence highlighting the benefits of our approach in the model misspecification regime compared to likelihood-based methods.