Create a task dataset using gpt

legent/dataset/controller.py

@@ -301,6 +301,8 @@ def parse_arg(input_string):
             elif solu.startswith("speak("):
                 actions = Speak(parse_arg(solu))
                 self.actions_queue.append(actions)
+            else:
+                raise
 
         self.actions = self.actions_queue.pop(0)
         self.actions.init_actions(self.env)

legent/dataset/task.py

@@ -1,25 +1,29 @@
-from legent.server.scene_generator import generate_scene, prefabs
-from legent.utils.config import TASKS_FOLDER
-from legent.utils.io import store_json, load_json_from_toolkit, time_string, scene_string, log_green, log
-from legent.utils.math import is_point_on_box
 import time
 import os
-from typing import Literal
 import random
-import numpy as np
 import re
+import time
+import uuid
+from typing import Literal
+
+import numpy as np
+
+from legent.server.scene_generator import generate_scene, prefabs
+from legent.utils.config import TASKS_FOLDER, OPENAI_API_KEY, OPENAI_BASE_URL, MODEL_CHAT
+from legent.utils.io import store_json, load_json_from_toolkit, time_string, scene_string, log_green, log
+from legent.utils.math import is_point_on_box
 
 
 class ChatBase:
-    def __init__(self, api_key=None, base_url=None) -> None:
+    def __init__(self, api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL) -> None:
         import openai
 
         if api_key:
             self.client = openai.OpenAI(api_key=api_key, base_url=base_url)
 
     def send_chat(self, messages):
         response = self.client.chat.completions.create(
-            model="gpt-4",  # 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-4', 'gpt-4-32k'
+            model=MODEL_CHAT,  # 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-4', 'gpt-4-32k'
             messages=messages,
             max_tokens=None,
             n=1,
@@ -150,21 +154,41 @@ def create_task_for_scene_by_prompting(self, task_type=Literal["come", "goto", "
         examples = [f"Task: {e['example']}; Plan: {e['plan']}; Solution: {e['solution']}" for e in task_prompt["examples"]]
         task_prompt["example"] = "\n".join(examples)
         # TODO: Add descriptions for functions goto_user(), goto(object_id), grab(), release().
-        task_description = f"""You need to {task_prompt['message']}
-You need to propose {sample_num} independent tasks and corresponding solutions, in the format of "Task: task; Plan: plan; Solution: solution.", with no line breaks between the three (use ';' to seperate). One sentence in Plan should match one function call in Solution.
+        task_description = \
+            f"""You need to {task_prompt['message']}
+You need to propose {sample_num} independent tasks and corresponding solutions, in the format of "Task: task; Plan: plan; Solution: solution.", with no line breaks between the three (use ';' to seperate). 
+One sentence in Plan should match one function call in Solution, and the Solution should contain only the following instructions and no other irrelevant output:
+1. goto_user()
+2. goto(object_id)
+3. find(object_id)
+4. grab()
+5. release()
+6. speak(text)
 For example (The examples are from other scenes. The number means object_id):
 {task_prompt['example']}
 """
         content = f"{scene_description}\n{task_description}"
         messages = [
-            {"role": "system", "content": system_message},
-            {"role": "assistant", "content": task_prompt["example"]},  # this message can ensure the format correct
-            {"role": "user", "content": content},
+            {
+                "role": "user",
+                "content": content
+            },
+            {
+                "role": "system",
+                "content": system_message
+            },
+            {  # this message can ensure the format correct
+                "role": "assistant",
+                "content": task_prompt['example']
+            },
         ]
+        id = uuid.uuid4()
+        log_green(f"\nid:{id} start send chat...")
 
         ret = self.send_chat(messages)
 
         log_green(f"<g>Send to ChatGPT<g/>:\n{content}\n<g>Received from ChatGPT<g/>:\n{ret}")
+        log_green(f"\nid:{id} end send chat...")
 
         task_lines = [task for task in ret.split("\n") if task]
         samples = []

legent/utils/config.py

@@ -17,6 +17,16 @@
 DATASET_FOLDER = f"{resource_path}/dataset"
 MODEL_FOLDER = f"{resource_path}/models"
 EVAL_FOLDER = f"{resource_path}/eval"
+
+# Default number of scene rooms
+ROOM_NUM = None
+
+# OpenAI
+OPENAI_API_KEY = None
+OPENAI_BASE_URL = 'https://api.keya.pw/v1'
+MODEL_CHAT = 'claude-3-haiku-20240307' # 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-4', 'gpt-4-32k','chatglm3-6b','claude-3-haiku-20240307'
+MODEL_VISION_PREVIEW = None # 'gpt-4-vision-preview'
+
 PACKED_FOLDER = f"{resource_path}/packed_scenes"
 
 

scripts/create_traj_prompt.py

@@ -0,0 +1,44 @@
+import random
+
+from legent import Environment, ResetInfo, TaskCreator, Controller, TrajectorySaver
+
+env = Environment(env_path="auto", use_animation=False, camera_resolution=448, camera_field_of_view=120,
+                  run_options={"port": 50099})  # significantly increase the sampling rate without using animations
+
+scene_num = 1000
+
+try:
+    saver = TrajectorySaver()
+    for i in range(scene_num):
+        while True:
+            try:
+                # The probability of occurrence of 1 is 1, the probability of 1 is 1, the probability of 2 is 2, the probability of 2 is 3.
+                room_num = random.choices([1, 2], weights=[1, 2])[0]
+
+                actions = ["come", "goto", "take", "bring", "put", "exist", "where"]
+                actions_weights = [1, 1, 1, 1, 1, 1, 1]  # Corresponding to the weight of each action
+                chosen_action = random.choices(actions, weights=actions_weights, k=1)[0]
+
+                task = TaskCreator().create_task_for_scene_by_prompting(task_type=chosen_action, sample_num=1)[0]
+
+                break
+            except Exception as e:
+                print("Task Exeption", e)
+                pass
+
+        try:
+            env.reset(ResetInfo(scene=task['scene']))
+            controller = Controller(env, task['solution'])
+            traj = controller.collect_trajectory(task)
+            if traj:
+                # The task has been completed successfully
+                saver.save_traj(traj=traj)
+                print(f'Complete task "{task["task"]}" in {traj.steps} steps.')
+            else:
+                print(f'Complete task "{task["task"]}" failed. Deserted.')
+        except Exception as e:
+            print("controller Exeption", e)
+            pass
+    saver.save()
+finally:
+    env.close()