Learning Policy Representations for Steerable Behavior Synthesis

TL;DR

This work develops a general policy representation framework that encodes policies as the expected state–action feature map under their occupancy measures, $h_\pi = \mathbb{E}_{d_\pi}[f(s,a)]$, enabling fixed-dimensional, behavior-centered embeddings. The authors combine a set-based estimator within a $\beta$-VAE with a value-aware contrastive objective to produce a smooth, $\delta$-ordered latent space aligned with multiple value functions, and they design a gradient-based optimization procedure to synthesize policies at test time under unseen value constraints. Their architecture couples a Transformer-based encoder, per-task projection heads, and differentiable value predictors, allowing zero-shot steering by optimizing in latent space while staying on the learned manifold. Experiments in Multi-Objective MuJoCo demonstrate accurate policy reconstruction, a structured latent geometry that correlates with returns, and successful constraint-satisfying behavior synthesis without additional environment interaction. Overall, the approach enables controllable, data-efficient behavior synthesis by manipulating policy representations rather than retraining policies.

Abstract

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.

Figures (21)

• Figure 1: The model consists of (1) an encoder $g_{\theta}$ that maps sets of state-action pairs to policy representations $\tilde{h}_\pi$; (2) a set of projectors $\{u_{\psi^{(k)}}\}_{k=1}^K$ that map $\tilde{h}_\pi$ to task-specific embeddings $\{z_\pi^{(k)}\}_{k=1}^K$ and a set of value regressors $\{v_{\xi^{(k)}}\}_{k=1}^K$; (3) a policy $\pi_\phi$ that outputs actions given $\tilde{h}_\pi$ and a query state $s$.
• Figure 2: (a) Distribution of policies in the training dataset, where color intensity represents the number of trajectories falling within each return bin. (b) Relative difference in returns between ground-truth and returns obtained by the latent-conditioned decoder.
• Figure 3: UMAP visualizations of the projected embeddings $\{z_i{^{(k)}}\}_{i=1}^I$ colored with corresponding velocity returns (Left) and energy penalty returns (Right). The learned embeddings exhibit a consistent ordering.
• Figure 4: UMAP visualizations of (a) policy representations trained with our full method, underlayed with contours of return values; and (b) policy representations trained with VAE baseline ($\alpha = 0$), showing a disordered landscape with no clear value-based geometry.
• Figure 5: We compare latent optimization trajectories for a constrained synthesis task using naive gradient descent (Red) versus projected gradient descent (Blue). The plots show the forward return (Left) and energy return (Right) of trajectories collected by the decoder (solid lines) and the corresponding value predictions (dashed lines) over optimization steps.

Theorems & Definitions (4)

• Theorem 2.1
• Proposition 3.1
• proof
• proof