Adapting Visual Policies via Predicted Rewards

Image-based reinforcement learning (RL) faces significant challenges in generalization when the visual environment undergoes substantial changes between the training and testing phases. Under such circumstances, learned policies may not perform well leading to degraded results. Previous approaches to this issue largely have aimed to broaden the training observation distribution, employing techniques like data augmentation and domain randomization. Nevertheless, given the sequential nature of the decision-making problem, and residual errors propagated by the policy model can accumulate throughout the trajectory, resulting in highly degraded performance. 

In this paper, we leverage the observation that a learned reward prediction function is often able to still reliably predict rewards under domain shift. We exploit this property to fine-tune the policy from reward prediction in the target domain. We have found that, even under significant domain shift, the predicted reward can still provide meaningful signal and fine-tuning substantially improves the policy. Our approach, termed Predicted Reward Fine-tuning (PRFT), improves performance across diverse tasks in both simulated benchmarks and real-world experiments.

We propose to jointly learn a reward prediction model to predict rewards from experience. The policy is then fine-tuned using such predicted rewards in the target testing environment. Through extensive experiments in both simulations and real-world experiments, we show that the reward prediction model generalizes well across visual effect shifts and significantly enhances policy performance through fine-tuning.