PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions

Bogdan Mazoure, Ilya Kostrikov, Ofir Nachum, Jonathan Tompson

arXiv:2111.14629v1cs.LGcs.AI

TL;DR

Offline zero-shot generalization is difficult when learning from limited, fixed datasets, while online self-supervised approaches rely on environment access. GSF learns observation similarity from generalized value functions and outperforms prior offline RL and representation-learning baselines across offline Procgen levels.

  • Problem

    Zero-shot generalization from offline data is challenging because high-dimensional observations require large datasets, while online self-supervised methods depend on environment access.

    Source: §1 · p. 2

  • Method

    GSF uses generalized value functions in a self-supervised framework to aggregate observations according to their expected future behavior.

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

  • Results

    GSF outperforms prior state-of-the-art offline RL and representation-learning baselines across the offline Procgen distribution of levels.

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

  • Takeaways & Limitations

    The framework can recover objectives similar to prior self-supervised methods by designing the cumulant function.

    Source: §5.4 · p. 8

    The paper leaves open whether GVF-based distances can avoid quantile discretization and whether cumulants can be chosen a priori to yield optimal representations beyond Procgen.

    Source: §7 · pp. 9–10

Abstract

from arXiv · show

Reinforcement learning (RL) agents are widely used for solving complex sequential decision making tasks, but still exhibit difficulty in generalizing to scenarios not seen during training. While prior online approaches demonstrated that using additional signals beyond the reward function can lead to better generalization capabilities in RL agents, i.e. using self-supervised learning (SSL), they struggle in the offline RL setting, i.e. learning from a static dataset. We show that performance of online algorithms for generalization in RL can be hindered in the offline setting due to poor estimation of similarity between observations. We propose a new theoretically-motivated framework called Generalized Similarity Functions (GSF), which uses contrastive learning to train an offline RL agent to aggregate observations based on the similarity of their expected future behavior, where we quantify this similarity using \emph{generalized value functions}. We show that GSF is general enough to recover existing SSL objectives while also improving zero-shot generalization performance on a complex offline RL benchmark, offline Procgen.

1 Introduction

Zero-shot generalization from offline, high-dimensional data is difficult because test tasks provide no training data and practical datasets can be limited in size and quality. The paper introduces GSF, using generalized value functions to learn observation similarity, and evaluates it on offline Procgen.

1 Introduction

  • Online self-supervised signals can improve generalization, but their reliance on environment access limits their applicability to offline learning.The introduction notes that online generalization methods can outperform PPO on Procgen by leveraging auxiliary signals.Source: §1 · p. 2
  • Offline zero-shot generalization is challenging because pixel-based datasets can be costly to generate, share, and sufficiently representative.A standard 25-million-frame Procgen PPO run generates more than 300 Gb of data.Source: §1 · p. 2
  • GSF uses generalized value functions to aggregate observations according to their expected future behavior, with the cumulant determining the similarity being encouraged.Reward cumulants yield reward-aware similarity, while future state-action counts yield reward-free similarity.Source: §1 · p. 2
  • GSF outperforms prior offline RL and representation-learning baselines across the offline Procgen distribution of levels.The benchmark contains 5M transitions from 200 related levels across 16 games.Source: §1 · p. 2

2 Related Works

Prior work studies generalization across rewards, observations, and transition dynamics, using methods such as exploration, auxiliary signals, and data augmentation. This paper situates its approach among representation-learning and offline-RL methods, noting the limits of handcrafted augmentation and conservative policy updates.

Generalization in reinforcement learning

  • Prior RL generalization work addresses unseen rewards, observations, and transition dynamics, with methods including exploration, auxiliary signals, and data augmentation.Data augmentation can require handcrafted domain information, and the paper reports that it alone performs poorly in its offline RL experiments.Source: §2 · p. 3

Representation learning

  • Self-supervised representation learning has supported generalization and sample-efficiency results in challenging RL settings, including Atari, DeepMind Control, and Procgen.The paper also cites heuristic-guided learning and random Fourier features as theoretically motivated representation-learning approaches.Source: §2 · p. 3

