Research map / Reinforcement learning / Visual and representation RL
Efficient Reinforcement Learning Through Adaptively Pretrained Visual Encoder
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes APE, which pretrains a ResNet18 visual encoder on ImageNet-100 with an adaptive closed-loop augmentation schedule (MoCo v2-based), then freezes most of it for DreamerV3 or DrQ-v2 policy learning on DMC, Atari, and Memory Maze.
Problem
Visual RL agents have poor sample efficiency and generalize poorly because task-dependent visual encoders extract task-specific rather than generalizable features; existing pretrained-encoder approaches neglect the pretraining augmentation schedule.
Approach
Pretrain a ResNet18 encoder with contrastive MoCo v2 on ImageNet-100 where each batch is split into sub-batches with different augmentation compositions, and the sampling probability of each composition is updated every epoch via softmax over pretext-task accuracy; the encoder (first three layers frozen) is then plugged into DreamerV3 or DrQ-v2 for policy learning.
Key results
- On DMC vision at 100K steps, APE (DreamerV3) reaches a mean of 836.7 across three tasks versus 458.6 for DreamerV3 and 488.2 for PIE-G.
- APE (DrQ-v2) achieves 560.2 mean at 100K versus 284.6 for DrQ-v2 on the same three tasks.
- On the full DMC vision suite at 1M steps, APE averages 828 versus 802.6 for DreamerV3 and 722.8 for DrQ-v2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use adaptively pretrained encoder instead of the standard DreamerV3 setup | DreamerV3 · DMC 1M mean | score: 802.6 -> 828 | Helped |
| Use adaptively pretrained encoder with DreamerV3 instead of DreamerV3 alone | DreamerV3 · DMC Mean 100K | score: 458.6 -> 836.7 | Helped |
| Use adaptively pretrained encoder with DrQ-v2 instead of DrQ-v2 alone | DrQ-v2 · DMC Mean 100K | score: 284.6 -> 560.2 | Helped |
| Use adaptively pretrained encoder with DreamerV3 instead of DreamerV3 alone | DreamerV3 · DMC Mean 500K | score: 749.1 -> 882.8 | Helped |
| Use adaptively pretrained encoder with DrQ-v2 instead of DrQ-v2 alone | DrQ-v2 · DMC Mean 500K | score: 773.1 -> 843.7 | Helped |
| Use adaptive gaussian-blur augmentation rather than adaptive color-jitter augmentation (ablation) | APE pretraining · ImageNet-100 validation | top-5 classification accuracy: 91.08 -> 91.7 | Helped |
| Use adaptive gaussian-blur augmentation instead of MoCo v2 pretraining | ResNet18 encoder · ImageNet-100 validation | top-5 classification accuracy: 90.84 -> 91.7 | Helped |
| Use a pretrained encoder rather than random initialization, frozen or trainable (ablation) | DreamerV3 visual encoder · DMC control tasks | score | Helped |
| Use ResNet18 instead of the deeper ResNet50 encoder (ablation) | APE · DMC control tasks | score | Helped |
| Increase encoder depth from ResNet18 to ResNet50 (ablation) | APE · DMC fine-grained control tasks | score | Hurt |
| Change the main augmentation strategy, including gaussian blur and color jitter frequency (ablation) | APE pretraining · DMC control tasks | sampling efficiency | Mixed |
| Use ImageNet-pretrained encoder rather than random initialization on multi-item Atari observations (ablation) | APE · Atari Boxing | performance | Hurt |
Surprising findings
- Deeper encoders (ResNet50) and the pretrained ImageNet-100 ResNet18 with a random-initialized frozen encoder show that deeper networks do not guarantee better features for RL.
- On some Atari tasks with multi-item observations (e.g., Boxing), APE performs slightly worse, which the authors attribute to ImageNet models being biased toward single-item classification.
Limitations
- Performance drops slightly on some Atari tasks with multi-item observations, where ImageNet-trained features are biased toward single items.
- Memory Maze runs were limited to 2M steps due to limited computational resources.
- Deeper encoders such as ResNet50 underperform on fine-grained control tasks.
Experimental setup
Datasets: ImageNet-100 DeepMind Control Suite Atari 100k Memory Maze
Benchmarks: DeepMind Control Suite Atari 100k Memory Maze ImageNet-100 linear probe
Metrics: episode return / score linear probe top-5 accuracy sample efficiency
Baselines: DreamerV3 DrQ-v2 SAC:state SAC:vision RRL PIE-G CURL MAE
Builds on: DreamerV3 MoCo v2 DrQ-v2 PIE-G AdDA
Closest papers by meaning
- Investigating Multi-task Pretraining and Generalization in Reinforcement Learning (ICLR 2023 · similarity 0.91)
- Pretraining Representations for Data-Efficient Reinforcement Learning (NeurIPS 2021 · similarity 0.91)
- Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning (NeurIPS 2022 · similarity 0.90)
- Investigating Pre-Training Objectives for Generalization in Vision-Based Reinforcement Learning (ICML 2024 · similarity 0.90)
- The Surprising Ineffectiveness of Pre-Trained Visual Representations for Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.90)
- Skill-aware Mutual Information Optimisation for Zero-shot Generalisation in Reinforcement Learning (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.