Back to work

rl

A Q-learning implementation to solve package collection tasks in the classic Four-Rooms grid world environment. The agent learns optimal policies across three complexity levels: single-package collection, multi-package unordered collection, and strict sequential multi-package collection (RGB order), supporting both deterministic and stochastic environment dynamics.

Mastering the Four-Rooms Environment with Q-Learning

This project focuses on applying the fundamental reinforcement learning algorithm, Q-learning, to solve dynamic navigation and collection tasks within the renowned Four-Rooms grid world environment. By developing a robust Q-Agent, we tackled three progressively challenging scenarios, incorporating state expansion and strategic reward functions to achieve optimal policies.

1. The Q-Learning Agent Implementation

The core of this project is the QAgent class, designed to learn optimal state-action policies using the $\epsilon$-greedy strategy.

State Space Representation

A crucial design decision was defining the state space to capture the necessary environmental context. Since the agent's objective depends not just on its physical location $(x, y)$ but also on the packages remaining, the state was expanded to a three-dimensional tuple: $$ S = (x, y, k) $$ where $x$ and $y$ are the coordinates on the $13 \times 13$ grid, and $k$ is the number of packages yet to be collected (ranging from 0 up to 4). This resulted in a Q-table dimension of $13 \times 13 \times 5 \times 4$ (for $4$ actions), ensuring all necessary states for the multi-package tasks could be differentiated.

The Q-Learning Update Rule

The agent updates its quality value $Q(s, a)$ using the standard Bellman equation for Q-learning: $$ Q(s, a) \leftarrow Q(s, a) + \alpha \left[ R + \gamma \max_{a'} Q(s', a') - Q(s, a) \right] $$ We utilized a learning rate ($\alpha$) of 0.1 and a discount factor ($\gamma$) of 0.99, prioritizing long-term rewards, which is essential for multi-step navigation problems.

Exploration Strategy

Exploration was managed using linear $\epsilon$-decay. The exploration rate $\epsilon$ gradually decreases from a starting value ($\epsilon_{start} = 1.0$) down to a minimum ($\epsilon_{end} = 0.01$ or $0.05$ for stochastic environments) over a defined number of training epochs.

2. Scenario Analysis and Reward Engineering

We successfully implemented and trained the Q-Agent across three distinct scenarios:

Scenario 1: Simple Collection

In this single-package environment, the agent only needs to find one target.

  • Reward: $-1$ per step, $+100$ upon collecting the single package.
  • Exploration Comparison: We compared a "Fast Decay/Low End" strategy ($\epsilon$ drops quickly to 0.01) against a "Slower Decay/High End" strategy ($\epsilon$ drops slowly to 0.10). The comparison plot (scenario1_exploration_comparison.png) demonstrated that while the fast decay strategy led to quicker convergence initially, the slower decay maintained higher performance stability, especially in stochastic environments, highlighting the trade-off between exploitation speed and maintaining generalization.

Scenario 2: Multi-Package Unordered Collection

The environment contains 4 packages, and the agent must collect all of them, regardless of order.

  • State Impact: The state $k$ (packages remaining) is critical here, as the optimal path changes based on which packages are still present.
  • Reward: $-1$ per step, $+50$ for collecting any package, and an additional $+100$ when the last package is collected (total $+150$ for the final collection). This immediate reward structure helps reinforce the behavior of actively seeking and collecting packages.

Scenario 3: Sequential Collection (RGB)

This scenario introduced significant complexity, requiring the agent to collect packages in a predefined order: P4 (initial random package) $\rightarrow$ Red $\rightarrow$ Blue $\rightarrow$ Green.

  • State Expansion & Reward Penalty: The use of the $k$ dimension in the state space implicitly tracks the progress (k=4 means looking for P4, k=3 means looking for Red, etc.). Crucially, if the agent lands on the wrong package for the current step (e.g., lands on Blue when $k=3$ and Red is expected), the episode terminates with a large penalty ($-100$), encouraging strict adherence to the sequence.
  • Performance: Due to the massive state space and the strict sequential constraints, Scenario 3 required significantly more training epochs (up to 75,000 in stochastic mode) to converge on a reliable policy. The generated path visualizations show the agent successfully navigating the complex sequence, demonstrating the Q-learning agent's capability in sequential decision-making tasks when supported by precise state tracking and reward engineering.

3. Conclusion

This project successfully leveraged Q-learning and a strategically expanded state space to solve increasingly complex package collection tasks in the Four-Rooms environment. Effective hyperparameter tuning, especially adjusting the $\epsilon$-decay rate based on environmental uncertainty (stochastic vs. deterministic), proved vital for achieving convergence and optimal policy performance. The project provides a strong practical foundation for understanding multi-goal reinforcement learning in grid worlds.

Built with

PythonReinforcement LearningQ-LearningGrid WorldNumPyMatplotlibMachine Learning