Offline reinforcement learning

  • Offline RL algorithms address extrapolation risk through behavior regularization or pessimistic models, but conservative updates may limit use of dataset information for generalization.The discussion specifically names CQL and MBS as methods that keep policy improvement supported by the batch.Source: §2 · p. 3

3 Problem setting

The problem setting trains from logged experience on a limited set of related POMDPs and evaluates on held-out tasks with different observation functions. Representation learning matters because decoding unseen observations into latent states affects the policy and its generalization.

3.1 Partially-observable Markov decision processes

  • Each task is a POMDP with states, actions, observations, transitions, rewards, and a discount factor, and the agent acts from observations.The agent receives observations generated from states and transitions according to the task’s observation and transition functions.Source: §3.1 · p. 3
  • Training uses a fixed historical dataset and a limited set of POMDPs; test tasks share transition dynamics and rewards but have different observation functions.Training and test POMDPs are assumed to come from a common task distribution.Source: §3.1 · p. 4

3.2 Representation learning

  • Prior action-based similarity objectives match observations with similar optimal-policy action sequences, but can fail when multiple optimal policies produce different actions over similar latent states.This motivates considering similarity measures beyond action sequences.Source: §3.2 · p. 4
  • The paper assumes value functions are linear in decoded latent states, so decoding unseen observations correctly directly affects policy generalization.A latent state decoder maps observations into the representation space used by the value function.Source: §3.2 · p. 4

4 Motivating example

Action-sequence similarity can misrepresent observations when multiple optimal policies yield different actions despite similar latent states and values. The Climber example motivates comparing observations through value-based or other future-behavior signals.

4 Motivating example

  • For two Climber levels, the true latent state and value function are near-identical (value-function average normalized mean-squared error 0.0398), but the same PPO policy yields very different action sequences.Across episodes, action-sequence average total variation distance is 0.4423.Source: Figure 1 · p. 4, §4 · p. 4
  • Action-based objectives can fail in POMDPs with multiple optimal policies, where trajectories visit the same latent states but take different actions.The paper proposes that value-function similarity can provide a better representation in this example, while action similarity may suit differently structured POMDPs.Source: §4 · pp. 4–5
  • The example motivates learning representations from discounted future behavior rather than relying on action sequences alone.The paper frames the goal as a general way to improve representation learning through self-supervised learning.Source: §4 · p. 5

5 Method

GSF compares observations by future behavior quantified through GVFs of chosen cumulants, then uses quantile-based labels for contrastive representation learning alongside offline RL.

5.1 Quantifying future behavior with GVFs

  • GVFs extend value functions to bounded cumulative signals, letting future behavior reflect selected signals such as rewards or observation visitation.The framework also uses vector-valued cumulants to learn successor features.Source: §5.1 · p. 5

5.2 Measuring distances between GVFs of different POMDPs

  • Because GVF magnitudes can differ sharply across POMDPs, GSF compares their order statistics and partitions each task’s GVF range into equal-size quantile bins.This avoids meaningless pointwise comparisons across POMDPs with different GVF scales.Source: §5.2 · p. 6

5.3 Self-supervised learning of GSFs

  • GVF quantile-bin labels supervise multiclass InfoNCE, bringing observations with identical labels closer in latent space during offline reinforcement learning.The labels are based on the quantile bin containing each observation’s GVF.Source: §5.2 · p. 6, §5.3 · p. 6
  • The authors report that their multiclass loss is more stable than other multiclass contrastive losses.The loss uses a linear classifier to predict quantile-bin labels from projected representations.Source: §5.3 · pp. 6–7

5.4 Algorithm

  • Implemented atop CQL, GSF adds representation learning, and its cumulant choice lets the framework recover objectives resembling prior self-supervised methods.The algorithm also relies on LearnGVF to estimate GVFs used in contrastive learning.Source: §5.4 · pp. 7–8

5.5 Choice of number of quantiles K

  • Increasing K reduces within-POMDP approximation error but raises finite-sample variance; too-small K groups unrelated observations and induces high bias.Large K approaches ranking states, but those rankings are estimated from data and each bin may contain only one sample.Source: §5.5 · p. 8

6 Experiments

On offline Procgen, GSF outperforms the other offline RL and representation-learning baselines. The experiments compare methods across the benchmark’s games and examine how GSF’s cumulant choice affects performance.

