Reproduce “Meta reinforcement learning as task inference” on OpenResearch

About this paper

Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A. Ortega, Yee Whye Teh, Nicolas Heess

The authors introduce a separate belief network that estimates task information from experience and supplies it to a policy. Privileged task information supervises the network during meta-training but is unavailable at meta-test. In Numpad, the design outperforms a baseline LSTM and auxiliary-head agent. In tested sparse 2D navigation, it adapts faster than PEARL and avoids revisiting previously tried places.

Key insights

• Leveraging readily available "privileged information" (true task specifications) during meta-training provides a powerful supervisory signal for learning robust task representations, significantly improving meta-RL performance.
• The task posterior's independence from the agent's policy allows the belief network to be trained efficiently using off-policy data from a replay buffer, enhancing sample efficiency.
• Information Bottleneck (IB) regularization is critical for enabling efficient off-policy training with recurrent architectures, promoting compact and robust task-relevant representations.

Summary by alphaXiv.