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

1 Introduction

Developing reinforcement learning agents capable of zeroshot adaptation to arbitrary objectives, ranging from simple navigation to complex, risk-sensitive constraints, remains a persistent challenge in the quest for generalist intelligence [Touati et al., 2023]. To achieve such versatility, a system must structurally decouple its knowledge of environmental dynamics from specific task incentives. While ModelBased RL (MBRL) offers a theoretical avenue for such decoupling, practical implementations often rely on autoregressive rollouts that suffer from compounding errors and prohibitive inference costs [Xu et al., 2022]. Successor Features (SFs) [Barreto et al., 2017; Barreto et al., 2018] offer a rigorous alternative by factorizing the value function into two distinct components: a predictive successor representation that summarizes future state occupancies, and a task-specific weight vector. This structural decomposition reduces the evaluation of novel rewards from a complex planning problem to a simple linear projection, enabling near-instantaneous task inference without the instability of learned simulators.

However, the computational efficiency of standard SFs relies on a restrictive linearity assumption. Conventional approaches squash the stochastic distribution of future outcomes into a single expectation vector [Barreto et al., 2017]. By doing so, they act as if the cumulative future is a deterministic quantity rather than a random one. While this reduction is sufficient for linear reward specifications, it fails to account for the higher-order statistics required to evaluate non-linear or safety-constrained objectives [Bellemare et al., 2017; Tamar et al., 2012]. As a second-order analysis reveals-and as is corroborated by recent theoretical findings on distributional bounds [Wang et al., 2024]-the true utility under such complex rewards depends crucially not only on the mean outcome but also on its dispersion. By discarding this variance information, mean-centric representations suffer from a systematic structural deficiency: they conflate highly volatile, dangerous trajectories with stable ones whenever their expected outcomes align, rendering the agent incapable of enforcing safety constraints during zero-shot adaptation.

To bridge this gap, recent research has diverged into two primary paradigms, yet each introduces significant computational or representational trade-offs. Discrete distributional methods, such as those extending C51 [Bellemare et al., 2017] to successor representations [Gimelfarb et al., 2021], attempt to capture uncertainty via categorical approximations. However, these approaches typically suffer from the curse of dimensionality: discretizing a high-dimensional outcome space incurs prohibitive parameter growth, scaling poorly to complex environments [Reinke and AlamedaPineda, 2023]. Conversely, generative approaches like DiSPO leverage diffusion models to represent the full outcome distribution [Zhu et al., 2024]. While theoretically expressive, these models rely on iterative sampling chains that impose substantial computational overhead [Janner et al., 2022]. Consequently, to maintain tractable planning speeds, such methods may use simplified reward signals during the policy improvement phase. This fallback partially undermines their distributional advantages, limiting their ability to perform true zero-shot optimization against non-linear risk metrics.

We introduce Gaussian Distributional Successor Features (GDSF) to address these limitations. Instead of relying on heavy generative priors or coarse discretization, GDSF models the successor measure using a state-conditioned Gaussian approximation learned through recursive moment matching. This explicit parameterization provides a lightweight mechanism to capture aleatoric uncertainty, avoiding the computational overhead of iterative sampling models. Crucially, this structure supports a dual-mode adaptation capability: it preserves the closed-form efficiency of standard SFs for linear reward specifications, while enabling tractable Latent Outcome Optimization for non-linear, risk-aware objectives. By conditioning a policy on these optimized latent outcomes, the agent can adapt to novel safety constraints zero-shot, ensuring precise adherence to the test-time reward landscape without additional fine-tuning.

Our main contributions are:

• Gaussian Distributional Successor Features (GDSF): A second-order Taylor analysis shows why mean-centric SFs are structurally misaligned with non-linear reward evaluation. GDSF augments SFs with an explicit parameterization of aleatoric uncertainty, correcting this bias without incurring the inference cost of full generative models.
• Dual-Mode Zero-Shot Adaptation: We combine closed-form recursive moment matching with latent outcome optimization. This mechanism traverses the uncertainty envelope to identify risk-compatible outcomes during zero-shot adaptation.
• Experimental Validation: Experiments on D4RL benchmarks [Younis et al., 2024] show that GDSF alleviates the structural bias of typical SFs. It also achieves strong performance on non-linear reward landscapes while substantially reducing inference time relative to heavier distributional or generative baselines.

