├── videos └── .gitkeep ├── data ├── ours │ ├── DBs │ │ └── .gitkeep │ ├── images │ │ └── .gitkeep │ ├── llm_responses │ │ └── .gitkeep │ └── old_images │ │ └── .gitkeep └── ours_objective │ ├── DBs │ └── .gitkeep │ ├── images │ └── .gitkeep │ └── llm_responses │ └── .gitkeep ├── images └── seif_avatar.jpeg ├── requirements.txt ├── README.md ├── db.py ├── LICENSE ├── config └── config.py ├── .gitignore ├── core.py ├── llm.py ├── gui.py ├── robot.py ├── controller.py ├── .zshrc ├── simulation_local.py ├── prompts └── prompts.py └── simulation_http.py /videos/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours/DBs/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours/images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours/llm_responses/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours/old_images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours_objective/DBs/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours_objective/images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/ours_objective/llm_responses/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /images/seif_avatar.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vateseif/narrate/HEAD/images/seif_avatar.jpeg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | tqdm 2 | numpy 3 | matplotlib 4 | gymnasium 5 | ipympl 6 | pydantic==2.6.0 7 | do-mpc==4.6.4 8 | onnx 9 | openai==1.11.1 10 | aiohttp==3.9.3 11 | asyncua==1.0.6 12 | langchain 13 | ipywidgets 14 | streamlit 15 | tiktoken 16 | opencv-python==4.9.0.80 17 | python-docx 18 | langchain-openai 19 | git+https://github.com/vateseif/Safe-panda-gym.git@safe_rl-panda-gym -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # NARRATE: Versatile Language Architeture for Optimal Control in Robotics 2 | This repo contains a reference implementation of the paper [NARRATE](https://narrate-mpc.github.io). 3 | We refer the reader to our project page to find the paper and details about the method. 4 | 5 | https://github.com/vateseif/l2o/assets/45405956/7a685788-5d3e-4f34-b1e0-21bc76f2ce6a 6 | 7 | ## Setup 8 | ### Env 9 | Create a python environment (i.e. with conda): 10 | ~~~ 11 | conda create --name narrate python=3.9 12 | conda activate narrate 13 | ~~~ 14 | ### Requirements 15 | Install requirements 16 | ~~~ 17 | pip install -r requirements.txt 18 | ~~~ 19 | ### OpenAI key 20 | You need to create the file `keys/gpt4.key` and put your OpenAI key. Make sure to have acces to GPT4. 21 | 22 | ## Run 23 | You will need to run 2 files in order to interact with the simulation environment. 24 | 25 | To start the chat interface you have to execute in your terminal: 26 | ~~~ 27 | streamlit run gui.py 28 | ~~~ 29 | 30 | To start the simulation you have to execute in your terminal 31 | ~~~ 32 | python simulation_http.py 33 | ~~~ 34 | -------------------------------------------------------------------------------- /db.py: -------------------------------------------------------------------------------- 1 | from config.config import DBConfig 2 | from sqlalchemy.orm import relationship 3 | from sqlalchemy.ext.declarative import declarative_base 4 | from sqlalchemy import Column, Integer, String, ForeignKey, JSON 5 | 6 | Base = declarative_base() 7 | 8 | 9 | class Episode(Base): 10 | __tablename__ = 'episodes' 11 | 12 | id = Column(Integer, primary_key=True) 13 | name = Column(String) # Optional, if you want to name or otherwise identify episodes 14 | state_trajectories = Column(JSON) # Store state trajectories as JSON 15 | mpc_solve_times = Column(JSON) # Store MPC solve times as JSON 16 | #n_collisions = Column(Integer) 17 | epochs = relationship("Epoch", backref="episode") 18 | 19 | class Epoch(Base): 20 | __tablename__ = 'epochs' 21 | 22 | id = Column(Integer, primary_key=True) 23 | episode_id = Column(Integer, ForeignKey('episodes.id')) 24 | time_step = Column(Integer, nullable=False) 25 | role = Column(String, nullable=False) 26 | content = Column(String, nullable=False) 27 | image = Column(String, nullable=False) # Store images as binary data 28 | 29 | 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Seif Ismail 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /config/config.py: -------------------------------------------------------------------------------- 1 | from prompts.prompts import PROMPTS 2 | from core import AbstractControllerConfig, AbstractLLMConfig, AbstractRobotConfig, AbstractSimulaitonConfig 3 | from typing import List 4 | 5 | 6 | class SimulationConfig(AbstractSimulaitonConfig): 7 | render: bool = True 8 | debug: bool = False 9 | logging: bool = False 10 | logging_video: bool = False 11 | logging_video_fps: int = 5 12 | logging_video_frequency: int = 10 13 | task: str = "Cubes" # [Cubes, CleanPlate, Sponge, CookSteak] 14 | save_video: bool = False 15 | fps: int = 20 # only used if save_video = True 16 | dt: float = 0.05 # simulation timestep. Must be equal to that of controller 17 | frame_width: int = 1024 18 | frame_height: int = 1024 19 | frame_target_position: List[float] = [0.0, -0.1, 0.] 20 | frame_distance: float = 1.6 21 | frame_yaw: int = -125 22 | frame_pitch: int = -30 23 | method:str = 'ours' 24 | 25 | 26 | class LLMConfig(AbstractLLMConfig): 27 | def __init__(self, avatar:str, task:str=None) -> None: 28 | self.avatar: str = avatar 29 | self.mock_task = None 30 | self.prompt: str = PROMPTS[avatar][task] 31 | model_name: str = "gpt-4-0125-preview" 32 | streaming: bool = False 33 | temperature: float = 0.9 34 | max_tokens: int = 1000 35 | 36 | 37 | class ControllerConfig(AbstractControllerConfig): 38 | def __init__(self, task:str=None) -> None: 39 | self.task: str = task 40 | nx: int = 3 41 | nu: int = 3 42 | T: int = 15 43 | dt: float = 0.05 44 | lu: float = -0.2 # lower bound on u 45 | hu: float = 0.2 # higher bound on u 46 | model_type: str = "discrete" 47 | penalty_term_cons: float = 1e7 48 | 49 | 50 | class RobotConfig(AbstractRobotConfig): 51 | def __init__(self, task:str=None) -> None: 52 | self.task: str = task 53 | open_gripper_time: int = 28 54 | method: str = "optimization" # ['optimization', 'objective'] 55 | COST_THRESHOLD: float = 1e-5 56 | COST_DIIFF_THRESHOLD: float = 1e-7 57 | GRIPPER_WIDTH_THRESHOLD: float = 4e-6 58 | TIME_THRESHOLD: float = 25 59 | MAX_OD_ATTEMPTS: int = 2 60 | 61 | 62 | class DBConfig: 63 | db_name: str = "data/DBs/cubes_objective.db" 64 | 65 | 66 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | 132 | # my files 133 | zzdlmpc.ipynb 134 | 135 | # trained model weights 136 | models/ 137 | results/ 138 | 139 | # tensorboard 140 | board/ 141 | 142 | # videos 143 | *.mp4 144 | *.webm 145 | 146 | # images 147 | *.png 148 | 149 | deletable.* 150 | 151 | keys/ 152 | *.key 153 | 154 | # Mac stuff 155 | *.DS_Store 156 | 157 | *.docx 158 | 159 | *.db 160 | 161 | *.json 162 | 163 | cache/* 164 | -------------------------------------------------------------------------------- /core.py: -------------------------------------------------------------------------------- 1 | import os 2 | import inspect 3 | from typing import List 4 | from abc import abstractmethod 5 | 6 | import gym 7 | import panda_gym 8 | import numpy as np 9 | 10 | BASE_DIR = os.path.dirname(__file__) 11 | # GPT4 api key. 12 | os.environ["OPENAI_API_KEY"] = open(BASE_DIR + '/keys/gpt4.key', 'r').readline().rstrip() 13 | 14 | 15 | class AbstractLLMConfig: 16 | prompt: str 17 | parsing: str 18 | model_name: str 19 | streaming: bool 20 | temperature: float 21 | 22 | class AbstractControllerConfig: 23 | T: int 24 | nx: int 25 | nu: int 26 | dt: float 27 | lu: float # lower bound on u 28 | hu: float # higher bound on u 29 | 30 | class AbstractRobotConfig: 31 | name: str 32 | controller_type: str 33 | 34 | class AbstractSimulaitonConfig: 35 | render: bool 36 | env_name: str 37 | task: str 38 | save_video: bool 39 | 40 | class ObjBase: 41 | ''' 42 | The object base that defines debugging tools 43 | ''' 44 | def initialize (self, **kwargs): 45 | pass 46 | 47 | def sanityCheck (self): 48 | # check the system parameters are coherent 49 | return True 50 | 51 | def errorMessage (self,msg): 52 | print(self.__class__.__name__+'-'+inspect.stack()[1][3]+': [ERROR] '+msg+'\n') 53 | return False 54 | 55 | def warningMessage (self,msg): 56 | print(self.__class__.__name__+'-'+inspect.stack()[1][3]+': [WARNING] '+msg+'\n') 57 | return False 58 | 59 | 60 | 61 | class AbstractController(ObjBase): 62 | 63 | def __init__(self, cfg: AbstractControllerConfig) -> None: 64 | self.cfg = cfg 65 | 66 | @abstractmethod 67 | def reset(self, x0:np.ndarray) -> None: 68 | return 69 | 70 | @abstractmethod 71 | def apply_gpt_message(self, gpt_message:str) -> None: 72 | return 73 | 74 | @abstractmethod 75 | def step(self) -> np.ndarray: 76 | return 77 | 78 | 79 | class AbstractLLM(ObjBase): 80 | 81 | def __init__(self, cfg:AbstractLLMConfig) -> None: 82 | self.cfg = cfg 83 | 84 | @abstractmethod 85 | def run(self): 86 | return 87 | 88 | 89 | 90 | class AbstractRobot(ObjBase): 91 | 92 | def __init__(self, cfg:AbstractRobotConfig) -> None: 93 | self.cfg = cfg 94 | 95 | # components 96 | self.TP: AbstractLLM # Task planner 97 | self.OD: AbstractLLM # Optimization Designer 98 | self.MPC: AbstractController # Controller 99 | 100 | 101 | @abstractmethod 102 | def reset_gpt(self): 103 | return 104 | 105 | class AbstractSimulation(ObjBase): 106 | def __init__(self, cfg: AbstractSimulaitonConfig) -> None: 107 | self.cfg = cfg 108 | # init env 109 | self.env = gym.make(f"Panda{cfg.env_name}-v2", render=cfg.render) 110 | # init robots 111 | self.robot 112 | # count number of tasks solved from a plan 113 | self.task_counter = 0 114 | 115 | def reset(self): 116 | """ Reset environment """ 117 | pass 118 | 119 | def create_plan(self): 120 | """ Triggers the Task Planner to generate a plan of subtasks""" 121 | pass 122 | 123 | def next_task(self): 124 | """ Tasks the Optimization Designer to carry out the next task in the plam""" 125 | pass 126 | 127 | def _solve_task(self): 128 | """ Applies the optimization designed by the Optimization Designer""" 129 | pass 130 | 131 | def _run(self): 132 | """ Start the simulation """ 133 | pass 134 | 135 | def run(self): 136 | """ Executes self._run() in a separate thread""" 137 | pass -------------------------------------------------------------------------------- /llm.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | from core import AbstractLLM, AbstractLLMConfig 3 | 4 | import os 5 | import json 6 | import requests 7 | import tiktoken 8 | from streamlit import empty, session_state 9 | from pydantic import BaseModel, Field 10 | from langchain_openai import ChatOpenAI 11 | from langchain.schema import HumanMessage, AIMessage 12 | from langchain.prompts.chat import SystemMessagePromptTemplate 13 | from langchain.output_parsers import PydanticOutputParser 14 | from langchain.callbacks.base import BaseCallbackHandler 15 | 16 | 17 | TOKEN_ENCODER = tiktoken.encoding_for_model("gpt-4") 18 | 19 | class Message: 20 | def __init__(self, text, base64_image=None, role="user"): 21 | self.role = role 22 | self.text = text 23 | self.base64_image = base64_image 24 | 25 | def to_dict(self): 26 | message = [{"type": "text", "text": self.text}] 27 | if self.base64_image: 28 | message.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{self.base64_image}", "detail": "high"}}) 29 | return {"role": self.role, "content": message} 30 | 31 | class StreamHandler(BaseCallbackHandler): 32 | 33 | def __init__(self, avatar:str, parser: PydanticOutputParser) -> None: 34 | super().__init__() 35 | self.avatar = avatar 36 | self.parser = parser 37 | 38 | def on_llm_start(self, serialized, prompts, **kwargs) -> None: 39 | """Run when LLM starts running.""" 40 | self.text = "" 41 | self.container = empty() 42 | 43 | def on_llm_new_token(self, token: str, *, chunk, run_id, parent_run_id=None, **kwargs): 44 | super().on_llm_new_token(token, chunk=chunk, run_id=run_id, parent_run_id=parent_run_id, **kwargs) 45 | self.text += token 46 | self.container.write(self.text + "|") 47 | 48 | def on_llm_end(self, response, **kwargs): 49 | pretty_text = self.parser.parse(self.text).pretty_print() 50 | self.container.markdown(pretty_text, unsafe_allow_html=False) 51 | session_state.messages.append({"type": self.avatar, "content": pretty_text}) 52 | 53 | class LLM(AbstractLLM): 54 | 55 | def __init__(self, cfg: AbstractLLMConfig) -> None: 56 | super().__init__(cfg) 57 | 58 | # init messages 59 | self.messages = [Message(text=self.cfg.prompt, role="system")] 60 | # request headers 61 | self.headers = { 62 | "Content-Type": "application/json", 63 | "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}" 64 | } 65 | 66 | def reset(self): 67 | self.messages = [Message(text=self.cfg.prompt, role="system")] 68 | 69 | def run(self, user_message:str, base64_image=None, short_history=False) -> dict: 70 | # add user message to chat history 71 | self.messages.append(Message(text=user_message, role="user", base64_image=base64_image)) 72 | # select the last 2 user messages and the last assistant message 73 | selected_messages = [self.messages[0]] + [m for m in self.messages[-1:] if m.role!="system"] if short_history else self.messages 74 | # send request to OpenAI API 75 | payload = { 76 | "model": self.cfg.model_name, 77 | "messages": [m.to_dict() for m in selected_messages], 78 | "max_tokens": self.cfg.max_tokens, 79 | "response_format": {"type": "json_object"} 80 | } 81 | #print([m.text for m in selected_messages]) 82 | t0 = time() 83 | response = requests.post("https://api.openai.com/v1/chat/completions", headers=self.headers, json=payload).json() 84 | solve_time = time() - t0 85 | # retrieve text response 86 | try: 87 | AI_response = response['choices'][0]['message']['content'] 88 | self.messages.append(Message(text=AI_response, role="assistant")) 89 | AI_response = json.loads(AI_response) 90 | except Exception as e: 91 | print(f"Error: {e}") 92 | AI_response = {"instruction": response['error']['message']} 93 | 94 | AI_response["solve_time"] = solve_time 95 | return AI_response 96 | 97 | 98 | 99 | 100 | -------------------------------------------------------------------------------- /gui.py: -------------------------------------------------------------------------------- 1 | import json 2 | import requests 3 | import streamlit as st 4 | from time import sleep 5 | 6 | # Init streamlit page 7 | st.title("Language to Optimization") 8 | 9 | # Create sidebar 10 | st.sidebar.title("Choose model") 11 | 12 | # Add a sidebar radio button to select the message type 13 | model = st.sidebar.radio("Select the model to talk to", ["Task Planner", "Optimization Designer"]) 14 | 15 | # init the avatars for the message icons 16 | avatars = {"human":"images/seif_avatar.jpeg", "OD":"images/wall-e.png", "TP":"images/eve.png"} 17 | 18 | # base url to reach sim 19 | base_url = 'http://localhost:8080/' 20 | 21 | # init robot simulation 22 | if "stage" not in st.session_state: 23 | # init state machine state: 24 | # 0 = There's no plan the OD can execute 25 | # 1 = There is a plan the TP can execute. A button pops up to allow the user to execute the plan 26 | # 2 = Trigger the execution of the plan 27 | st.session_state.stage = 0 28 | # init state machine state: 29 | # 0 = You can press start to start recording frames 30 | # 1 = You can press stop to save the recording or cacel the recording 31 | # 2 = Saves the recording and stops saving frames 32 | # 3 = cancels the recording and stops saving frames 33 | st.session_state.recording = 0 34 | # init number of tasks solved from a plan 35 | st.session_state.task = None 36 | 37 | def set_state(i): 38 | # Function to update the state machine stage 39 | st.session_state.stage = i 40 | 41 | def set_recording_state(i): 42 | # Function to update the recording state machine stage 43 | st.session_state.recording = i 44 | 45 | def append_message(message:dict): 46 | # Function to append a message to the chat history 47 | message_type = message["type"] 48 | if message_type == "image": 49 | with st.chat_message("human", avatar=avatars["human"]): 50 | st.image(message["content"], width=400, caption="Current scene") 51 | else: 52 | with st.chat_message(message_type, avatar=avatars[message_type]): 53 | st.markdown(message["content"]) 54 | 55 | # Initialize chat history 56 | if "messages" not in st.session_state: 57 | st.session_state.messages = [] 58 | 59 | # Display chat messages from history on app rerun 60 | for message in st.session_state.messages: 61 | append_message(message) 62 | 63 | # Accept user input 64 | if prompt := st.chat_input("What should the robot do?"): 65 | # Add user message to chat history 66 | st.session_state.messages.append({"type": "human", "content":prompt}) 67 | # Display user message in chat message container 68 | with st.chat_message("human", avatar=avatars["human"]): 69 | st.markdown(prompt) 70 | 71 | # Display assistant response in chat message container 72 | if model == "Task Planner": 73 | response = requests.post(base_url+'make_plan', json={"content": prompt}).json() 74 | st.session_state.messages += response 75 | for m in response: append_message(m) 76 | #st.session_state.task = response[-2]["content"] 77 | set_state(1) 78 | elif model == "Optimization Designer": 79 | response = requests.post(base_url+'solve_task', json={"content": prompt}).json() 80 | st.session_state.messages += response 81 | for m in response: append_message(m) 82 | 83 | if st.session_state.stage == 1: 84 | st.button(f'Execute Plan', on_click=set_state, args=[2]) 85 | 86 | if st.session_state.stage == 2: 87 | st.button(f'Stop Plan', on_click=set_state, args=[0]) 88 | response = requests.get(base_url+'next_task').json() 89 | if response[-1]['content'] == "finished": 90 | set_state(0) 91 | elif response[-1]['content'] is not None: 92 | for m in response: append_message(m) 93 | st.session_state.messages += response 94 | sleep(3) 95 | st.rerun() 96 | 97 | # Reset the simulation 98 | st.sidebar.button('Reset', on_click=requests.get, args=[base_url+'reset']) 99 | 100 | # Close the simulation 101 | st.sidebar.button('Close', on_click=requests.get, args=[base_url+'close']) 102 | 103 | if st.session_state.recording == 0: 104 | st.sidebar.button('Start recording', on_click=set_recording_state, args=[1]) 105 | 106 | if st.session_state.recording == 1: 107 | _ = requests.get(base_url+'start_recording') 108 | st.sidebar.button('Stop recording', on_click=set_recording_state, args=[2]) 109 | st.sidebar.button('Cancel recording', on_click=set_recording_state, args=[3]) 110 | 111 | if st.session_state.recording == 2: 112 | _ = requests.get(base_url+'save_recording') 113 | set_recording_state(0) 114 | 115 | if st.session_state.recording == 3: 116 | _ = requests.get(base_url+'cancel_recording') 117 | set_recording_state(0) -------------------------------------------------------------------------------- /robot.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from time import time 3 | 4 | from llm import LLM 5 | from core import AbstractRobot 6 | from controller import Controller 7 | from typing import Tuple, List, Dict, Optional 8 | from config.config import RobotConfig, LLMConfig, ControllerConfig 9 | 10 | 11 | class Robot(AbstractRobot): 12 | def __init__(self, env_info:Tuple[List], cfg=RobotConfig()) -> None: 13 | self.cfg = cfg 14 | self.robots_info, self.objects_info = env_info 15 | 16 | self.gripper = 1. # 1 means the gripper is open 17 | self.gripper_timer = 0 18 | self.gripper_is_moving = False 19 | self.t_prev_task = time() 20 | self.TP = LLM(LLMConfig("TP", self.cfg.task)) 21 | self.OD = LLM(LLMConfig("OD", self.cfg.task)) 22 | 23 | self.MPC = Controller(env_info, ControllerConfig(self.cfg.task)) 24 | 25 | def init_states(self, observation:Dict[str, np.ndarray], t:float): 26 | """ Update simulation time and current state of MPC controller""" 27 | self.MPC.init_states(observation, t, self.gripper==-0.02) 28 | 29 | def pretty_print(self, response:dict): 30 | if "instruction" in response.keys(): 31 | pretty_msg = "**Reasoning:**\n" 32 | pretty_msg += f"{response['reasoning']}\n" 33 | pretty_msg += "**Instruction:**\n" 34 | pretty_msg += f"{response['instruction']}\n" 35 | else: 36 | pretty_msg = "```\n" 37 | pretty_msg += f"min {response['objective']}\n" 38 | pretty_msg += f"s.t.\n" 39 | for c in response['equality_constraints']: 40 | pretty_msg += f"\t {c} = 0\n" 41 | for c in response['inequality_constraints']: 42 | pretty_msg += f"\t {c} <= 0\n" 43 | pretty_msg += "```\n" 44 | 45 | return pretty_msg 46 | 47 | def _get_instruction(self, query:str) -> str: 48 | instruction = f"objects = {[o['name'] for o in self.objects_info]}\n" 49 | instruction += f"# Query: {query}" 50 | return instruction 51 | 52 | def _open_gripper(self): 53 | self.gripper = -0.01 54 | self.gripper_timer = 0 55 | self.gripper_is_moving = True 56 | 57 | def _close_gripper(self): 58 | self.gripper = -0.02 59 | self.gripper_timer = 0 60 | self.gripper_is_moving = True 61 | 62 | def reset(self): 63 | # open grfipper 64 | self.gripper = 1. 65 | # reset llms 66 | self.TP.reset() 67 | self.OD.reset() 68 | # reset mpc 69 | self.MPC.reset() 70 | 71 | def plan_task(self, user_message:str, base64_image=None) -> str: 72 | """ Runs the Task Planner by passing the user message and the current frame """ 73 | plan = self.TP.run(self._get_instruction(user_message), base64_image, short_history=True) 74 | print(f"\33[92m {plan} \033[0m \n") 75 | return plan 76 | 77 | def is_robot_busy(self): 78 | return not ((np.abs(self.MPC.prev_cost - self.MPC.cost) <= self.cfg.COST_DIIFF_THRESHOLD or 79 | self.MPC.cost <= self.cfg.COST_THRESHOLD or 80 | time()-self.t_prev_task>=self.cfg.TIME_THRESHOLD) and 81 | not self.gripper_is_moving) 82 | 83 | def update_gripper(self, plan:str) -> Optional[str]: 84 | if "open_gripper" in plan.lower(): 85 | self._open_gripper() 86 | #simulate_stream("OD", "\n```\n open_gripper()\n```\n") 87 | return "\n```\n open_gripper()\n```\n" 88 | elif "close_gripper" in plan.lower(): 89 | self._close_gripper() 90 | #simulate_stream("OD", "\n```\n close_gripper()\n```\n") 91 | return "\n```\n close_gripper()\n```\n" 92 | else: 93 | return None 94 | 95 | def solve_task(self, query:str, optimization:Optional[dict]=None) -> str: 96 | """ Applies and returns the optimization designed by the Optimization Designer """ 97 | # if custom function is called apply that 98 | 99 | if self.is_robot_busy(): 100 | return None 101 | 102 | self.t_prev_task = time() 103 | 104 | gripper_update = self.update_gripper(query) 105 | if gripper_update is not None: 106 | return gripper_update 107 | 108 | if optimization is not None: 109 | # apply optimization functions to MPC 110 | try: 111 | self.MPC.setup_controller(optimization) 112 | return self.pretty_print(optimization) 113 | except Exception as e: 114 | print(f"Error with Open Loop Optimization: {e}") 115 | 116 | # catch if reply cannot be parsed. i.e. when askin the LLM a question 117 | for i in range(self.cfg.MAX_OD_ATTEMPTS): 118 | try: 119 | # design optimization functions 120 | if i == 0: 121 | query += "The previous optimization was not feasible. Please try again with a simpler formulation. You can assume the size of all objects is the same." 122 | optimization = self.OD.run(self._get_instruction(query), short_history=True) 123 | print(f"\33[92m {optimization} \033[0m \n") 124 | if self.cfg.method == "objective": 125 | optimization["equality_constraints"] = [] 126 | optimization["inequality_constraints"] = [] 127 | # apply optimization functions to MPC 128 | self.MPC.setup_controller(optimization) 129 | return self.pretty_print(optimization) 130 | except Exception as e: 131 | print(f"Error: {e}") 132 | 133 | return None 134 | 135 | def step(self): 136 | action = [] 137 | # compute actions from controller (single or dual robot) 138 | control: List[np.ndarray] = self.MPC.step() 139 | for u in control: 140 | action.append(np.hstack((u, self.gripper))) 141 | 142 | if self.gripper_timer >= self.cfg.open_gripper_time: 143 | self.gripper_is_moving = False 144 | else: 145 | self.gripper_timer += 1 146 | self.gripper_is_moving = True 147 | if self.gripper_timer == int(self.cfg.open_gripper_time/2) and self.gripper == -0.01: 148 | self.gripper = 1. 149 | 150 | return action 151 | 152 | def retrieve_trajectory(self): 153 | return self.MPC.retrieve_trajectory() 154 | 155 | 156 | -------------------------------------------------------------------------------- /controller.py: -------------------------------------------------------------------------------- 1 | import do_mpc 2 | import numpy as np 3 | import casadi as ca 4 | from time import time 5 | from itertools import chain 6 | from typing import Dict, List, Optional, Tuple 7 | 8 | from core import AbstractController 9 | from config.config import ControllerConfig 10 | 11 | 12 | class Object: 13 | def __init__(self, name:str, position:np.ndarray=np.zeros((3,1)), psi:float=0., size:float=0.) -> None: 14 | self.name = name 15 | self.position = position 16 | self.psi = psi 17 | self.size = size 18 | 19 | class Controller(AbstractController): 20 | 21 | def __init__(self, env_info:Tuple[List], cfg=ControllerConfig) -> None: 22 | super().__init__(cfg) 23 | 24 | # init info of robots and objects 25 | self.robots_info, self.objects_info = env_info 26 | 27 | 28 | 29 | # gripper fingers offset for constraints 30 | self.gripper_offsets = [(np.array([0., -0.048, 0.003]), 0.013), (np.array([0., 0.048, 0.003]), 0.013), 31 | (np.array([0., 0.0, 0.06]), 0.025), (np.array([0., -0.045, 0.06]), 0.025), 32 | (np.array([0., 0.045, 0.06]), 0.025), (np.array([0., -0.09, 0.06]), 0.025), (np.array([0., 0.09, 0.06]), 0.025) 33 | ] 34 | 35 | self.gripper_offsets_load = [(np.array([0., 0., 0.]), 0.03)] 36 | 37 | 38 | self.gripper_closed = False 39 | 40 | # init controller 41 | self.setup_controller() 42 | 43 | def init_model(self): 44 | # inti do_mpc model 45 | self.model = do_mpc.model.Model(self.cfg.model_type) 46 | 47 | # simulation time 48 | self.t = self.model.set_variable('parameter', 't') 49 | # position of objects 50 | self.objects = {} 51 | for o in self.objects_info: 52 | position = self.model.set_variable(var_type='_p', var_name=o['name']+'_position', shape=(3,1)) 53 | psi = self.model.set_variable(var_type='_p', var_name=o['name']+'_psi') 54 | size = self.model.set_variable(var_type='_p', var_name=o['name']+'_size') 55 | obj = Object(o['name'], position, psi, size) 56 | self.objects[o['name']] = obj 57 | 58 | # gripper pose [x, y, z, theta, gamma, psi] 59 | self.pose = [] 60 | 61 | self.x = [] # gripper position (x,y,z) 62 | self.psi = [] # gripper psi (rotation around z axis) 63 | self.dx = [] # gripper velocity (vx, vy, vz) 64 | self.dpsi = [] # gripper rotational speed 65 | self.u = [] # gripper control (=velocity) 66 | self.u_psi = [] # gripper rotation control (=rotational velocity) 67 | self.cost = 1. # cost function 68 | self.prev_cost = float('inf') # previous cost function 69 | self.solve_time = 0. # time to solve the optimization problem 70 | for i, r in enumerate(self.robots_info): 71 | # position (x, y, z) 72 | self.x.append(self.model.set_variable(var_type='_x', var_name=f'x{r["name"]}', shape=(self.cfg.nx,1))) 73 | self.psi.append(self.model.set_variable(var_type='_x', var_name=f'psi{r["name"]}', shape=(1,1))) 74 | self.dx.append(self.model.set_variable(var_type='_x', var_name=f'dx{r["name"]}', shape=(self.cfg.nx,1))) 75 | self.dpsi.append(self.model.set_variable(var_type='_x', var_name=f'dpsi{r["name"]}', shape=(1,1))) 76 | self.u.append(self.model.set_variable(var_type='_u', var_name=f'u{r["name"]}', shape=(self.cfg.nu,1))) 77 | self.u_psi.append(self.model.set_variable(var_type='_u', var_name=f'u_psi{r["name"]}', shape=(1,1))) 78 | # system dynamics 79 | self.model.set_rhs(f'x{r["name"]}', self.x[i] + self.dx[i] * self.cfg.dt) 80 | self.model.set_rhs(f'psi{r["name"]}', self.psi[i] + self.dpsi[i] * self.cfg.dt) 81 | self.model.set_rhs(f'dx{r["name"]}', self.u[i]) 82 | self.model.set_rhs(f'dpsi{r["name"]}', self.u_psi[i]) 83 | 84 | def setup_controller(self, optimization={"objective":None, "equality_constraints":[], "inequality_constraints":[]}): 85 | self.init_model() 86 | # init cost function 87 | self.model.set_expression('cost', self._eval(optimization["objective"])) 88 | # setup model 89 | self.model.setup() 90 | # init variables and expressions 91 | self.init_expressions() 92 | # init 93 | self.init_mpc() 94 | # set functions 95 | self.set_objective(self._eval(optimization["objective"])) 96 | # set base constraint functions 97 | constraints = [] 98 | # positive equality constraint 99 | constraints += [self._eval(c) for c in optimization["equality_constraints"]] 100 | # negative equality constraint 101 | constraints += [-self._eval(c) for c in optimization["equality_constraints"]] 102 | # inequality constraints 103 | gripper_offsets = self.get_gripper_offsets() 104 | inequality_constraints = [[*map(lambda const: self._eval(c, const), gripper_offsets)] for c in optimization["inequality_constraints"]] 105 | constraints += list(chain(*inequality_constraints)) 106 | # set constraints 107 | self.set_constraints(constraints) 108 | # setup 109 | self.mpc.set_uncertainty_values(t=np.array([0.])) # init time to 0 110 | self.mpc.setup() 111 | self.mpc.set_initial_guess() 112 | 113 | def _normalize_angle(self, angle): 114 | """ 115 | Normalize an angle to be within the range [-pi, pi]. 116 | 117 | Parameters: 118 | angle (float): The angle in radians to be normalized. 119 | 120 | Returns: 121 | float: The normalized angle within the range [-pi, pi]. 122 | """ 123 | normalized_angle = np.arctan2(np.sin(angle), np.cos(angle)) 124 | # Check if the angle is outside the range [-pi/2, pi/2] and adjust 125 | if normalized_angle > np.pi/2: 126 | normalized_angle -= np.pi 127 | elif normalized_angle < -np.pi/2: 128 | normalized_angle += np.pi 129 | return normalized_angle 130 | 131 | def set_objective(self, mterm: ca.SX=ca.DM([[0]])): 132 | # reegularization term for numerical stability 133 | regularization = 0 134 | for i, r in enumerate(self.robots_info): 135 | regularization += .1 * ca.norm_2(self.dx[i])**2 136 | regularization += .0002 * ca.norm_2(ca.sin(self.psi[i]) * ca.cos(self.psi[i]))**2 137 | mterm = mterm + regularization 138 | lterm = 2*mterm 139 | # state objective 140 | self.mpc.set_objective(mterm=mterm, lterm=lterm) 141 | # input objective 142 | u_kwargs = {f'u{r["name"]}':0.5 for r in self.robots_info} | {f'u_psi{r["name"]}':1e-5 for r in self.robots_info} 143 | self.mpc.set_rterm(**u_kwargs) 144 | 145 | def set_constraints(self, nlp_constraints: Optional[List[ca.SX]] = None): 146 | 147 | for r in self.robots_info: 148 | # base constraints (state) 149 | self.mpc.bounds['lower','_x', f'x{r["name"]}'] = np.array([-3., -3., 0.0]) # stay above table 150 | self.mpc.bounds['upper','_x', f'psi{r["name"]}'] = np.pi * 0.55 * np.ones((1, 1)) # rotation upper bound 151 | self.mpc.bounds['lower','_x', f'psi{r["name"]}'] = -np.pi * 0.55 * np.ones((1, 1)) # rotation lower bound 152 | 153 | # base constraints (input) 154 | self.mpc.bounds['upper','_u', f'u{r["name"]}'] = self.cfg.hu * np.ones((self.cfg.nu, 1)) # input upper bound 155 | self.mpc.bounds['lower','_u', f'u{r["name"]}'] = self.cfg.lu * np.ones((self.cfg.nu, 1)) # input lower bound 156 | self.mpc.bounds['upper','_u', f'u_psi{r["name"]}'] = np.pi * np.ones((1, 1)) # input upper bound 157 | self.mpc.bounds['lower','_u', f'u_psi{r["name"]}'] = -np.pi * np.ones((1, 1)) # input lower bound 158 | 159 | if nlp_constraints == None: 160 | return 161 | # soft constraints used as a logarithmic barrier for numerical stability 162 | for i, constraint in enumerate(nlp_constraints): 163 | self.mpc.set_nl_cons(f'const{i}', expr=constraint, ub=0., 164 | soft_constraint=True, 165 | penalty_term_cons=self.cfg.penalty_term_cons) 166 | 167 | def init_mpc(self): 168 | # init mpc model 169 | self.mpc = do_mpc.controller.MPC(self.model) 170 | # setup params 171 | setup_mpc = {'n_horizon': self.cfg.T, 't_step': self.cfg.dt, 'store_full_solution': False} 172 | # setup mpc 173 | self.mpc.set_param(**setup_mpc) 174 | self.mpc.settings.supress_ipopt_output() # => verbose = False 175 | 176 | 177 | def init_expressions(self): 178 | # init variables for python evaluation 179 | self.eval_variables = {"ca":ca, "np":np} # python packages 180 | 181 | self.R = [] # rotation matrix for angle around z axis 182 | for i in range(len(self.robots_info)): 183 | # rotation matrix 184 | self.R.append(np.array([[ca.cos(self.psi[i]), -ca.sin(self.psi[i]), 0], 185 | [ca.sin(self.psi[i]), ca.cos(self.psi[i]), 0], 186 | [0, 0, 1.]])) 187 | 188 | def _quaternion_to_euler_angle_vectorized2(self, quaternion): 189 | x, y, z, w = quaternion 190 | ysqr = y * y 191 | 192 | t0 = +2.0 * (w * x + y * z) 193 | t1 = +1.0 - 2.0 * (x * x + ysqr) 194 | X = np.arctan2(t0, t1) 195 | 196 | t2 = +2.0 * (w * y - z * x) 197 | 198 | t2 = np.clip(t2, a_min=-1.0, a_max=1.0) 199 | Y = np.arcsin(t2) 200 | 201 | t3 = +2.0 * (w * z + x * y) 202 | t4 = +1.0 - 2.0 * (ysqr + z * z) 203 | Z = np.arctan2(t3, t4) 204 | 205 | return np.array([X, Y, Z]) 206 | 207 | def _set_x0(self, observation: Dict[str, np.ndarray]): 208 | x0 = [] 209 | self.pose = [] 210 | for r in self.robots_info: 211 | obs = observation[f'robot{r["name"]}'] # observation of each robot 212 | x = obs[:3] 213 | psi = self._normalize_angle(np.array([obs[5]])) 214 | dx = obs[6:9] 215 | x0.append(np.concatenate((x, psi, dx, [0]))) 216 | self.pose.append(obs[:6]) 217 | # set x0 in MPC 218 | self.mpc.x0 = np.concatenate(x0) 219 | 220 | def init_states(self, observation:Dict[str, np.ndarray], t:float, gripper_closed:bool=False): 221 | """ Set the values the MPC initial states and variables """ 222 | self.gripper_closed = gripper_closed 223 | self.observation = observation 224 | # set mpc x0 225 | self._set_x0(observation) 226 | # set variable parameters 227 | parameters = {'t': [t]} 228 | parameters = parameters | {o['name']+'_position': [observation[o['name']]['position']] for o in self.objects_info} 229 | parameters = parameters | {o['name']+'_size': [observation[o['name']]['size']] for o in self.objects_info} 230 | parameters = parameters | {o['name']+'_psi': [self._quaternion_to_euler_angle_vectorized2(observation[o['name']]['orientation'])[-1]] for o in self.objects_info if o["name"].endswith("_orientation")} 231 | #print(parameters) 232 | self.mpc.set_uncertainty_values(**parameters) 233 | 234 | def get_gripper_offsets(self): 235 | if self.gripper_closed: 236 | gripper_offset = self.gripper_offsets + self.gripper_offsets_load 237 | else: 238 | gripper_offset = self.gripper_offsets 239 | return gripper_offset 240 | 241 | def reset(self) -> None: 242 | """ 243 | observation: robot observation from simulation containing position, angle and velocities 244 | """ 245 | self.setup_controller() 246 | return 247 | 248 | def _eval(self, code_str: str, offset=(np.zeros(3), 0.)): 249 | # put together variables for python code evaluation: 250 | if code_str == None: return ca.SX(0) 251 | 252 | # parse offset 253 | collision_xyz, collision_radius = offset 254 | # initial state of robots before applying any action 255 | x0 = {f'x0{r["name"]}': self.observation[f'robot{r["name"]}'][:3] for r in self.robots_info} 256 | # robot variable states (decision variables in the optimization problem) 257 | robots_states = {} 258 | for i, r in enumerate(self.robots_info): 259 | robots_states[f'x{r["name"]}'] = self.x[i] + self.R[i]@collision_xyz 260 | robots_states[f'dx{r["name"]}'] = self.dx[i] 261 | robots_states[f'psi{r["name"]}'] = self.psi[i] 262 | 263 | eval_variables = self.eval_variables | robots_states | self.objects | x0 | {'t': self.t} 264 | # evaluate code 265 | evaluated_code = eval(code_str, eval_variables) + collision_radius 266 | return evaluated_code 267 | 268 | def _solve(self) -> List[np.ndarray]: 269 | """ Returns a list of conntrols, 1 for each robot """ 270 | # solve mpc at state x0 271 | t0 = time() 272 | u0 = self.mpc.make_step(self.mpc.x0).squeeze() 273 | self.solve_time = time() - t0 274 | # compute action for each robot 275 | action = [] 276 | for i in range(len(self.robots_info)): 277 | ee_displacement = u0[4*i:4*i+3] # positon control 278 | theta_regularized = self.pose[i][3] if self.pose[i][3]>=0 else self.pose[i][3] + 2*np.pi 279 | theta_rotation = [(np.pi - theta_regularized)*1.5] 280 | gamma_rotation = [-self.pose[i][4] * 1.5] # P control for angle around y axis 281 | psi_rotation = [u0[4*i+3]] # rotation control 282 | action.append(np.concatenate((ee_displacement, theta_rotation, gamma_rotation, psi_rotation))) 283 | 284 | self.prev_cost = self.cost 285 | self.cost = self.mpc.data['_aux'][-1][-1] 286 | 287 | return action 288 | 289 | def step(self): 290 | if not self.mpc.flags['setup']: 291 | return [np.zeros(6) for i in range(len(self.robots_info))] 292 | return self._solve() 293 | 294 | def retrieve_trajectory(self): 295 | trajectory = [] 296 | try: 297 | for _x in self.mpc.opt_x_num['_x', :, 0, 0]: 298 | _x = _x.toarray().flatten() 299 | trajectory.append(_x[:3]) 300 | except: 301 | pass 302 | return trajectory 303 | 304 | -------------------------------------------------------------------------------- /.zshrc: -------------------------------------------------------------------------------- 1 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 2 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 3 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 4 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 5 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 6 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 7 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 8 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 9 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 10 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 11 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 12 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 13 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 14 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 15 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 16 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 17 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 18 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 19 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 20 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 21 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 22 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 23 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 24 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 25 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 26 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 27 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 28 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 29 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 30 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 31 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 32 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 33 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 34 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 35 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 36 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 37 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 38 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 39 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 40 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 41 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 42 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 43 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin 44 | export PATH=/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/opt/homebrew/bin 45 | -------------------------------------------------------------------------------- /simulation_local.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import gym 4 | import json 5 | import panda_gym 6 | import numpy as np 7 | from PIL import Image 8 | from tqdm import tqdm 9 | from typing import List 10 | from datetime import datetime 11 | from sqlalchemy import create_engine 12 | from sqlalchemy.orm import sessionmaker 13 | 14 | 15 | from robot import Robot 16 | from db import Base, Episode, Epoch 17 | from core import AbstractSimulation, BASE_DIR 18 | from config.config import SimulationConfig, RobotConfig 19 | 20 | 21 | 22 | class Simulation(AbstractSimulation): 23 | def __init__(self, cfg=SimulationConfig()) -> None: 24 | #super().__init__(cfg) 25 | 26 | self.cfg = cfg 27 | # init env 28 | self.env = gym.make(f"Panda{cfg.task}-v2", render=cfg.render, debug=cfg.debug) 29 | # init robots 30 | # count number of tasks solved from a plan 31 | self.plan = None 32 | self.optimizations = [] 33 | self.task_counter = 0 34 | self.prev_instruction = "None" 35 | 36 | # simulation time 37 | self.t = 0. 38 | self.epoch = 0 39 | env_info = (self.env.robots_info, self.env.objects_info) 40 | self.robot = Robot(env_info,RobotConfig(self.cfg.task)) 41 | # count number of tasks solved from a plan 42 | self.task_counter = 0 43 | # bool for stopping simulation 44 | self.stop_thread = False 45 | # whether to save frame (initialized to false) 46 | self.save_video = self.cfg.save_video 47 | # init list of RGB frames if wanna save video 48 | self.frames_list = [] 49 | self.frames_list_logging = [] 50 | self.video_name = f"{self.cfg.task}_{datetime.now().strftime('%d-%m-%Y_%H:%M:%S')}" 51 | self.video_path = os.path.join(BASE_DIR, f"videos/{self.video_name}.mp4") 52 | # 53 | self.state_trajectories = [] 54 | self.mpc_solve_times = [] 55 | self.session = None 56 | 57 | if self.cfg.logging: 58 | engine = create_engine(f'sqlite:///data/{self.cfg.method}/DBs/{cfg.task}.db') 59 | Base.metadata.create_all(engine) 60 | self.Session = sessionmaker(bind=engine) 61 | 62 | 63 | def _round_list(self, l, n=2): 64 | """ round list and if the result is -0.0 convert it to 0.0 """ 65 | return [r if (r:=round(x, n)) != -0.0 else 0.0 for x in l] 66 | 67 | def _append_robot_info(self): 68 | for r in self.env.robots_info: 69 | obs = self.observation[f'robot{r["name"]}'] # observation of each robot 70 | self.state_trajectories.append(obs.tolist()) 71 | self.mpc_solve_times.append(self.robot.MPC.solve_time) 72 | 73 | def _create_scene_description(self): 74 | """ Look at the observation and create a string that describes the scene to be passed to the task planner """ 75 | ri = 0 76 | description = "The following is the description of the current scene:\n" 77 | for name in self.observation.keys(): 78 | if name.startswith("robot"): 79 | robot_xyz = self._round_list(self.observation[name][:3]) 80 | description += f"- The gripper of the {name} is located at {robot_xyz}.\n" 81 | if self.robot.gripper==-1: 82 | if round(self.env.robots[ri].get_fingers_width(),2) <= 0.01: 83 | description += f"- The gripper fingers have closed but they are grasping no object.\n" 84 | else: 85 | distances = {cube_name: np.linalg.norm(np.array(robot_xyz)-np.array(self.observation[cube_name])) for cube_name in self.observation.keys() if cube_name.endswith("_cube")} 86 | closest_cube = min(distances, key=distances.get)[:-5] 87 | description += f"- The gripper fingers are closed and they are firmly grasping the {closest_cube} cube.\n" 88 | else: 89 | description += f"- The gripper fingers are open.\n" 90 | ri += 1 91 | elif name.endswith("_cube"): 92 | description += f"- The center of the {name[:-5]} cube is located at {self._round_list(self.observation[name])}\n" 93 | else: 94 | pass 95 | 96 | description += """Please carefully analyze the scene description and decide what to do next. Some helpful tips are: 97 | (1) If the gripper is not at the location where it should be it is surely because of collisions. Specify in your instruction to the robot about collision avoidance and constraints. 98 | (2) Be careful when placing a cube on top of another one that you leave some clearance between those 2 cubes. Be very careful and explain how much space should be left between. 99 | (a) It's ok if the cube is not at the same exact x and y position as the cube below. 100 | (3) Make sure that the cube you've put on the stack has not fallen. Always check every cube to understand if it is on the ground or on top of another cube. 101 | (a) A cube is on the ground if it's height is 0.02m. 102 | (b) If you stacked a cube and need to go to another one, make sure to instruct the robot to avoid collisions with the cubes in the stack. 103 | (4) The description of scene is ALWAYS correct, the instructions you give may be wrong or at times mis-interepted by the robot. Always try to fix this when it happens. 104 | (5) Make sure the gripper is open before it needs to go to an object to grasp it. 105 | """ 106 | 107 | return description 108 | 109 | def _uplaod_image(self, rgba_image:np.ndarray) -> str: 110 | # Convert the NumPy array to a PIL Image object 111 | image = Image.fromarray(rgba_image, 'RGBA') 112 | image_path = f"{self.episode_folder}/{datetime.now().strftime('%d-%m-%Y_%H:%M:%S')}.png" # Specify your local file path here 113 | image.save(image_path, 'PNG') 114 | return image_path 115 | 116 | def _retrieve_image(self) -> np.ndarray: 117 | frame_np = np.array(self.env.render("rgb_array", 118 | width=self.cfg.frame_width, height=self.cfg.frame_height, 119 | target_position=self.cfg.frame_target_position, 120 | distance=self.cfg.frame_distance, 121 | yaw=self.cfg.frame_yaw, 122 | pitch=self.cfg.frame_pitch)) 123 | frame_np = frame_np.reshape(self.cfg.frame_width, self.cfg.frame_height, 4).astype(np.uint8) 124 | 125 | return frame_np 126 | 127 | def _store_epoch_db(self, episode_id, role, content, image_url): 128 | session = self.Session() 129 | 130 | # Find the last epoch number for this episode 131 | last_epoch = session.query(Epoch).filter_by(episode_id=episode_id).order_by(Epoch.time_step.desc()).first() 132 | if last_epoch is None: 133 | next_time_step = 1 # This is the first epoch for the episode 134 | else: 135 | next_time_step = last_epoch.time_step + 1 136 | 137 | # Create and insert the new epoch 138 | epoch = Epoch(episode_id=episode_id, time_step=next_time_step, role=role, content=content, image=image_url) 139 | session.add(epoch) 140 | session.commit() 141 | session.close() 142 | 143 | def _make_plan(self, user_message:str="") -> str: 144 | self.plan:dict = self.robot.plan_task(user_message) 145 | self.task_counter = 0 146 | pretty_msg = "Tasks:\n" 147 | pretty_msg += "".join([f"{i+1}. {task}\n" for i, task in enumerate(self.plan["tasks"])]) 148 | if self.cfg.logging: 149 | image = self._retrieve_image() 150 | image_url = self._uplaod_image(image) 151 | self._store_epoch_db(self.episode.id, "human", self.robot._get_instruction(user_message), image_url) 152 | self._store_epoch_db(self.episode.id, "TP", pretty_msg, image_url) 153 | return pretty_msg 154 | 155 | def _solve_task(self, task:str, optimization:dict=None) -> str: 156 | AI_response = self.robot.solve_task(task, optimization) if task != "finished" else task 157 | if self.cfg.logging and AI_response is not None: 158 | image = self._retrieve_image() 159 | image_url = self._uplaod_image(image) 160 | self._store_epoch_db(self.episode.id, "OD", AI_response, image_url) 161 | return AI_response 162 | 163 | def reset(self): 164 | # reset pand env 165 | self.observation = self.env.reset() 166 | # reset robot 167 | self.robot.reset() 168 | # reset controller 169 | self.robot.init_states(self.observation, self.t) 170 | # reset task counter 171 | self.plan = None 172 | self.optimizations = [] 173 | self.task_counter = 0 174 | # init list of RGB frames if wanna save video 175 | if self.save_video: 176 | self._save_video() 177 | if self.cfg.logging: 178 | if self.session is not None: 179 | self.episode.state_trajectories = json.dumps(self.state_trajectories) 180 | self.episode.mpc_solve_times = json.dumps(self.mpc_solve_times) 181 | if self.cfg.logging_video: 182 | self._save_video() 183 | self.session.commit() 184 | self.state_trajectories = [] 185 | self.mpc_solve_times = [] 186 | self.session.close() 187 | self.session = self.Session() 188 | self.episode = Episode() # Assuming Episode has other fields you might set 189 | self.session.add(self.episode) 190 | self.session.commit() 191 | n_episodes = len(os.listdir(f"data/{self.cfg.method}/images")) 192 | self.episode_folder = f"data/{self.cfg.method}/images/{n_episodes}" 193 | os.mkdir(self.episode_folder) 194 | self.video_path = os.path.join(BASE_DIR, f"data/{self.cfg.method}/videos/{self.cfg.task}_{n_episodes}_full.mp4") 195 | self.video_path_logging = os.path.join(BASE_DIR, f"data/{self.cfg.method}/videos/{self.cfg.task}/{self.episode.id}.mp4") 196 | 197 | # init list of RGB frames if wanna save video 198 | if self.save_video: 199 | self._save_video() 200 | self.frames_list = [] 201 | self.frames_list_logging = [] 202 | self.t = 0. 203 | self.epoch = 0 204 | 205 | 206 | def step(self): 207 | # increase timestep 208 | self.t += self.cfg.dt 209 | self.epoch += 1 210 | # update controller (i.e. set the current gripper position) 211 | self.robot.init_states(self.observation, self.t) 212 | # compute action 213 | action = self.robot.step() 214 | # apply action 215 | self.observation, _, done, _ = self.env.step(action) 216 | # add states to state_trajectories 217 | if self.cfg.logging: 218 | self._append_robot_info() 219 | if self.cfg.logging_video and self.epoch%20 == 0: 220 | self.frames_list_logging.append(self._retrieve_image()) 221 | # visualize trajectory 222 | if self.cfg.debug: 223 | trajectory = self.robot.retrieve_trajectory() 224 | self.env.visualize_trajectory(trajectory) 225 | # store RGB frames if wanna save video 226 | if self.save_video: 227 | frame = np.array(self.env.render("rgb_array", width=self.cfg.frame_width, height=self.cfg.frame_height)) 228 | frame = frame.reshape(self.cfg.frame_width, self.cfg.frame_height, 4).astype(np.uint8) 229 | self.frames_list.append(frame) 230 | 231 | return done 232 | 233 | def close(self): 234 | # init list of RGB frames if wanna save video 235 | if self.save_video: 236 | self._save_video() 237 | # store state_trajectories and mpc_solve_times 238 | if self.cfg.logging: 239 | if self.cfg.logging_video: 240 | self._save_video() 241 | self.episode.state_trajectories = json.dumps(self.state_trajectories) 242 | self.episode.mpc_solve_times = json.dumps(self.mpc_solve_times) 243 | self.session.commit() 244 | self.session.close() 245 | # close env 246 | self.env.close() 247 | 248 | 249 | def _save_video(self): 250 | for s, p, l in [(self.cfg.save_video, self.video_path, self.frames_list), (self.cfg.logging_video, self.video_path_logging, self.frames_list_logging)]: 251 | if not s: continue 252 | # Define the parameters 253 | fourcc = cv2.VideoWriter_fourcc(*'mp4v') 254 | # Create a VideoWriter object 255 | out = cv2.VideoWriter(p, fourcc, self.cfg.fps, (self.cfg.frame_width, self.cfg.frame_height)) 256 | # Write frames to the video 257 | for frame in tqdm(l): 258 | # Ensure the frame is in the correct format (RGBA) 259 | if frame.shape[2] == 3: 260 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2RGBA) 261 | # Convert the frame to BGR format (required by VideoWriter) 262 | frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR) 263 | out.write(frame_bgr) 264 | # Release the VideoWriter 265 | out.release() 266 | 267 | 268 | def run(self, query:str, plan:dict, optimizations:List[dict]): 269 | self.task_counter = 0 270 | if plan is not None: 271 | self.plan = plan 272 | else: 273 | self._make_plan(query) 274 | if optimizations is not None: 275 | self.optimizations = optimizations 276 | pretty_msg = "Tasks:\n" 277 | pretty_msg += "".join([f"{i+1}. {task}\n" for i, task in enumerate(self.plan["tasks"])]) 278 | if self.cfg.logging: 279 | image = self._retrieve_image() 280 | image_url = self._uplaod_image(image) 281 | self._store_epoch_db(self.episode.id, "human", query, image_url) 282 | self._store_epoch_db(self.episode.id, "TP", pretty_msg, image_url) 283 | 284 | is_plan_unfinished = True 285 | while is_plan_unfinished: 286 | is_plan_unfinished = self.task_counter < len(self.plan["tasks"]) 287 | task = self.plan["tasks"][self.task_counter] if is_plan_unfinished else "finished" 288 | optimization = self.optimizations[self.task_counter] if (self.optimizations and is_plan_unfinished) else None 289 | _ = self._solve_task(task, optimization) 290 | self.task_counter += 1 291 | 292 | while self.robot.is_robot_busy(): 293 | self.step() 294 | 295 | 296 | if __name__=="__main__": 297 | # init sim 298 | s = Simulation() 299 | for _ in range(2): 300 | s.reset() 301 | # load data 302 | task_folder = f'data/{s.cfg.method}/llm_responses/{s.cfg.task}' 303 | # run sim 304 | s.run("use right robot to move container to sink and left robot to move sponge to the sink. the sponge is wet so keep it above the container to avoid water dropping on the floor", None, None) 305 | 306 | s.close() -------------------------------------------------------------------------------- /prompts/prompts.py: -------------------------------------------------------------------------------- 1 | TP_PROMPT = """ 2 | You are a helpful assistant in charge of controlling a robot manipulator. 3 | The user will give you a goal and you have to formulate a plan that the robot will follow to achieve the goal. 4 | 5 | You can control the robot in the following way: 6 | (1) Instructions in natural language to move the gripper and follow constriants. 7 | (2) open_gripper() 8 | (3) close_gripper() 9 | (a) you can firmly grasp an object only if the gripper is at the same position of the center of the object and the gripper is open. 10 | 11 | Rules: 12 | (1) You MUST ALWAYS specificy which objects specifically the gripper has to avoid collisions with in your instructions. 13 | (2) NEVER avoid collisions with an object you are gripping. 14 | (3) Use these common sense rules for spatial reasoning: 15 | (a) 'in front of' and 'behind' for positive and negative x-axis directions. 16 | (b) 'to the left' and 'to the right' for positive and negative y-axis directions. 17 | (c) 'above' for positive z-axis directions. 18 | 19 | 20 | You MUST always respond with a json following this format: 21 | { 22 | "tasks": ["task1", "task2", "task3", ...] 23 | } 24 | 25 | Here are some general examples: 26 | 27 | objects = ['coffee pod', 'coffee machine'] 28 | # Query: put the coffee pod into the coffee machine 29 | { 30 | "tasks": ["move gripper to the coffee pod and avoid collisions with the coffee machine", "close_gripper()", "move the gripper above the coffee machine", "open_gripper()"] 31 | } 32 | 33 | objects = ['blue block', 'yellow block', 'mug'] 34 | # Query: stack the blue block on the yellow block, and avoid the mug at all time. 35 | { 36 | "tasks": ["move gripper to the blue block and avoid collisions with the yellow block and the mug", "close_gripper()", "move the gripper above the yellow block and avoid collisions with the yellow block and the mug", "open_gripper()"] 37 | } 38 | 39 | objects = ['apple', 'drawer handle', 'drawer'] 40 | # Query: put apple into the drawer. 41 | { 42 | "tasks": ["move gripper to drawer handle and avoid collisions with apple and drawer", "close_gripper()", "move gripper 0.25m in the y direction", "open_gripper()", "move gripper to the apple and avoid collisions with the drawer and its handle", "close_gripper()", "move gripper above the drawer and avoid collisions with the drawer", "open_gripper()"] 43 | } 44 | 45 | objects = ['plate', 'fork', 'knife', 'glass] 46 | # Query: Order the kitchen objects flat on the table in the x-y plane. 47 | { 48 | "tasks": ["move gripper to the fork and avoid collisions with plate, knife, glass", "close_gripper()", "move gripper to the left side of the plate avoiding collisions with plate, knife, glass", "open_gripper()", "move gripper to the glass and avoid collisions with fork, plate, knife", "close_gripper()", "move gripper in front of the plate avoiding collisions with fork, plate and knife", "open_gripper()", "move gripper to the knife and avoid collisions with fork, plate, glass", "close_gripper()", "move gripper to the right side of the plate avoiding collisions with fork, plate and glass", "open_gripper()"] 49 | } 50 | 51 | """ 52 | 53 | 54 | OD_PROMPT = """ 55 | You are a helpful assistant in charge of designing the optimization problem for an MPC controller that is controlling a robot manipulator. 56 | At each step, I will give you a task and you will have to return the objective and (optionally) the constraint functions that need to be applied to the MPC controller. 57 | 58 | This is the scene description: 59 | (1) Casadi is used to program the MPC. 60 | (2) The variable `x` represents the gripper position of the gripper in 3D, i.e. (x, y, z). 61 | (2) The variable `x0` represents the initial gripper position at the current time step before any action is applied i.e. (x, y, z). 62 | (3) The orientation of the gripper around the z-axis is defined by variable `psi`. 63 | (4) The variable `t` represents the simulation time. 64 | (5) Each time I will also give you a list of objects you can interact with (i.e. objects = ['peach', 'banana']). 65 | (a) The position of each object is an array [x, y, z] obtained by adding `.position` (i.e. 'banana.position'). 66 | (b) The size of each cube is a float obtained by adding '.size' (i.e. 'banana.size'). 67 | (c) The rotaton around the z-axis is a float obtained by adding '.psi' (i.e. 'banana.psi'). 68 | (6) 69 | (a) 'in front of' and 'behind' for positive and negative x-axis directions. 70 | (b) 'to the left' and 'to the right' for positive and negative y-axis directions. 71 | (c) 'above' and 'below' for positive and negative z-axis directions. 72 | 73 | 74 | Rules: 75 | (1) You MUST write every equality constraints such that it is satisfied if it is = 0: 76 | (a) If you want to write "ca.norm_2(x) = 1" write it as "1 - ca.norm_2(x)" instead. 77 | (2) You MUST write every inequality constraints such that it is satisfied if it is <= 0: 78 | (a) If you want to write "ca.norm_2(x) >= 1" write it as "1 - ca.norm_2(x)" instead. 79 | (3) You MUST avoid colliding with an object IFF you're moving the gripper specifically to that object or nearby it (i.e. above the object), even if not specified in the query. 80 | (4) NEVER avoid collisions with an object you're not moving to or nearby if not specified in the query. 81 | (4) Use `t` in the inequalities especially when you need to describe motions of the gripper. 82 | 83 | You must format your response into a json. Here are a few examples: 84 | 85 | objects = ['object_1', 'object_2'] 86 | # Query: move the gripper to [0.2, 0.05, 0.2] and avoid collisions with object_2 87 | { 88 | "objective": "ca.norm_2(x - np.array([0.2, 0.05, 0.2]))**2", 89 | "equality_constraints": [], 90 | "inequality_constraints": ["object_2.size - ca.norm_2(x - object_2.position)"] 91 | } 92 | Notice how the inequality constraint holds if <= 0. 93 | 94 | objects = ['red_cube', 'yellow_cube'] 95 | # Query: move the gripper to red cube and avoid colliding with the yellow cube 96 | { 97 | "objective": "ca.norm_2(x - red_cube.position)**2", 98 | "equality_constraints": [], 99 | "inequality_constraints": ["red_cube.size*0.85 - ca.norm_2(x - red_cube.position)", "yellow_cube.size - ca.norm_2(x - yellow_cube.position)"] 100 | } 101 | Notice the collision avoidance constraint with the red_cube despite not being specified in the query because the gripper has to go to the red cube. 102 | 103 | objects = ['coffee_pod', 'coffee_machine'] 104 | # Query: move gripper above the coffe pod and keep gripper at a height higher than 0.1m 105 | { 106 | "objective": "ca.norm_2(x - (coffee_pod.position + np.array([0, 0, coffee_pod.size])))**2", 107 | "equality_constraints": [], 108 | "inequality_constraints": ["coffee_pod.size - ca.norm_2(x - coffee_pod.position)", "0.1 - x[2]"] 109 | } 110 | Notice that there's no collision avoidance constraint with the coffee_machine because it is not in the query and because gripper is not moving to or nearby it. 111 | 112 | 113 | objects = ['blue_container', 'yellow_container', 'green_container'] 114 | # Query: Move gripper above stack composed by blue, yellow, and green container 115 | { 116 | "objective": "ca.norm_2(x - (blue_container.position + np.array([0, 0, blue_container.size + yellow_container.size + green_container.size])))**2", 117 | "equality_constraints": [], 118 | "inequality_constraints": ["blue_container.size*0.85 - ca.norm_2(x - blue_container.position)", "yellow_container.size*0.85 - ca.norm_2(x - yellow_container.position)", "green_container.size*0.85 - ca.norm_2(x - green_container.position)"] 119 | } 120 | 121 | objects = ['mug'] 122 | # Query: Move the gripper 0.1m upwards 123 | { 124 | "objective": "ca.norm_2(x - (x0 + np.array([0, 0, 0.1])))**2", 125 | "equality_constraints": [], 126 | "inequality_constraints": [] 127 | } 128 | 129 | objects = ['apple', 'pear'] 130 | # Query: move the gripper to apple and stay 0.04m away from pear 131 | { 132 | "objective": "ca.norm_2(x - apple.position)**2", 133 | "equality_constraints": [], 134 | "inequality_constraints": ["apple.size*0.85 - ca.norm_2(x - apple.position)", "0.04 - ca.norm_2(x - pear.position)"] 135 | } 136 | 137 | objects = ['joystick', 'remote'] 138 | # Query: Move the gripper at constant speed along the x axis while keeping y and z fixed at 0.2m 139 | { 140 | "objective": "ca.norm_2(x_left[0] - t)**2", 141 | "equality_constraints": ["np.array([0.2, 0.2]) - x[1:]"], 142 | "inequality_constraints": [] 143 | } 144 | 145 | objects = ['fork', 'spoon', 'plate'] 146 | # Query: Move the gripper behind fork and avoid collisions with spoon 147 | { 148 | "objective": "ca.norm_2(x - (fork.position + np.array([-fork.size, 0, 0])))**2", 149 | "equality_constraints": [], 150 | "inequality_constraints": ["fork.size*0.85 - ca.norm_2(x - fork.position)", "spoon.size - ca.norm_2(x - spoon.position)"] 151 | } 152 | """ 153 | 154 | TP_PROMPT_COLLAB = """ 155 | You are a helpful assistant in charge of controlling 2 robot manipulators. 156 | The user will give you a goal and you have to formulate a plan that the robots will follow to achieve the goal. 157 | 158 | You can control the robot in the following way: 159 | (1) Instructions in natural language to move the robot grippers and follow constriants. 160 | (2) open_gripper() 161 | (3) close_gripper() 162 | (a) you can firmly grasp an object only if the gripper is at the same position of the center of the object and the gripper is open. 163 | 164 | Rules: 165 | (1) You MUST ALWAYS specificy which objects specifically the grippers have to avoid collisions with in your instructions. 166 | (2) NEVER avoid collisions with an object you are gripping. 167 | (3) Use these common sense rules for spatial reasoning: 168 | (a) 'in front of' and 'behind' for positive and negative x-axis directions. 169 | (b) 'to the left' and 'to the right' for positive and negative y-axis directions. 170 | (c) 'above' for positive z-axis directions. 171 | 172 | 173 | You MUST always respond with a json following this format: 174 | { 175 | "tasks": ["task1", "task2", "task3", ...] 176 | } 177 | 178 | Here are some general examples: 179 | 180 | objects = ['coffee pod 1', 'coffee pod 2', 'coffee machine'] 181 | # Query: put the coffee pod into the coffee machine 182 | { 183 | "tasks": ["left robot: move gripper to the coffee pod 1 and avoid collisions with coffe pod 1 and coffee machine. right robot: move gripper to the coffee pod 2 and avoid collisions with coffee pod 2 and coffee machine", "left robot: close_gripper(). right robot: do nothing", "left robot: move the gripper above the coffee machine. right robot: do nothing", "left robot: open_gripper(). right robot: do nothing", "left robot: move the gripper behind the coffee machine. right robot: move the robot above the coffee machine and avoid collisions with the coffee machine. keep the robots at a distance greater than 0.1m". "left robot: do nothing. right robot: open_gripper()"] 184 | } 185 | 186 | objects = ['rope left', 'rope right', 'rope'] 187 | # Query: move the rope 0.1m to the left 188 | { 189 | "tasks": ["left robot: move gripper to rope left and avoid collisions with it. right robot: move gripper to rope left and avoid collisions with it", "left robot: close_gripper(). right robot: close_gripper()", "left robot: move gripper 0.1m to the left. right robot: move gripper 0.1m to the left. keep the distance of the robots equal to their current distance", "left robot: open_gripper(). right robot: open_gripper()"] 190 | } 191 | 192 | objects = ['apple', 'drawer handle', 'drawer'] 193 | # Query: put apple into the drawer. 194 | { 195 | "tasks": ["left robot: move gripper to drawer handle and avoid collisions with handle and drawer. right robot: move gripper to the apple and avoid collisions with apple", "left robot: close_gripper(). right robot: close gripper.", "left robot: move gripper 0.25m in the y direction. right robot: do nothing.", "left robot: do nothing. right robot: move gripper above the drawer.", "left robot: do nothing. right robot: open_gripper()"] 196 | } 197 | """ 198 | 199 | 200 | OD_PROMPT_COLLAB = """ 201 | You are a helpful assistant in charge of designing the optimization problem for an MPC controller that is controlling 2 robot manipulators. 202 | At each step, I will give you a task and you will have to return the objective and (optionally) the constraint functions that need to be applied to the MPC controller. 203 | 204 | This is the scene description: 205 | (1) Casadi is used to program the MPC. 206 | (2) The variables `x_left` and `x_right` represent the position of the gripper in 3D, i.e. (x, y, z) of the 2 robots. 207 | (2) The variables `x0_left` and `x0_right` represent the initial gripper position before any action is applied i.e. (x, y, z) of the 2 robots. 208 | (3) The orientation of the grippers around the z-axis is defined by variables `psi_left`, `psi_right`. 209 | (4) The variable `t` represents the simulation time. 210 | (5) Each time I will also give you a list of objects you can interact with (i.e. objects = ['peach', 'banana']). 211 | (a) The position of each object is an array [x, y, z] obtained by adding `.position` (i.e. 'banana.position'). 212 | (b) The size of each cube is a float obtained by adding '.size' (i.e. 'banana.size'). 213 | (c) The rotaton around the z-axis is a float obtained by adding '.psi' (i.e. 'banana.psi'). 214 | (6) 215 | (a) 'in front of' and 'behind' for positive and negative x-axis directions. 216 | (b) 'to the left' and 'to the right' for positive and negative y-axis directions. 217 | (c) 'above' and 'below' for positive and negative z-axis directions. 218 | 219 | Rules: 220 | (1) A robot MUST avoid colliding with an object if it is moving the gripper specifically to that object. You MUST do this even when not specified in the query. 221 | (2) You MUST write every equality constraints such that it is satisfied if it is = 0: 222 | (a) If you want to write "ca.norm_2(x_left) = 1" write it as "1 - ca.norm_2(x_left)" instead. 223 | (3) You MUST write every inequality constraints such that it is satisfied if it is <= 0: 224 | (a) If you want to write "ca.norm_2(x_right) >= 1" write it as "1 - ca.norm_2(x_right)" instead. 225 | (4) NEVER avoid collisions with an object you're not moving to or nearby if not specified in the query. 226 | 227 | You must format your response into a json. Here are a few examples: 228 | 229 | objects = ['object_1', 'object_2'] 230 | # Query: left robot: move the gripper to [0.2, 0.05, 0.2] and avoid collisions with object_2. right robot: do nothing. 231 | { 232 | "objective": "ca.norm_2(x_left - np.array([0.2, 0.05, 0.2])**2 + ca.norm_2(x_right - x0_right)**2)**2", 233 | "equality_constraints": [], 234 | "inequality_constraints": ["object_2.size - ca.norm_2(x_left - object_2.position)"] 235 | } 236 | Notice how the inequality constraint holds if <= 0. 237 | 238 | objects = ['red_cube', 'yellow_cube'] 239 | # Query: left robot: move the gripper to red cube and avoid colliding with the yellow cube. right robot: move the gripper to yellow cube and avoid colliding with the red cube. 240 | { 241 | "objective": "ca.norm_2(x_left - red_cube.position)**2 + ca.norm_2(x_right - yellow_cube.position)**2", 242 | "equality_constraints": [], 243 | "inequality_constraints": ["red_cube.size - ca.norm_2(x_left - red_cube.position)", "yellow_cube.size - ca.norm_2(x_left - yellow_cube.position)", "yellow_cube.size - ca.norm_2(x_right - yellow_cube.position)", "red_cube.size - ca.norm_2(x_right - red_cube.position)"] 244 | } 245 | Notice the collision avoidance constraint with the red_cube despite not being specified in the query because the left gripper has to go to the red cube. 246 | 247 | objects = ['coffee_pod', 'coffee_machine'] 248 | # Query: left robot: move gripper above the coffe pod. right robot: move gripper above the coffe machine. keep the 2 grippers at a distance greater than 0.1m. 249 | { 250 | "objective": "ca.norm_2(x_left - (coffee_pod.position + np.array([0, 0, coffee_pod.size])))**2 + ca.norm_2(x_right - (coffee_machine.position + np.array([0, 0, coffee_machine.size])))**2", 251 | "equality_constraints": [], 252 | "inequality_constraints": ["coffee_pod.size - ca.norm_2(x_left - coffee_pod.position)", "coffee_machine.size - ca.norm_2(x_right - coffee_machine.position)", "0.1**2 - ca.norm_2(x_left - x_right)**2"] 253 | } 254 | Notice that there's no collision avoidance constraint with the coffee_machine because it is not in the query and because gripper is not moving to or nearby it. 255 | 256 | 257 | objects = ['mug'] 258 | # Query: left robot: Move the gripper 0.1m upwards. right robot: move the gripper 0.1m to the right. 259 | { 260 | "objective": "ca.norm_2(x_left - (x0_left + np.array([0, 0, 0.1])))**2 + ca.norm_2(x_right - (x0_right + np.array([0, -0.1, 0])))**2", 261 | "equality_constraints": [], 262 | "inequality_constraints": [] 263 | } 264 | 265 | objects = ['joystick', 'remote'] 266 | # Query: left robot: Move the gripper at constant speed along the x axis. right robot: Move the gripper at constant speed along the x axis. keep the distance of the robots equal to their current distance. 267 | { 268 | "objective": "ca.norm_2(x_left[0] - t)**2 + ca.norm_2(x_right[0] - t)**2", 269 | "equality_constraints": ["ca.norm_2(x0_left-x0_right)**2 - ca.norm_2(x_left-x_right)**2"], 270 | "inequality_constraints": [] 271 | } 272 | """ 273 | 274 | 275 | PROMPTS = { 276 | "TP": { 277 | "Cubes": TP_PROMPT, 278 | "CleanPlate": TP_PROMPT, 279 | "Sponge": TP_PROMPT_COLLAB, 280 | "CookSteak": TP_PROMPT_COLLAB, 281 | }, 282 | "OD": { 283 | "Cubes": OD_PROMPT, 284 | "CleanPlate": OD_PROMPT, 285 | "Sponge": OD_PROMPT_COLLAB, 286 | "CookSteak": OD_PROMPT_COLLAB, 287 | } 288 | } 289 | -------------------------------------------------------------------------------- /simulation_http.py: -------------------------------------------------------------------------------- 1 | import os 2 | import io 3 | import sys 4 | import cv2 5 | import gym 6 | import json 7 | import base64 8 | import asyncio 9 | import requests 10 | import panda_gym 11 | import numpy as np 12 | from PIL import Image 13 | from tqdm import tqdm 14 | from aiohttp import web 15 | from datetime import datetime 16 | from sqlalchemy import create_engine 17 | from sqlalchemy.orm import sessionmaker 18 | 19 | 20 | from robot import Robot 21 | from db import Base, Episode, Epoch 22 | from core import AbstractSimulation, BASE_DIR 23 | from config.config import SimulationConfig, RobotConfig 24 | 25 | 26 | class Simulation(AbstractSimulation): 27 | def __init__(self, cfg=SimulationConfig()) -> None: 28 | #super().__init__(cfg) 29 | 30 | self.cfg = cfg 31 | # init env 32 | self.env = gym.make(f"Panda{cfg.task}-v2", render=cfg.render, debug=cfg.debug) 33 | # init robots 34 | # count number of tasks solved from a plan 35 | self.plan = None 36 | self.optimizations = [] 37 | self.task_counter = 0 38 | self.prev_instruction = "None" 39 | 40 | # simulation time 41 | self.t = 0. 42 | self.epoch = 0 43 | env_info = (self.env.robots_info, self.env.objects_info) 44 | self.robot = Robot(env_info,RobotConfig(self.cfg.task)) 45 | # count number of tasks solved from a plan 46 | self.task_counter = 0 47 | # bool for stopping simulation 48 | self.stop_thread = False 49 | # whether to save frame (initialized to false) 50 | self.save_video = self.cfg.save_video 51 | # init list of RGB frames if wanna save video 52 | self.frames_list = [] 53 | self.frames_list_logging = [] 54 | # 55 | self.state_trajectories = [] 56 | self.mpc_solve_times = [] 57 | self.session = None 58 | 59 | self.video_path = os.path.join(BASE_DIR, f"data/videos/{self.cfg.task}_{datetime.now()}_full.mp4") 60 | 61 | if self.cfg.logging: 62 | engine = create_engine(f'sqlite:///data/{self.cfg.method}/DBs/{cfg.task}.db') 63 | Base.metadata.create_all(engine) 64 | self.Session = sessionmaker(bind=engine) 65 | 66 | 67 | def _round_list(self, l, n=2): 68 | """ round list and if the result is -0.0 convert it to 0.0 """ 69 | return [r if (r:=round(x, n)) != -0.0 else 0.0 for x in l] 70 | 71 | def _append_robot_info(self): 72 | for r in self.env.robots_info: 73 | obs = self.observation[f'robot{r["name"]}'] # observation of each robot 74 | self.state_trajectories.append(obs.tolist()) 75 | self.mpc_solve_times.append(self.robot.MPC.solve_time) 76 | 77 | def _create_scene_description(self): 78 | """ Look at the observation and create a string that describes the scene to be passed to the task planner """ 79 | ri = 0 80 | description = "The following is the description of the current scene:\n" 81 | for name in self.observation.keys(): 82 | if name.startswith("robot"): 83 | robot_xyz = self._round_list(self.observation[name][:3]) 84 | description += f"- The gripper of the {name} is located at {robot_xyz}.\n" 85 | if self.robot.gripper==-1: 86 | if round(self.env.robots[ri].get_fingers_width(),2) <= 0.01: 87 | description += f"- The gripper fingers have closed but they are grasping no object.\n" 88 | else: 89 | distances = {cube_name: np.linalg.norm(np.array(robot_xyz)-np.array(self.observation[cube_name])) for cube_name in self.observation.keys() if cube_name.endswith("_cube")} 90 | closest_cube = min(distances, key=distances.get)[:-5] 91 | description += f"- The gripper fingers are closed and they are firmly grasping the {closest_cube} cube.\n" 92 | else: 93 | description += f"- The gripper fingers are open.\n" 94 | ri += 1 95 | elif name.endswith("_cube"): 96 | description += f"- The center of the {name[:-5]} cube is located at {self._round_list(self.observation[name])}\n" 97 | else: 98 | pass 99 | 100 | description += """Please carefully analyze the scene description and decide what to do next. Some helpful tips are: 101 | (1) If the gripper is not at the location where it should be it is surely because of collisions. Specify in your instruction to the robot about collision avoidance and constraints. 102 | (2) Be careful when placing a cube on top of another one that you leave some clearance between those 2 cubes. Be very careful and explain how much space should be left between. 103 | (a) It's ok if the cube is not at the same exact x and y position as the cube below. 104 | (3) Make sure that the cube you've put on the stack has not fallen. Always check every cube to understand if it is on the ground or on top of another cube. 105 | (a) A cube is on the ground if it's height is 0.02m. 106 | (b) If you stacked a cube and need to go to another one, make sure to instruct the robot to avoid collisions with the cubes in the stack. 107 | (4) The description of scene is ALWAYS correct, the instructions you give may be wrong or at times mis-interepted by the robot. Always try to fix this when it happens. 108 | (5) Make sure the gripper is open before it needs to go to an object to grasp it. 109 | """ 110 | 111 | return description 112 | 113 | def _uplaod_image(self, rgba_image:np.ndarray) -> str: 114 | # Convert the NumPy array to a PIL Image object 115 | image = Image.fromarray(rgba_image, 'RGBA') 116 | image_path = f"{self.episode_folder}/{datetime.now().strftime('%d-%m-%Y_%H:%M:%S')}.png" # Specify your local file path here 117 | image.save(image_path, 'PNG') 118 | return image_path 119 | 120 | def _retrieve_image(self) -> np.ndarray: 121 | frame_np = np.array(self.env.render("rgb_array", 122 | width=self.cfg.frame_width, height=self.cfg.frame_height, 123 | target_position=self.cfg.frame_target_position, 124 | distance=self.cfg.frame_distance, 125 | yaw=self.cfg.frame_yaw, 126 | pitch=self.cfg.frame_pitch)) 127 | frame_np = frame_np.reshape(self.cfg.frame_width, self.cfg.frame_height, 4).astype(np.uint8) 128 | 129 | return frame_np 130 | 131 | def _store_epoch_db(self, episode_id, role, content, image_url): 132 | session = self.Session() 133 | 134 | # Find the last epoch number for this episode 135 | last_epoch = session.query(Epoch).filter_by(episode_id=episode_id).order_by(Epoch.time_step.desc()).first() 136 | if last_epoch is None: 137 | next_time_step = 1 # This is the first epoch for the episode 138 | else: 139 | next_time_step = last_epoch.time_step + 1 140 | 141 | # Create and insert the new epoch 142 | epoch = Epoch(episode_id=episode_id, time_step=next_time_step, role=role, content=content, image=image_url) 143 | session.add(epoch) 144 | session.commit() 145 | session.close() 146 | 147 | def _make_plan(self, user_message:str="") -> str: 148 | self.plan:dict = self.robot.plan_task(user_message) 149 | self.task_counter = 0 150 | pretty_msg = "Tasks:\n" 151 | pretty_msg += "".join([f"{i+1}. {task}\n" for i, task in enumerate(self.plan["tasks"])]) 152 | if self.cfg.logging: 153 | image = self._retrieve_image() 154 | image_url = self._uplaod_image(image) 155 | self._store_epoch_db(self.episode.id, "human", self.robot._get_instruction(user_message), image_url) 156 | self._store_epoch_db(self.episode.id, "TP", pretty_msg, image_url) 157 | return pretty_msg 158 | 159 | def _solve_task(self, task:str, optimization:dict=None) -> str: 160 | AI_response = self.robot.solve_task(task, optimization) if task != "finished" else task 161 | if self.cfg.logging and AI_response is not None: 162 | image = self._retrieve_image() 163 | image_url = self._uplaod_image(image) 164 | self._store_epoch_db(self.episode.id, "OD", AI_response, image_url) 165 | return AI_response 166 | 167 | def reset(self): 168 | # reset pand env 169 | self.observation = self.env.reset() 170 | # reset robot 171 | self.robot.reset() 172 | # reset controller 173 | self.robot.init_states(self.observation, self.t) 174 | # reset task counter 175 | self.plan = None 176 | self.optimizations = [] 177 | self.task_counter = 0 178 | if self.cfg.logging: 179 | if self.session is not None: 180 | self.episode.state_trajectories = json.dumps(self.state_trajectories) 181 | self.episode.mpc_solve_times = json.dumps(self.mpc_solve_times) 182 | if self.cfg.logging_video: 183 | self._save_video() 184 | self.session.commit() 185 | self.state_trajectories = [] 186 | self.mpc_solve_times = [] 187 | self.session.close() 188 | self.session = self.Session() 189 | self.episode = Episode() # Assuming Episode has other fields you might set 190 | self.session.add(self.episode) 191 | self.session.commit() 192 | n_episodes = len(os.listdir(f"data/{self.cfg.method}/images")) 193 | self.episode_folder = f"data/{self.cfg.method}/images/{n_episodes}" 194 | os.mkdir(self.episode_folder) 195 | self.video_path_logging = os.path.join(BASE_DIR, f"data/{self.cfg.method}/videos/{self.cfg.task}_{n_episodes}.mp4") 196 | self.video_path = os.path.join(BASE_DIR, f"data/videos/{self.cfg.task}_{datetime.now()}_full.mp4") 197 | # init list of RGB frames if wanna save video 198 | if self.save_video: 199 | self._save_video() 200 | self.frames_list = [] 201 | self.frames_list_logging = [] 202 | self.t = 0. 203 | self.epoch = 0 204 | 205 | 206 | def step(self): 207 | # increase timestep 208 | self.t += self.cfg.dt 209 | self.epoch += 1 210 | # update controller (i.e. set the current gripper position) 211 | self.robot.init_states(self.observation, self.t) 212 | # compute action 213 | action = self.robot.step() 214 | # apply action 215 | self.observation, _, done, _ = self.env.step(action) 216 | # add states to state_trajectories 217 | if self.cfg.logging: 218 | self._append_robot_info() 219 | if self.cfg.logging_video and self.epoch%20 == 0: 220 | self.frames_list_logging.append(self._retrieve_image()) 221 | # visualize trajectory 222 | if self.cfg.debug: 223 | trajectory = self.robot.retrieve_trajectory() 224 | self.env.visualize_trajectory(trajectory) 225 | # store RGB frames if wanna save video 226 | if self.save_video: 227 | self.frames_list.append(self._retrieve_image()) 228 | 229 | return done 230 | 231 | def close(self): 232 | # close environment 233 | #self.thread.join() 234 | self.stop_thread = True 235 | # init list of RGB frames if wanna save video 236 | if self.save_video: 237 | self._save_video() 238 | 239 | if self.cfg.logging: 240 | if self.cfg.logging_video: 241 | self._save_video() 242 | self.episode.state_trajectories = json.dumps(self.state_trajectories) 243 | self.episode.mpc_solve_times = json.dumps(self.mpc_solve_times) 244 | self.session.commit() 245 | self.session.close() 246 | # exit 247 | #sys.exit() 248 | 249 | def _save_video(self): 250 | for s, p, l in [(self.save_video, self.video_path, self.frames_list), (self.cfg.logging_video, self.video_path_logging, self.frames_list_logging)]: 251 | if not s: continue 252 | # Define the parameters 253 | fourcc = cv2.VideoWriter_fourcc(*'mp4v') 254 | # Create a VideoWriter object 255 | out = cv2.VideoWriter(p, fourcc, self.cfg.fps, (self.cfg.frame_width, self.cfg.frame_height)) 256 | # Write frames to the video 257 | for frame in tqdm(l): 258 | # Ensure the frame is in the correct format (RGBA) 259 | if frame.shape[2] == 3: 260 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2RGBA) 261 | # Convert the frame to BGR format (required by VideoWriter) 262 | frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR) 263 | out.write(frame_bgr) 264 | # Release the VideoWriter 265 | out.release() 266 | 267 | async def http_close(self, request): 268 | self.close() 269 | return web.json_response({"content": "Simulation closed"}) 270 | 271 | async def http_reset(self, request): 272 | self.reset() 273 | return web.json_response({"content": "Simulation reset"}) 274 | 275 | async def http_solve_task(self, request): 276 | data = await request.json() 277 | user_task = data.get('content') 278 | AI_response = self._solve_task(user_task) 279 | return web.json_response([{"type": "OD", "content": AI_response}]) 280 | 281 | async def http_make_plan(self, request): 282 | data = await request.json() 283 | user_message = data.get('content') 284 | pretty_msg = self._make_plan(user_message) 285 | return web.json_response([{"type": "TP", "content": pretty_msg}]) 286 | 287 | async def http_next_task(self, request): 288 | is_plan_unfinished = self.task_counter < len(self.plan["tasks"]) 289 | task = self.plan["tasks"][self.task_counter] if is_plan_unfinished else "finished" 290 | optimization = self.optimizations[self.task_counter] if (self.optimizations and is_plan_unfinished) else None 291 | AI_response = self._solve_task(task, optimization) 292 | if AI_response is not None: self.task_counter += 1 293 | return web.json_response([{"type": "OD", "content": AI_response}]) 294 | 295 | async def http_upload_plan(self, request): 296 | data = await request.json() 297 | query = data['query'] 298 | plan = data['plan'] 299 | self.task_counter = 0 300 | self.plan = plan 301 | pretty_msg = "Tasks:\n" 302 | pretty_msg += "".join([f"{i+1}. {task}\n" for i, task in enumerate(self.plan["tasks"])]) 303 | if self.cfg.logging: 304 | image = self._retrieve_image() 305 | image_url = self._uplaod_image(image) 306 | self._store_epoch_db(self.episode.id, "human", query, image_url) 307 | self._store_epoch_db(self.episode.id, "TP", pretty_msg, image_url) 308 | return web.json_response({"response": "Plan uploaded"}) 309 | 310 | async def http_upload_optimizations(self, request): 311 | data = await request.json() 312 | self.task_counter = 0 313 | self.optimizations = data.get('content') 314 | return web.json_response({"response": "Optimizations uploaded"}) 315 | 316 | async def http_save_recording(self, request): 317 | self._save_video() 318 | self.save_video = False 319 | return web.json_response({"response": "Recording saved"}) 320 | 321 | async def http_start_recording(self, request): 322 | self.save_video = True 323 | return web.json_response({"response": "Recording started"}) 324 | 325 | async def http_cancel_recording(self, request): 326 | self.save_video = False 327 | self.frames_list = [] 328 | return web.json_response({"response": "Recording cancelled"}) 329 | 330 | async def main(self, app): 331 | runner = web.AppRunner(app) 332 | await runner.setup() 333 | site = web.TCPSite(runner, 'localhost', 8080) 334 | await site.start() 335 | 336 | await self._run() 337 | 338 | async def _run(self): 339 | while not self.stop_thread: 340 | done = self.step() 341 | await asyncio.sleep(0.05) 342 | if done: 343 | break 344 | self.env.close() 345 | 346 | def run(self): 347 | app = web.Application() 348 | app.add_routes([ 349 | web.post('/make_plan', self.http_make_plan), 350 | web.post('/solve_task', self.http_solve_task), 351 | web.post('/upload_plan', self.http_upload_plan), 352 | web.post('/upload_optimizations', self.http_upload_optimizations), 353 | web.get('/close', self.http_close), 354 | web.get('/reset', self.http_reset), 355 | web.get('/next_task', self.http_next_task), 356 | web.get('/save_recording', self.http_save_recording), 357 | web.get('/start_recording', self.http_start_recording), 358 | web.get('/cancel_recording', self.http_cancel_recording) 359 | ]) 360 | 361 | asyncio.run(self.main(app)) 362 | 363 | 364 | if __name__=="__main__": 365 | s = Simulation() 366 | s.reset() 367 | s.run() --------------------------------------------------------------------------------