Optimizing Monetary Policy with Q-Learning

import numpy as np
import gym
from gym import spaces
from gym.utils import seeding
import matplotlib.pyplot as plt

# ==============================================================================
# PART 1: THE ECONOMIC ENVIRONMENT
# ==============================================================================

class EconomyEnv(gym.Env):
    """
    A custom Gym environment simulating a simple economy based on the Phillips Curve.
    """
    metadata = {'render.modes': ['human']}

    def __init__(self,
                 inflation_persistence=1.0,
                 rate_sensitivity=0.2,
                 policy_change_cost=0.5,
                 max_episode_steps=100,
                 potential_scaling_factor=10
                ):
        super().__init__()

        # --- Economic Model Parameters ---
        self.beta = inflation_persistence    # How much past inflation affects current inflation
        self.sigma = rate_sensitivity      # How much interest rates affect inflation
        self.alpha = policy_change_cost      # Penalty for volatile policy changes
        self.kappa = potential_scaling_factor

        # --- Environment Configuration ---
        self.shock_std_dev = 0.5
        self.goal_inflation = 2.0
        self._max_episode_steps = max_episode_steps

        # --- State Space Boundaries ---
        self.min_inflation = -2.0
        self.max_inflation = 8.0
        self.min_interest_rate = 0.0  # Enforce the Zero Lower Bound (ZLB)
        self.max_interest_rate = 10.0 # 

        # --- Gym Setup ---
        self.state = None
        self._current_step = 0

        # Define observation space (a continuous range of values)
        obs_low = np.array([self.min_inflation, self.min_interest_rate], dtype=np.float32)
        obs_high = np.array([self.max_inflation, self.max_interest_rate], dtype=np.float32)
        self.observation_space = spaces.Box(low=obs_low, high=obs_high, dtype=np.float32)

        # Define action space (a discrete set of 5 actions)
        self.action_space = spaces.Discrete(5)
        self._action_map = {
            0: -50.0, # Large Cut
            1: -25.0, # Standard Cut
            2: 0.0,   # Hold
            3: 25.0,  # Standard Hike
            4: 50.0   # Large Hike
        }
        self.seed()

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        """Execute one time step within the environment."""
        current_inflation, current_interest_rate = self.state

        # --- Action ---
        # Map the discrete action (0-4) to a change in basis points
        bps_change = self._action_map[action]
        new_interest_rate = np.clip(
            current_interest_rate + (bps_change / 100.0),
            self.min_interest_rate,
            self.max_interest_rate
        )

        # --- Reward Calculation (via Reward Shaping) ---
        # This "dense" reward function guides the agent at every step.
        # 1. Calculate the deviation from the inflation target in the current state.
        current_deviation = np.abs(current_inflation - self.goal_inflation)

        # 2. Calculate the expected deviation in the next state.
        expected_next_inflation = (self.beta * current_inflation) - (self.sigma * new_interest_rate)
        next_deviation = np.abs(expected_next_inflation - self.goal_inflation)

        # 3. The primary reward is the *improvement* in deviation. A positive
        #    reward is given for getting closer to the target.
        reward_for_progress = self.kappa * (current_deviation - next_deviation)

        # 4. A small penalty is applied for making large, volatile policy changes.
        policy_penalty = self.alpha * (bps_change / 100.0)**2
        reward = reward_for_progress - policy_penalty

        # --- State Transition ---
        # The actual next state is the expected state plus a random economic shock.
        economic_shock = self.np_random.normal(0, self.shock_std_dev)
        new_inflation = expected_next_inflation + economic_shock

        self._current_step += 1

        # --- Terminal Conditions ---
        is_unstable = not (self.min_inflation <= new_inflation <= self.max_inflation)
        is_over_time = self._current_step >= self._max_episode_steps
        is_successful = bool(np.abs(new_inflation - self.goal_inflation) < 0.2)
        done = is_unstable or is_over_time or is_successful

        # Add large terminal rewards/penalties to strongly guide final outcomes.
        if is_unstable:
            reward -= 100.0
        if is_successful:
            reward += 100.0

        # Update the state and return results
        self.state = np.array([
            np.clip(new_inflation, self.min_inflation, self.max_inflation),
            new_interest_rate
        ], dtype=np.float32)

        return self.state, reward, done, {}

    def reset(self):
        """Reset the environment to a random initial state."""
        self._current_step = 0
        # Start each episode in a random valid state to ensure broad learning
        self.state = self.observation_space.sample()
        return self.state

