Research map / Reinforcement learning / Visual and representation RL

What Effects the Generalization in Visual Reinforcement Learning: Policy Consistency with Truncated Return Prediction

AAAI 2024 · Reinforcement learning › Visual and representation RL · approach: Robust visual encoders for RL · 9 citations (OpenAlex)

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 Truncated Return Prediction (TRP), an auxiliary task with a Transformer predictor that keeps state-action value predictions consistent across augmented and unaugmented observations for multiple training-stage policies, backed by a theoretical generalization upper bound, to improve generalization in visual RL.

Problem

Generalization of visual reinforcement learning to unseen visual environments (changed backgrounds/colors/videos) when training on a single environment, where policies overfit to training visuals.

Approach

Derives an upper bound on the generalization gap made of policy divergence and Bellman error, then maintains cross-domain consistency for multiple policies in the policy space by predicting truncated Monte Carlo returns of historical trajectories from the first observation (plus action sequence) before and after augmentation, using a Transformer predictor.

Key results

• TRP achieves the best generalization on 6 of 10 DMControl-GB settings (color-hard and video-easy across five tasks), per Table 1.
• On Robotic Manipulation, TRP achieves state-of-the-art in 6 of 8 test environments, e.g., PegBox Test3 162.5 ± 23 vs. SGQN 154.6 ± 7.
• TRP improves sample efficiency in the training environment over SVEA and SGQN, especially on walker run and cheetah run.

Measured effects reported in the paper

Surprising findings

• Replacing the Transformer predictor with an MLP or GRU yields no obvious performance gap, which the authors interpret as the gains coming from policy consistency rather than architecture.
• TRP underperforms on Robotic Manipulation Test2 for both Reach and PegBox, and on finger spin and cartpole swingup in generalization.

Limitations

• Underperforms on some tasks (finger spin, cartpole swingup) where the replay buffer contains transitions from similar policies, limiting policy diversity.
• Theoretical analysis relies on an assumption on bounded marginal distribution ratio and the policy-divergence step uses a heuristic (expert policy approximated by the learned policy).
• Results only on SAC and simulated benchmarks; underperforms on Robotic Manipulation Test2.

Experimental setup

Datasets: DMControl-GB Places Robotic Manipulation tasks (Jangir et al. 2022)

Benchmarks: DMControl-GB Robotic Manipulation (Reach, PegBox)

Metrics: average episode return standard deviation of return

Baselines: SAC CURL RAD PAD DrQ SODA SVEA SGQN

Builds on: SAC SVEA DrQ SODA RAD

Closest papers by meaning

• Focus On What Matters: Separated Models For Visual-Based RL Generalization (NeurIPS 2024 · similarity 0.90)
• Generalizing Consistency Policy to Visual RL with Prioritized Proximal Experience Regularization (NeurIPS 2024 · similarity 0.89)
• VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning (NeurIPS 2026 · similarity 0.88)
• The Surprising Ineffectiveness of Pre-Trained Visual Representations for Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.88)
• Zero Shot Generalization of Vision-Based RL Without Data Augmentation (ICML 2025 · similarity 0.87)
• A Simple Framework for Generalization in Visual RL under Dynamic Scene Perturbations (NeurIPS 2024 · similarity 0.87)