🐍🔥 Lava Escape RL

A stochastic reinforcement learning research environment where an autonomous snake agent learns to navigate a dynamic lava maze and reach the exit — while avoiding Goodhart's Law.

---

📋 Table of Contents

1. Project Overview
2. Motivation
3. Reinforcement Learning Concepts
4. Mathematical Formulation
5. Goodhart's Law Discussion
6. Environment Design
7. Training Pipeline
8. Results
9. Future Improvements
10. Screenshots

---

1. Project Overview

Lava Escape RL transforms a classic Snake game into a professional reinforcement learning research environment. The snake is no longer controlled by a human — it is an autonomous learning agent trained via Q-Learning and Deep Q-Networks (DQN).

The agent must:

• 🧭 Navigate a procedurally-generated lava maze
• 🍎 Collect food for energy bonuses
• 🚪 Reach the exit tile to win the episode
• 🔥 Avoid dynamically-spreading lava
• 🔄 Learn to avoid reward hacking (Goodhart's Law)

Repository Structure

Snake/
├── src/                        # Original Java Snake game (unchanged)
├── python_rl/
│   ├── config.yaml             # Central hyperparameter configuration
│   ├── env/
│   │   ├── lava_escape_env.py  # Gymnasium-compatible MDP environment
│   │   ├── maze.py             # Wilson's algorithm maze generator
│   │   ├── lava.py             # Stochastic lava cellular automaton
│   │   └── renderer.py         # Pygame visual renderer
│   ├── agents/
│   │   ├── q_learning.py       # Tabular Q-learning (3 exploration modes)
│   │   ├── dqn.py              # Deep Q-Network (PyTorch)
│   │   └── replay_buffer.py    # Experience replay buffer
│   ├── rewards/
│   │   └── __init__.py         # NaiveReward + ShapedReward (anti-Goodhart)
│   ├── training/
│   │   ├── train_q.py          # Q-learning training loop
│   │   ├── train_dqn.py        # DQN training loop
│   │   └── logger.py           # CSV episode logger
│   ├── visualization/
│   │   ├── plot_curves.py      # Learning curves + Goodhart plots
│   │   └── policy_heatmap.py   # Q-value & policy arrow heatmap
│   └── evaluate.py             # Evaluation with greedy policy
└── results/
    ├── models/                 # Saved model checkpoints
    ├── logs/                   # CSV training logs
    └── plots/                  # Generated figures

---

2. Motivation

"When a measure becomes a target, it ceases to be a good measure." — Charles Goodhart, 1975

Most toy RL environments have trivially aligned reward functions — if the agent maximises reward, it automatically solves the task. The real world is different: optimising a proxy metric often leads to unintended behaviour.

This project uses a Snake game to demonstrate this concretely:

Reward Design	Agent Behaviour	True Task Solved?
Food count only	Circles food endlessly	❌ Never exits
Survival time only	Actively avoids exit	❌ Exit = death
Shaped reward	Navigates to exit	✅ Yes

This makes Lava Escape RL both an RL research tool and a Goodhart's Law demonstration platform.

---

3. Reinforcement Learning Concepts

3.1 Agent and Environment

An agent interacts with an environment in discrete time steps. At each step $t$:

1. The agent observes state $s_t \in \mathcal{S}$
2. The agent selects action $a_t \in \mathcal{A}$ according to policy $\pi$
3. The environment transitions to $s_{t+1} \sim P(\cdot \mid s_t, a_t)$
4. The agent receives scalar reward $r_t = R(s_t, a_t)$

3.2 Exploration Strategies

Three strategies are implemented and compared:

ε-greedy: $$a_t = \begin{cases} \text{random} & \text{with probability } \varepsilon \ \arg\max_a Q(s_t, a) & \text{with probability } 1 - \varepsilon \end{cases}$$

Decaying ε-greedy: $$\varepsilon_t = \max!\left(\varepsilon_{\min},; \varepsilon_0 \cdot d^t\right)$$ where $d \in (0,1)$ is the decay rate (configured as epsilon_decay in config.yaml).

Softmax (Boltzmann) Exploration: $$\pi(a \mid s) = \frac{\exp!\left(Q(s,a) / \tau\right)}{\sum_{a'} \exp!\left(Q(s,a') / \tau\right)}$$ Higher temperature $\tau$ → more uniform exploration. Lower $\tau$ → exploitation.

---

4. Mathematical Formulation

4.1 Markov Decision Process (MDP)

The environment is formalised as an MDP $\mathcal{M} = \langle \mathcal{S}, \mathcal{A}, P, R, \gamma \rangle$:

