chenbao-dev

app.py

@@ -3,13 +3,13 @@ import numpy as np
 from tempfile import NamedTemporaryFile
 import copy
 import shapely
+from hydra.core.global_hydra import GlobalHydra
 from shapely.geometry import *
 from shapely.affinity import *
 from omegaconf import OmegaConf
 from moviepy.editor import ImageSequenceClip
 import gradio as gr
 
-
 from consts import ALL_BLOCKS, ALL_BOWLS
 from md_logger import MarkdownLogger
 
@@ -37,12 +37,17 @@ from gensim.sim_runner import SimulationRunner
 from gensim.memory import Memory
 from gensim.utils import set_gpt_model, clear_messages
 
+
 class DemoRunner:
-        
+
     def __init__(self):
         self._env = None
+        GlobalHydra.instance().clear()
+        hydra.initialize(version_base="1.2", config_path='cliport/cfg')
+        self._cfg = hydra.compose(config_name="data")
 
     def setup(self, api_key):
+        cfg = self._cfg
         openai.api_key = api_key
         cfg['model_output_dir'] = 'temp'
         cfg['prompt_folder'] = 'topdown_task_generation_prompt_simple_singleprompt'
@@ -56,48 +61,43 @@ class DemoRunner:
         self.simulation_runner = SimulationRunner(cfg, agent, critic, memory)
 
         info = '### Build'
-        img = np.zeros((720, 640, 0))
+        img = np.zeros((720, 640, 3))
 
         return info, img
 
     def run(self, instruction):
+        cfg = self._cfg
         cfg['target_task_name'] = instruction
 
-        self._env.cache_video = []
-        self._md_logger.clear()
-
-        try:
-            self.simulation_runner.task_creation()
-            self.simulation_runner.simulate_task()
-        except Exception as e:
-            return f'Error: {e}', None, None
-
-        video_file_name = None
-        if self._env.cache_video:
-            rendered_clip = ImageSequenceClip(self._env.cache_video, fps=25)
-            video_file_name = NamedTemporaryFile(suffix='.mp4').name
-            rendered_clip.write_videofile(video_file_name, fps=25)
+        # self._env.cache_video = []
+        self.simulation_runner._md_logger = ''
 
-        return self.simulation_runner.chat_log, self.simulation_runner.env.curr_video, video_file_name
+        self.simulation_runner.task_creation()
+        self.simulation_runner.simulate_task()
+        print("self.video_path = ", self.simulation_runner.video_path)
+        return self.simulation_runner._md_logger, self.simulation_runner.video_path
 
 
 def setup(api_key):
     if not api_key:
         return 'Please enter your OpenAI API key!', None, None
-    
+
     demo_runner = DemoRunner()
-    
+
     info, img = demo_runner.setup(api_key)
     return info, img, demo_runner
 
 
 def run(instruction, demo_runner):
     if demo_runner is None:
-        return 'Please run setup first!', None, None
+        return 'Please run setup first!', None
+
+    # return None, "/home/baochen/Desktop/projects/GenSim2/data/assemble-pallet-ball-train/videos/000001.mp4"
     return demo_runner.run(instruction)
 
 
 if __name__ == '__main__':
+    os.environ['GENSIM_ROOT'] = os.getcwd()
     with open('README.md', 'r') as f:
         for _ in range(12):
             next(f)
@@ -120,22 +120,21 @@ if __name__ == '__main__':
 
         with gr.Row():
             with gr.Column():
-                
                 inp_instruction = gr.Textbox(label='Task Name', lines=1)
                 btn_run = gr.Button("Run (this may take 30+ seconds)")
                 info_run = gr.Markdown(label='Generated Code')
             with gr.Column():
                 video_run = gr.Video(label='Video of Last Instruction')
-        
+
         btn_setup.click(
-            setup, 
+            setup,
             inputs=[inp_api_key],
             outputs=[info_setup, img_setup, state]
         )
         btn_run.click(
-            run, 
-            inputs=[inp_instruction, state], 
-            outputs=[info_run, img_setup, video_run]
+            run,
+            inputs=[inp_instruction, state],
+            outputs=[info_run, video_run]
         )
-        
+
     demo.queue().launch(show_error=True)

cliport/cfg/config.yaml

@@ -5,7 +5,7 @@ tag: default
 debug: False
 gpt_temperature: 0.8 # GPT-4 response temperature. higher means more diversity
 prompt_folder: vanilla_task_generation_prompt # the prompt folder that stores the prompt chain
-max_env_run_cnt: 3 # maximum number of runs for each environment
+max_env_run_cnt: 1 # maximum number of runs for each environment
 trials: 10 # how many times of spawning each environment generated
 output_folder: 'output/output_stats'
 model_output_dir: '' # to be filled in with date

cliport/cfg/data.yaml

@@ -25,7 +25,7 @@ dataset:
 
 # record videos (super slow)
 record:
-  save_video: False
+  save_video: True
   save_video_path: ${data_dir}/${task}-${mode}/videos/
   add_text: False
   add_task_text: True

gensim/sim_runner.py

@@ -22,8 +22,10 @@ from gensim.utils import (
 )
 import pybullet as p
 
+
 class SimulationRunner:
     """ the main class that runs simulation loop """
+
     def __init__(self, cfg, agent, critic, memory):
         self.cfg = cfg
         self.agent = agent
@@ -46,19 +48,25 @@ class SimulationRunner:
         self.generated_task_programs = []
         self.generated_task_names = []
         self.generated_tasks = []