2 RelatedWork

2.1 Zero-Shot Transfer with Successor Features

SFs allow for zero-shot adaptation by splitting environment dynamics from rewards [Barreto et al., 2017]. Yet, the standard toolkit-including Universal SFs [Borsa et al., 2019] and Forward-Backward representations [Touati and Ollivier, 2021]-suffers from a mean-centric bias. By betting on reward linearity, these methods effectively blind the agent to aleatoric uncertainty, leading to errors whenever risk matters. Newer distributional variants (like SFR [Reinke and Alameda-Pineda, 2023] or generative approaches [Zhu et al., 2024]) do try to map the full outcome landscape. But they come with a heavy price tag: computational drag from complex parameterizations like diffusion models. This leaves a specific gap: we need a method that captures just enough uncertainty to handle risk, without carrying the computational baggage of full distributional modeling.

2.2 Distributional Reinforcement Learning

Distributional RL approximates the full return distribution rather than just its expectation [Bellemare et al., 2017], leveraging algorithms like QR-DQN [Dabney et al., 2018] to improve performance. Crucially, recent theory proves that learning distributions enables tighter second-order bounds that scale with the variance of the return [Wang et al., 2024]. This insight strongly motivates capturing outcome dispersion. While methods like RaSF [Gimelfarb et al., 2021] and DiSPO [Zhu et al., 2024] extend this to successor features, they often suffer from restrictive assumptions or high sampling latency due to diffusion models. In contrast, GDSF leverages the benefits of second-order information via a lightweight Gaussian approximation, ensuring analytic efficiency for zero-shot adaptation.

3 Preliminaries

Environmental and Outcome Indicators. We view the controlled Markov process as an MDP without rewards, described by the tuple (S,A, P, γ). It is natural to associate a vector-valued signal ϕ(s, a) ∈ Rd with each transition. To this end, we define the cumulative discounted outcome of a policy π when starting in state s as the following random vari-

The most straightforward approach to represent SFs is to interpret the expectation of this random variable, denoted as ψπ (s) ≜ E[Zπ (s)].

Zero-Shot Reward Specification. In zero-shot transfer, the goal is to generalize to tasks that are unseen during the training phase. At evaluation time, a task is typically specified through a utility function f : Rd → R. The objective then becomes maximizing the expected utility:

However, the majority of previous SF methods require that f needs to be linear (that is, f(z) = w⊤z). Under this linearity assumption, the evaluation reduces to a simple dot product J π (s) = w⊤ψπ (s). This kind of representation is typically not applicable when you have non-linear reward specifications, because it relies only on expectations and ignores the distributional information that comes from cumulative outcomes.

4 Methodology

Our goal in this work is to achieve zero-shot adaptation to non-linear rewards: once the policy π(a|s, z) is trained offline, no further retraining or finetuning is required when a new task is specified at test time. We avoid modeling full distributional dynamics; instead, we preserve only the rewardrelevant variance needed for risk-aware selection at inference time. Injecting this lightweight signal into standard successor features yields a representation that is both risk-attentive and computationally tractable.

Overview of GDSF. GDSF works in two stages. During offline learning, we train networks to predict the mean µθ(s) and variance σ2θ (s) of cumulative outcomes using the offline dataset. At test time, when given a new reward function f, we sample outcomes from the learned Gaussian and select the one that maximizes f while staying in high-density regions. A goal-conditioned policy then executes this selected outcome. Figure 1 shows this process.

4.1 Motivation

Mean-centric SFs work well when the environment has linear rewards, but they become inadequate once the reward function f at test time is non-linear. Rather than reasoning about the full outcome distribution, they simply apply the reward function to the expected cumulative outcome:

This kind of simplification has drawbacks. It forces the stochastic cumulative outcome Zπ (s) to be mapped, in effect, to a deterministic point-which is just its mean. To tackle this bias properly, we need a more honest representation of which outcomes are actually attained.

Second-Order Taylor Analysis. The specific cause of this error can be identified by looking at the second-order Taylor expansion of f around the mean outcome µ = E[Zπ (s)]:

