# Learning Task-relevant Representations for Generalization via Characteristic Functions of Reward Sequence Distributions

arXiv: https://arxiv.org/abs/2205.10218 · HTML: https://mlsanity.com/paper/2205.10218/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 36 references that are in the map.
- First posted on arXiv 20 May 2022.

## Builds on (most cited first, 29 of 29)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [Self-Supervised Policy Adaptation during Deployment](https://mlsanity.com/paper/2007.04309/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Generalization and Regularization in DQN](https://mlsanity.com/paper/1810.00123/) (2018)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1910.05396/) (2019)
- [Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck](https://mlsanity.com/paper/1910.12911/) (2019)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Observational Overfitting in Reinforcement Learning](https://mlsanity.com/paper/1912.02975/) (2019)
- [Automatic Data Augmentation for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/2006.12862/) (2020)
- [Towards Generalization and Simplicity in Continuous Control](https://mlsanity.com/paper/1703.02660/) (2017)
- [SECANT: Self-Expert Cloning for Zero-Shot Generalization of Visual Policies](https://mlsanity.com/paper/2106.09678/) (2021)
- [Invariant Policy Optimization: Towards Stronger Generalization in Reinforcement Learning](https://mlsanity.com/paper/2006.01096/) (2020)
- [A Characteristic Function Approach to Deep Implicit Generative Modeling](https://mlsanity.com/paper/1909.07425/) (2019)
- [Rotation, Translation, and Cropping for Zero-Shot Generalization](https://mlsanity.com/paper/2001.09908/) (2020)
- [Sample-Efficient Reinforcement Learning via Conservative Model-Based Actor-Critic](https://mlsanity.com/paper/2112.10504/) (2021)
- [Invariant Policy Learning: A Causal Perspective](https://mlsanity.com/paper/2106.00808/) (2021)

## Built on by (most cited first, 15 of 15)

- [De Novo Molecular Generation via Connection-aware Motif Mining](https://mlsanity.com/paper/2302.01129/) (2023)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [KiGRAS: Kinematic-Driven Generative Model for Realistic Agent Simulation](https://mlsanity.com/paper/2407.12940/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [Apollo-MILP: An Alternating Prediction-Correction Neural Solving Framework for Mixed-Integer Linear Programming](https://mlsanity.com/paper/2503.01129/) (2025)
- [Scalable and Effective Arithmetic Tree Generation for Adder and Multiplier Designs](https://mlsanity.com/paper/2405.06758/) (2024)
- [Learning to Stop Cut Generation for Efficient Mixed-Integer Linear Programming](https://mlsanity.com/paper/2401.17527/) (2024)
- [Road Planning for Slums via Deep Reinforcement Learning](https://mlsanity.com/paper/2305.13060/) (2023)
- [Accelerate Presolve in Large-Scale Linear Programming via Reinforcement Learning](https://mlsanity.com/paper/2310.11845/) (2023)
- [State Sequences Prediction via Fourier Transform for Representation Learning](https://mlsanity.com/paper/2310.15888/) (2023)
- [Generalization in Visual Reinforcement Learning with the Reward Sequence Distribution](https://mlsanity.com/paper/2302.09601/) (2023)
- [Similarity-based Knowledge Transfer for Cross-Domain Reinforcement Learning](https://mlsanity.com/paper/2312.03764/) (2023)
- [Sequential Action-Induced Invariant Representation for Reinforcement Learning](https://mlsanity.com/paper/2309.12628/) (2023)
- [DRARL: Disengagement-Reason-Augmented Reinforcement Learning for Efficient Improvement of Autonomous Driving Policy](https://mlsanity.com/paper/2506.16720/) (2025)
- [Intrinsic Dynamics-Driven Generalizable Scene Representations for Vision-Oriented Decision-Making Applications](https://mlsanity.com/paper/2405.19736/) (2024)

## Abstract (arXiv)

Generalization across different environments with the same tasks is critical for successful applications of visual reinforcement learning (RL) in real scenarios. However, visual distractions -- which are common in real scenes -- from high-dimensional observations can be hurtful to the learned representations in visual RL, thus degrading the performance of generalization. To tackle this problem, we propose a novel approach, namely Characteristic Reward Sequence Prediction (CRESP), to extract the task-relevant information by learning reward sequence distributions (RSDs), as the reward signals are task-relevant in RL and invariant to visual distractions. Specifically, to effectively capture the task-relevant information via RSDs, CRESP introduces an auxiliary task -- that is, predicting the characteristic functions of RSDs -- to learn task-relevant representations, because we can well approximate the high-dimensional distributions by leveraging the corresponding characteristic functions. Experiments demonstrate that CRESP significantly improves the performance of generalization on unseen environments, outperforming several state-of-the-arts on DeepMind Control tasks with different visual distractions.