-        self.passed_tasks = [] # accepted ones
+        self.passed_tasks = []  # accepted ones
+
+        self.video_path = None
+        self._md_logger = ''
 
     def print_current_stats(self):
         """ print the current statistics of the simulation design """
         print("=========================================================")
-        print(f"{self.cfg['prompt_folder']} Trial {self.curr_trials} SYNTAX_PASS_RATE: {(self.syntax_pass_rate / (self.curr_trials)) * 100:.1f}% RUNTIME_PASS_RATE: {(self.runtime_pass_rate / (self.curr_trials)) * 100:.1f}% ENV_PASS_RATE: {(self.env_pass_rate / (self.curr_trials)) * 100:.1f}%")
+        print(
+            f"{self.cfg['prompt_folder']} Trial {self.curr_trials} SYNTAX_PASS_RATE: {(self.syntax_pass_rate / (self.curr_trials)) * 100:.1f}% RUNTIME_PASS_RATE: {(self.runtime_pass_rate / (self.curr_trials)) * 100:.1f}% ENV_PASS_RATE: {(self.env_pass_rate / (self.curr_trials)) * 100:.1f}%")
         print("=========================================================")
 
     def save_stats(self):
         """ save the final simulation statistics """
         self.diversity_score = compute_diversity_score_from_assets(self.task_asset_logs, self.curr_trials)
-        save_stat(self.cfg, self.cfg['model_output_dir'], self.generated_tasks, self.syntax_pass_rate / (self.curr_trials),
-                self.runtime_pass_rate / (self.curr_trials), self.env_pass_rate / (self.curr_trials), self.diversity_score)
+        save_stat(self.cfg, self.cfg['model_output_dir'], self.generated_tasks,
+                  self.syntax_pass_rate / (self.curr_trials),
+                  self.runtime_pass_rate / (self.curr_trials), self.env_pass_rate / (self.curr_trials),
+                  self.diversity_score)
         print("Model Folder: ", self.cfg['model_output_dir'])
         print(f"Total {len(self.generated_tasks)} New Tasks:", [task['task-name'] for task in self.generated_tasks])
         try:
@@ -66,7 +74,6 @@ class SimulationRunner:
         except:
             pass
 
-
     def task_creation(self):
         """ create the task through interactions of agent and critic """
         self.task_creation_pass = True
@@ -93,16 +100,15 @@ class SimulationRunner:
         # self.curr_task_name = self.generated_task['task-name']
         print("task creation time {:.3f}".format(time.time() - start_time))
 
-
     def setup_env(self):
         """ build the new task"""
         env = Environment(
-                self.cfg['assets_root'],
-                disp=self.cfg['disp'],
-                shared_memory=self.cfg['shared_memory'],
-                hz=480,
-                record_cfg=self.cfg['record']
-            )
+            self.cfg['assets_root'],
+            disp=self.cfg['disp'],
+            shared_memory=self.cfg['shared_memory'],
+            hz=480,
+            record_cfg=self.cfg['record']
+        )
 
         task = eval(self.curr_task_name)()
         task.mode = self.cfg['mode']
@@ -118,15 +124,15 @@ class SimulationRunner:
         print(f"Mode: {task.mode}")
 
         # Start video recording
-        if record:
-            env.start_rec(f'{dataset.n_episodes+1:06d}')
+        # if record:
+        #     env.start_rec(f'{dataset.n_episodes+1:06d}')
 
         return task, dataset, env, expert
 
     def run_one_episode(self, dataset, expert, env, task, episode, seed):
         """ run the new task for one episode """
         add_to_txt(
-                self.chat_log, f"================= TRIAL: {self.curr_trials}", with_print=True)
+            self.chat_log, f"================= TRIAL: {self.curr_trials}", with_print=True)
         record = self.cfg['record']['save_video']
         np.random.seed(seed)
         random.seed(seed)
@@ -138,6 +144,12 @@ class SimulationRunner:
         reward = 0
         total_reward = 0
 
+        save_data = self.cfg['save_data']
+        # Start recording video (NOTE: super slow)
+        if record:
+            video_name = f'{dataset.n_episodes + 1:06d}'
+            env.start_rec(video_name)
+
         # Rollout expert policy
         for _ in range(task.max_steps):
             act = expert.act(obs, info)
@@ -149,6 +161,12 @@ class SimulationRunner:
             if done:
                 break
 
+        # End recording video
+        if record:
+            env.end_rec()
+            self.video_path = os.path.join(self.cfg['record']['save_video_path'],
+                                           f"{video_name}.mp4")
+
         episode.append((obs, None, reward, info))
         return total_reward
 
@@ -159,10 +177,10 @@ class SimulationRunner:
         env_success_cnt = 0.
         seed = 123
         self.curr_trials += 1
-        
+
         if p.isConnected():
             p.disconnect()
-        
+
         if not self.task_creation_pass:
             print("task creation failure => count as syntax exceptions.")
             return
@@ -179,6 +197,7 @@ class SimulationRunner:
             save_text(self.cfg['model_output_dir'], self.generated_task_name + '_error', str(traceback.format_exc()))
             print("========================================================")
             print("Syntax Exception:", to_print)
+            self._md_logger = str(traceback.format_exc())
             return
 
         try:
@@ -201,4 +220,6 @@ class SimulationRunner:
             save_text(self.cfg['model_output_dir'], self.generated_task_name + '_error', str(traceback.format_exc()))
             print("========================================================")
             print("Runtime Exception:", to_print)
+            self._md_logger = str(traceback.format_exc())
+
         self.memory.save_run(self.generated_task)