Qlearning.py

'''
*
*   ===================================================
*       CropDrop Bot (CB) Theme [eYRC 2025-26]
*   ===================================================
*
*  This script is intended to be an Boilerplate for 
*  Task 1B of CropDrop Bot (CB) Theme [eYRC 2025-26].
*
*  Filename:		Qlearning.py
*  Created:		    24/08/2025
*  Last Modified:	24/08/2025
*  Author:		    e-Yantra Team
*  Team ID:		    [ CB_2202 ]
*  This software is made available on an "AS IS WHERE IS BASIS".
*  Licensee/end user indemnifies and will keep e-Yantra indemnified from
*  any and all claim(s) that emanate from the use of the Software or
*  breach of the terms of this agreement.
*  
*  e-Yantra - An MHRD project under National Mission on Education using ICT (NMEICT)
*
*****************************************************************************************
'''
'''You can Modify the this file,add more functions According to your usage.
   You are not allowed to add any external packges,Beside the included Packages.You can use Built-in Python modules.
'''
import numpy as np
import random
import pickle
import os

class QLearningController:
    def __init__(self, n_states=34, n_actions=5, filename="q_table.pkl"): 
        """
        Initialize the Q-learning controller.

        Parameters:
        - n_states (int): Total number of discrete states the agent can be in.
        - n_actions (int): Total number of discrete actions the agent can take.
        - filename (str): Filename to save/load the Q-table (persistent learning).
        """

        self.n_states = n_states
        self.n_actions = n_actions

        # === Configure your learning rate (alpha) and exploration rate (epsilon) here ===
        self.lr = 0.1  # Learning rate: how much new information overrides old
        self.epsilon = 0.3  # Exploration rate: chance of choosing a random action
        self.gamma = 0.9

        self.filename = filename

        # Initialize Q-table with zeros; dimensions: [states x actions]
        self.q_table = np.zeros((n_states, n_actions))

        # Action list: should be populated with your lineFollowers valid actions.The Below is just an Example.
        self.action_list = [0, 1, 2, 3, 4]  # Example: 0 = left, 1 = forward, 2 = right

        self.base_speed = 2.5

        # Mapping of action index to specific commands (e.g., motor speeds)
        self.actions = {
            0: (self.base_speed / 2, self.base_speed * 2),
            1: (self.base_speed / 1.25, self.base_speed * 1.25),
            2: (self.base_speed, self.base_speed),
            3: (self.base_speed * 1.25, self.base_speed / 1.25),
            4: (self.base_speed * 2, self.base_speed / 2),
        }

        self.avg_reward = 0
        self.episodes = 1
        self.iterations = 0

        self.prev_state = None

    def Get_state(self, sensor_data):
        """
        Convert raw sensor data into a discrete state index.
        
        Parameters:
        - sensor_data: Any sensor readings from your environment (e.g., distance sensors).
        
        Returns:
        - state (int): A unique index representing the current state.

        === TODO: Implement your logic to convert sensor_data to a discrete state ===
        """
        keys = ["left_corner", "left", "middle", "right", "right_corner"]
        threshold = .5
        state_array = [1 if sensor_data[key] > threshold else 0 for key in keys]
        state_str = "".join(map(str, state_array))
        state_int = int(state_str, 2)

        if state_int == 0:
            if self.prev_state in [1, 3, 7, 32]: # 32 ->  Left
                state_int = 32
            elif self.prev_state in [16, 24, 28, 33]: # 33 -> Right
                state_int = 33
            else:
                state_int = 0
        
        self.prev_state = state_int
        return state_int

    def Calculate_reward(self, state):
        """
        Calculate the reward based on the State.

        Parameters:
        - state: Current State of the Linefollower.

        Returns:
        - reward : The reward for the last action.

        === TODO: Implement your reward function here. 
                  For example, give +1 for going straight, -1 for hitting a wall, etc. ===
        """
        reward = 0
        if state in [4, 14]:
            reward = 5
        elif state in [6, 12]:
            reward = 4
        elif state in [2, 8]: 
            reward = 3
        elif state in [3, 24]:
            reward = 2  
        elif state in [1, 16]:
            reward = 1
        elif state in [28, 7]:
            reward = 0
        elif state in [0, 32, 33]:
            reward = -3
        else:
            reward = -1

        self.iterations += 1
        self.avg_reward += (reward - self.avg_reward) / self.iterations 

        return reward

    def update_q_table(self, state, action, reward, next_state):
        """
        Parameters:
        - state (int): Current state.
        - action (int): Action taken.
        - reward (float): Reward received.
        - next_state (int): State reached after taking the action.

        === TODO: Implement the Q-learning update rule here ===
        """
        a_idx = self.action_list.index(action)

        self.q_table[state][a_idx] = self.q_table[state][a_idx] + self.lr * (reward + self.gamma * np.max(self.q_table[next_state]) - self.q_table[state][a_idx])

        if self.epsilon > 0.1: self.epsilon *= 0.9999

    def choose_action(self, state):
        """
        Choose an action to perform based on the current state.

        Uses an epsilon-greedy strategy:
        - With probability epsilon: choose a random action (exploration)
        - Otherwise: choose the best known action (exploitation)

        Parameters:
        - state (int): Current state index.

        Returns:
        - action : The action chosen (from action_list).

        === TODO: Replace with your epsilon-greedy selection logic ===
        """
        if random.uniform(0, 1) < self.epsilon:
            return random.choice(self.action_list)
        else:
            return int(np.argmax(self.q_table[state]))
        
    def perform_action(self, action):
        """
        Translate an action into motor commands or robot movement.

        Parameters:
        - action : Action selected by the agent.

        Returns:
        - left_speed : Speed for the left motor.
        - right_speed : Speed for the right motor.

        === TODO: Implement action-to-motor translation logic based on your robot ===
        """
        left_motor_speed, right_motor_speed = self.actions.get(action, (self.base_speed, self.base_speed))
        return left_motor_speed, right_motor_speed

    def save_q_table(self):
        """
        Save the current Q-table and parameters to a file.

        Useful for keeping learned behavior between runs.

        === INSTRUCTIONS: You may Save Additional Thing while saving but do not Remove the the following Parameters ===
        """
        with open(self.filename, 'wb') as f:
            pickle.dump({
                'q_table': self.q_table,
                'epsilon': self.epsilon,
                'n_action': self.n_actions,
                'n_states': self.n_states,
                # Add any additional data you want to save
                'episodes': self.episodes,
                
            }, f)

        print(f"Avg Q: {np.mean(self.q_table)} | Avg Reward: {self.avg_reward} | Epsilon: {self.epsilon} | LR: {self.lr} | Episode: {self.episodes} | Iterations {self.iterations}")

    def load_q_table(self):
        """
        Load the Q-table and parameters from file, if it exists.

        Returns:
        - True if data was loaded successfully, False otherwise.
        """
        if os.path.exists(self.filename):
            with open(self.filename, 'rb') as f:
                data = pickle.load(f)

            self.q_table = data.get('q_table', self.q_table)
            self.epsilon = data.get('epsilon', self.epsilon) 
            self.n_actions = data.get('n_action', self.n_actions)
            self.n_states = data.get('n_states', self.n_states)
            # Load other data here if needed
            self.episodes = data.get('episodes', self.episodes) + 1
            return True
        return False