Abstract:Representation learning based on multi-task pretraining has become a powerful approach in many domains. In particular, task-aware representation learning aims to learn an optimal representation for a specific target task by sampling data from a set of source tasks, while task-agnostic representation learning seeks to learn a universal representation for a class of tasks. In this paper, we propose a general and versatile algorithmic and theoretic framework for \textit{active representation learning}, where the learner optimally chooses which source tasks to sample from. This framework, along with a tractable meta algorithm, allows most arbitrary target and source task spaces (from discrete to continuous), covers both task-aware and task-agnostic settings, and is compatible with deep representation learning practices. We provide several instantiations under this framework, from bilinear and feature-based nonlinear to general nonlinear cases. In the bilinear case, by leveraging the non-uniform spectrum of the task representation and the calibrated source-target relevance, we prove that the sample complexity to achieve $\varepsilon$-excess risk on target scales with $ (k^*)^2 \|v^*\|_2^2 \varepsilon^{-2}$ where $k^*$ is the effective dimension of the target and $\|v^*\|_2^2 \in (0,1]$ represents the connection between source and target space. Compared to the passive one, this can save up to $\frac{1}{d_W}$ of sample complexity, where $d_W$ is the task space dimension. Finally, we demonstrate different instantiations of our meta algorithm in synthetic datasets and robotics problems, from pendulum simulations to real-world drone flight datasets. On average, our algorithms outperform baselines by $20\%-70\%$.

Task Aware Dreamer for Task Generalization in Reinforcement Learning

Reward Informed Dreamer for Task Generalization in Reinforcement Learning

Leveraging the Efficiency of Multi-Task Robot Manipulation Via Task-Evoked Planner and Reinforcement Learning

HarmonyDream: Task Harmonization Inside World Models

Mastering Diverse Domains through World Models

A Task-Agnostic Regularizer for Diverse Subpolicy Discovery in Hierarchical Reinforcement Learning

Dream to Adapt: Meta Reinforcement Learning by Latent Context Imagination and MDP Imagination

Sample Efficient Myopic Exploration Through Multitask Reinforcement Learning with Diverse Tasks

Decompose a Task into Generalizable Subtasks in Multi-Agent Reinforcement Learning.

World Models with Hints of Large Language Models for Goal Achieving

Efficient Multi-Task Reinforcement Learning via Task-Specific Action Correction

Continual Deep Reinforcement Learning with Task-Agnostic Policy Distillation

DREAM: Adaptive Reinforcement Learning based on Attention Mechanism for Temporal Knowledge Graph Reasoning

Synthesizing Programmatic Policy for Generalization Within Task Domain

Mastering Atari with Discrete World Models

TransDreamer: Reinforcement Learning with Transformer World Models

Understanding What Affects the Generalization Gap in Visual Reinforcement Learning: Theory and Empirical Evidence

Powderworld: A Platform for Understanding Generalization via Rich Task Distributions

Active Representation Learning for General Task Space with Applications in Robotics

Towards Understanding How to Reduce Generalization Gap in Visual Reinforcement Learning.