Zero-Shot Reinforcement Learning via Function Encoders

2024-01-30

International Conference on Machine Learning

AI Summary

This work addresses the challenge of zero-shot cross-task transfer in reinforcement learning. We propose the Function Encoder framework, which maps reward and state-transition functions into low-dimensional, semantically consistent task embeddings via weighted combinations of nonlinear basis functions—enabling task alignment and immediate transfer without online fine-tuning. The framework is modular and seamlessly integrates with mainstream RL algorithms including PPO, SAC, and DQN. Experiments across multiple benchmark domains demonstrate substantial improvements in zero-shot generalization, achieving state-of-the-art performance in data efficiency, asymptotic policy quality, and training stability. Our core contribution is the first explicit encoding of task-level functional representations (i.e., reward and dynamics functions) into transferable vector embeddings—departing from conventional paradigms that rely solely on policy- or value-function-based transfer. This paradigm shift enables more principled and scalable cross-task knowledge reuse.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Multiagent Learning
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Although reinforcement learning (RL) can solve many challenging sequential decision making problems, achieving zero-shot transfer across related tasks remains a challenge. The difficulty lies in finding a good representation for the current task so that the agent understands how it relates to previously seen tasks. To achieve zero-shot transfer, we introduce the function encoder, a representation learning algorithm which represents a function as a weighted combination of learned, non-linear basis functions. By using a function encoder to represent the reward function or the transition function, the agent has information on how the current task relates to previously seen tasks via a coherent vector representation. Thus, the agent is able to achieve transfer between related tasks at run time with no additional training. We demonstrate state-of-the-art data efficiency, asymptotic performance, and training stability in three RL fields by augmenting basic RL algorithms with a function encoder task representation.

Problem

• Achieving zero-shot transfer across related RL tasks
• Finding good task representations for agent understanding
• Enhancing RL algorithms with function encoder representation

Innovation

• Function encoder for zero-shot transfer learning
• Learned non-linear basis functions combination
• Coherent vector representation for task relation