Layout Generalization and Reward Shaping in Cooperative Overcooked-AI

This project studies whether self-play PPO policies can generalize across kitchen layouts in the cooperative Overcooked-AI environment. The experimental setup combines reward shaping with a stepwise decay schedule and a deep policy network, and examines the effect of entropy regularization and reward shaping on both learning efficiency and final policy performance.

The full write-up is in report.pdf.

Example replay

Two self-play agents cooperating in the cramped_room layout:

Generalization experiment: learning curves

Learning curves across the three training layouts (cramped_room, asymmetric_advantages, forced_coordination):

Project layout

.
├── configs/                    # YAML experiment configs (cramped_room, generalization)
├── src/overcooked_rl/
│   ├── config.py                # TrainConfig / EvalConfig dataclasses + YAML loading
│   ├── rewards.py                # Shaping decay + scale schedule
│   ├── envs/generalized.py       # Multi-layout Overcooked-AI wrapper
│   ├── agents/                   # PolicyNetwork, ValueNetwork, PPOAgent
│   ├── training/                  # Training loop, checkpoint helpers
│   ├── evaluation/                # Single-episode rollout, multi-layout generalization eval
│   ├── logging_utils.py          # CSV + TensorBoard logging
│   └── cli.py                    # `overcooked-train` / `overcooked-eval` entry points
├── notebooks/demo.ipynb          # Evaluation notebook with rollout GIFs
├── checkpoints/                  # Tracked: the final policy/value pair used by the demo
├── assets/                       # Tracked: README/notebook media
├── results/                      # Untracked: logs, extra checkpoints, TensorBoard runs, GIFs
├── scripts/setup_overcooked.sh   # Clones + installs the Overcooked-AI dependency
└── report.pdf

Setup (Linux)

Overcooked-AI is a required external dependency and is not vendored in this repo, so it's installed separately. Python 3.10+.

# 1. Create and activate a virtual environment
python3.10 -m venv .venv
source .venv/bin/activate

# 2. Install Overcooked-AI and this package (editable)
./scripts/setup_overcooked.sh

# Optional: notebook + test extras
pip install -e ".[dev]"

scripts/setup_overcooked.sh clones overcooked_ai next to this repo (or into a directory you pass as its first argument) and installs it with the [harl] extra (gym, pygame, and the rest of the RL stack), then installs this repo in editable mode.

Conda works the same way — create the env with conda create -n overcooked-rl python=3.10 before running the script.

Windows

The code itself is platform-independent (paths are handled with pathlib), but scripts/setup_overcooked.sh is a bash script. On Windows, either run it under WSL/Git Bash, or do the two steps manually:

git clone https://github.com/HumanCompatibleAI/overcooked_ai.git ../overcooked_ai
pip install -e "../overcooked_ai[harl]"
pip install -e .

Training

Training is config-driven; the two committed experiments each have a YAML file under configs/:

overcooked-train --config configs/cramped_room.yaml
overcooked-train --config configs/generalization.yaml

Any field can be overridden from the command line without editing the YAML:

overcooked-train --config configs/generalization.yaml \
    --total-episodes 5000 --seed 0 \
    --set ppo.lr=0.0003 --set reward_shaping.decay_total=10000

Progress is logged to a CSV file (log_dir in the config) and to TensorBoard:

tensorboard --logdir results/tensorboard_logs

Checkpoints are saved every checkpoint_every episodes, plus a best (on greedy-eval improvement) and final pair at the end of the run, under checkpoint_dir.

Evaluation

overcooked-eval loads a checkpoint and reports greedy + sampling reward across one or more layouts, saving a rollout GIF per layout:

overcooked-eval --config configs/eval_generalization.yaml

Demo notebook

notebooks/demo.ipynb reproduces the report's final results using the checkpoints tracked under checkpoints/, and displays rollout GIFs inline:

jupyter notebook notebooks/demo.ipynb

Repository data

Only a curated set of binaries is tracked in git: the final policy/value checkpoints the demo notebook loads (checkpoints/) and the media referenced by this README (assets/). Everything else training/evaluation produces (intermediate checkpoints, CSV/TensorBoard logs, per-run GIFs) is written under results/, which is gitignored. The full set of checkpoints and logs from the original experiments remains available in this repository's git history if needed.