Toward Task Generalization via Memory Augmentation in Meta-Reinforcement Learning
This addresses the challenge of deploying RL in diverse and dynamic settings, though it appears incremental as it builds on existing memory-based RL approaches.
The paper tackles the problem of reinforcement learning agents struggling with tasks that differ from training by introducing memory augmentation to improve task generalization, achieving zero-shot generalization to unseen tasks in simulation and real-world legged locomotion experiments.
Agents trained via reinforcement learning (RL) often struggle to perform well on tasks that differ from those encountered during training. This limitation presents a challenge to the broader deployment of RL in diverse and dynamic task settings. In this work, we introduce memory augmentation, a memory-based RL approach to improve task generalization. Our approach leverages task-structured augmentations to simulate plausible out-of-distribution scenarios and incorporates memory mechanisms to enable context-aware policy adaptation. Trained on a predefined set of tasks, our policy demonstrates the ability to generalize to unseen tasks through memory augmentation without requiring additional interactions with the environment. Through extensive simulation experiments and real-world hardware evaluations on legged locomotion tasks, we demonstrate that our approach achieves zero-shot generalization to unseen tasks while maintaining robust in-distribution performance and high sample efficiency.