Module 7: Reinforcement Learning

from geometry_msgs.msg import Point
import unittest

UP = 0
RIGHT = 1
DOWN = 2
LEFT = 3

HEADINGS = [math.pi/2, 0, -math.pi/2, -math.pi]

class GridConverter:
    def __init__(self, width: float, height: float, cell_size: float):
        self.cols = math.ceil(width / cell_size)
        self.rows = math.ceil(height / cell_size)
        
    def num_states(self) -> int:
        return self.cols * self.rows
	
    def grid_to_state(self, grid_square: tuple[int, int]) -> int:
        col, row = grid_square
        return # YOUR CODE HERE 
    
    def state_to_grid(self, state_num: int) -> tuple[int, int]:
        return # YOUR CODE HERE
        
    def forbidden_actions(self) -> list[tuple[int, int]]:
        forbidden = []
        for i in range(self.cols):
            forbidden.append((self.grid_to_state((i, 0)), DOWN))
            forbidden.append((self.grid_to_state((i, self.rows - 1)), UP))
        for i in range(self.rows):
            forbidden.append((self.grid_to_state((0, i)), LEFT))
            forbidden.append((self.grid_to_state((self.cols - 1, i)), RIGHT))
        return forbidden

class GridConverterTest(unittest.TestCase):
    def test_state_to_square(self):
        conv = GridConverter(1.5, 1.5, 0.5)
        self.assertEqual(conv.cols, 3)
        self.assertEqual(conv.rows, 3)
        expected = [
            (0, 0), (1, 0), (2, 0), 
            (0, 1), (1, 1), (2, 1), 
            (0, 2), (1, 2), (2, 2)
        ]
        for state in range(9):
            grid_square = conv.state_to_grid(state)
            self.assertEqual(expected[state], grid_square)
            self.assertEqual(conv.grid_to_state(grid_square), state)
			

if __name__ == '__main__':
    unittest.main()

    def __init__(self, width: float, height: float, cell_size: float):
        self.cols = int(width / cell_size)
        self.rows = int(height / cell_size)
        self.cell_size = cell_size
        self.max_odom_x = width / 2.0
        self.max_odom_y = height / 2.0

    def odom_to_grid(self, position: Point) -> tuple[int, int]:
        col = # YOUR CODE HERE
        row = # YOUR CODE HERE
        return col, row
    
    def grid_to_odom(self, grid_square: tuple[int, int]) -> Point:
        col, row = grid_square
        p = Point()
        p.x = # YOUR CODE HERE
        p.y = # YOUR CODE HERE
        return p

    def test_from_odom(self):
        conv = GridConverter(10.0, 10.0, 2.0)
        for (odom_x, odom_y, expected_col, expected_row) in [
            (4.5, 2.0, 4, 3), 
            (-3.1, -0.5, 0, 2), 
            (0.0, 0.0, 2, 2), 
            (0.8, -0.8, 2, 2), 
            (1.0, 1.0, 3, 3), 
            (-1.0, -1.0, 2, 2)]:
            p = Point()
            p.x = odom_x
            p.y = odom_y
            col, row = conv.odom_to_grid(p)
            self.assertEqual(expected_col, col)
            self.assertEqual(expected_row, row)

    def test_from_grid(self):
        conv = GridConverter(10.0, 10.0, 2.0)
        for (col, row, expected_odom_x, expected_odom_y) in [
            (4, 3, 4.0, 2.0),
            (0, 2, -4.0, 0.0),
            (2, 2, 0.0, 0.0),
            (3, 3, 2.0, 2.0)
        ]:
            odom = conv.grid_to_odom((col, row))
            self.assertAlmostEqual(odom.x, expected_odom_x, 5)
            self.assertAlmostEqual(odom.y, expected_odom_y, 5)

import random
from typing import Sequence, Tuple


class QParameters:
    def __init__(self, target_visits: int=2, discount: float=0.5, num_states: int=2, num_actions: int=2):
        self.target_visits = target_visits
        self.discount = discount
        self.num_states = num_states
        self.num_actions = num_actions
        self.state_actions = {state: [act for act in range(num_actions)] for state in range(num_states)}

    def forbid_state_actions(self, forbidden_state_actions: Sequence[Tuple[int,int]]):
        for (state, action) in forbidden_state_actions:
            if state in self.state_actions and action in self.state_actions[state]:
                self.state_actions[state].remove(action)


