# Unified State Representation Learning under Data Augmentation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 9 references that are in the map.
- First posted on arXiv 12 Sep 2022.

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

- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [Disentangling Factors of Variation with Cycle-Consistent Variational Auto-Encoders](https://mlsanity.com/paper/1804.10469/) (2018)
- [PlayVirtual: Augmenting Cycle-Consistent Virtual Trajectories for Reinforcement Learning](https://mlsanity.com/paper/2106.04152/) (2021)
- [Domain Adaptation In Reinforcement Learning Via Latent Unified State Representation](https://mlsanity.com/paper/2102.05714/) (2021)

## Abstract (arXiv)

The capacity for rapid domain adaptation is important to increasing the applicability of reinforcement learning (RL) to real world problems. Generalization of RL agents is critical to success in the real world, yet zero-shot policy transfer is a challenging problem since even minor visual changes could make the trained agent completely fail in the new task. We propose USRA: Unified State Representation Learning under Data Augmentation, a representation learning framework that learns a latent unified state representation by performing data augmentations on its observations to improve its ability to generalize to unseen target domains. We showcase the success of our approach on the DeepMind Control Generalization Benchmark for the Walker environment and find that USRA achieves higher sample efficiency and 14.3% better domain adaptation performance compared to the best baseline results.