Reinforcement Learning on Tower of Hanoi

import numpy as np

class TowerOfHanoiEnv:
    def __init__(self, n_discs=3, max_steps=50):
        self.n_discs = n_discs
        self.max_steps = max_steps
        self.actions = [(1,2),(1,3),(2,1),(2,3),(3,1),(3,2)] #action space
        self.reset()

    def reset(self):
        # Towers: top of stack on right
        self.towers = [list(range(self.n_discs, 0, -1)), [], []]
        self.steps = 0
        return self._get_state()

    def _get_state(self):
        # Flattened state (pad with 0s)
        state = []
        for peg in self.towers:
            s = peg + [0]*(self.n_discs - len(peg))
            state.extend(s)
        return np.array(state, dtype=int)

    def _is_valid_move(self, src, dest):
        # Check for illegal move
        if not self.towers[src]:
            return False
        if self.towers[dest] and self.towers[src][-1] > self.towers[dest][-1]:
            return False
        return True

    def step(self, action):
        """
        Action: integer 0-5 mapping to move (src, dest)
        Mapping: [(1->2),(1->3),(2->1),(2->3),(3->1),(3->2)]
        """
        src, dest = self.actions[action]
        src -= 1
        dest -= 1
        reward = 0
        done = False

        if self._is_valid_move(src, dest):
            self.towers[dest].append(self.towers[src].pop())
        else:
            reward = -1  # penalty for illegal move

        self.steps += 1

        if self.towers[2] == list(range(self.n_discs, 0, -1)):
            reward = 1  # solved
            done = True
        elif self.steps >= self.max_steps:
            done = True

        return self._get_state(), reward, done

    def render(self):
        # Render state to command line for human interpretability
        max_height = self.n_discs
        for level in range(max_height, 0, -1):
            row = []
            for peg in self.towers:
                if len(peg) >= level:
                    row.append(str(peg[level-1]).center(5))
                else:
                    row.append(" ".center(5))
            print(" ".join(row))
        print("__________________")
        print("  P1    P2    P3  ")
        print("-" * 18)

import random
from TowerOfHanoiEnv import *

env = TowerOfHanoiEnv(n_discs=3)

# Hyperparameters
episodes = 5000
alpha = 0.1 # Learning Rate - extent to which the newly acquired information overrides the old information (the current Q-value).
gamma = 0.99 # Discount Factor - determines the importance of future rewards vs. immediate rewards.
epsilon = 0.2 # Epsilon Greedy (20% chance to explore new action)

algorithm_choice = 0
algorithm_choices = {0: 'sarsa', 1: 'q_learning', 2: 'expected_sarsa'} #q_learning a.k.a. sarsamax
current_algorithm = algorithm_choices[algorithm_choice]
print(f"{current_algorithm} RL algorithm is used.")

Q = {}
ep_steps = []

def select_action(state, Q, epsilon, action_space_size):
    """Epsilon-greedy action selection."""
    if state not in Q:
        # Add unseen state to Q-table
        Q[state] = [0.0] * action_space_size
        
    if random.random() < epsilon:
        # Explore: Choose a random action
        return random.randint(0, action_space_size - 1)
    else:
        # Exploit: Choose the best-known action
        return int(np.argmax(Q[state]))

for ep in range(episodes):
    print(f"Episode: {ep}\n")
    state = tuple(env.reset())
    done = False
    steps = 0

    action = select_action(state, Q, epsilon, len(env.actions))

    while not done:
        current_action = action
        
        next_state, reward, done = env.step(action)
        next_state = tuple(next_state)
        
        if next_state not in Q:
            Q[next_state] = [0]*len(env.actions)
        
        # td_target - Temporal Difference Target
        td_target = 0

        if done:
            td_target = reward
        else:
            greedy_action_idx = int(np.argmax(Q[next_state]))

            if current_algorithm == 'sarsa':
                action_for_next_step = select_action(next_state, Q, epsilon, len(env.actions))
                target_q_value = Q[next_state][action_for_next_step]
            elif current_algorithm == 'q_learning': #sarsamax
                # Greedy Policy - Use max (best) Q i.e. rewards from previous to select next action
                action_for_next_step = select_action(next_state, Q, epsilon, len(env.actions)) # Still need next step action for behavior (Off-Policy)
                target_q_value = max(Q[next_state])
            elif current_algorithm == 'expected_sarsa':
                # Expected SARSA: A' is the average value over the behavior policy
                action_for_next_step = select_action(next_state, Q, epsilon, len(env.actions)) # Still need next step action for behavior (Off-Policy)
                
                expected_q_value = 0
                prob_greedy = 1 - epsilon + (epsilon / len(env.actions))
                prob_non_greedy = epsilon / len(env.actions)

                for a_prime in range(len(env.actions)):
                    q_value = Q[next_state][a_prime]
                    if a_prime == greedy_action_idx:
                        expected_q_value += prob_greedy * q_value
                    else:
                        expected_q_value += prob_non_greedy * q_value
                target_q_value = expected_q_value

            # Calculate the target Q - only differences between the three methods
            td_target = reward + gamma * target_q_value
        # Update Q values
        old_q_value = Q[state][current_action]
        Q[state][current_action] = old_q_value + alpha * (td_target - old_q_value)
        
        state = next_state
        action = action_for_next_step 

        steps += 1
        env.render()
        print(f"Action: {env.actions[action]} Reward: {reward}\n")

    print(f"Episode {ep} Finished! Total steps: {steps}")
    ep_steps.append(steps)