class QTable:
    def __init__(self, params: QParameters):
        self.q = [[0.0] * params.num_actions for i in range(params.num_states)]
        self.visits = [[0] * params.num_actions for i in range(params.num_states)]
        self.target_visits = params.target_visits
        self.discount = params.discount
        self.last_state = 0
        self.last_action = 0
        self.state_actions = params.state_actions
        self.total_updates = 0

    def num_states(self) -> int:
        return len(self.q)
    
    def num_actions(self) -> int:
        return len(self.q[0])

    def sense_act_learn(self, new_state: int, reward: float) -> int:
        assert 0 <= new_state < self.num_states()
        alpha = self.learning_rate(self.last_state, self.last_action)
        update = alpha * (self.discount * self.q[new_state][self.best_action(new_state)] + reward)
        self.q[self.last_state][self.last_action] *= 1.0 - alpha
        self.q[self.last_state][self.last_action] += update

        self.visits[self.last_state][self.last_action] += 1
        if self.is_exploring(new_state):
            new_action = self.least_visited_action(new_state)
        else:
            new_action = self.best_action(new_state)

        self.last_state = new_state
        self.last_action = new_action
        self.total_updates += 1
        return new_action

    def learning_rate(self, state: int, action: int) -> float:
        return 1 / (1 + self.visits[state][action])

    def best_action(self, state: int) -> int:
        best = self.state_actions[state][0]
        for action in self.state_actions[state]:
            if self.q[state][best] < self.q[state][action]:
                best = action
        assert 0 <= best < len(self.q[state])
        return best

    def is_exploring(self, state: int) -> bool:
        below_target = False
        for action in self.state_actions[state]:
            if self.visits[state][action] < self.target_visits:
                below_target = True
        return below_target

    def least_visited_action(self, state: int) -> int:
        least_visited = self.state_actions[state][0]
        for action in self.state_actions[state]:
            if self.visits[state][least_visited] > self.visits[state][action]:
                least_visited = action
        return least_visited

import sys, curses, math, random, copy
import rclpy
from rclpy.node import Node
from rclpy.qos import qos_profile_sensor_data
from geometry_msgs.msg import TwistStamped
from robot_pose import RobotPose, odom2pose, turn_from_to
from nav_msgs.msg import Odometry
from irobot_create_msgs.msg import InterfaceButtons
from grid import GridConverter, UP, DOWN, LEFT, RIGHT, HEADINGS
from fuzzy import fuzzify, defuzzify, f_and, f_or, f_not


ANGLE_TOLERANCE = math.pi/24
DISTANCE_TOLERANCE = 0.05
MAX_FORWARD = 0.25
MAX_TURN = 0.5
MAX_ALIGN = math.pi/8


class GridNode(Node):
    def __init__(self, robot_name: str, width: float, height: float, cell_size: float):
        super().__init__(f"{robot_name}_GridNode")
        self.create_subscription(Odometry, f"{robot_name}/odom", self.odom_callback, qos_profile_sensor_data)
        self.motors = self.create_publisher(TwistStamped, f"{robot_name}/cmd_vel_stamped", qos_profile_sensor_data)
        self.grid = GridConverter(width, height, cell_size) 

        self.current_action = None
        self.pose = None
        self.goal_pose = None

        self.last_key = None

    def grid_col_row(self) -> tuple[int, int]:
        if self.pose is not None:
            return self.grid.odom_to_grid(self.pose.point())

    def state(self):
        col_row = self.grid_col_row()
        if col_row is not None:
            return self.grid.grid_to_state(col_row)

    def odom_callback(self, odom: Odometry):    
        self.pose = odom2pose(odom)
        if self.goal_pose is not None:
            self.head_for_goal()

    def make_twist(self) -> TwistStamped:
        t = TwistStamped()
        t.header.frame_id = "base_link"
        t.header.stamp = self.get_clock().now().to_msg()
        return t

    def update_goal(self, action: int):
        self.current_action = action
        self.goal_pose = copy.copy(self.pose)
        self.goal_pose.theta = HEADINGS[self.current_action]        
        self.goal_pose.move(self.grid.cell_size)

    def head_for_goal(self):
        target_angle = self.pose.turn_to(self.goal_pose.x, self.goal_pose.y)
        distance = self.pose.distance_to(self.goal_pose.x, self.goal_pose.y)
        if distance < DISTANCE_TOLERANCE:
            self.final_heading_align()        
        else:
            self.go_to_goal(target_angle, distance)    

    def final_heading_align(self):
        offset = turn_from_to(self.pose.theta, self.goal_pose.theta)
        if abs(offset) < ANGLE_TOLERANCE:
            self.goal_pose = None
        else:
            t = self.make_twist()
            offset_left, offset_right = fuzzy_angle_left_right(offset)
            t.twist.angular.z = defuzzy_angle(offset_left, offset_right)
            self.motors.publish(t)

    def go_to_goal(self, target_angle: float, distance: float):
        fuzzy_left, fuzzy_right = fuzzy_angle_left_right(target_angle)
        fuzzy_distance = f_and(fuzzify(distance, 0.0, self.grid.cell_size), 
                               f_not(f_or(fuzzy_left, fuzzy_right)))
        t = self.make_twist()
        t.twist.angular.z = defuzzy_angle(fuzzy_left, fuzzy_right)
        t.twist.linear.x = defuzzify(fuzzy_distance, 0.0, MAX_FORWARD)
        self.motors.publish(t)

    def process_keystroke(self, k: str):
        self.last_key = k
        if k.isdigit():
            self.update_goal(int(k))

    def running(self) -> bool:
        return self.last_key != 'q'

    def display(self, stdscr):
        stdscr.addstr(0, 0, "0: up, 1: right, 2: down, 3: left")
        if self.last_key is not None:
            stdscr.addstr(1, 0, self.last_key)
        display_padded(stdscr, 2, f"pose     : {self.pose}")
        display_padded(stdscr, 3, f"goal_pose: {self.goal_pose}")
        display_padded(stdscr, 4, f"action   : {self.current_action}")
        if self.pose is not None:
            display_padded(stdscr, 5, f"state    : {self.state()}; ({self.grid_col_row()})")


