Generative Actor Critic

TL;DR

The paper addresses offline-to-online RL by moving from expected-return optimization to learning a joint distribution over trajectories and returns, p(τ,y), and performing test-time inference for policy improvement. It introduces Generative Actor-Critic (GAC), a latent-variable framework with a continuous latent plan z and factorization p(τ,y,z)=p(τ|z)p(y|z)p(z), enabling exploitation through gradient-based latent optimization and exploration through sampling z conditioned on optimistic targets y^+. Trained with an ELBO objective using variational Bayes, GAC achieves strong offline performance and substantially better offline-to-online improvement on Gym-MuJoCo and Maze2D compared with baselines, even when step-wise rewards are absent. The work reveals emergent latent-space structure, including implicit world-modeling and cognitive-map-like organization, suggesting that modeling the full trajectory-return distribution can underpin robust planning in complex environments.

Abstract

Conventional Reinforcement Learning (RL) algorithms, typically focused on estimating or maximizing expected returns, face challenges when refining offline pretrained models with online experiences. This paper introduces Generative Actor Critic (GAC), a novel framework that decouples sequential decision-making by reframing \textit{policy evaluation} as learning a generative model of the joint distribution over trajectories and returns, $p(τ, y)$, and \textit{policy improvement} as performing versatile inference on this learned model. To operationalize GAC, we introduce a specific instantiation based on a latent variable model that features continuous latent plan vectors. We develop novel inference strategies for both \textit{exploitation}, by optimizing latent plans to maximize expected returns, and \textit{exploration}, by sampling latent plans conditioned on dynamically adjusted target returns. Experiments on Gym-MuJoCo and Maze2D benchmarks demonstrate GAC's strong offline performance and significantly enhanced offline-to-online improvement compared to state-of-the-art methods, even in absence of step-wise rewards.

Figures (8)

• Figure 1: Conceptual illustration of exploitation and exploration. After modeling the data distribution (gray), GAC supports distinct test-time inference queries for various decision-making objectives. For exploitation (blue), the objective is to maximize the expected return, leading to a focused, low-variance policy that targets high-certainty outcomes. For exploration (green), the generative model is conditioned on a target distribution shifted towards higher returns, guiding the search for novel and potentially superior trajectories.
• Figure 2: Illustration of return distributions in Mujoco. The left column displays the return distributions achieved by various inference objectives after offline pre-training, highlighting the explorative and exploitative behavior. The middle column presents a split violinplot for the inference objectives of GAC-$\mathbb{E}[y]$ and GAC-$p(y^+)$, showing the progressive performance gains and the explorative and exploitative behaviors at three sequential fine-tuning stages. The right column visualizes the progressive improvement in the datasets' quality, as evidenced by the upward shift of the dataset's return distribution over the fine-tuning stages.
• Figure 3: Illustration of online improvement and emergent properties in Maze. The top row heatmaps visualize performance from different starting cells. The first two panels compare inference strategies at Stage 1, demonstrating the superiority of the exploitative strategy GAC-$\mathbb{E}[y]$ (Stage 1 Score) over sampling from the prior (Stage 1 Prior Score). The progression from Stage 1 to Stage 3 (second to fourth panels) then illustrates a monotonic improvement in online fine-tuning. Subplots in the bottom row visualize the model's emergent capabilities. The model is trained on suboptimal, axis-aligned trajectories (far left). After fine-tuning, GAC generates novel, efficient diagonal shortcuts (second from left). This advanced planning is likely supported by an emergent cognitive map: a t-SNE visualization shows the latent space of plans, $z$, self-organizing into distinct clusters (third from left), which directly correspond to spatially coherent regions of starting states, $s_0$ (far right).
• Figure 4: Analysis of actor and critic. The left two panels illustrate the actor's consistency. For each target return (x-axis), we infer 50 latent plans to generate corresponding trajectories, and plot the mean of their evaluated returns (y-axis). The right two panels demonstrate the critic's consistency. For each trajectory from the dataset with a ground truth return (x-axis), we infer 50 latent plans and use their average predicted return as the y-axis value. Points with higher ELBO values are less transparent and closer to the ideal $y=x$ line, indicating more reliable predictions. The Stage 3 plot distinguishes original (blue) from newly collected online data (red).
• Figure 5: Experiment on replanning. While the initial plan from $A$ rarely succeeds, with a replanning at the critical intermediate state $B$, the agent can correct its course towards the goal. In the left we visualize a failure and a correction starting from $A$ and a reference from $B$. In the right we report the success rate.