runner.py

from utils.socket_utils import send_array, recv_array
import numpy as np
import zmq
from robots.tiago import Tiago
import math
from utils.llm_utils import get_env_variable
from LMs.gpt import GPTAssistant, GPT
from LMs.gemini import GeminiPro
from pddls.problem import generate_pddl_problem
import cv2
import base64
import json
from pddls.predicates import Predicators


class Runner:
    def __init__(self):
        # source spatial target
        self.source = None
        self.spatial = None
        self.target = None
        self.holding = None
        self.at = "floor"
        self.goal = None
        self.preds = None

    def _init_predicate(self, robot, instruction, random=False):
        self.predicators = Predicators(robot, self.gpt, instruction, random)
        # self.predicators._init_server(
        #     self.nav_socket, self.manip_socket, self.theta_offset
        # )

    def init_assistant(self):
        azure_openai_key = get_env_variable("AZURE_OPENAI_KEY")
        azure_endpoint = get_env_variable("AZURE_OPENAI_ENDPOINT")
        self.gpt = GPTAssistant(
            azure_openai_key=azure_openai_key,
            endpoint=azure_endpoint,
            model_name="gpt-4o",
        )
        with open("prompts/assistant/system.txt", "r") as f:
            system_msg = f.read()
        self.gpt.create_prompt(system_message=system_msg)

    def init_real(self, x1, y1, x2, y2, nav_ip, nav_port, manip_ip, manip_port):
        """
        Take coordinates of two tapes: the robot stands on tape (x1, y1) and look at tape (x2, y2)
        Compute two rotation matrices r2n_matrix and n2r_matrix that can transform coordinates
        from robot hector slam system coordinate system to record3d coordinate system and
        from record3d coordinate system to robot hector slam system coordinate system respectively
        """
        x_offset = x1
        y_offset = y1
        self.theta_offset = np.arctan2((y2 - y1), (x2 - x1))
        print(f"offsets: {x_offset}, {y_offset}, {self.theta_offset}")
        self.r2n_matrix = np.array(
            [[1, 0, x_offset], [0, 1, y_offset], [0, 0, 1]]
        ) @ np.array(
            [
                [np.cos(self.theta_offset), -np.sin(self.theta_offset), 0],
                [np.sin(self.theta_offset), np.cos(self.theta_offset), 0],
                [0, 0, 1],
            ]
        )
        self.n2r_matrix = np.array(
            [
                [np.cos(self.theta_offset), np.sin(self.theta_offset), 0],
                [-np.sin(self.theta_offset), np.cos(self.theta_offset), 0],
                [0, 0, 1],
            ]
        ) @ np.array([[1, 0, -x_offset], [0, 1, -y_offset], [0, 0, 1]])
        # init socket
        context = zmq.Context()
        self.nav_socket = context.socket(zmq.REQ)
        self.nav_socket.connect("tcp://" + nav_ip + ":" + str(nav_port))
        self.manip_socket = context.socket(zmq.REQ)
        self.manip_socket.connect("tcp://" + manip_ip + ":" + str(manip_port))

    def instruct_augment(self, instruction, robot, domain="room", fb=False):
        fb_pred = []
        # get A, B
        with open("prompts/assistant/find_object_and_related.txt", "r") as f:
            query = f.read()
        query = query.replace("[INSTRUCTION]", instruction)
        r = self.gpt.generate(message=query, clear_conversation=True)
        print(r)
        # r is json format string, to dict
        json_r = json.loads(r)
        A = json_r["object_name"]
        B = json_r["related_object_name"]
        others = json_r["other_object_names"]
        self.source = A
        self.target = B

        # generate pddl problem
        name = instruction
        domain = domain
        objects = ["floor", "tiago - robot", A, B] + others
        goals = [""]
        initial_states = []
        # add current object status
        relation = self.predicators.relation(A, B)
        at_ = self.predicators.at(["floor", A, B] + others)
        initial_states.append(at_)

        holding_ = self.predicators.holding(A, None)
        initial_states.append(holding_)

        if holding_[1:].split(" ")[0] == "not":
            initial_states.append(relation)
            find_A = self.predicators.find(A)
            find_B = self.predicators.find(B)
            if fb:
                fb_pred.append(find_A)
                fb_pred.append(find_B)
            else:
                initial_states.append(find_A)
                initial_states.append(find_B)
        else:
            for obj_name in others:
                find_others = self.predicators.find(obj_name)
                if fb:
                    fb_pred.append(find_others)
                else:
                    initial_states.append(find_others)

        if self.preds is None:
            with open("prompts/assistant/generate_init.txt", "r") as f:
                query = f.read()
            query = query.replace("[INSTRUCTION]", instruction)
            query = query.replace("[OBJECTS]", ", ".join(objects))
            r = self.gpt.generate(message=query, clear_conversation=True)
            print(r)
            self.preds = r
        # r is lines of conditions, split by \n
        for line in self.preds.split("\n"):
            pred = line[1:-1].split(" ")[0]
            if pred in ["at", "on", "in", "find", "holding"]:
                continue
            elif pred == "detected":
                obj = line[1:-1].split(" ")[1]
                detected_ = self.predicators.detected(obj, None)
                if fb:
                    fb_pred.append(detected_)
                else:
                    initial_states.append(detected_)
            elif pred == "blocked":
                obj = line[1:-1].split(" ")[1]
                blocked_ = self.predicators.blocked(obj, None)
                if fb:
                    fb_pred.append(blocked_)
                else:
                    initial_states.append(blocked_)
            elif pred == "opened":
                obj = line[1:-1].split(" ")[1]
                opened_ = self.predicators.opened(obj, None)
                initial_states.append(opened_)
            elif pred == "graspable":
                obj = line[1:-1].split(" ")[1]
                graspable_ = self.predicators.graspable(obj, None)
                if fb:
                    fb_pred.append(graspable_)
                else:
                    initial_states.append(graspable_)
            elif pred == "reachable":
                obj = line[1:-1].split(" ")[1]
                reachable_ = self.predicators.reachable(obj, None)
                if fb:
                    fb_pred.append(reachable_)
                else:
                    initial_states.append(reachable_)
            elif pred == "placeable":
                if len(line[1:-1].split(" ")) == 2:
                    obj1 = line[1:-1].split(" ")[1]
                    placeable_ = self.predicators.placeable(obj1, None)
                    if fb:
                        fb_pred.append(placeable_)
                    else:
                        initial_states.append(placeable_)

        # blocked_ = self.predicators.blocked(A, None)
        # initial_states.append(blocked_)

        # opened_ = self.predicators.opened(A, None)
        # initial_states.append(opened_)

        # closed_ = self.predicators.closed(A, None)
        # initial_states.append(closed_)

        # detected_ = self.predicators.detected(A, None)
        # initial_states.append(detected_)

        # graspable_ = self.predicators.graspable(A, None)
        # initial_states.append(graspable_)

        # reachable_ = self.predicators.reachable(A, None)
        # initial_states.append(reachable_)

        # placeable_ = self.predicators.placeable(A, None)
        # initial_states.append(placeable_)

        # image, depth, points = robot.camera.capture_image()
        # fx, fy, cx, cy = (
        #     robot.camera.fx,
        #     robot.camera.fy,
        #     robot.camera.cx,
        #     robot.camera.cy,
        # )
        # image, depth, points = [1], [1], [1]
        # fx, fy, cx, cy, head_tilt, head_pane = (0, 0, 0, 0, 0, 0)
        # is_detected_A = self.check_status_with_server(
        #     robot,
        #     "is_detected",
        #     image,
        #     depth,
        #     [fx, fy, cx, cy, head_tilt, head_pane],
        #     A,
        # )
        # detected_ = self.predicators.detected(A, None)
        # initial_states.append(detected_)
        # # if A is not detected, scan and detect it
        # if is_detected_A is None:
        #     if self.holding is not None:
        #         for obj_name in others:
        #             is_detected_obj = self.check_status_with_server(
        #                 robot,
        #                 "is_detected",
        #                 obj_name,
        #                 image,
        #                 depth,
        #                 [fx, fy, cx, cy],
        #                 obj_name,
        #             )
        #             if is_detected_obj is None:
        #                 initial_states.append(f"(not (is_detected {obj_name}))")
        #             else:
        #                 initial_states.append(f"(is_detected {obj_name})")
        #                 for obj_name in others:
        #                     is_graspable_obj = self.check_status_with_server(
        #                         robot,
        #                         "is_placeable",
        #                         obj_name,
        #                         image,
        #                         depth,
        #                         [fx, fy, cx, cy],
        #                         obj_name,
        #                     )
        #                     if is_graspable_obj is None:
        #                         initial_states.append(
        #                             f"(not (is_graspable {obj_name}))"
        #                         )
        #                     else:
        #                         initial_states.append(f"(is_graspable {obj_name})")
        # else:
        #     initial_states.append(f"(is_detected {A})")
        #     is_graspable_A = self.check_status_with_server(
        #         robot, "is_graspable", A, image, depth, [fx, fy, cx, cy], A
        #     )
        #     if is_graspable_A is None:
        #         initial_states.append(f"(not (is_graspable {A}))")
        #     else:
        #         initial_states.append(f"(is_graspable {A})")
        #         self.holding = A

        # goal
        if self.goal is None:
            with open("prompts/assistant/generate_goal_condition.txt", "r") as f:
                query = f.read()
            query = query.replace("[INSTRUCTION]", instruction)
            r = self.gpt.generate(message=query, clear_conversation=True)
            self.goal = r
        goals[0] = self.goal

        pddl_problem = generate_pddl_problem(
            name, domain, objects, initial_states, goals
        )
        return pddl_problem, fb_pred

    def navigate(self, robot, A, B):
        # start_xy = np.array([0, 0, -0.816174])
        start_xy = robot._get_base_pose()
        print(start_xy)
        transformed_start_xy = self.r2n_matrix @ np.array([start_xy[0], start_xy[1], 1])
        start_xy[0], start_xy[1] = transformed_start_xy[0], transformed_start_xy[1]
        start_xy[2] += self.theta_offset
        print(start_xy)

        # Send start_xy, A, B to the workstation and receive planned paths
        send_array(self.nav_socket, start_xy)
        print(self.nav_socket.recv_string())
        self.nav_socket.send_string(A)
        print(self.nav_socket.recv_string())
        self.nav_socket.send_string(B)
        print(self.nav_socket.recv_string())
        self.nav_socket.send_string("Waiting for path")
        paths = recv_array(self.nav_socket)
        print(paths)
        self.nav_socket.send_string("Path received")
        end_xyz = recv_array(self.nav_socket)
        z = end_xyz[2]
        end_xyz = self.n2r_matrix @ np.array([end_xyz[0], end_xyz[1], 1])
        end_xyz[2] = z

        final_paths = []
        for path in paths:
            transformed_path = self.n2r_matrix @ np.array([path[0], path[1], 1])
            transformed_path[2] = path[2] - self.theta_offset
            print(transformed_path)
            final_paths.append(transformed_path)
            robot.move(transformed_path)

    def navigate_by_point(self, robot, x, y, theta):
        path = np.array([x, y, theta])
        robot.move(path)

    def manipulate(self, robot: Tiago, target: str, mode="pick", grasp_axis="y"):
        retry = 1
        retry_flag = True
        while retry <= 2:
            # get head tilt angle
            # head tilt is joint value, ranged from [-1.53, 0.79] in stretch robot
            # TODO: do we need to convert it to tiago?
            head_pan, head_tilt = robot.head_joint_values
            # capture image from current camera
            image, depth, points = robot.camera.capture_image()
            ## communicate with server
            # send RGB, depth and camera intrinsics data
            print("Sending image and depth to workstation")
            send_array(self.manip_socket, image)
            print(self.manip_socket.recv_string())
            send_array(self.manip_socket, depth)
            print(self.manip_socket.recv_string())
            send_array(
                self.manip_socket,
                np.array(
                    [
                        robot.camera.fy,
                        robot.camera.fx,
                        robot.camera.cy,
                        robot.camera.cx,
                        head_tilt,
                        head_pan,
                    ]
                ),
            )
            print(self.manip_socket.recv_string())
            # Sending Object text and Manipulation mode
            self.manip_socket.send_string(target)
            print(self.manip_socket.recv_string())
            self.manip_socket.send_string(mode)
            print(self.manip_socket.recv_string())
            # Waiting for the base and camera transforms to center the object vertically and horizontallsockety
            # get translation and rotation based on the camera frame
            self.manip_socket.send_string(
                "Waiting for gripper pose/ base and head trans from workstation"
            )
            translation = recv_array(self.manip_socket)
            self.manip_socket.send_string("translation received by robot")
            rotation = recv_array(self.manip_socket)
            self.manip_socket.send_string("rotation received by robot")
            add_data = recv_array(self.manip_socket)
            self.manip_socket.send_string("Additional data received by robot")
            # additional data
            depth = add_data[0]
            cropped = add_data[1]
            retry_flag = add_data[2]
            print(f"Additional data received - {add_data}")
            print("translation: ")
            print(translation)
            print("rotation: ")
            print(rotation)
            print(self.manip_socket.recv_string())

            if retry == 1:
                robot._center_object(translation, rotation)
                retry += 1
                continue

            ## grasp the object
            if mode == "pick":
                retry += 1
                robot.grasp(translation.tolist(), rotation.tolist(), axis=grasp_axis)

            if mode == "place":
                retry += 1
                robot.place(translation.tolist(), rotation.tolist())


if __name__ == "__main__":
    instruction = "take the pill box in the drawer to me"
    # instruction = "take the paper box on the wooden table and place it on the black table"
    # instruction = "take the blue jacket on the cloth hanger to me"
    # instruction = "take the book on the black table and insert it into the bookshelf"
    # instruction = "lift the bucket on the floor and place it on the white table"

    azure_openai_key = get_env_variable("AZURE_OPENAI_KEY")
    azure_endpoint = get_env_variable("AZURE_OPENAI_ENDPOINT")
    brain = GPT(
        azure_openai_key=azure_openai_key,
        endpoint=azure_endpoint,
        model_name="gpt-4o",
    )
    brain.create_prompt()

    runner = Runner()
    runner.init_assistant()
    runner._init_predicate(None, instruction)

    while True:
        state = runner.instruct_augment(instruction, None)
        actions, scores = brain.generate(instruction, state, "")

        # save state and action to txt file
        with open("pddls/take_pill/wo_prec/1.txt", "a") as f:
            f.write(f"{state}\n")
            f.write(f"{actions}\n")
            f.write(f"{scores}\n")