State Chrono Representation for Enhancing Generalization in Reinforcement Learning

TL;DR

This paper introduces the State Chrono Representation (SCR), a novel method that enhances reinforcement learning by incorporating long-term temporal information into state representations, improving generalization in complex environments.

Contribution

SCR extends metric-based state representations by integrating extensive temporal information without adding many parameters, addressing limitations in long-term dependency modeling.

Findings

SCR outperforms recent metric-based methods in generalization tasks.

SCR demonstrates improved performance in DeepMind Control and Meta-World environments.

The approach effectively captures long-term dynamics and rewards in state representations.

Abstract

In reinforcement learning with image-based inputs, it is crucial to establish a robust and generalizable state representation. Recent advancements in metric learning, such as deep bisimulation metric approaches, have shown promising results in learning structured low-dimensional representation space from pixel observations, where the distance between states is measured based on task-relevant features. However, these approaches face challenges in demanding generalization tasks and scenarios with non-informative rewards. This is because they fail to capture sufficient long-term information in the learned representations. To address these challenges, we propose a novel State Chrono Representation (SCR) approach. SCR augments state metric-based representations by incorporating extensive temporal information into the update step of bisimulation metric learning. It learns state distances within a temporal framework that considers both future dynamics and cumulative rewards over current and long-term future states. Our learning strategy effectively incorporates future behavioral information into the representation space without introducing a significant number of additional parameters for modeling dynamics. Extensive experiments conducted in DeepMind Control and Meta-World environments demonstrate that SCR achieves better performance comparing to other recent metric-based methods in demanding generalization tasks. The codes of SCR are available in https://github.com/jianda-chen/SCR.