[Paper Review] Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning

Mhairi Dunion, Trevor McInroe+3 more|arXiv (Cornell University)|Jul 12, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper introduces Temporal Disentanglement (TED), a self-supervised auxiliary task that learns disentangled representations in reinforcement learning by leveraging temporal consistency in sequential observations. TED improves generalization to unseen environment variables—both task-irrelevant (e.g., background color) and task-relevant (e.g., goal positions)—enabling faster adaptation and better performance recovery than state-of-the-art methods, especially in image-based RL settings.

ABSTRACT

Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training. This issue is especially problematic for image-based RL, where a change in just one variable, such as the background colour, can change many pixels in the image. The changed pixels can lead to drastic changes in the agent's latent representation of the image, causing the learned policy to fail. To learn more robust representations, we introduce TEmporal Disentanglement (TED), a self-supervised auxiliary task that leads to disentangled image representations exploiting the sequential nature of RL observations. We find empirically that RL algorithms utilising TED as an auxiliary task adapt more quickly to changes in environment variables with continued training compared to state-of-the-art representation learning methods. Since TED enforces a disentangled structure of the representation, our experiments also show that policies trained with TED generalise better to unseen values of variables irrelevant to the task (e.g. background colour) as well as unseen values of variables that affect the optimal policy (e.g. goal positions).

Motivation & Objective

• Address the poor generalization of image-based RL agents to unseen environment variations, such as changes in background color or goal positions.
• Overcome limitations of domain randomization and invariant representation learning, which are sample-inefficient or fail to generalize to task-relevant variables.
• Develop a self-supervised, online method that learns disentangled representations using temporal data from consecutive timesteps, avoiding reliance on i.i.d. data.
• Enable lifelong learning in RL by allowing agents to quickly adapt to new environment values without catastrophic forgetting.
• Improve robustness and performance recovery after distributional shift by enforcing a disentangled structure in the latent representation.

Proposed method

• Propose TED, a self-supervised auxiliary loss that encourages disentanglement by contrasting representations of temporally adjacent observations.
• Use two types of non-temporal samples: (1) from the same episode (non-stationary features), and (2) from different episodes (stationary features), to separate agent-controlled and environment-controlled factors.
• Train a contrastive classifier with a disentangled structure that learns to distinguish between changes due to agent actions and those due to environmental changes.
• Integrate TED as an auxiliary loss into existing RL algorithms (e.g., RAD, SAC, PPO) with minimal architectural changes and no need for a decoder.
• Optimize the TED loss with a hyperparameter α to balance disentanglement and policy performance, tuning it empirically for optimal generalization.
• Use a disentanglement metric (based on factor-VAE) to evaluate the quality of learned representations, showing improved disentanglement with TED.

Experimental results

Research questions

• RQ1 Can a self-supervised auxiliary task based on temporal consistency improve generalization in image-based RL to unseen values of environment variables?
• RQ2 Does learning disentangled representations through temporal contrastive learning lead to faster adaptation and better performance recovery after distributional shifts?
• RQ3 How does TED compare to state-of-the-art representation learning methods in terms of generalization to both task-irrelevant and task-relevant environment variables?
• RQ4 What is the impact of using only same-episode or only different-episode non-temporal samples on generalization performance?
• RQ5 How critical is the disentangled classifier structure in TED compared to a standard linear classifier?

Key findings

• TED significantly improves generalization of three base RL algorithms (RAD, SAC, PPO) on unseen values of environment variables, including both irrelevant (e.g., background color) and relevant (e.g., goal position) factors.
• Agents trained with TED adapt more quickly to new environment values during continued training, recovering optimal policy performance faster than baselines.
• The ablation study shows that using both same-episode and different-episode non-temporal samples yields the best generalization, outperforming versions using only one type.
• Replacing the disentangled classifier with a standard linear classifier reduces generalization performance, demonstrating the importance of the disentangled structure in TED.
• The TED loss coefficient α must be carefully tuned: α=200 achieves optimal performance, while α=50 or α=500 leads to worse generalization due to insufficient or excessive disentanglement.
• TED improves disentanglement scores on the factor-VAE metric, confirming that the learned representations are more disentangled than those from baselines, even without explicit supervision.