Offline Procgen benchmark

  • Offline Procgen was devised to test offline RL generalization across observation functions, using procedurally generated levels from 16 games.The evaluation uses Procgen's easy distribution; the dataset procedure begins with PPO pretraining on 200 easy-distribution levels per environment.Source: §6 · p. 8

Results

  • The baseline comparison includes behavioral cloning, CQL, CURL, CTRL, DeepMDP, VPN, CSSC, and Policy Similarity Embeddings.These methods assess policy quality, offline RL performance, and alternative representation-learning objectives.Source: §6 · p. 9
  • GSF performs better than the other offline RL and representation-learning baselines on offline Procgen.Figure 3 summarizes returns over five seeds and 16 games, normalized by each game's median CQL performance.Source: Figure 3 · p. 9, §6 · p. 9
  • Reward and successor-feature cumulants yield similar performance, while action-based distance produces larger variance.The three cumulants correspond to learning the policy's Q-value, an induced distribution over states, and the discounted policy, respectively.Source: §6 · p. 9

7 Discussion

The discussion presents GSF as combining RL and representation learning to improve zero-shot generalization on pixel-based control tasks. It identifies open questions about quantile discretization and selecting cumulants beyond Procgen, and notes a theoretical trade-off between approximation and extrapolation error.

Discussion

  • GSF combines reinforcement learning and representation learning to improve zero-shot generalization on challenging pixel-based control tasks.The method compares observation pairs using similarity defined by an instantaneous accumulated signal, and improves empirical performance on offline Procgen.Source: §7 · p. 9
  • GSF's hyperparameter choice involves a trade-off between finite-sample approximation error and extrapolation error.The discussion presents this as a theoretical result.Source: §7 · p. 9

Open questions

  • The authors leave open whether GVF-based distances can avoid quantile discretization and whether cumulants can be chosen a priori beyond Procgen to yield optimal representations.For continuous observations, the described implementation learns successor features via TD and computes quantiles over their L1 norm.Source: §7 · pp. 9–10

Appendix

The appendix details GVF estimator training and theoretical properties, then reports additional Procgen analyses of GSF’s similarity choices and contrastive-loss and hyperparameter variants.

7.1 Experimental details

  • GVF estimators are jointly trained for 50k iterations, with PopArt re-normalization used to prevent gradient instability from differing cumulant scales across POMDPs.A shared representation and network output split into POMDP-specific chunks support the joint training.Source: §7.1 · p. 14
  • Each method was allowed to tune one hyperparameter because of computational-budget constraints; GSF tuned the number of quantiles K.The dataset was collected from PPO rollouts with a decaying ε-greedy strategy to address insufficient action coverage for CQL.Source: §7.1 · p. 14

7.2 Proofs

  • For state-indicator cumulants, observations with similar representations have similar state-visitation frequencies; the highest quantile bin has the lowest interpolation error and the first bin the highest.Theorem 2 relates the quantile bins’ interpolation errors to state visitation counts.Source: §7.2 · p. 16

7.3 Additional results

  • In the Procgen trajectory example, value sequences are more similar than discrete action sequences, supporting grouping observations by value magnitude to ignore background and small positional differences.The passage reports value sequences of 8.20, 8.05, 8.74, 9.06, 9.29 and 8.35, 8.30, 8.26, 8.55, 8.40.Source: §7.3 · p. 16
  • Table 2 reports average performance across 16 Procgen games and five seeds, standardized by CQL performance on zero-shot generalization across the easy-level distribution.The supplied caption describes the table’s scope and normalization but gives no method scores.Source: Table 2 · p. 17
  • Categorical cross-entropy yields a more stable and lower error than the alternative multi-class contrastive loss, so it is used in all experiments.Figure 5 compares the categorical-cross-entropy and pairwise-InfoNCE versions.Source: §7.3 · p. 17, Figure 5 · p. 18
  • The selected GSF hyperparameter combination has one of the highest inter-game median values and the lowest inter-quartile range among the tested combinations.The ablation varies temperature τ, bin count K, and projection-network structure h.Source: §7.3 · p. 18, Figure 6 · p. 19

Related papers

Loading 2111.14629v1…