Tr(·) denotes the trace operator. Here, ∇2 f(µ) represents the Hessian of the reward function, which captures its curvature and Cov(Z) is the covariance in cumulative outcomes. This expansion serves as a diagnostic tool for understanding locally where mean-based evaluation fails, rather than being a theoretical guarantee that holds globally.

Implications of Outcome Dispersion. Eq. (4) reveals the weakness of mean-only representations. In the linear setting (∇2 f = 0), there is no second-order term and meanbased evaluation is exact, explaining why standard SFs perform well under linear reward specifications. More generally, for non-linear reward functions, the expected utility also depends on the variation across outcomes through the covariance term. Mean-centric approaches ignore this dependence, leading to misaligned evaluation when outcome variability is high. This motivates us to augment successor representations with lightweight uncertainty information for improved zeroshot evaluation.

4.2 Gaussian Distributional Successor Features

A simple generalization of Standard Successor Features in order to account for the dispersion in non-linear reward evaluation is provided as follows: We approximate the uncertainty over outcomes with a state-conditioned diagonal Gaussian:

This Gaussian captures the uncertainty needed for riskaware evaluation without modeling full multimodal dynamics. We use a state-conditioned distribution to separate reward evaluation from action-level control-it acts as a sufficient statistic for outcome dispersion while remaining lightweight.

Control is handled by a separate goal-conditioned policy π(a|s, z), which we train on the offline dataset using standard hindsight relabeling techniques. This separation means our uncertainty model only needs to capture risk-relevant distinctions for evaluation, not generative control information.

4.3 Learning viaMoment-Matching Regression

The model is trained so that it fits the first two moments of the cumulative outcome distribution. This is done via temporaldifference learning.

First Moment (Mean). The mean µθ(s) is trained by minimizing the squared Bellman error:

SecondMoment (Outcome Dispersion). We focus on dispersion along the sampled trajectory rather than branching across next states, since the latter is orthogonal to our riskaware selection objective and tends to underestimate uncertainty at poor states. Variance is therefore propagated recursively along observed transitions.

The way we implement this is by forcing the variance network σ2θ (s) to align with a bootstrapped target distribution:

Here, KL(·∥·) denotes the Kullback-Leibler divergence. where yµ = ϕ(s, a) + γµ¯ θ(s′ ) and

Here ϵ absorbs the irreducible transition noise. The objective yields a conservative proxy for long-horizon dispersion that is sufficient for identifying reward-sensitive regions and more stable than exhaustive distributional modeling.

4.4 Inference-Time Evaluation and Latent Outcome Optimization

When we get to evaluation time and are given a non-linear reward function f, we estimate the expected utility by sampling plausible cumulative outcomes. These outcomes need to be consistent with the uncertainty representation that we learned:

The way we view inference here is as a selection problem over feasible outcomes. This is different from treating it as a planning procedure or as policy optimization. To make zero-shot adaptation possible, what we do is perform latent outcome optimization. This filters for the best outcomes that are feasible:

What the constraint does is guarantee that the outcomes we choose are restricted to regions that were commonly reached in the offline dataset. This way, we avoid out-of-distribution hallucinations when doing inference. What we're essentially doing is swapping test-time policy optimization for a policy that is constrained over outcome selection instead.

This enables instantaneous zero-shot adaptation to new reward specifications without environment interaction.

4.5 Theoretical Analysis

In this section, we analyze the approximations inherent in GDSF. Our goal is to justify why the proposed uncertainty envelope, although simplified, can serve as a sufficient and conservative statistic for zero-shot reward evaluation.

Definition 1 (Conservative Outcome Envelope). Let's consider the learned Gaussian approximation ξθ(z|s) ≜ N(µθ(s), diag(σ2θ (s))). What we define as the Conservative Outcome Envelope is the set of outcomes that reside within k standard deviations of the mean:

In practical terms, this envelope delimits the specific region of the outcome space that we consider to be reliably supported by the offline dataset that we have. This elementwise constraint arises naturally from the diagonal Gaussian structure: since each outcome dimension is modeled independently, each dimension i is bounded by its own σθ,i(s) without conflating dimensions of different scales. In our implementation, the latent outcome optimization in Sec. 4.4 restricts z∗ to Ek(s) as an axis-aligned, tractable surrogate for the density level set Z(s); both characterize the high-density support of ξθ.