Component	Symbol	Description
State space	$\mathcal{S}$	Flattened grid channels: cell type, lava age, BFS distance-to-exit
Action space	$\mathcal{A}$	${$UP, DOWN, LEFT, RIGHT$}$, $|\mathcal{A}| = 4$
Transition	$P(s' \mid s, a)$	Stochastic — lava spread makes $s'$ random even for same $(s, a)$
Reward	$R(s, a)$	Shaped multi-objective signal (see §4.3)
Discount	$\gamma \in [0,1)$	$\gamma = 0.99$ — agent is nearly far-sighted

4.2 Stochastic Transitions: Lava Spread Model

Lava cells spread probabilistically, making the transition function non-deterministic:

$$P(\text{cell}_{r,c} = \text{LAVA} \mid k \text{ lava neighbours}) = 1 - (1 - p_{\text{spread}})^k$$

This is derived from the independent competitor model: each lava neighbour independently tries to ignite the cell with probability $p_{\text{spread}}$. The probability that at least one succeeds is $1 - (1 - p)^k$.

Lava state	Transition	Probability
EMPTY → LAVA	Spread from $k$ neighbours	$1-(1-p_s)^k$
LAVA → COOLED	Cooling	$p_c = 0.03$
COOLED → LAVA	Reignition	$p_r = 0.02$
COOLED → EMPTY	Full cooling	$1 - p_r = 0.98$

4.3 Return and Bellman Equation

The discounted return from time $t$: $$G_t = \sum_{k=0}^{\infty} \gamma^k r_{t+k+1} = r_{t+1} + \gamma r_{t+2} + \gamma^2 r_{t+3} + \cdots$$

The action-value function (Q-function) under policy $\pi$: $$Q^\pi(s, a) = \mathbb{E}_\pi!\left[G_t \mid s_t = s,, a_t = a\right]$$

The Bellman optimality equation: $$Q^(s, a) = \mathbb{E}!\left[r + \gamma \max_{a'} Q^(s', a') \mid s, a\right]$$

4.4 Q-Learning Update Rule

Q-learning estimates $Q^*$ via the temporal difference (TD) update:

$$Q(s_t, a_t) ;\leftarrow; Q(s_t, a_t) ;+; \alpha \underbrace{\left[, r_t + \gamma \max_{a'} Q(s_{t+1}, a') - Q(s_t, a_t),\right]}_{\delta_t ;:=; \text{TD error}}$$

Where:

• $\alpha \in (0,1]$ — learning rate (step size)
• $\delta_t$ — the TD error: how wrong our current estimate was

4.5 Deep Q-Network (DQN) Loss

DQN parameterises $Q(s,a;\theta)$ with a neural network. The loss is:

$$\mathcal{L}(\theta) = \mathbb{E}_{(s,a,r,s') \sim \mathcal{D}}!\left[\left(y - Q(s,a;\theta)\right)^2\right]$$

where the Bellman target uses a frozen target network $\theta^-$:

$$y = r + \gamma \max_{a'} Q(s', a';, \theta^-)$$

The target network $\theta^-$ is updated from $\theta$ every $N$ steps (hard copy), providing a stable regression target and preventing oscillations.

4.6 Shaped Reward Function

The anti-Goodhart reward is a linear combination of aligned objectives:

$$R(s,a) = w_{\text{exit}} \cdot \mathbb{1}[\text{exit}] \cdot B_{\text{exit}}(\text{coverage}) \

• w_{\text{food}} \cdot \mathbb{1}[\text{food}] \cdot r_f \
• w_{\text{surv}} \cdot r_s \
• w_{\text{expl}} \cdot \mathbb{1}[\text{new cell}] \cdot r_e \

• w_{\text{lava}} \cdot \Phi_{\text{lava}}(s) \
• w_{\text{loop}} \cdot \mathbb{1}[\text{revisit within } W] \
• w_{\text{stag}} \cdot \mathbb{1}[\text{no progress} > T]$$

The exit bonus is scaled by grid coverage: $$B_{\text{exit}}(\text{coverage}) = 100 \cdot \left(0.5 + 0.5 \cdot \frac{|\mathcal{V}|}{|\mathcal{C}|}\right)$$

where $\mathcal{V}$ is visited cells and $\mathcal{C}$ is total navigable cells. This incentivises thorough exploration before exiting.

The lava proximity penalty uses a distance-weighted sum: $$\Phi_{\text{lava}}(s) = \sum_{(r',c') \in N_2(s)} \frac{\mathbb{1}[\text{grid}[r',c'] = \text{LAVA}]}{d(s, (r',c')) + 1}$$

where $N_2(s)$ is the 2-step Manhattan neighbourhood of the snake's head.

---

5. Goodhart's Law Discussion

The Problem with Naive Rewards

Goodhart's Law states that when a measure becomes a target, it ceases to be a good measure. In RL, this manifests as reward hacking: the agent finds unexpected ways to maximise the proxy reward without achieving the intended goal.

Case 1: Food-only reward

# NaiveReward(mode='food_only')
# Exit gives NO reward → agent ignores the goal
reward = food_reward if ate_food else 0.0

The agent learns to circulate around food tiles indefinitely, never attempting to reach the exit. It achieves high reward while completely failing the task.

Case 2: Survival-only reward

# NaiveReward(mode='survival_only')
# Exit TERMINATES the episode → agent avoids it!
reward = survival_bonus  # per step alive

The agent learns that reaching the exit kills it (from a reward perspective). It actively avoids the exit and survives as long as possible — maximum reward, zero success.

The Fix: Multi-Objective Shaped Reward

The shaped reward ensures that the globally optimal policy under the reward function is also the intended behaviour:

1. Exit dominates: $B_{\text{exit}} = 100 \gg r_f = 5 \gg r_s = 0.1$ — no other strategy can match cumulative exit reward
2. Food is instrumental: food helps survival but alone can never match exit reward
3. Stagnation penalty: breaks food-farming loops
4. Loop penalty: prevents circular motion for proximity survival

This design makes reward hacking mathematically suboptimal.

---

6. Environment Design

Maze Generation (Wilson's Algorithm)

The maze is generated using Wilson's loop-erased random walk, which samples uniformly from all possible spanning trees of the grid graph. This guarantees:

• Every cell is reachable (no isolated pockets)
• No closed loops — the agent must navigate properly
• Unbiased structure — no directional preference

Observation Space

The agent receives a flat vector of dimensionality $H \times W \times 3 + 2$:

Channel	Content	Range
0	Cell type (normalised)	[0, 1]
1	Lava age (how long cell has been lava)	[0, 1]
2	BFS distance to exit (inverted, closer=higher)	[0, 1]
+2 scalars	Snake head position (normalised)	[0, 1]

Action Space

Discrete: $\mathcal{A} = {0=\text{UP},; 1=\text{DOWN},; 2=\text{LEFT},; 3=\text{RIGHT}}$

Stochasticity Sources

1. Lava spread (probabilistic neighbour model)
2. Food respawn locations (random empty cell)
3. Episode-to-episode maze variation (different random seeds)
4. Lava cool/reignite transitions

---

7. Training Pipeline

Quick Start

cd python_rl
pip install -r requirements.txt

Train Q-Learning (shaped reward)

python training/train_q.py --reward shaped --strategy epsilon_greedy --episodes 5000

Train Q-Learning (Goodhart demo — food only)

python training/train_q.py --reward naive_food --episodes 5000

Train DQN

python training/train_dqn.py --reward shaped --episodes 3000 --device cpu

Compare exploration strategies

python training/train_q.py --strategy epsilon_greedy
python training/train_q.py --strategy softmax

Generate all plots

python visualization/plot_curves.py --compare

Evaluate a trained agent

python evaluate.py --agent q --model ../results/models/q_shaped_final.pkl
python evaluate.py --agent dqn --model ../results/models/dqn_shaped_final.pt

Configuration

All hyperparameters live in python_rl/config.yaml — edit them without touching code.

---

8. Results

Training logs are saved as CSV in results/logs/. Plots are generated in results/plots/.

Expected Metrics After Training

Agent	Reward Mode	Exit Success Rate	Mean Return
Q-Learning	Shaped	~35–60%	+40 to +80
Q-Learning	Food Only	~0–2%	high proxy ❌
Q-Learning	Survival Only	~0%	moderate proxy ❌
DQN	Shaped	~50–75%	+60 to +110

Results vary by seed and episode count. Run evaluate.py after training for exact numbers.

---

9. Future Improvements

• Prioritised Experience Replay — sample high-TD-error transitions more often
• Dueling DQN — separate value and advantage streams
• PPO / A3C — policy gradient methods for comparison
• Curriculum learning — start with small mazes, gradually increase difficulty
• Multi-agent — competitive/cooperative multi-snake scenarios
• LSTM-DQN — partial observability with memory
• Intrinsic motivation — curiosity-driven exploration (ICM)
• Hyperparameter sweep — automated tuning with Optuna
• Web dashboard — real-time training metrics via Streamlit

---

10. Screenshots

Add your screenshots here after running training and evaluation.

Environment Visualization

[SCREENSHOT_1 — Lava maze with snake, food, and exit tiles]

Learning Curves

[SCREENSHOT_2 — Reward vs Episode and exit success rate over training]

Goodhart's Law Comparison

[SCREENSHOT_3 — Side-by-side: naive vs shaped reward agent behaviour]

---

Citation

If you use this project in academic work, please cite:

@misc{lava_escape_rl_2026,
  title   = {Lava Escape RL: Stochastic Reinforcement Learning with Anti-Goodhart Reward Design},
  year    = {2026},
  url     = {https://github.com/houssam-05-ctrl/Snake}
}

License

MIT — see LICENSE for details.