[Paper Review] Memory Augmented Policy Optimization for Program Synthesis with Generalization

Liang Chen, Mohammad Norouzi+3 more|arXiv (Cornell University)|Jul 6, 2018

Reinforcement Learning in Robotics41 references16 citations

TL;DR

This paper proposes Memory Augmented Policy Optimization (MAPO), a policy gradient method that uses a memory buffer of high-reward trajectories to reduce variance and improve sample efficiency in deterministic environments with discrete actions. MAPO achieves state-of-the-art performance on program synthesis benchmarks—improving accuracy to 46.2% on WikiTableQuestions and 74.9% on WikiSQL—outperforming fully supervised baselines with only weak supervision.

ABSTRACT

This paper presents Memory Augmented Policy Optimization (MAPO): a novel policy optimization formulation that incorporates a memory buffer of promising trajectories to reduce the variance of policy gradient estimates for deterministic environments with discrete actions. The formulation expresses the expected return objective as a weighted sum of two terms: an expectation over a memory of trajectories with high rewards, and a separate expectation over the trajectories outside the memory. We propose 3 techniques to make an efficient training algorithm for MAPO: (1) distributed sampling from inside and outside memory with an actor-learner architecture; (2) a marginal likelihood constraint over the memory to accelerate training; (3) systematic exploration to discover high reward trajectories. MAPO improves the sample efficiency and robustness of policy gradient, especially on tasks with a sparse reward. We evaluate MAPO on weakly supervised program synthesis from natural language with an emphasis on generalization. On the WikiTableQuestions benchmark we improve the state-of-the-art by 2.5%, achieving an accuracy of 46.2%, and on the WikiSQL benchmark, MAPO achieves an accuracy of 74.9% with only weak supervision, outperforming several strong baselines with full supervision. Our code is open sourced at this https URL

Motivation & Objective

• Address the challenge of sample inefficiency and high variance in policy gradient methods for program synthesis tasks with sparse rewards.
• Improve generalization in weakly supervised program synthesis by leveraging memory of high-reward trajectories.
• Develop a training algorithm that efficiently explores and exploits promising trajectories in deterministic environments with discrete actions.
• Achieve superior performance on program synthesis benchmarks using only weak supervision, reducing reliance on expensive full supervision.
• Enhance robustness and convergence speed in policy optimization through memory-based variance reduction and systematic exploration.

Proposed method

• Formulate the expected return as a weighted sum of expectations over trajectories inside and outside a memory buffer of high-reward trajectories.
• Implement distributed sampling using an actor-learner architecture to efficiently sample from within and outside the memory buffer.
• Apply a marginal likelihood constraint on the memory to stabilize and accelerate training by preserving high-quality trajectory distributions.
• Integrate systematic exploration strategies to proactively discover high-reward trajectories not yet in the memory.
• Use a deterministic policy gradient framework enhanced with memory-augmented objective to reduce variance and improve learning stability.
• Maintain a dynamic memory buffer that is updated based on trajectory returns, prioritizing high-performing sequences.

Experimental results

Research questions

• RQ1 Can memory-augmented policy optimization reduce variance and improve sample efficiency in program synthesis with sparse rewards?
• RQ2 How does incorporating a memory of high-reward trajectories affect generalization in weakly supervised program synthesis?
• RQ3 To what extent can MAPO outperform fully supervised baselines when trained with only weak supervision on program synthesis tasks?
• RQ4 What impact does systematic exploration have on discovering high-reward trajectories in sparse-reward environments?
• RQ5 How does the marginal likelihood constraint over the memory buffer influence training convergence and stability?

Key findings

• MAPO achieves a state-of-the-art accuracy of 46.2% on the WikiTableQuestions benchmark, improving the prior SOTA by 2.5%.
• On the WikiSQL benchmark, MAPO attains 74.9% accuracy using only weak supervision, surpassing several strong baselines that rely on full supervision.
• The memory-augmented objective significantly reduces policy gradient variance, leading to more stable and efficient training in sparse-reward environments.
• Systematic exploration enables the agent to discover high-reward trajectories more effectively, especially in early training phases.
• The marginal likelihood constraint over the memory accelerates convergence and improves sample efficiency without compromising performance.
• MAPO demonstrates strong generalization capabilities, generalizing effectively to unseen test instances in program synthesis tasks.