Proposition 1 (Variance Lower Bound Property). It's worth recalling that for the true cumulative outcome variable Zπ (s), the variance decomposes in a natural way via the Law of Total Variance:

Here, P π (s′ |s) = Ea∼π(·|s)[P(s′ |s, a)] denotes the policyinduced transition kernel. Our recursive objective for σ2θ (Eq. 7) has been designed to explicitly capture the first term, which is the accumulated dispersion. At the same time, it intentionally omits the second term, which is the cross-state branching. What this means is that the learned variance constitutes a theoretical lower bound on the total outcome variance: E[σ2θ (s)] ≤ Var(Zπ (s)).

This property is important because it shows that our Gaussian envelope provides a conservative abstraction of outcome uncertainty under non-linear rewards. Here "conservative" refers to anchoring inference to the high-density support (Remark 2), rather than inflating the envelope with branching variance.

To unpack this decomposition intuitively: the first term captures the variance attributable to stochastic execution conditional on reaching each next state-this irreducible noise is what our recursive objective (Eq. 7) is designed to track. The second term captures additional variance from stochastic branching across different next states-this cross-state variability is intentionally omitted by our model. Because we retain only the first term, the learned σ2θ (s) constitutes a lower bound on total outcome variance under the true model, making our envelope conservative by construction.

Remark 1 (Implications for Evaluation). What this lower bound property implies is that GDSF constructs an envelope that is tight around the anticipated trajectory. It isolates the irreducible execution noise from the multi-modal branching dynamics that can occur. When it comes to zero-shot evaluation, this is actually a feature rather than a bug. It mitigates the "average-case" bias that you get from mean-only methods. At the same time, it circumvents the excessive pessimism that would arise if we were to mistake unique branching routes for chaos that has high variance.

Remark 2 (Safeguard Against Hallucination). An important point is that the latent outcome optimization z∗ = argmaxf(z) gets strictly constrained to z ∈ Ek(s). What Ek(s) does is anchor the search to the high-density region that comes from the behavior policy. Because of this, the selected target z∗ is in-distribution by definition. If we assume that the goal-conditioned policy π(a|s, z) generalizes reasonably well within the data support (which is a standard assumption in offline RL), then this constraint effectively prevents the agent from pursuing "hallucinated" high-reward outcomes. These would be outcomes that are physically unreachable in practice.

5 Experiments

We assess GDSF on zero-shot adaptation for non-linear rewards on the D4RL benchmark. Our experiments examine the behavior of different successor representations under an evaluation-time reward shift, particularly for non-linear rewards. We consider:

(RQ1) In comparison to other methods that are mainly designed for linear reward specifications, how well can GDSF generalize to holdout non-linear rewards without retraining?

• (RQ2) To what extent does adding a Gaussian outcome envelope over the successor representation improve policy evaluation when the reward function is sensitive to outcome dispersion?

(RQ3) How does GDSF perform with zero-shot reward shift, compared to offline RL and successor feature baselines?

5.1 Experimental Setup

Environments and Outcome Space. We use two continuous-control domains: Navigation (AntMaze umaze, medium, large) and Locomotion (MuJoCo HalfCheetah, Hopper, Walker2d). For the Navigation domain, we define the outcome ϕ(s, a) as the (x, y) coordinates that specify position. For Locomotion, the outcome includes both the forward velocity and the control costs that are incurred. Thus, ϕ(s, a) ∈ R2 in both domains. Figure 2 gives a visual overview of these environments.

Test-Time Reward Specifications. To evaluate how well zero-shot adaptation works under reward transformations that happen at evaluation time, we look at three different families of test-time objectives. These families differ in what their functional form is and how they depend on cumulative outcomes:

1. Linear Reward (Base): f(z) = w⊤z. This serves as a reference case. In this case, linear composition of suc-
2. cessor features is exact, and this approach is commonly adopted when doing standard SF-based evaluation.
3. ThresholdReward (Const): f(z) =w⊤z·⊮[∥z∥∞ < δ]. Thus, the reward becomes zero whenever ∥z∥∞ ≥ δ. Here ⊮[·] denotes the indicator function, equal to 1 when its argument is true and 0 otherwise. What these objectives do is impose hard constraints on the cumulative outcomes. This happens via a discontinuous feasibility region.
4. Gaussian Reward (Shape): f(z) = exp(-∥z - g∥2 /β). These objectives are for precision-sensitive tasks that emphasize how close you are to a target outcome g. The decay is smooth and is governed by the parameter β.

