MountainCar Q-Learning Pipeline

A step-by-step guide to applying Q-learning to the classic MountainCar environment from Gymnasium.

Libraries Used

This project uses a few key Python libraries. Here is a quick overview of their roles:

• numpy: Used for numerical operations and managing arrays.
• gymnasium: The library that provides the classic MountainCar-v0 environment.
• collections: Used for the `defaultdict` to store the Q-table.

1. Define the Discretization and Q-Learning Agent

• What the code does: It defines a `Discretizer` class to convert the continuous state space of MountainCar into a discrete grid. It also defines a `QLearningAgent` class, which is a standard implementation for tabular Q-learning.
• Behind the scenes: The MountainCar environment's state space (position and velocity) is continuous, but `Q-learning` requires a discrete state space to work with a Q-table. The `Discretizer` handles this by mapping continuous values to integer bins.
• Abstraction: The `Discretizer` and `QLearningAgent` classes abstract the complex logic for managing the state space and the learning algorithm.

import numpy as np
from collections import defaultdict

class Discretizer:
    def __init__(self, bins):
        self.bins = bins

    def __call__(self, obs) -> tuple[int, ...]:
        obs = np.array(obs, dtype=float).ravel()
        return tuple(int(np.digitize(o, b)) for o, b in zip(obs, self.bins))

class QLearningAgent:
    def __init__(self, n_actions, alpha=0.1, gamma=0.99, epsilon=1.0, epsilon_min=0.05, epsilon_decay=0.995):
        self.n_actions = n_actions
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min
        self.epsilon_decay = epsilon_decay
        self.Q = defaultdict(lambda: np.zeros(n_actions, dtype=np.float32))

    def act(self, state):
        if np.random.rand() < self.epsilon:
            return np.random.randint(self.n_actions)
        return int(np.argmax(self.Q[state]))

    def learn(self, s, a, r, s_next, done):
        best_next_q = 0.0 if done else np.max(self.Q[s_next])
        td_target = r + self.gamma * best_next_q
        td_error = td_target - self.Q[s][a]
        self.Q[s][a] += self.alpha * td_error
        
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

2. Create the Environment and Run the Experiment

• What the code does: It creates the MountainCar environment, the agent, and the discretizer. It then runs a training loop for thousands of episodes, allowing the agent to learn the optimal policy.
• Behind the scenes: In each episode, the agent interacts with the environment, using the `Discretizer` to convert the continuous state into a discrete state for the Q-table. It takes actions, receives rewards, and learns. Over time, the agent learns how to gain momentum to reach the goal.
• Abstraction: The training loop abstracts the complex, long-term process of learning and policy improvement. The `env.step()` function abstracts the physics of the MountainCar environment.

import gymnasium as gym

env = gym.make("MountainCar-v0")
agent = QLearningAgent(n_actions=env.action_space.n)

bins = [
    np.linspace(env.observation_space.low[0], env.observation_space.high[0], 10),
    np.linspace(env.observation_space.low[1], env.observation_space.high[1], 10)
]
discretizer = Discretizer(bins)

n_episodes = 20000
for i in range(n_episodes):
    state = discretizer(env.reset()[0])
    done = False
    episode_reward = 0
    while not done:
        action = agent.act(state)
        next_obs, reward, done, _, _ = env.step(action)
        next_state = discretizer(next_obs)
        agent.learn(state, action, reward, next_state, done)
        state = next_state
        episode_reward += reward
    
    if (i + 1) % 1000 == 0:
        print(f"Episode {i + 1}, Epsilon: {agent.epsilon:.4f}, Total Reward: {episode_reward:.2f}")

print("\nTraining finished. You can now test the policy.")