# ==============================================================================
# PART 2: THE Q-LEARNING AGENT 
# ==============================================================================

class QLearningAgent:
    """A standard agent that learns via the Q-learning algorithm."""
    def __init__(self, state_bins, num_actions, alpha, gamma, epsilon):
        self.num_actions = num_actions
        self.alpha = alpha      # Learning rate: How much to update Q-values
        self.gamma = gamma      # Discount factor: Importance of future rewards
        self.epsilon = epsilon    # Exploration rate: Probability of taking a random action
        self.q_table = np.zeros(state_bins + (num_actions,))

    def choose_action(self, state_tuple):
        """Chooses an action using an epsilon-greedy policy."""
        if np.random.uniform(0, 1) < self.epsilon:
            # Exploration: Take a random action
            return np.random.choice(self.num_actions)
        else:
            # Exploitation: Choose the action with the highest Q-value
            q_values = self.q_table[state_tuple]
            return np.argmax(q_values)

    def update_q_table(self, state, action, reward, next_state):
        """Updates the Q-value for a state-action pair using the Bellman equation."""
        old_value = self.q_table[state][action]
        next_max = np.max(self.q_table[next_state])

        # The core of Q-learning
        new_value = old_value + self.alpha * (reward + self.gamma * next_max - old_value)
        self.q_table[state][action] = new_value

# ==============================================================================
# PART 3: TRAINING AND VISUALIZATION
# ==============================================================================

def discretize_state(state, bins):
    """Converts a continuous state vector into a discrete tuple for Q-table indexing."""
    inflation_val, interest_rate_val = state
    # Find which bin each continuous value falls into
    inflation_bin = np.digitize(inflation_val, bins[0]) - 1
    interest_rate_bin = np.digitize(interest_rate_val, bins[1]) - 1

    # Clip values to ensure they are valid indices for the Q-table
    return (
        np.clip(inflation_bin, 0, len(bins[0]) - 2),
        np.clip(interest_rate_bin, 0, len(bins[1]) - 2)
    )

if __name__ == '__main__':
    # --- Hyperparameters ---
    NUM_EPISODES = 3_000_000
    NUM_INFLATION_BINS = 30
    NUM_INTEREST_BINS = 40
    LEARNING_RATE = 0.1
    DISCOUNT_FACTOR = 0.99
    EPSILON_START = 1.0
    EPSILON_END = 0.01
    # Epsilon decays over the first 80% of episodes
    EPSILON_DECAY_RATE = (EPSILON_START - EPSILON_END) / (NUM_EPISODES * 0.8)

    # --- Initialization ---
    env = EconomyEnv()
    state_discretization_bins = [
        np.linspace(env.min_inflation, env.max_inflation, NUM_INFLATION_BINS),
        np.linspace(env.min_interest_rate, env.max_interest_rate, NUM_INTEREST_BINS)
    ]
    q_table_shape = (NUM_INFLATION_BINS - 1, NUM_INTEREST_BINS - 1)

    agent = QLearningAgent(
        state_bins=q_table_shape,
        num_actions=env.action_space.n,
        alpha=LEARNING_RATE,
        gamma=DISCOUNT_FACTOR,
        epsilon=EPSILON_START
    )

    rewards_per_episode = []
    print("--- Starting Training ---")

    # --- Training Loop ---
    for episode in range(NUM_EPISODES):
        state_continuous = env.reset()
        state_discrete = discretize_state(state_continuous, state_discretization_bins)
        done = False
        total_episode_reward = 0

        while not done:
            action = agent.choose_action(state_discrete)
            next_state_continuous, reward, done, _ = env.step(action)
            next_state_discrete = discretize_state(next_state_continuous, state_discretization_bins)
            agent.update_q_table(state_discrete, action, reward, next_state_discrete)
            state_discrete = next_state_discrete
            total_episode_reward += reward

        # Decay epsilon after each episode
        if agent.epsilon > EPSILON_END:
            agent.epsilon -= EPSILON_DECAY_RATE

        rewards_per_episode.append(total_episode_reward)

        # Print progress update periodically
        if (episode + 1) % 50000 == 0:
            avg_reward = np.mean(rewards_per_episode[-5000:])
            print(f"Episode: {episode + 1:>7}/{NUM_EPISODES} | Avg Reward (last 5k): {avg_reward:8.2f} | Epsilon: {agent.epsilon:.4f}")

    print("--- Training Finished ---")