Baselines. We compare GDSF against baselines that fall into three categories. This comparison lets us rigorously evaluate whether modeling outcome dispersion is necessary and whether our Gaussian representation is efficient:

• Mean-Centric SFs: USFA [Borsa et al., 2019] and FB [Touati and Ollivier, 2021], which compress the outcome distribution into a single expectation E[ϕ], isolating the gain from explicitly capturing outcome dispersion.
• Distributional & Generative SFs: SFR [Reinke and Alameda-Pineda, 2023] (categorical discretization) and DiSPO [Zhu et al., 2024] (diffusion-based planning); these test whether our lightweight Gaussian matches their expressivity at lower compute cost.
• Standard Offline RL: IQL [Kostrikov et al., 2022], MOPO [Yu et al., 2020], and COMBO [Yu et al., 2021]. These model-free baselines reference standard offline performance and highlight SF decomposition flexibility.

Implementation and Inference Protocol. GDSF models the outcome distribution using a dual-head MLP that outputs the mean µθ(s) and variance σ2θ (s) for each state.

At inference, we adapt the evaluation to the task type:

• For linear rewards, we perform standard analytic evaluation using the predicted mean.
• For non-linear rewards, we apply latent outcome optimization: we sample K = 64 candidate outcomes from the Gaussian envelope N(µ, σ2 ) and select the outcome z∗ = argmaxf(z). This explicitly leverages the learned uncertainty to filter out risky trajectories.

We report mean episodic returns over 10 episodes per seed, averaged across 5 random seeds.

5.2 Comparative Experiment

Table 1 reports quantitative results on the navigation and locomotion domains. We structure the analysis according to the three test-time reward families (Sec. 5.1), ordered by increasing sensitivity to outcome dispersion.

Performance on Linear Rewards (Mean-Sufficient Baseline). We first consider linear rewards, where the reward function is flat (∇2 f = 0) and mean-based evaluation is exact. This setting verifies that GDSF introduces no unnecessary conservatism when uncertainty is irrelevant.

As expected, GDSF achieves performance comparable to strong offline RL baselines (IQL, COMBO) and mean-centric SF methods (USFA). For instance, on AntMaze-Large, GDSF obtains 394±43, closely matching DiSPO (402±45). This confirms that the dual-mode planner correctly defaults to mean-based evaluation in linear regimes, introducing neither bias nor instability.

Threshold Non-Linearity (Safety Constraints). Threshold rewards impose sharp discontinuities whenever a constraint is violated. This makes the setting highly sensitive to outcome dispersion: two policies with the same expected outcome can yield widely different utility. We find that meancentric SFs (USFA, FB) collapse in this setting because their structural evaluation bias makes them blind to the risk of threshold violations. For example, on AntMaze-Large, the performance of USFA plummets from 264 ± 42 to 98 ± 35. In contrast, GDSF achieves significantly higher returns (312 ± 52). It even scores higher than the generative baseline DiSPO (608 vs. 534 on AntMaze-Umaze, Threshold regime). We attribute this gain to our explicit Gaussian parameterization. While DiSPO relies on costly iterative sampling to estimate constraint violation probabilities, GDSF directly leverages the predicted variance σ2θ in latent outcome w/o Latent Outcome Optimization 77.81 ± 1.11 -9.4% Validates Active Outcome Selection optimization, enforcing constraints with greater precision and computational efficiency.

Table 3: Ablation of GDSF Mechanisms. We isolate the impact of explicit variance modeling and latent outcome optimization. Removing the variance constraint (reverting to a Mean-Only baseline) causes catastrophic failure on risk-sensitive tasks, validating the structural bias analysis in Sec. 3.

Gaussian Non-Linearity (Precision-Sensitive Objectives). Gaussian-shaped rewards impose smooth, exponential penalties for deviations from a target outcome, testing the agent's ability to minimize outcome dispersion. This setting highlights the challenges generative SFs face in high-precision, zero-shot optimization. While DiSPO models the full distribution, its sampling-based inference struggles to identify high-precision outcomes. In contrast, GDSF explicitly separates mean and variance, allowing latent outcome optimization to actively select the outcome closest to the target g within the learned envelope σ2θ , thereby enabling precisionaware adaptation. This results in higher returns, e.g., 2418 ± 348 on Hopper-Medium versus 2186 ± 405 for DiSPO.

