PaperlayerView exact revision on arXiv

Source-linked AI summary

Robust Successor Features

Erik Nikulski, Yamen Habib, Vicenç Gomez, Anders Jonsson, Rubén Moreno-Bote, Javier Segovia-Aguas

arXiv:2609.31016v1cs.LG

TL;DR

Transfer RL generally handles reward changes and robust RL handles transition uncertainty, but the successor-feature and GPI setting lacked a guarantee spanning both. This paper introduces robust successor features for linear MDPs and extends GPI to account for transition mismatch. On grid-based experiments, RSF generalizes to unseen reward-dynamics combinations with higher policy accuracy than the reported baseline and ablations.

  • Problem

    Transfer methods typically assume shared transition kernels and robust RL methods fixed rewards, leaving no successor-feature and GPI suboptimality guarantee for tasks differing in both.

    Source: §1 · p. 2

  • Method

    Robust successor features factorize value into a task-agnostic representation and task-specific linear weights for reward and transition dynamics in feature-based linear MDPs.

    Source: §1 · p. 2

  • Results

    93.6% policy accuracy on unseen reward-dynamics combinations, above the USFA baseline (86.7%) and reward-only (86.8%) and dynamics-only (91.3%) ablations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The GPI bound explicitly accounts for transition-kernel mismatch and recovers the reward-only successor-feature bound as dynamics become shared.

    Source: §1 · p. 2

    The work assumes linear MDPs and leaves learning the reward and dynamics feature maps beyond its scope.

    Source: §1 · p. 2, §4 · p. 7

Abstract

from arXiv · show

Generalization in Reinforcement Learning (RL) refers to the ability to execute close-to-optimal policies in unseen tasks after the agent has been trained on a different set of tasks. Building on the seminal work of the successor representation and further adaptations with function approximation, Transfer in RL has traditionally focused on generalizing to tasks that only differ in the reward function. A decade after the introduction of the successor representation, Robust RL emerged simultaneously from several articles in the field of operations research. In Robust RL, the transition kernel is unknown, and the goal is to maximize the expected reward under this uncertainty. Our work unifies these two paradigms through robust successor features, which generalize across both the reward function and the transition kernel, under the assumption that tasks are linear Markov Decision Processes. We derive a bound on Generalized Policy Improvement (GPI) that explicitly quantifies how performance degrades with the mismatch between transition kernels, recovering existing successor-feature guarantees when dynamics are shared. Finally, the generalization capabilities of robust successor features are validated on several grid-based benchmarks and compared to previous alternatives that focus solely on either the reward or the transition kernel.

1 Introduction

Transfer RL typically generalizes across rewards, while robust RL addresses uncertainty in dynamics; the paper combines both in robust successor features for linear MDPs. It extends the GPI guarantee to dynamics mismatch and reports grid-based generalization across unseen reward-dynamics combinations.

1 Introduction

  • Prior transfer methods assume shared transition kernels and robust RL methods assume fixed rewards; the successor-feature and GPI setting lacked a guarantee when both differ.The paper frames this as the gap its unified approach addresses.Source: §1 · pp. 1–2
  • RSF unifies transfer and robust RL for feature-based linear MDPs, factorizing value into a task-agnostic representation and task-specific reward and transition-kernel weights.Combined with GPI, the representation generalizes zero-shot across changes in both reward and dynamics.Source: §1 · p. 2
  • The extended GPI bound makes dynamics mismatch explicit, and its mismatch-dependent term vanishes as reference dynamics approach target dynamics, recovering the reward-only bound.The recovered special case is the bound of Barreto et al. [2017].Source: §1 · p. 2
  • 93.6% policy accuracy on unseen reward-dynamics combinations, exceeding the USFA baseline (86.7%) and reward-only (86.8%) and dynamics-only (91.3%) ablations.The grid-based experiments validate the RSF algorithm against alternatives focused on either reward or dynamics variation.Source: §1 · p. 2

2 Preliminaries

The preliminaries establish the MDP and successor-feature formalisms, then introduce robust RL as optimization under transition uncertainty. Together, these provide the context for the paper's later framework.

2.1 Markov Decision Process (MDP)

  • A discounted MDP specifies states, actions, transition kernel, reward, and discount factor; its policy seeks to maximize expected discounted reward.The transition kernel maps state-action pairs to distributions over subsequent states.Source: §2.1 · pp. 2–3

2.2 Universal Successor Features (USFs)

  • Successor features aggregate future discounted features, allowing task rewards to be represented through linear weights; universal successor features also encode policies.This framework supports zero-shot task solving when combined with Generalized Policy Improvement.Source: §2.2 · pp. 3–4

