[Paper Review] BabyAI 1.1

David Yu-Tung Hui, Maxime Chevalier-Boisvert+2 more|arXiv (Cornell University)|Jul 24, 2020

Spectroscopy Techniques in Biomedical and Chemical Research4 citations

TL;DR

BabyAI 1.1 introduces three architectural and representational improvements to enhance sample efficiency in training agents to follow grounded-language instructions: removing max-pooling in the visual encoder, adding residual connections around FiLM layers, and adopting a Bag-of-Words (BOW) representation for visual inputs. These changes increase reinforcement learning sample efficiency by up to 3× and boost imitation learning success on the hardest level (BossLevel) from 77% to 90.4%.

ABSTRACT

The BabyAI platform is designed to measure the sample efficiency of training an agent to follow grounded-language instructions. BabyAI 1.0 presents baseline results of an agent trained by deep imitation or reinforcement learning. BabyAI 1.1 improves the agent's architecture in three minor ways. This increases reinforcement learning sample efficiency by up to 3 times and improves imitation learning performance on the hardest level from 77 % to 90.4 %. We hope that these improvements increase the computational efficiency of BabyAI experiments and help users design better agents.

Motivation & Objective

• Improve sample efficiency in training agents to follow grounded-language instructions in the BabyAI environment.
• Address limitations in the original BabyAI 1.0 agent architecture, particularly in information flow and training efficiency.
• Enhance imitation learning performance on the most challenging levels, especially BossLevel, where baseline success was only 77%.
• Increase computational efficiency of experiments by optimizing network architecture and visual representation.
• Provide a more robust and scalable foundation for future research in vision-language instruction following and zero-shot generalization.

Proposed method

• Remove max-pooling layers in the early stages of the visual encoder to preserve spatial resolution and improve information flow to higher layers.
• Introduce residual connections around FiLM layers to stabilize training and improve gradient flow in the visual-linguistic fusion module.
• Replace the raw integer-based tile representation with a learned Bag-of-Words (BOW) embedding, where each tile attribute (type, color, state) is embedded via a lookup table and averaged.
• Use a 3-channel RGB image representation as an alternative visual input, converting the 7×7 grid into a 56×56×3 tensor for end-to-end CNN processing.
• Train agents using both deep reinforcement learning (PPO) and imitation learning (behavior cloning), with ablation studies on architecture and visual representation variants.
• Evaluate performance across multiple levels, including both single-room (small) and multi-room (big) environments, with success rate and training sample efficiency as primary metrics.

Experimental results

Research questions

• RQ1 How do architectural modifications—specifically removing max-pooling and adding residual connections—affect sample efficiency in reinforcement learning on the BabyAI platform?
• RQ2 To what extent does adopting a Bag-of-Words (BOW) visual representation improve imitation learning performance compared to the original integer-based representation?
• RQ3 What is the impact of visual input representation (BOW vs. pixels vs. original) on training efficiency and final performance across different levels?
• RQ4 Can the combination of architectural and representational improvements significantly boost performance on the most challenging level (BossLevel), where baseline success was only 77%?
• RQ5 How do these changes affect computational efficiency, particularly in terms of frames per second during RL training?

Key findings

• Removing max-pooling in the visual encoder improved reinforcement learning sample efficiency by up to 3×, significantly reducing the number of training episodes required.
• The addition of residual connections around FiLM layers improved sample efficiency on the most difficult tasks, particularly in GoToLocal and PutNextLocal, despite mixed results across all tasks.
• Adopting a Bag-of-Words (BOW) visual representation increased imitation learning success on the hardest level (BossLevel) from 77% to 90.4%, demonstrating a substantial improvement in zero-shot generalization.
• The BOW representation achieved near-perfect imitation learning performance (≥99.5% success rate) on the six easiest levels with only 10,000 demonstrations, outperforming both original and pixel-based representations.
• The BOW-based architecture with residual connections (bow_endpool_res) achieved the best overall performance in both RL and IL, with the highest success rates and fastest convergence.
• Training with pixel-based visual inputs required a reduced learning rate (5×10⁻⁵) to stabilize training, especially on the most complex levels, indicating higher training difficulty compared to BOW and original representations.