Comparison with Fixed-Policy Baselines. Standard offline RL algorithms (IQL, MOPO) optimize policies with respect to the original training reward, and their performance under linear shifts reflects only passive re-evaluation of the original behavior (non-adaptation) under the new reward - such algorithms cannot achieve success on a new objective that requires an adapted policy. In contrast, GDSF performs active inference-time adaptation: it exploits the learned dynamics and selects the lowest cost outcome according to the Gaussian envelope, directly steering behavior to achieve higher expected rewards.

5.3 Ablation Study

We perform ablation studies to quantify the contribution of each key GDSF component. Specifically, we examine the impact of explicit Gaussian parameterization (supported by Prop. 1) and latent outcome optimization (enabling the dualmode inference mechanism). Results (Tables 2 and 3) are obtained on PointMaze-UMaze-v3 under a composite nonlinear reward (ϕ-quadratic plus cross-dimension interaction terms) designed to stress-test the mean and variance channels of the architecture.

Feature Encoder Architecture (left side of Table 2). We analyze the effect of the state encoder. Consistent with prior work in kernel-based RL, Random Fourier Features (RFF) provide the best balance between performance (85.8) and training stability (W1 error 0.13). Learnable MLPs achieve a lower mean return (81.4 vs. 85.8 for RFF), possibly due to spectral bias when fitting high-frequency reward landscapes. Hence, RFF is used as the default encoder.

Efficiency of Distributional Representation (right side of Table 2). A core hypothesis of GDSF is that a lightweight Gaussian approximation suffices for control, making highdimensional categorical parameterizations unnecessary. The right side of Table 2 provides clear empirical support: the 2parameter Gaussian representation (µ, σ2 ) outperforms C51 (51d) and Quantile Regression (51d) baselines (85.8 vs. 79.8/81.5). This suggests that, in continuous control, the first two moments constitute a sufficient statistic for risk-sensitive decision-making, whereas high-dimensional categorical distributions may overfit to sampling noise and reduce generalization.

Impact of Variance and Optimization Mechanisms. Table 3 presents an ablation study quantifying the specific contributions of GDSF components. We first examine the necessity of variance modeling. Removing the uncertainty envelope (w/o Variance) causes returns to precipitate from 85.84 to 70.12, a drop of 18.3%. This result empirically validates the structural bias analysis in Sec. 3, demonstrating that without second-moment correction, the agent fails to distinguish between stable and volatile trajectories, leading to suboptimal risk assessment.

We further analyze the benefit of latent outcome selection. Disabling this inference-time optimization (w/o Latent Opt.) and relying on passive sampling results in a 9.4% performance decline (77.81). This confirms that the learned Gaussian envelope functions as a navigable feasible region, enabling the agent to actively identify and select high-utility outcomes within the support, rather than settling for the average outcome of the underlying behavior policy.

6 Conclusion

We addressed the structural bias of mean-centric successor features by showing that they misvalue policies under non-linear reward specifications. To resolve this, we introduced GDSF, which augments standard representations with a lightweight, state-conditioned uncertainty envelope via recursive moment matching, formalized as a lower bound on total variance that isolates accumulated execution noise from branching dynamics.

This formulation supports a dual-mode inference mechanism: it preserves closed-form efficiency for linear rewards while enabling tractable Latent Outcome Optimization for risk-sensitive objectives, bypassing the computational overhead of iterative planning or generative sampling.

Future directions include extending the explicit parameterization to Mixture-of-Gaussians for branching dynamics, and leveraging the learned envelope for active risk-aware exploration in online settings.

Acknowledgements

This work was supported in part by RC Grant No. EP/Y027787/1, in part by the National Natural Science Foundation of China under Grant 62272391 and Grant 62402391, in part by the Fundamental Research Funds for the Central Universities under Grant No. D5000250275, and in part by the Fundamental and Interdisciplinary Disciplines Breakthrough Plan of the Ministry of Education of China under Grant JYB2025XDXM116.