2.3 Robust Reinforcement Learning

  • Robust RL optimizes expected discounted reward against worst-case transition probabilities, while ambiguity sets restrict plausible kernels around estimates using a robustness budget.The cited work also describes less conservative approaches that optimize in expectation over an uncertainty-set subset.Source: §2.3 · pp. 4–5

3 Robust Successor Features

Robust successor features extend successor-feature transfer to tasks with changing reward functions and transition kernels. They use independent linear feature representations for reward and dynamics, enabling zero-shot generalization to novel representable dynamics.

3 Robust Successor Features

  • RSF relaxes the shared-transition assumption, allowing tasks to differ in both reward and transition kernels when both have linear feature representations.The reward and transition kernel use independent feature spaces.Source: §3 · p. 5
  • The transition model represents each next-state probability as transition-triplet features weighted by task-specific parameters.The weights are restricted to the probability simplex without loss of generality because rescaling can be absorbed into the feature function.Source: §3 · p. 5
  • Universal successor features are expressed as robust successor features weighted by task-specific dynamics parameters, separating dynamics from task and policy.The paper writes this as ψ(s, a, z) = τ(s, a, z, q)⊤q.Source: §3 · p. 5
  • When q matches the exact linearization weights, τ weighted by q recovers universal successor features and can generalize zero-shot to novel representable transition dynamics.The dynamics weights p are task-specific, with the probability-simplex restriction stated to be without loss of representational capacity.Source: §3 · pp. 5–6

4 Learning Robust Successor Features

The section extends the GPI guarantee to tasks with different transition kernels and describes learning approximate robust successor features. It also states that learning the feature maps is beyond this paper’s scope.

4 Learning Robust Successor Features

  • The robust GPI bound includes a transition-kernel mismatch term κpq, whose variation is scaled by the chosen policy encoding z.The bound uses κpq to bound the ℓ1 distance between task transition kernels.Source: §4 · p. 6
  • As q approaches p, κpq approaches zero and the bound reduces to the reward-only successor-feature bound.This recovers the shared-dynamics case.Source: §4 · p. 7
  • The method learns approximate robust successor features with a one-step TD error under GPI, using sampled training tasks, nearby task encodings, and ϵ-greedy action selection.The implementation uses Double DQN and an off-policy replay buffer for updates.Source: §4 · p. 7
  • Learning both reward and transition feature maps from MDPs is beyond the paper’s scope and is identified as future work for continuous RL.The stated focus is comparing robust-SF generalization with alternatives addressing only reward or transition variation.Source: §4 · p. 7

5 Experiments

The experiments assess zero-shot generalization across reward and transition changes in stochastic Gridworld, comparing RSF with USFA and single-axis ablations.

5.1 Experimental Setup

  • The setup uses a stochastic 3 × 4 Gridworld; RSF and USFA share four training environments, while each ablation varies only reward or dynamics.Policies face different trade-offs as step penalties and transition stochasticity change.Source: §5.1 · p. 8

5.2 Environment-aware zero-shot generalization

  • Models were tested zero-shot across unseen step probabilities and rewards with the true reward weights and dynamics parameters supplied.Evaluation measured Q-value MAE and policy accuracy.Source: §5.2 · p. 9
  • RSF led in Q-function mean error (0.063) and policy accuracy (93.6%), outperforming USFA and both single-axis ablations on the stochastic grid.Source: Figure 2 · p. 9, §5.2 · pp. 9–10, Figure 3 · p. 10
  • Q-value MAEs were 0.186 ± 0.094 for RSF-reward-only and 0.263 ± 0.150 for RSF-prob-only, each limited to training variation on one axis.Source: §5.2 · p. 9
  • RSF-prob-only had the worst Q-value error but second-best policy accuracy, since inaccurate Q magnitudes need not change the argmax action.Reward-only had lower MAE than dynamics-only, yet lower policy accuracy.Source: §5.2 · p. 10

6 Discussion

The discussion identifies representation learning, exploration, and linear-model expressivity as important boundaries and extensions for robust successor features.

  • The representation predicts outcomes for task weights w and p; few-shot tests estimate these weights, while jointly learning feature maps ϕ and φ remains open.Finite MDPs are trivially linearizable, but succinct joint representations remain a central question.Source: §6 · p. 10
  • GPI training may explore poorly because it relies on an ϵ-greedy mechanism; entropy-regularized objectives and maximum-occupancy regularization are potential alternatives.Source: §6 · p. 10
  • The TD-based loss may cause representation collapse when representations are learned jointly, although collapse was not observed here because the representation was given.Contrastive losses are cited as an alternative proposed to mitigate this effect.Source: §6 · p. 10
  • The linear reward and transition-kernel assumption may not hold in complex environments with continuous state or action spaces, such as robotic control.The discussion suggests integrating Hidden-Parameter MDPs as a future direction.Source: §6 · p. 10