done = False
step = 0

# Setup environment
env = TowerOfHanoiEnv(n_discs=3)
state = tuple(env.reset()) #initial state
print(f"step {step}:\n")
env.render()

while not done:
    action = Q[state].index(max(Q[state])) # Choose best action given by Q-Table
    step += 1
    print(f"\nstep {step}: {env.actions[action][0]}->{env.actions[action][1]}")
    next_state, reward, done = env.step(action)
    state = tuple(next_state)
    env.render()

import numpy as np
import random
from collections import deque, namedtuple

from TowerOfHanoiEnv import *

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

class QNetwork(nn.Module):
    def __init__(self, state_size, action_size, seed):
        """Initializes parameters and builds model.
        
        Args:
            state_size (int): Dimension of each state (e.g., 9 for 3 discs).
            action_size (int): Dimension of action space (6 possible moves).
            seed (int): Random seed for reproducibility.
        """
        super(QNetwork, self).__init__()
        # Set a seed for reproducibility
        self.seed = torch.manual_seed(seed)
        
        # Simple feedforward network
        self.fc1 = nn.Linear(state_size, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, action_size)

    def forward(self, state):
        """Maps state input to action Q-values."""
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        return self.fc4(x)        

# Hyperparameters for the Agent
BUFFER_SIZE = int(1e5)  # Replay buffer size
BATCH_SIZE = 64         # Minibatch size for sampling
GAMMA_DQN = 0.99        # Discount factor
LR = 2e-4               # Learning rate of the optimizer
UPDATE_EVERY = 4        # How often to update the network (step count)
TARGET_UPDATE_FREQ = 50 # How often to update the target network (step count)

# Define the device for PyTorch operations
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Replay Buffer named tuple
Experience = namedtuple('Experience', field_names=['state', 'action', 'reward', 'next_state', 'done'])

class ReplayBuffer:
    """Fixed-size buffer to store experience tuples."""
    def __init__(self, buffer_size, batch_size, seed):
        self.memory = deque(maxlen=buffer_size)  
        self.batch_size = batch_size
        self.seed = random.seed(seed)
    
    def add(self, state, action, reward, next_state, done):
        """Add a new experience to memory."""
        e = Experience(state, action, reward, next_state, done)
        self.memory.append(e)
    
    def sample(self):
        """Randomly sample a batch of experiences from memory."""
        experiences = random.sample(self.memory, k=self.batch_size)

        # Convert experiences to Tensors
        states = torch.from_numpy(np.vstack([e.state for e in experiences if e is not None])).float().to(device)
        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(device)
        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(device)
        next_states = torch.from_numpy(np.vstack([e.next_state for e in experiences if e is not None])).float().to(device)
        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(device)
  
        return (states, actions, rewards, next_states, dones)

    def __len__(self):
        """Return the current size of internal memory."""
        return len(self.memory)

