# TOM: Learning Policy-Aware Models for Model-Based Reinforcement Learning via Transition Occupancy Matching

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 35 references that are in the map.
- First posted on arXiv 22 May 2023.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (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)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training](https://mlsanity.com/paper/2210.00030/) (2022)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [Deep Dynamics Models for Learning Dexterous Manipulation](https://mlsanity.com/paper/1909.11652/) (2019)
- [Reinforcement Learning via Fenchel-Rockafellar Duality](https://mlsanity.com/paper/2001.01866/) (2020)
- [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)
- [How Far I'll Go: Offline Goal-Conditioned Reinforcement Learning via f-Advantage Regression](https://mlsanity.com/paper/2206.03023/) (2022)
- [Investigating Compounding Prediction Errors in Learned Dynamics Models](https://mlsanity.com/paper/2203.09637/) (2022)
- [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)
- [Live in the Moment: Learning Dynamics Model Adapted to Evolving Policy](https://mlsanity.com/paper/2207.12141/) (2022)
- [Self-Consistent Models and Values](https://mlsanity.com/paper/2110.12840/) (2021)

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

- [Query-Policy Misalignment in Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2305.17400/) (2023)
- [A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2310.06253/) (2023)
- [OMPO: A Unified Framework for RL under Policy and Dynamics Shifts](https://mlsanity.com/paper/2405.19080/) (2024)
- [Off-Policy Maximum Entropy RL with Future State and Action Visitation Measures](https://mlsanity.com/paper/2412.06655/) (2024)
- [Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2505.13709/) (2025)
- [Climate Surrogates for Scalable Multi-Agent Reinforcement Learning: A Case Study with CICERO-SCM](https://mlsanity.com/paper/2510.07971/) (2025)
- [Maximum-Entropy Exploration with Future State-Action Visitation Measures](https://mlsanity.com/paper/2603.18965/) (2026)
- [Are Good Generators Good Decision-Makers? Policy Learning for General Interventions via Retargeted Counterfactual Generation](https://mlsanity.com/paper/2606.07399/) (2026)
- [DT2: Decision-Targeted Digital Twins](https://mlsanity.com/paper/2606.25923/) (2026)
- [Beyond Conservatism: Recoverability-Conditioned Exploration for Model-Based Imitation Learning](https://mlsanity.com/paper/2609.33336/) (2026)

## Abstract (arXiv)

Standard model-based reinforcement learning (MBRL) approaches fit a transition model of the environment to all past experience, but this wastes model capacity on data that is irrelevant for policy improvement. We instead propose a new "transition occupancy matching" (TOM) objective for MBRL model learning: a model is good to the extent that the current policy experiences the same distribution of transitions inside the model as in the real environment. We derive TOM directly from a novel lower bound on the standard reinforcement learning objective. To optimize TOM, we show how to reduce it to a form of importance weighted maximum-likelihood estimation, where the automatically computed importance weights identify policy-relevant past experiences from a replay buffer, enabling stable optimization. TOM thus offers a plug-and-play model learning sub-routine that is compatible with any backbone MBRL algorithm. On various Mujoco continuous robotic control tasks, we show that TOM successfully focuses model learning on policy-relevant experience and drives policies faster to higher task rewards than alternative model learning approaches.