# Learning Policy Representations for Steerable Behavior Synthesis

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 40 references that are in the map.
- First posted on arXiv 29 Jan 2026.

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Diversity is All You Need: Learning Skills without a Reward Function](https://mlsanity.com/paper/1802.06070/) (2018)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Deep metric learning using Triplet network](https://mlsanity.com/paper/1412.6622/) (2014)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [A Survey of Multi-Objective Sequential Decision-Making](https://mlsanity.com/paper/1402.0590/) (2014)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [DART: Noise Injection for Robust Imitation Learning](https://mlsanity.com/paper/1703.09327/) (2017)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Autonomous Agents Modelling Other Agents: A Comprehensive Survey and Open Problems](https://mlsanity.com/paper/1709.08071/) (2017)
- [Robust Imitation of Diverse Behaviors](https://mlsanity.com/paper/1707.02747/) (2017)
- [Rank-N-Contrast: Learning Continuous Representations for Regression](https://mlsanity.com/paper/2210.01189/) (2022)
- [Control functionals for Monte Carlo integration](https://mlsanity.com/paper/1410.2392/) (2014)
- [PD-MORL: Preference-Driven Multi-Objective Reinforcement Learning Algorithm](https://mlsanity.com/paper/2208.07914/) (2022)
- [RSL-RL: A Learning Library for Robotics Research](https://mlsanity.com/paper/2509.10771/) (2025)
- [Learning Policy Representations in Multiagent Systems](https://mlsanity.com/paper/1806.06464/) (2018)
- [Policy Evaluation Networks](https://mlsanity.com/paper/2002.11833/) (2020)
- [A Deep Policy Inference Q-Network for Multi-Agent Systems](https://mlsanity.com/paper/1712.07893/) (2017)
- [Fast Adaptation via Policy-Dynamics Value Functions](https://mlsanity.com/paper/2007.02879/) (2020)

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

- [Unsupervised Behavioral Compression: Learning Low-Dimensional Policy Manifolds through State-Occupancy Matching](https://mlsanity.com/paper/2603.27044/) (2026)

## Abstract (arXiv)

Given a Markov decision process (MDP), we seek to learn representations for a range of policies to facilitate behavior steering at test time. As policies of an MDP are uniquely determined by their occupancy measures, we propose modeling policy representations as expectations of state-action feature maps with respect to occupancy measures. We show that these representations can be approximated uniformly for a range of policies using a set-based architecture. Our model encodes a set of state-action samples into a latent embedding, from which we decode both the policy and its value functions corresponding to multiple rewards. We use variational generative approach to induce a smooth latent space, and further shape it with contrastive learning so that latent distances align with differences in value functions. This geometry permits gradient-based optimization directly in the latent space. Leveraging this capability, we solve a novel behavior synthesis task, where policies are steered to satisfy previously unseen value function constraints without additional training.