class DQNAgent:
    """Interacts with and learns from the environment."""
    def __init__(self, state_size, action_size, seed):
        self.state_size = state_size
        self.action_size = action_size
        self.seed = random.seed(seed)

        # Q-Network (Local and Target)
        self.qnetwork_local = QNetwork(state_size, action_size, seed).to(device)
        self.qnetwork_target = QNetwork(state_size, action_size, seed).to(device)
        self.optimizer = optim.Adam(self.qnetwork_local.parameters(), lr=LR)

        # Replay memory
        self.memory = ReplayBuffer(BUFFER_SIZE, BATCH_SIZE, seed)
        # Initialize time step for learning updates
        self.t_step = 0
    
    def step(self, state, action, reward, next_state, done):
        # Save experience in replay memory
        self.memory.add(state, action, reward, next_state, done)
        
        # Learn every UPDATE_EVERY time steps.
        self.t_step = (self.t_step + 1) % UPDATE_EVERY
        if self.t_step == 0:
            # If enough samples are available in memory, get a random subset and learn
            if len(self.memory) > BATCH_SIZE:
                experiences = self.memory.sample()
                self.learn(experiences, GAMMA_DQN)

        # Update target network every TARGET_UPDATE_FREQ steps
        if self.t_step % TARGET_UPDATE_FREQ == 0:
            self.update_target_network()

    def act(self, state, epsilon, inference=False):
        """Returns actions for given state as per current policy (epsilon-greedy)."""
        # Convert state from numpy array to PyTorch tensor
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        
        # Get Q-values from local network
        self.qnetwork_local.eval() # Set network to evaluation mode (no gradient tracking)
        with torch.no_grad():
            # Note: Using torch.no_grad() is crucial during action selection to prevent training.
            action_values = self.qnetwork_local(state)

        if not inference:
            self.qnetwork_local.train() # Set network back to training mode

        # Epsilon-greedy action selection
        if random.random() > epsilon:
            return np.argmax(action_values.cpu().data.numpy()) # Best action
        else:
            return random.choice(np.arange(self.action_size)) # Random action
    
    def learn(self, experiences, gamma):
        """Update value parameters using given batch of experience tuples."""
        states, actions, rewards, next_states, dones = experiences

        # Get max predicted Q values (for next states) from target model
        # The Q-Learning update rule uses the max Q value for the next state
        Q_targets_next = self.qnetwork_target(next_states).detach().max(1)[0].unsqueeze(1)
        
        # Compute Q targets for current states: Q_target = R + gamma * max_a(Q_target(s', a))
        Q_targets = rewards + (gamma * Q_targets_next * (1 - dones))

        # Get expected Q values from local model: Q_expected = Q_local(s, a)
        # Gather the Q-value for the action that was actually taken
        Q_expected = self.qnetwork_local(states).gather(1, actions)

        # Compute loss (Mean Squared Error)
        loss = F.mse_loss(Q_expected, Q_targets)
        
        # Minimize the loss
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def update_target_network(self):
        """Soft update model parameters: $\\theta_{\\text{target}} = \\tau*\\theta_{\\text{local}} + (1 - \\tau)*\\theta_{\\text{target}}$"""
        # A full copy is often used in basic DQN, meaning $\\tau=1$.
        # This is a full copy update.
        for target_param, local_param in zip(self.qnetwork_target.parameters(), self.qnetwork_local.parameters()):
            target_param.data.copy_(local_param.data)

    def save_agent(self, path):
        torch.save(self.qnetwork_local.state_dict(), path)
    
    def load_agent(self, path):
        self.qnetwork_local.load_state_dict(torch.load(path))
        self.qnetwork_target.load_state_dict(torch.load(path))

# --- DQN Specific Hyperparameters ---
EPISODES_DQN = 40000 
EPS_START = 1.0       # Starting value of epsilon
EPS_END = 0.001       # Minimum value of epsilon
EPS_DECAY = 0.9998    # Decay rate for epsilon - lower epsilon after each episode (i.e. less likely to choose random actions later in the training)
SEED = 0
best_agent_path = "TowerHanoi_DQN.pth"

# --- Environment Setup (Using the existing TowerOfHanoiEnv) ---
env = TowerOfHanoiEnv(n_discs=3, max_steps=50)
state_size = env.n_discs * 3 # 3 pegs * n_discs (for padding)
action_size = len(env.actions)

# --- Initialize the DQNAgent ---
agent = DQNAgent(state_size=state_size, action_size=action_size, seed=SEED)

# --- DQN Training Loop ---
eps = EPS_START             # Initialize epsilon
ep_steps_DQN = []
solved_in_episodes = None
min_steps = 2**env.n_discs - 1 # Optimal moves for N discs

print(f"--- Starting DQN Training for {env.n_discs} Discs ---")

