robot.py

# Imports from external libraries
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib import pyplot as plt
import math

# Imports from this project
# You should not import any other modules, including config.py
# If you want to create some configuration parameters for your algorithm, keep them within this robot.py file
import config
import constants
from graphics import VisualisationLine

# Configure matplotlib for interactive mode
plt.ion()


# CONFIGURATION PARAMETERS. Add whatever configuration parameters you like here.
# Remember, you will only be submitting this robot.py file, no other files.
SEED = 5
################### State Estimator Class ###################
class StateEstimator:
    def __init__(self):
        self.observation_params = [None, None, None, None, None]
        self.calibrated = False
        self.initial_state = None

    def calibrate_from_known_x(self, obs, known_x=0.05):
        self.observation_params[4] = obs[4]
        param3 = math.atan(obs[2]) - known_x - obs[4]
        self.observation_params[3] = param3
        self.calibrated = True
        self.initial_state = self.recover_state(obs)
        print("[INFO] Calibration complete.")
        return (param3, obs[4])
    
    def recover_state(self, obs):
        """Recover the original state from an observation"""
        if not self.calibrated:
            print("Error: Calibration not done.")
            return None
        
        # Recover x from tan function
        # We know x should be in range 0-200, so adjust if needed
        x_raw = math.atan(obs[2]) - self.observation_params[3] - self.observation_params[4]
        
        # Adjust x to ensure it's in the expected range (0-200)
        # Since tangent repeats every π, we may need to add multiples of π
        while x_raw < 0:
            x_raw += math.pi
        
        # If x is still too large, wrap it into the valid range
        x = x_raw % 200
        
        # Recover y from exponential function (this should be correct already)
        y = obs[3] / math.exp(self.observation_params[3])
        
        state = np.array([x, y], dtype=np.float32)
        return state

    



