Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning

Ruijie Zheng, Xiyao Wang, Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé, Furong Huang

TL;DR

TACO addresses sample inefficiency in visual reinforcement learning by learning concurrent state and action representations via a temporal contrastive objective that maximizes the mutual information $J_{ ext{TACO}}=\mathcal{I}(Z_{t+K}; [Z_t, U_t, ..., U_{t+K-1}])$. It provides a theoretically sufficient objective for control and is designed as a plug-in module compatible with online and offline RL pipelines using a DrQ-v2 backbone, with horizon $K$ typically $1$ or $3$. Empirically, TACO achieves about a 40% average performance boost after 1M environment steps across nine DM Control Suite tasks in online RL and sets new state-of-the-art results for offline visual RL when paired with strong baselines like TD3+BC and CQL across diverse datasets. The work highlights the importance of learning action representations in continuous control and demonstrates that a stable, MI-based objective can effectively capture dynamics for improved sample efficiency and generalization.

Abstract

Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO: Temporal Action-driven Contrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.

TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning

TL;DR

TACO addresses sample inefficiency in visual reinforcement learning by learning concurrent state and action representations via a temporal contrastive objective that maximizes the mutual information . It provides a theoretically sufficient objective for control and is designed as a plug-in module compatible with online and offline RL pipelines using a DrQ-v2 backbone, with horizon typically or . Empirically, TACO achieves about a 40% average performance boost after 1M environment steps across nine DM Control Suite tasks in online RL and sets new state-of-the-art results for offline visual RL when paired with strong baselines like TD3+BC and CQL across diverse datasets. The work highlights the importance of learning action representations in continuous control and demonstrates that a stable, MI-based objective can effectively capture dynamics for improved sample efficiency and generalization.

Abstract

Despite recent progress in reinforcement learning (RL) from raw pixel data, sample inefficiency continues to present a substantial obstacle. Prior works have attempted to address this challenge by creating self-supervised auxiliary tasks, aiming to enrich the agent's learned representations with control-relevant information for future state prediction. However, these objectives are often insufficient to learn representations that can represent the optimal policy or value function, and they often consider tasks with small, abstract discrete action spaces and thus overlook the importance of action representation learning in continuous control. In this paper, we introduce TACO: Temporal Action-driven Contrastive Learning, a simple yet powerful temporal contrastive learning approach that facilitates the concurrent acquisition of latent state and action representations for agents. TACO simultaneously learns a state and an action representation by optimizing the mutual information between representations of current states paired with action sequences and representations of the corresponding future states. Theoretically, TACO can be shown to learn state and action representations that encompass sufficient information for control, thereby improving sample efficiency. For online RL, TACO achieves 40% performance boost after one million environment interaction steps on average across nine challenging visual continuous control tasks from Deepmind Control Suite. In addition, we show that TACO can also serve as a plug-and-play module adding to existing offline visual RL methods to establish the new state-of-the-art performance for offline visual RL across offline datasets with varying quality.
Paper Structure (28 sections, 3 theorems, 17 equations, 14 figures, 6 tables)

This paper contains 28 sections, 3 theorems, 17 equations, 14 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Visual reinforcement learning
  4. Contrastive learning and the InfoNCE objective
  5. TACO: temporal action-driven contrastive Loss
  6. Temporal contrastive learning objectives and analysis
  7. TACO implementation
  8. Experiments and results
  9. Comparison between TACO and strong baselines in online RL tasks
  10. Combining TACO with offline RL algorithms
  11. Related work
  12. Contrastive learning in visual reinforcement learning
  13. Action representation in reinforcement learning
  14. Conclusion
  15. Acknowledgement
  16. ...and 13 more sections

Key Result

Theorem 3.1

Let $K\in \mathbb{N}^+$, and $\mathbb{J}_\text{TACO\xspace}=\mathcal{I}(Z_{t+K}; [Z_t,U_t,..., U_{t+K-1}])$. If for a given state and action representation $\phi_Z, \psi_U$, $\mathbb{J}_\text{TACO\xspace}$ is maximized, then for arbitrary state-action pairs $(s_1,a_1), (s_2, a_2)$ such that $\phi(s_

Figures (14)

  • Figure 1: Comparison of average episode reward across nine challenging tasks in Deepmind Control Suite after one million environment steps.
  • Figure 2: A demonstration of our temporal contrastive loss: Given a batch of state-action transition triples $\{(s_{t}^{(i)},[a^{(i)}_t,..., a^{(i)}_{t+K-1}],s^{(i)}_{t+K})\}_{i=1}^{N}$, we first apply the state encoder and action encoder to get latent state-action encodings: $\{(z^{(i)}_t,[u^{(i)}_t,..., u^{(i)}_{t+K-1}],z^{(i)}_{t+K})\}_{i=1}^{N}$. Then we apply two different projection layers to map $(z^{(i)}_t,[u^{(i)}_t,..., u^{(i)}_{t+K-1}])$ and $z^{(i)}_{t+K}$ into the shared contrastive embedding space. Finally, we learn to predict the correct pairings between $(z_t,[u_t,..., u_{t+K-1}])$ and $z_{t+K}$ using an InfoNCE loss.
  • Figure 3: 1M Performance of TACO with and without action representation
  • Figure 4: (Deepmind Control Suite) Performance of TACO against two strongest model-free visual RL baselines. Results of DrQ-v2 and A-LIX are reproduced from their open-source implementations, and all results are averaged over 6 random seeds.
  • Figure 5: (Meta-world) Performance of TACO against DrQ-v2 and A-LIX. All results are averaged over 6 random seeds.
  • ...and 9 more figures

Theorems & Definitions (4)

  • Theorem 3.1
  • Proposition H.1
  • proof
  • Proposition H.2