[Paper Review] Learning Invariances for Policy Generalization

Rémi Tachet, Philip Bachman+1 more|arXiv (Cornell University)|Sep 7, 2018

Reinforcement Learning in Robotics17 references3 citations

TL;DR

This paper investigates methods to improve policy generalization in reinforcement learning by learning invariances, particularly translation invariance in a simple jumping task. It finds that data augmentation effectively induces invariance, enabling 99.8% generalization across 1,265 unseen obstacle positions, while meta-learning and adversarial training fail to outperform the baseline despite theoretical promise.

ABSTRACT

While recent progress has spawned very powerful machine learning systems, those agents remain extremely specialized and fail to transfer the knowledge they gain to similar yet unseen tasks. In this paper, we study a simple reinforcement learning problem and focus on learning policies that encode the proper invariances for generalization to different settings. We evaluate three potential methods for policy generalization: data augmentation, meta-learning and adversarial training. We find our data augmentation method to be effective, and study the potential of meta-learning and adversarial learning as alternative task-agnostic approaches.

Motivation & Objective

• To address the challenge of policy generalization in reinforcement learning, where agents trained on limited tasks fail to generalize to unseen variations.
• To investigate whether invariance learning—specifically translation invariance—can be induced through different training paradigms.
• To evaluate data augmentation, meta-learning, and adversarial training as task-agnostic methods for improving zero-shot generalization.
• To determine whether learning disentangled or invariant features leads to better transferability across unseen task variations.
• To identify which method most effectively enables agents to generalize beyond memorization of task-specific features.

Proposed method

• Data augmentation is applied by randomly shifting the game screen embedding during training, forcing the agent to learn translation-invariant representations.
• Meta-learning is implemented via a first-order approximation of minimizing the sum of losses on two random tasks after a gradient update, encouraging generalizable parameter updates.
• Adversarial training employs a discriminator that classifies the task identity from a hidden layer’s activation, pushing the agent to produce task-agnostic representations.
• The agent is trained using Double-DQN and A2C algorithms, with and without the proposed invariance-inducing techniques.
• The performance is evaluated on 1,271 total tasks, with 6 used for training and 1,265 for zero-shot generalization testing.
• t-SNE visualization is used to analyze hidden layer representations, comparing task clustering with and without adversarial training.

Experimental results

Research questions

• RQ1 Can data augmentation effectively induce translation invariance in a reinforcement learning policy, leading to improved zero-shot generalization?
• RQ2 Does meta-learning, through gradient alignment across tasks, lead to better generalization than standard RL training?
• RQ3 Can adversarial training successfully remove task-specific information from hidden representations while preserving task-solving capability?
• RQ4 Is learning independent or disentangled features sufficient for generalization, or are invariances more critical?
• RQ5 Why do meta-learning and adversarial training fail to generalize better than standard RL despite their theoretical advantages?

Key findings

• Data augmentation achieves 99.8% generalization performance, solving 1,263 out of 1,265 unseen obstacle positions, significantly outperforming the baseline.
• The baseline agent without data augmentation generalizes to only 2.8% of the test tasks (36 out of 1,265), indicating poor transferability.
• The agent trained with data augmentation generalizes to two consecutive obstacles, a scenario never seen during training, demonstrating robust invariance learning.
• Meta-learning and adversarial training do not improve generalization beyond the standard RL baseline, despite their theoretical potential for invariance learning.
• t-SNE visualization confirms that adversarial training successfully removes task-specific clustering in hidden representations, yet generalization remains poor.
• The results show that invariance must be explicitly encouraged—simply learning disentangled features is insufficient for generalization, as shown by the failure of meta-learning and adversarial training to generalize effectively.