# Latent State Marginalization as a Low-cost Approach for Improving Exploration

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 94 references that are in the map.
- First posted on arXiv 3 Oct 2022.

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

- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (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)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Residual Flows for Invertible Generative Modeling](https://mlsanity.com/paper/1906.02735/) (2019)
- [Provably Efficient Maximum Entropy Exploration](https://mlsanity.com/paper/1812.02690/) (2018)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [Bridging the Gap Between Value and Policy Based Reinforcement Learning](https://mlsanity.com/paper/1702.08892/) (2017)
- [A unified view of entropy-regularized Markov decision processes](https://mlsanity.com/paper/1705.07798/) (2017)
- [Maximum Entropy RL (Provably) Solves Some Robust RL Problems](https://mlsanity.com/paper/2103.06257/) (2021)
- [Deep Kalman Filters](https://mlsanity.com/paper/1511.05121/) (2015)
- [Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data](https://mlsanity.com/paper/1605.06432/) (2016)
- [Discovering and Achieving Goals via World Models](https://mlsanity.com/paper/2110.09514/) (2021)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Deep Variational Reinforcement Learning for POMDPs](https://mlsanity.com/paper/1806.02426/) (2018)
- [Sequential Neural Models with Stochastic Layers](https://mlsanity.com/paper/1605.07571/) (2016)
- [Generative Flow Networks for Discrete Probabilistic Modeling](https://mlsanity.com/paper/2202.01361/) (2022)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)

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

- [Better Training of GFlowNets with Local Credit and Incomplete Trajectories](https://mlsanity.com/paper/2302.01687/) (2023)
- [Generative Augmented Flow Networks](https://mlsanity.com/paper/2210.03308/) (2022)
- [Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets](https://mlsanity.com/paper/2305.17010/) (2023)
- [Distributional GFlowNets with Quantile Flows](https://mlsanity.com/paper/2302.05793/) (2023)
- [Improving GFlowNets for Text-to-Image Diffusion Alignment](https://mlsanity.com/paper/2406.00633/) (2024)
- [Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow](https://mlsanity.com/paper/2405.13629/) (2024)
- [Structure in Deep Reinforcement Learning: A Survey and Open Problems](https://mlsanity.com/paper/2306.16021/) (2023)
- [Value Gradient Guidance for Flow Matching Alignment](https://mlsanity.com/paper/2512.05116/) (2025)
- [Sequential Monte Carlo for Policy Optimization in Continuous POMDPs](https://mlsanity.com/paper/2505.16732/) (2025)
- [DODT: Enhanced Online Decision Transformer Learning through Dreamer's Actor-Critic Trajectory Forecasting](https://mlsanity.com/paper/2410.11359/) (2024)
- [Revisiting Mixture Policies in Entropy-Regularized Actor-Critic](https://mlsanity.com/paper/2605.09157/) (2026)
- [On the Parallel Complexity of Multilevel Monte Carlo in Stochastic Gradient Descent](https://mlsanity.com/paper/2310.02402/) (2023)
- [Learning Temporal Abstractions via Variational Homomorphisms in Option-Induced Abstract MDPs](https://mlsanity.com/paper/2507.16473/) (2025)

## Abstract (arXiv)

While the maximum entropy (MaxEnt) reinforcement learning (RL) framework -- often touted for its exploration and robustness capabilities -- is usually motivated from a probabilistic perspective, the use of deep probabilistic models has not gained much traction in practice due to their inherent complexity. In this work, we propose the adoption of latent variable policies within the MaxEnt framework, which we show can provably approximate any policy distribution, and additionally, naturally emerges under the use of world models with a latent belief state. We discuss why latent variable policies are difficult to train, how naive approaches can fail, then subsequently introduce a series of improvements centered around low-cost marginalization of the latent state, allowing us to make full use of the latent state at minimal additional cost. We instantiate our method under the actor-critic framework, marginalizing both the actor and critic. The resulting algorithm, referred to as Stochastic Marginal Actor-Critic (SMAC), is simple yet effective. We experimentally validate our method on continuous control tasks, showing that effective marginalization can lead to better exploration and more robust training. Our implementation is open sourced at https://github.com/zdhNarsil/Stochastic-Marginal-Actor-Critic.
