Beyond the Mean: Gaussian Distributional Successor Features for Zero-Shot Non-Linear Reward Adaptation

Abstract

Zero-shot transfer for offline reinforcement learning involves generalizing to a wide range of tasks that are often risk-sensitive, without new interaction. We show here that although Successor Features (SFs) provide a principled framework for transfer through abstracting dynamics from rewards, their typical form is restricted to expected cumulative returns and is thus indifferent to variance in returns. We contend that such a mean-centric approach introduces bias by conflating risky and safe trajectories with identical expected outcomes, as it ignores outcome dispersion critical for risk-aware evaluation. Additionally, we note that distributional methods can be employed to model outcome distributions, but frequently do so at a high computational cost due to the need for discretization or elaborate sampling. To overcome this problem, we introduce a form of successor representation that incorporates an uncertainty envelope over states: Gaussian Distributional Successor Features (GDSF). Rather than inferring an implicit generative process, GDSF learns to approximate the distribution of cumulative outcomes with a Gaussian through recursive moment-matching. We present this framework for efficient representation of aleatoric uncertainty, without resorting to computationally expensive existing distributional methods. We also introduce Latent Outcome Optimization, an inference procedure which finds high-reward outcomes in the data support and drives a goal-conditioned policy. This dual-mode design allows the agent to transfer to non-linear objectives in a zero-shot way while preserving closed-form efficiency for linear tasks. Experiments on the D4RL benchmark show that GDSF yields better performance in risk-sensitive tasks while achieving at least competitive results in standard linear regimes.