A Theoretical Results

The appendix develops theoretical bounds for tasks with differing dynamics and defines expected features and a uniform feature bound for reward analysis.

A Theoretical Results

  • The generalized value-gap bound first accounts for differing transition kernels, then bounds reward differences via ϕmax, with the combined steps proving Theorem 1.Source: A Theoretical Results · p. 14

Definition 1 (Expected immediate features)

Definition 2 (Expected immediate reward)

  • ϕmax is taken over raw features rather than expectations under one fixed kernel, and Lemma 1 establishes that it uniformly bounds each task’s expected features.This preserves its bounding role across a family of transition kernels.Source: Definition 2 (Expected immediate reward) · p. 14

Proof of Prop. 1

These proof passages establish bounds for transition mismatch, value gaps, and reward differences across linear-MDP tasks.

  • The proof represents transition probabilities with feature vectors and learns weights p to linearize the transition kernel.Source: Proof of Prop. 1 · p. 15
  • A distribution-and-value bound can be tightened by exchanging the value functions and taking the smaller of the two resulting bounds.Source: Proof of Prop. 1 · pp. 15–16
  • The transition-kernel mismatch coefficient κ_pq approaches zero for shared dynamics and reaches the model-agnostic bound of 2 in the worst case.The specialized bound extends the shared-dynamics successor-feature result to differing dynamics.Source: Proof of Prop. 1 · p. 16
  • The propositions bound policy-value approximation error and reward differences for tasks with differing dynamics and linear reward functions.Source: Proof of Prop. 1 · pp. 16–18

Proof of Theorem 1

The proof bounds the gap between the GPI policy and the optimal policy for a new task using three terms: task distance, transition-kernel mismatch, and approximation error.

Proof of Theorem 1

  • The proof uses i for the new task (w, p) and j for a reference task (z, q) in C.Source: Proof of Theorem 1 · p. 19
  • The GPI-to-optimal-policy gap for a new task is bounded by task distance, transition-kernel mismatch scaled by reference-task size, and approximation error.The approximation error includes errors in both reward and transition-kernel approximation.Source: Proof of Theorem 1 · p. 20
  • The bound compares the GPI policy with the new task’s optimal policy, rather than only comparing it with previously computed policies.Source: Proof of Theorem 1 · p. 20

B Related Work

Prior transfer methods commonly handle reward variation with shared dynamics, while methods addressing both reward and dynamics use other representations or assumptions. The related work positions robust successor features against these approaches.

B Related Work

  • SF and GPI methods generalize across rewards but rely on shared dynamics, limiting their use when transition kernels shift.The passage contrasts these methods with the paper’s focus on simultaneous reward and dynamics variation.Source: B Related Work · p. 20
  • Function encoders support zero-shot capabilities, but their cited evaluations do not explicitly study combined perturbations in rewards and transition functions.Source: B Related Work · p. 20
  • Multitask offline RL methods address reward and transition variation, while SFs are described as more efficient to train but assume linear rewards.The passage attributes the other methods’ higher computational expense to approaches such as low-rank factorization or latent predictive models.Source: B Related Work · p. 20

C Few-Shot Learning

In few-shot adaptation, robust successor features rapidly adapt across unseen evaluation environments by estimating unknown reward and dynamics parameters from limited interactions. The paper also identifies learning the representations themselves as a remaining challenge for complex, high-dimensional tasks.

C Few-Shot Learning

  • The evaluation compares USFA, Robust RL, and Robust SF settings with scratch-trained DDQN using cumulative average return over the first 50 steps.This comparison isolates the sample efficiency of leveraging a generalized model rather than learning a distinct model for each new task.Source: C Few-Shot Learning · p. 20
  • Robust successor features adapt rapidly across all unseen evaluation environments, while USFA and Robust RL adapt only in some and scratch-trained DDQN improves minimally.Figure 4 reports cumulative average returns during few-shot adaptation in unseen environments.Source: C Few-Shot Learning · p. 20, Figure 4 · p. 21
  • By leveraging the generalized model captured during initial training, the agent can estimate new dynamics and reward structures and achieve near-optimal returns in a fraction of DDQN’s time.Source: C Few-Shot Learning · p. 21
  • Applying the approach to complex, high-dimensional tasks such as robotics requires learning the underlying representations ϕ and φ simultaneously.Source: C Few-Shot Learning · p. 21
  • All reported results use scalar TD-error to learn τ/ψ; extending the approach to vector or matrix TD-error remains ongoing work.Source: C Few-Shot Learning · p. 21
Loading 2609.31016v1…