# Model-Based Reinforcement Learning via Latent-Space Collocation

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 68 references that are in the map.
- First posted on arXiv 24 Jun 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [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)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Learning Complex Dexterous Manipulation with Deep Reinforcement Learning and Demonstrations](https://mlsanity.com/paper/1709.10087/) (2017)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [A Recurrent Latent Variable Model for Sequential Data](https://mlsanity.com/paper/1506.02216/) (2015)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Unsupervised Learning for Physical Interaction through Video Prediction](https://mlsanity.com/paper/1605.07157/) (2016)
- [Stochastic Video Generation with a Learned Prior](https://mlsanity.com/paper/1802.07687/) (2018)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [Semi-parametric Topological Memory for Navigation](https://mlsanity.com/paper/1803.00653/) (2018)
- [Deep Dynamics Models for Learning Dexterous Manipulation](https://mlsanity.com/paper/1909.11652/) (2019)
- [Search on the Replay Buffer: Bridging Planning and Reinforcement Learning](https://mlsanity.com/paper/1906.05253/) (2019)
- [PRECOG: PREdiction Conditioned On Goals in Visual Multi-Agent Settings](https://mlsanity.com/paper/1905.01296/) (2019)
- [Planning with Goal-Conditioned Policies](https://mlsanity.com/paper/1911.08453/) (2019)
- [End-to-End Robotic Reinforcement Learning without Reward Engineering](https://mlsanity.com/paper/1904.07854/) (2019)
- [Sequential Neural Models with Stochastic Layers](https://mlsanity.com/paper/1605.07571/) (2016)
- [Temporal Difference Models: Model-Free Deep RL for Model-Based Control](https://mlsanity.com/paper/1802.09081/) (2018)
- [Hierarchical Foresight: Self-Supervised Learning of Long-Horizon Tasks via Visual Subgoal Generation](https://mlsanity.com/paper/1909.05829/) (2019)
- [Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors](https://mlsanity.com/paper/2006.13205/) (2020)
- [Classical Planning in Deep Latent Space: Bridging the Subsymbolic-Symbolic Boundary](https://mlsanity.com/paper/1705.00154/) (2017)
- [Learning Plannable Representations with Causal InfoGAN](https://mlsanity.com/paper/1807.09341/) (2018)
- [Learning and Querying Fast Generative Models for Reinforcement Learning](https://mlsanity.com/paper/1802.03006/) (2018)
- [From Pixels to Torques: Policy Learning with Deep Dynamical Models](https://mlsanity.com/paper/1502.02251/) (2015)

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

- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [Reinforcement Learning with Action-Free Pre-Training from Videos](https://mlsanity.com/paper/2203.13880/) (2022)
- [Latent Plans for Task-Agnostic Offline Reinforcement Learning](https://mlsanity.com/paper/2209.08959/) (2022)
- [Visuo-Tactile Transformers for Manipulation](https://mlsanity.com/paper/2210.00121/) (2022)
- [Planning Goals for Exploration](https://mlsanity.com/paper/2303.13002/) (2023)
- [HARP: Autoregressive Latent Video Prediction with High-Fidelity Image Generator](https://mlsanity.com/paper/2209.07143/) (2022)
- [RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability](https://mlsanity.com/paper/2309.00082/) (2023)
- [Fighting Uncertainty with Gradients: Offline Reinforcement Learning via Diffusion Score Matching](https://mlsanity.com/paper/2306.14079/) (2023)
- [Horizon Generalization in Reinforcement Learning](https://mlsanity.com/paper/2501.02709/) (2025)
- [Free from Bellman Completeness: Trajectory Stitching via Model-based Return-conditioned Supervised Learning](https://mlsanity.com/paper/2310.19308/) (2023)
- [Self-Supervised Reinforcement Learning that Transfers using Random Features](https://mlsanity.com/paper/2305.17250/) (2023)
- [Distributional Successor Features Enable Zero-Shot Policy Optimization](https://mlsanity.com/paper/2403.06328/) (2024)
- [Semantic World Models](https://mlsanity.com/paper/2510.19818/) (2025)
- [Parallel Stochastic Gradient-Based Planning for World Models](https://mlsanity.com/paper/2602.00475/) (2026)
- [Variational Latent Branching Model for Off-Policy Evaluation](https://mlsanity.com/paper/2301.12056/) (2023)
- [Off-Policy Evaluation for Human Feedback](https://mlsanity.com/paper/2310.07123/) (2023)
- [Analytical Survey of Learning with Low-Resource Data: From Analysis to Investigation](https://mlsanity.com/paper/2510.08962/) (2025)
- [Just Round: Quantized Observation Spaces Enable Memory Efficient Learning of Dynamic Locomotion](https://mlsanity.com/paper/2210.08065/) (2022)
- [Grounding Generated Videos in Feasible Plans via World Models](https://mlsanity.com/paper/2602.01960/) (2026)
- [Learning Concept-Based Causal Transition and Symbolic Reasoning for Visual Planning](https://mlsanity.com/paper/2310.03325/) (2023)
- [Context-aware Diversity Enhancement for Neural Multi-Objective Combinatorial Optimization](https://mlsanity.com/paper/2405.08604/) (2024)
- [Deep Generative Models for Decision-Making and Control](https://mlsanity.com/paper/2306.08810/) (2023)
- [Learning Tactile Insertion in the Real World](https://mlsanity.com/paper/2405.00383/) (2024)
- [Off-Policy Selection for Initiating Human-Centric Experimental Design](https://mlsanity.com/paper/2410.20017/) (2024)
- [MInCo: Mitigating Information Conflicts in Distracted Visual Model-based Reinforcement Learning](https://mlsanity.com/paper/2504.04164/) (2025)
- [Wavelet Predictive Representations for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2510.04507/) (2025)

## Abstract (arXiv)

The ability to plan into the future while utilizing only raw high-dimensional observations, such as images, can provide autonomous agents with broad capabilities. Visual model-based reinforcement learning (RL) methods that plan future actions directly have shown impressive results on tasks that require only short-horizon reasoning, however, these methods struggle on temporally extended tasks. We argue that it is easier to solve long-horizon tasks by planning sequences of states rather than just actions, as the effects of actions greatly compound over time and are harder to optimize. To achieve this, we draw on the idea of collocation, which has shown good results on long-horizon tasks in optimal control literature, and adapt it to the image-based setting by utilizing learned latent state space models. The resulting latent collocation method (LatCo) optimizes trajectories of latent states, which improves over previously proposed shooting methods for visual model-based RL on tasks with sparse rewards and long-term goals. Videos and code at https://orybkin.github.io/latco/.