# The Robot class (which could be called "Agent") is the "brain" of the robot, and is used to decide what action to execute in the environment
class Robot:

    # Initialise a new robot
    def __init__(self):
        # The environment (only available during development mode)
        self.environment = None
        
        # Create the state estimator
        self.state_estimator = StateEstimator()
        self.estimator_calibrated = False
        
        # A list of visualisations which will be displayed on the bottom half of the window
        self.visualisation_lines = []
        
        # Initialize resistance map (200×100 grid)
        self.resistance_map = np.ones((200, 100))  * 0.9
        self.resistance_visits = np.zeros((200, 100))    # No visits initially
        
        # Q-learning parameters
        self.num_actions = 5  # We'll discretize actions into 8 directions
        self.q_values = np.zeros((200, 100, self.num_actions))  # State-action values
        self.actions = [  # Discrete action set
            [1.0, 0.0],    # Right
            [0.7, 0.7],    # Up-right
            [0.0, 1.0],    # Up
            [0.0, -1.0],   # Down
            [0.7, -0.7]    # Down-right
        ]
        
        self.epsilon = 0.3 # Initial exploration rate
        
        # For stuck detection
        self.position_history = []
        self.stuck_detection_window = 30
        self.stuck_threshold = 0.01
            
        # Budget management
        self.demo_received = True
        
        # Goal tracking
        self.goal_reached = False
        self.goal_bonus_reward = 50.0  # Large bonus for reaching the goal
        
        # Visualize the initial map
        self.update_visualization()
    
    def track_position(self, state):
        """Track position to detect if robot is stuck"""
        self.position_history.append(state.copy())
        
        # Keep history limited to window size
        if len(self.position_history) > self.stuck_detection_window:
            self.position_history.pop(0)

    def check_if_stuck(self, current_state):
        # Need enough history to make determination
        if len(self.position_history) < self.stuck_detection_window:
            return False
        
        # Check if robot has moved significantly in recent history
        start_pos = self.position_history[0]
        
        # Check x progress specifically
        x_progress = current_state[0] - start_pos[0]
        
        # Stuck condition: minimal x-axis progress
        is_stuck = x_progress < 0.005
        
        if is_stuck:
            print(f"Stuck detection: x_progress={x_progress:.3f}")
        
        return is_stuck
    # Get the next training action
    def training_action(self, obs, money):
        print("Money", money)
        # Make sure state estimator is calibrated
        if not self.estimator_calibrated:
            self.state_estimator.calibrate_from_known_x(obs)
            self.estimator_calibrated = True
            print("State estimator calibrated!")
        
        # Check if we have enough money for any action
        if money < min(constants.COST_PER_STEP, constants.COST_PER_RESET):
            print("Insufficient funds to continue training. Ending training.")
            return 4, None  # Finish training
        
        
        # Get current state
        state = self.state_estimator.recover_state(obs)
        if state is None:
            # Fallback if state estimation fails
            return 1, np.array([1.0, 0.0])
        
        # Check if we've reached the goal line
        if state[0] == constants.GOAL_LINE_X:
            self.goal_reached = True
            self.position_history = []
            # Reset to starting position
            reset_state = np.array([0.05, np.random.uniform(0, constants.ENVIRONMENT_HEIGHT)])
            # Reset goal flag
            return 2, reset_state
        
        # Check if robot is stuck
        is_stuck = self.check_if_stuck(state)
        
        # If stuck and not enough money for reset, end training
        if is_stuck and money < constants.COST_PER_RESET:
            print("Robot is stuck and insufficient funds to reset. Ending training.")
            return 4, None  # Finish training
        
        # If stuck and we have enough budget for a reset, do it
        if is_stuck:
            print("Robot appears stuck. Performing reset to a new position.")
            # Reset position history when resetting
            self.position_history = []
            
            # Reset to x=0.05 with a random y position
            reset_state = np.array([0.05, np.random.uniform(0, constants.ENVIRONMENT_HEIGHT)])
            return 2, reset_state
        
        # Convert to grid coordinates
        grid_x = int(state[0] * 100)
        grid_y = int(state[1] * 100)
        grid_x = max(0, min(grid_x, 199))
        grid_y = max(0, min(grid_y, 99))
        
        # Proper epsilon-greedy exploration/exploitation
        if np.random.random() < self.epsilon:
            # EXPLORATION: Random action with right bias
            right_actions = [idx for idx, action in enumerate(self.actions) if action[0] > 0]
            action_idx = np.random.choice(right_actions if right_actions else range(self.num_actions))
        else:
            # EXPLOITATION: Use resistance-based approach (like testing)
            # Calculate resistance for each possible action
            action_resistances = []
            for action_idx, action in enumerate(self.actions):
                # Predict next position
                next_x = min(199, max(0, grid_x + int(action[0] * 10)))
                next_y = min(99, max(0, grid_y + int(action[1] * 10)))
                
                # Get resistance at predicted position
                resistance = self.resistance_map[next_x, next_y]
                action_resistances.append((action_idx, resistance))
            
            # Sort actions by lowest resistance
            sorted_actions = sorted(action_resistances, key=lambda x: x[1])
            
            # Choose action with lowest resistance that moves forward
            right_biased_actions = [
                action for action in sorted_actions 
                if self.actions[action[0]][0] > 0  
            ]
            
            # If no right-moving low resistance actions, fall back to lowest resistance
            action_idx = right_biased_actions[0][0] if right_biased_actions else sorted_actions[0][0]
        
        # Get the actual action vector
        action_value = np.array(self.actions[action_idx]) * constants.MAX_ACTION_MAGNITUDE
        
        # Decrease epsilon over time (properly implementing epsilon decay)
        self.epsilon = max(0.1, self.epsilon - 0.0001)
        
        # Store for later Q-value updates
        self.last_state = (grid_x, grid_y)
        self.last_action_idx = action_idx
        
        # Track this position for stuck detection
        self.track_position(state)
        
        return 1, action_value

    # Get the next testing action
    def testing_action(self, obs):
        # Get current state
        state = self.state_estimator.recover_state(obs)
        if state is None:
            # Fallback if state estimation fails
            return np.array([1.0, 0.0])
        
        # Convert to grid coordinates
        grid_x = int(state[0] * 100)
        grid_y = int(state[1] * 100)
        grid_x = max(0, min(grid_x, 199))
        grid_y = max(0, min(grid_y, 99))
        
        # Check if we've reached or passed goal line
        if state[0] >= constants.GOAL_LINE_X:
            return np.array([0.0, 0.0])  # Stop moving
        
        # Calculate resistance for each possible action
        action_resistances = []
        for action_idx, action in enumerate(self.actions):
            # Predict next position
            next_x = min(199, max(0, grid_x + int(action[0] * 10)))
            next_y = min(99, max(0, grid_y + int(action[1] * 10)))
            
            # Get resistance at predicted position
            resistance = self.resistance_map[next_x, next_y]
            action_resistances.append((action_idx, resistance))
        
        # Sort actions by lowest resistance
        sorted_actions = sorted(action_resistances, key=lambda x: x[1])
        
        # Choose action with lowest resistance that moves forward
        right_biased_actions = [
            action for action in sorted_actions 
            if self.actions[action[0]][0] > 0  
        ]
        
        # If no right-moving low resistance actions, fall back to lowest resistance
        best_action_idx = right_biased_actions[0][0] if right_biased_actions else sorted_actions[0][0]
        
        # Get the actual action vector
        action = np.array(self.actions[best_action_idx]) * constants.MAX_ACTION_MAGNITUDE
        
        return action

    # Receive a transition
    def receive_transition(self, obs, action, next_obs, reward):

        # Only process if estimator is calibrated
        if not self.estimator_calibrated:
            return
                
        # Get states from observations
        state = self.state_estimator.recover_state(obs)
        next_state = self.state_estimator.recover_state(next_obs)
        
        if state is not None and next_state is not None:
            # Update resistance map
            self.update_resistance_map(state, action, next_state)
      
            # Occasionally update visualization with more frequency
            if np.random.random() < 0.2:
                self.update_visualization()

    
    # Helper function to find closest action index
    def find_closest_action_idx(self, normalized_action):
        """Find index of the discrete action closest to the given continuous action"""
        distances = [np.linalg.norm(np.array(a) - normalized_action) for a in self.actions]
        return np.argmin(distances)
    
    # Helper function for exploration action selection
    def get_exploration_action_idx(self, grid_x, grid_y):
        """Choose an action that biases toward unexplored areas and rightward movement"""
        # Initialize weights for each action
        action_weights = np.ones(self.num_actions)
        
        # Strongly bias toward rightward movement (action index 0)
        action_weights[0] *= 3.0  # Right
        action_weights[1] *= 2.0  # Up-right
        action_weights[4] *= 2.0  # Down-right
        
        # Check visited status of adjacent cells
        for i, (dx, dy) in enumerate([
            (1, 0),    # Right
            (1, 1),    # Up-right
            (0, 1),    # Up
            (0, -1),   # Down
            (1, -1)    # Down-right
        ]):
            # Check if this direction leads to a valid cell
            nx, ny = grid_x + dx, grid_y + dy
            if 0 <= nx < 200 and 0 <= ny < 100:
                # If cell has low visit count, increase probability
                if self.resistance_visits[nx, ny] < 1:
                    action_weights[i] *= 3.0
                
                # If cell has high resistance, decrease probability
                if self.resistance_map[nx, ny] > 0.7:
                    action_weights[i] *= 0.1
        
        # Normalize weights to form a probability distribution
        action_weights /= np.sum(action_weights)
        
        # Sample an action based on weights
        return np.random.choice(self.num_actions, p=action_weights)
    
    # Update resistance map based on observed transition
    def update_resistance_map(self, state, action, next_state):
        # Calculate expected distance without resistance
        expected_distance = 3 * np.linalg.norm(action)
        
        # Calculate actual distance traveled
        actual_distance = np.linalg.norm(next_state - state)
        
        # Calculate resistance (prevent division by zero)
        if expected_distance > 0.001:
            # Higher resistance = slower movement = lower ratio
            movement_ratio = actual_distance / expected_distance
            # Convert to resistance (0 = no resistance, 1 = full resistance)
            resistance = 1.0 - min(1.0, movement_ratio)
        else:
            resistance = 0.5  # Default if action was negligible
            
        # Convert state to grid coordinates
        center_x = int(state[0] * 100)
        center_y = int(state[1] * 100)
        
        # Define the neighborhood size (how many cells around the center to update)
        neighborhood_size = 20
        
        # Update the center cell and surrounding cells
        for dx in range(-neighborhood_size, neighborhood_size + 1):
            for dy in range(-neighborhood_size, neighborhood_size + 1):
                grid_x = center_x + dx
                grid_y = center_y + dy
                
                # Ensure coordinates are within bounds
                if 0 <= grid_x < 200 and 0 <= grid_y < 100:
                    # Calculate distance from center (for weighting)
                    distance = np.sqrt(dx**2 + dy**2)
                    
                    # Apply weight based on distance (higher weight for closer cells)
                    # Use Gaussian-like weighting
                    weight = np.exp(-distance**2 / (2 * 3.0**2))  # 1.0 is like the standard deviation
                    
                    # The further from center, the less we change the value
                    effective_resistance = weight * resistance + (1 - weight) * self.resistance_map[grid_x, grid_y]
                    
                    # Update resistance using running average
                    visits = self.resistance_visits[grid_x, grid_y]
                    current_value = self.resistance_map[grid_x, grid_y]
                    
                    # More weight to new observations when we have few visits
                    # But less weight for cells farther from the center
                    learning_rate = weight / (visits + 1)
                    
                    # Update the resistance value
                    self.resistance_map[grid_x, grid_y] = (1 - learning_rate) * current_value + learning_rate * effective_resistance
                    
                    # Only increment visit count for the center cell to maintain proper learning rates
                    if dx == 0 and dy == 0:
                        self.resistance_visits[grid_x, grid_y] += 1
                    # For surrounding cells, increment by a fraction based on distance
                    elif visits < 5:  # Only for cells with few visits
                        self.resistance_visits[grid_x, grid_y] += weight * 0.5
    
    def update_visualization(self):
        def get_heat_map_color(value):
            """Convert a value from 0-1 to a more distinctive color gradient"""
            # Blue (0.0) -> Cyan -> Green -> Yellow -> Red (1.0)
            if value < 0.25:
                # Blue to Cyan (0.0-0.25)
                r = 0
                g = int((value / 0.25) * 255)
                b = 255
            elif value < 0.5:
                # Cyan to Green (0.25-0.5)
                r = 0
                g = 255
                b = int(255 - ((value - 0.25) / 0.25) * 255)
            elif value < 0.75:
                # Green to Yellow (0.5-0.75)
                r = int(((value - 0.5) / 0.25) * 255)
                g = 255
                b = 0
            else:
                # Yellow to Red (0.75-1.0)
                r = 255
                g = int(255 - ((value - 0.75) / 0.25) * 255)
                b = 0
            
            return (r, g, b)
    
        # Clear previous visualization lines
        self.visualisation_lines = []
        
        # Sample points from the resistance map
        step_x, step_y = 5, 5
        
        # Add offsets to shift the grid up and left
        offset_x = -0.0  # Shift left (negative value)
        offset_y = 0.0   # Shift up (positive value)
        
        # Draw grid points with color based on resistance
        for x in range(0, 200, step_x):
            for y in range(0, 100, step_y):
                # Convert grid coordinates to environment coordinates with offset
                env_x = (x / 100.0) + offset_x  # Scale to [0,2] and shift left
                env_y = (y / 100.0) + offset_y  # Scale to [0,1] and shift up
                
                # Size of point based on confidence
                size = 0.01
                if self.resistance_visits[x, y] > 0:
                    size = min(0.02, 0.01 + 0.002 * self.resistance_visits[x, y])
                
                # Color based on resistance (blue for low, red for high)
                resistance = self.resistance_map[x, y]
                # RGB color: Convert from 0-1 range to 0-255 range for PyGame
                color = get_heat_map_color(resistance)
                
                # Draw a point (small square)
                self.visualisation_lines.append(VisualisationLine(
                    env_x, env_y, env_x + size, env_y, color
                ))
                self.visualisation_lines.append(VisualisationLine(
                    env_x + size, env_y, env_x + size, env_y + size, color
                ))
                self.visualisation_lines.append(VisualisationLine(
                    env_x + size, env_y + size, env_x, env_y + size, color
                ))
                self.visualisation_lines.append(VisualisationLine(
                    env_x, env_y + size, env_x, env_y, color
                ))
        
        # Add policy visualization (show arrows for learned policy)
        step = 20  # Visualize policy at coarser grid to avoid crowding
        for x in range(0, 200, step):
            for y in range(0, 100, step):
                if np.any(self.q_values[x, y] != 0):  # Only show where we have data
                    # Get best action
                    action_idx = np.argmax(self.q_values[x, y])
                    best_action = self.actions[action_idx]
                    
                    # Convert to environment coordinates
                    env_x = (x / 100.0) + offset_x
                    env_y = (y / 100.0) + offset_y
                    
                    # Calculate arrow endpoint
                    arrow_length = 0.05
                    end_x = env_x + best_action[0] * arrow_length
                    end_y = env_y + best_action[1] * arrow_length
                    
                    # Draw policy arrow (white)
                    self.visualisation_lines.append(VisualisationLine(
                        env_x, env_y, end_x, end_y, (255, 255, 255), 0.01
                    ))
        
        # Draw demonstration path if available
        if hasattr(self, 'demo_path') and self.demo_path:
            for state, next_state, _ in self.demo_path:
                # Convert states to environment coordinates
                start_x = state[0] + offset_x
                start_y = state[1] + offset_y
                end_x = next_state[0] + offset_x
                end_y = next_state[1] + offset_y
                
                # Draw demonstration path as green lines
                self.visualisation_lines.append(VisualisationLine(
                    start_x, start_y, end_x, end_y, (0, 255, 0), 0.01
                ))
        # Count actions with non-zero Q-values
        non_zero_q_actions = np.count_nonzero(self.q_values != 0)