Synthetic Data is Sufficient for Zero-Shot Visual Generalization from Offline Data

2025-08-17

AI Summary

Offline visual reinforcement learning suffers from poor zero-shot generalization, primarily due to insufficient diversity in real-world offline datasets, leading to overfitting. This paper proposes a two-stage synthetic data augmentation method that requires no modification to existing algorithms: first, applying geometric and color-robust transformations to original visual trajectories; second, generating semantically consistent, action-observation aligned synthetic data within the latent space of diffusion models. The approach is universally compatible with both continuous and discrete action spaces. Evaluated on visual D4RL and Procgen benchmarks, it significantly narrows the generalization gap between training and test environments, enhances zero-shot transfer performance, and maintains computational efficiency and plug-and-play compatibility. Its core contribution lies in the first integration of latent-space diffusion-based generation with trajectory-level data augmentation to improve generalization in offline visual RL.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Offline reinforcement learning (RL) offers a promising framework for training agents using pre-collected datasets without the need for further environment interaction. However, policies trained on offline data often struggle to generalise due to limited exposure to diverse states. The complexity of visual data introduces additional challenges such as noise, distractions, and spurious correlations, which can misguide the policy and increase the risk of overfitting if the training data is not sufficiently diverse. Indeed, this makes it challenging to leverage vision-based offline data in training robust agents that can generalize to unseen environments. To solve this problem, we propose a simple approach generating additional synthetic training data. We propose a two-step process, first augmenting the originally collected offline data to improve zero-shot generalization by introducing diversity, then using a diffusion model to generate additional data in latent space. We test our method across both continuous action spaces (Visual D4RL) and discrete action spaces (Procgen), demonstrating that it significantly improves generalization without requiring any algorithmic changes to existing model-free offline RL methods. We show that our method not only increases the diversity of the training data but also significantly reduces the generalization gap at test time while maintaining computational efficiency. We believe this approach could fuel additional progress in generating synthetic data to train more general agents in the future.

Problem

• Overcoming limited state diversity in offline RL training
• Reducing visual data noise and spurious correlations
• Improving zero-shot generalization for unseen environments

Innovation

• Generates synthetic data for offline RL
• Uses diffusion model in latent space
• Improves zero-shot generalization diversity