def display_padded(stdscr, row: int, s: str):
    height, width = stdscr.getmaxyx()
    stdscr.addstr(row, 0, f"{s}{' ' * (width - len(s) - 1)}")


def fuzzy_angle_left_right(angle: float) -> tuple[float, float]:
    fuzzy_left = fuzzify(angle, ANGLE_TOLERANCE, MAX_ALIGN)
    fuzzy_right = fuzzify(-angle, ANGLE_TOLERANCE, MAX_ALIGN)
    return fuzzy_left, fuzzy_right
    

def defuzzy_angle(fuzzy_left: float, fuzzy_right: float) -> float:
    if fuzzy_left > fuzzy_right:
        return defuzzify(fuzzy_left, 0.0, MAX_TURN)
    else:
        return defuzzify(fuzzy_right, 0.0, -MAX_TURN)


def main(stdscr):
    rclpy.init()
    node = GridNode(sys.argv[1], float(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]))
    curses.cbreak()
    stdscr.nodelay(True)
    stdscr.clear()
    while node.running():
        try:
            k = stdscr.getch()
            if k != -1:
                node.process_keystroke(chr(k))
            node.display(stdscr)
            rclpy.spin_once(node, timeout_sec=0.0)
        except curses.error as e:
            if str(e) != 'no input':
                stdscr.addstr(0, 0, traceback.format_exc())

    rclpy.shutdown()
    node.destroy_node()
    curses.nocbreak()
    curses.echo()
    stdscr.refresh()


if __name__ == '__main__':
    if len(sys.argv) < 5:
        print("Usage: python3 grid_node.py robot_name width height cell_size")
    else:
        curses.wrapper(main)

import sys, curses, math, random, copy
import rclpy
from rclpy.qos import qos_profile_sensor_data
from robot_pose import RobotPose, odom2pose, turn_from_to
from irobot_create_msgs.msg import InterfaceButtons
from qlearning import QParameters, QTable
from grid import HEADINGS
import grid_node


class QNode(grid_node.GridNode):
    def __init__(self, robot_name: str, target_visits: int, discount: float):
        super().__init__(robot_name, 1.2, 1.2, 0.4)
        self.create_subscription(InterfaceButtons, f"{robot_name}/interface_buttons", self.button_callback, qos_profile_sensor_data)

        params = QParameters(target_visits, discount, num_states=self.grid.num_states(), num_actions=len(HEADINGS))
        params.forbid_state_actions(self.grid.forbidden_actions())
        self.qtable = QTable(params)
        self.started = False

    def button_callback(self, buttons: InterfaceButtons):
        if self.started:
            if self.goal_pose is None:
                reward = None
                if buttons.button_1.is_pressed:
                    reward = -1.0
                elif buttons.button_2.is_pressed:
                    reward = 0.25
                elif buttons.button_power.is_pressed:
                    reward = 0.0
                if reward is not None:
                    action = self.qtable.sense_act_learn(self.state(), reward)
                    self.update_goal(action)
        elif self.pose is not None and (buttons.button_1.is_pressed or buttons.button_2.is_pressed or buttons.button_power.is_pressed):
            self.started = True
            self.update_goal(random.randint(0, len(HEADINGS) - 1)) 

    def display(self, stdscr):
        super().display(stdscr)
        for state in range(self.qtable.num_states()):
            col, row = self.grid.state_to_grid(state)    
            values = [round(v, 2) for v in self.qtable.q[state]]
            stdscr.addstr(6 + self.grid.rows + state, 0, f"state {state} ({col}, {row}): qs: {values} visits:  {self.qtable.visits[state]}")


def main(stdscr):
    rclpy.init()
    node = QNode(sys.argv[1], int(sys.argv[2]), float(sys.argv[3]))
    curses.cbreak()
    stdscr.nodelay(True)
    stdscr.clear()
    while node.running():
        try:
            k = stdscr.getch()
            if k != -1:
                node.process_keystroke(chr(k))
            node.display(stdscr)
            rclpy.spin_once(node, timeout_sec=0.0)
        except curses.error as e:
            if str(e) != 'no input':
                stdscr.addstr(0, 0, traceback.format_exc())

    rclpy.shutdown()
    node.destroy_node()
    curses.nocbreak()
    curses.echo()
    stdscr.refresh()


if __name__ == '__main__':
    if len(sys.argv) < 4:
        print("Usage: python3 qnode.py robot_name target_visits discount")
    else:
        curses.wrapper(main)