for ep in range(EPISODES_DQN):
    state = env.reset() # State is a numpy array
    done = False
    steps = 0
    
    while not done:
        # Agent selects an action using epsilon-greedy policy
        action = agent.act(state, eps)
        
        # Environment takes a step
        next_state, reward, done = env.step(action)
        
        # Agent stores experience and learns
        # Note: We pass the action index (int) and the reward/done must be float/int
        agent.step(state, action, reward, next_state, done)
        
        state = next_state
        steps += 1
        
        # Print rendering for every 100 episodes
        if ep % 1000 == 0:
             env.render()
             print(f"Action: {env.actions[action]} Reward: {reward}")

    eps = max(EPS_END, eps * EPS_DECAY)
    
    ep_steps_DQN.append(steps)
    
    # Check for successful completion
    if env.towers[2] == list(range(env.n_discs, 0, -1)):
        status = "**SOLVED!**"

        # Cheating:
        if steps == min_steps and solved_in_episodes is None:
            # if reached optimal solution (mininum steps)
            # However this doesn't mean anything when epsilon is still high (i.e. best step is still by random chance)
            solved_in_episodes = ep
    else:
        status = "Max steps reached"
        
    if ep % 100 == 0 or steps == min_steps:
        print(f"\nEpisode {ep}/{EPISODES_DQN} | Steps: {steps} | Status: {status} | Epsilon: {eps:.4f} | Avg Steps (last 100 episodes): {np.mean(ep_steps_DQN[-100:]):.1f}")
        
    if solved_in_episodes is not None and eps <= EPS_END: # <-- Only stop if epsilon has reached its minimum
        print("Stop early since epsilon has decayed and no significant improvement found.")
        break

# Save final agent**
agent.save_agent(best_agent_path)
print(f"\nFinal Agent Saved!")

print("\n--- DQN Training Finished ---")
if solved_in_episodes is not None:
    print(f"Optimal (?) solution found at episode: {solved_in_episodes} with {ep_steps_DQN[solved_in_episodes]} steps.")
print(f"Final average steps (last 100 episodes): {np.mean(ep_steps_DQN[-100:]):.1f}")

done = False
step = 0

# Setup environment
env = TowerOfHanoiEnv(n_discs=3)
state = env.reset()
print(f"step {step}:\n")
env.render()

best_agent = DQNAgent(state_size=state_size, action_size=action_size, seed=0) # Seed won't be needed if epsilon is decay to 0.0
best_agent.load_agent(best_agent_path)

while not done:
    action = agent.act(state, 0.001, inference=True) # Turn on inference and set epsilon to the final epsilon value
    step += 1
    next_state, reward, done = env.step(action)
    print(f"\nstep {step}: {env.actions[action]} Reward: {reward}")
    state = next_state
    env.render()
print(f"Total steps: {step}")

def hanoi_solver(n, src, dest, aux, moves_list):
    """
    Finds the optimal sequence of moves to solve the Tower of Hanoi puzzle.

    Args:
        n (int): The number of discs to move.
        src (int): The source peg (1, 2, or 3).
        dest (int): The destination peg (1, 2, or 3).
        aux (int): The auxiliary (temporary) peg (1, 2, or 3).
        moves_list (list): A list to append the required moves (src, dest) to.
    """
    if n == 1:
        # Base case: Move the single disc from source to destination
        moves_list.append((src, dest))
        return

    # 1. Move n-1 discs from source to auxiliary peg
    hanoi_solver(n - 1, src, aux, dest, moves_list)
    # 2. Move the largest disc (n) from source to destination peg
    moves_list.append((src, dest))
    # 3. Move n-1 discs from auxiliary to destination peg
    hanoi_solver(n - 1, aux, dest, src, moves_list)

def solve_hanoi_with_env(env):
    """
    Uses the recursive solver to get the optimal moves and applies them 
    to the provided TowerOfHanoiEnv instance.

    Args:
        env (TowerOfHanoiEnv): The environment instance (must be reset).
    
    Returns:
        list: The sequence of moves (tuples) found by the solver.
    """
    n_discs = env.n_discs
    
    # Pegs are 1, 2, 3 (consistent with environment's action definition)
    src_peg = 1
    dest_peg = 3
    aux_peg = 2
    
    optimal_moves = []
    
    # Generate the optimal sequence of moves
    hanoi_solver(n_discs, src_peg, dest_peg, aux_peg, optimal_moves)
    
    print(f"Optimal Moves Generated: {len(optimal_moves)} (2^{n_discs} - 1)")
       
    # Convert (src, dest) tuple to the environment's action index (0-5)
    action_map = {move: i for i, move in enumerate(env.actions)}
    
    print("\nStarting Board State:")
    env.render()
    
    for move in optimal_moves:
        # Get the action index for the move
        action_index = action_map[move]
        
        # Take a step in the environment
        _, _, done = env.step(action_index)
        
        # Render the board
        env.render()
        
        if done:
            break
            
    print("\n--- Solving Complete ---")
    print(f"Final Steps: {env.steps}")
    print(f"Solved: {env.towers[2] == list(range(n_discs, 0, -1))}")
    env.render()
    
    return optimal_moves

# %%
env_3discs = TowerOfHanoiEnv(n_discs=3, max_steps=800)
solve_hanoi_with_env(env_3discs)